perf: first-token latency is dominated by lazy weight materialization
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Context
Measured under #1545 on a Tesla V100 (sm_70), but the finding is general and not architecture specific, which is why it is filed outside epic #1536. The full record is `docs/benchmark_results/volta-ttft-fixed-cost-2026-09-01.md`.
On `qwen3.8-27B-4bit` at a 2-token prompt, with a warm PTX cache and a warm file page cache, `mlxcel generate --profile` reports a first token 15.54 s after the model is reported loaded (10 repetitions, 6.5% spread). Running the same prefill a second time in the same process costs 0.546 s, so 14.99 s of it (96.5%) is one-time process cost. Of that one-time cost, 12.08 s (77.8% of the whole first token) is materializing the language model's 15.13 GB of weights.
MLX's `load_safetensors()` returns unevaluated `Load` arrays, so the CLI's `Model loaded in 0.98s` line is followed by `resident: 0.00 GB` (printed at `src/commands/generate.rs:166`) and the whole read plus host-to-device copy is charged to prefill. `Load::eval_gpu` in `mlx/backend/cuda/load.cpp` does, per tensor: one device allocation, one `malloc` of the full tensor, one synchronous host read into it, one pageable `cudaMemcpyAsync`, and a `cudaLaunchHostFunc` to free the staging buffer. The count is structural rather than approximate: the checkpoint's index holds 1,847 `language_model.*` tensors and every profile of a text-only run shows exactly 1,847 `cudaMemcpyAsync` calls, flat in the number of generated tokens.
The measured legs, from #1545: host staging (`malloc`, read, `free`) 7.40 s at 2.04 GB/s, measured by a standalone microbenchmark that replicates the host path and touches no CUDA; host-to-device copy 4.19 s at 3.61 GB/s, which is the pageable-memory rate over PCIe Gen3 x16; device allocation 0.49 s over 2,255 `cudaMallocAsync` calls.
## Why this is worth an issue
Three separable candidates, each with a measured size attached:
1. **The host-to-device leg runs at the pageable rate.** Staging through pinned memory, or copying straight from the mmap so no staging buffer exists at all, would attack the 4.19 s leg. This is a change to `Load::eval_gpu`, which mlxcel already patches under `src/lib/mlx-cpp/patches/` (the CUDA backend patch directory already carries `jit_module.cpp` and `matmul.cpp`).
2. **The 7.40 s host staging leg is a `malloc` plus a page-cache read plus a `free` that unmaps 15 GB.** Copying from the mmap would remove most of it. Note that `free` alone is 0.63 s.
3. **Which phase is charged is a loader property, not a device property.** `load_gemma4_family_weights_with_backing` in `src/models/sanitize.rs` (defined at `src/models/sanitize.rs:1442`) ends with `eval_all` then `detach_all` (`src/models/sanitize.rs:1484-1485`), so a Gemma 4 checkpoint reports `resident: 6.28 GB` after load and a 2.4 s first token, while the Qwen 3.5 loader leaves everything lazy. Making that a deliberate, uniform decision would at minimum stop `--profile` misreporting 12 s of weight loading as prefill, and would move a server's cost from its first request to its startup.
Relabelling is not obviously free. Per byte of language model, the eager Gemma path measured 0.60 to 0.76 GB/s against the lazy Qwen path's 1.02 GB/s, so it needs its own measurement, and eager materialization raises peak memory for loaders that repack.
## Scope
**In scope:** the three candidates above, each measured before and after: the host-to-device leg in `Load::eval_gpu`, the host staging leg in `Load::eval_gpu`, and the eager-versus-lazy materialization decision in mlxcel's loaders (`src/lib/mlxcel-core/src/weights.rs`, `src/models/sanitize.rs`).
**Out of scope:** CUDA graph construction, JIT module loading, and graph cache capacity. #1545 measured all three and none is where the time is: graph instantiation is 0.10 s and saturates at 196 distinct graphs against a 2,000-entry cache, `MLX_USE_CUDA_GRAPHS=0` changes the first token by less than the repeat spread, and warm `cuModuleLoadDataEx` is 0.06 s over 6 calls.
## Implementation notes
- **Reuse**: the eager precedent is already in tree at `src/models/sanitize.rs:1484-1485`. Any uniform policy should go through the existing `mlxcel_core::eval_all` / `detach_all` pair rather than a second mechanism, and should live at the `load_weights_from_dir` layer (`src/lib/mlxcel-core/src/weights.rs:310`) so every loader inherits it instead of each family deciding for itself.
- **Reuse**: `MLXCEL_PROFILE_TTFT`, the pre-first-token instrumentation #1545 added, already breaks the pre-first-token phase into `setup`, `build`, `sample`, `eval` and `post`. Extend it rather than adding a parallel timer.
- **Constraints**: `Load::eval_gpu` changes belong under `src/lib/mlx-cpp/patches/mlx/backend/cuda/` following the existing patch convention, not as edits to the vendored MLX checkout.
- **Constraints**: eager materialization raises peak memory for loaders that repack, since the source arrays and the repacked results are live at the same time. Any eager change must report peak alongside wall clock.
- **Edge cases**: multimodal checkpoints hold tensors a text-only run never touches (`qwen3.8-27B-4bit` has 2,180 tensors, 1,847 under `language_model.` and 333 under `vision_tower.`), so an eager policy must not force the vision tower resident on a text-only run.
- **Edge cases**: a cold page cache turns the first token into a disk benchmark, so every arm must state its page-cache state or the comparison is meaningless.
- **Error handling**: pinned-memory allocation can fail on a host under memory pressure. Fall back to the current pageable path rather than aborting the load, and do not silently change the result.
## Acceptance criteria
- [ ] First-token time on `qwen3.8-27B-4bit` at a short prompt, before and after, with page-cache and PTX-cache state stated on both arms (method rule 3 and rule 7 of the baseline record program).
- [ ] The host staging leg and the host-to-device leg reported separately, since only one of them is a CUDA API and a CUDA profiler sees only that one.
- [ ] Steady-state decode rate unchanged, measured as a slope over two `-n` values with both runs reaching the token budget.
- [ ] Peak memory reported for any change that materializes weights earlier, on the largest checkpoint available.
- [ ] A decision recorded on whether weight materialization is charged to model load or to the first token, applied uniformly across loaders rather than per family.
- [ ] The change is integrated into the real load and generate path, not a standalone benchmark or an opt-in flag nobody sets.
## Validation
```bash
# Warm the page cache first. A TTFT number quoted without page-cache state can be wrong by 5x.
cat ./models/mlx-community/qwen3.8-27B-4bit/model*.safetensors > /dev/null
# First token, with the pre-first-token phase breakdown.
MLXCEL_PROFILE_TTFT=1 ./target/release/mlxcel generate \
-m ./models/mlx-community/qwen3.8-27B-4bit --no-chat-template -p "Hi." -n 1 --profile
# The same prefill a second time in one process. The difference is the one-time cost.
./target/release/mlxcel-bench-decode \
-m ./models/mlx-community/qwen3.8-27B-4bit --no-chat-template -p "Hi." \
--warmup-tokens 1 -n 1
# Attribution. --cuda-graph-trace=node is mandatory. Reconcile the profiled prefill
# against the unprofiled wall clock before reading any absolute number.
nsys profile -t cuda,nvtx --cuda-graph-trace=node -o ttft_warm_n1 \
./target/release/mlxcel generate -m ./models/mlx-community/qwen3.8-27B-4bit \
--no-chat-template -p "Hi." -n 1 --profile
nsys stats --report cuda_api_sum --report cuda_gpu_kern_sum --format csv ttft_warm_n1.nsys-rep
```
A pass is: the first-token time falls against the 15.54 s baseline by an amount that reconciles with the leg the change targets; `cudaMemcpyAsync` still shows 1,847 calls, since that count is a property of the checkpoint and not of the transfer path; the decode slope is unchanged; and peak memory is reported. The host staging leg is invisible to `nsys` and must be timed in-process or by the standalone microbenchmark described in the reference record, since it makes no CUDA call.
## References
- #1545, and `docs/benchmark_results/volta-ttft-fixed-cost-2026-09-01.md`.
- `src/lib/mlxcel-core/src/weights.rs` (`load_weights_from_dir`, `src/lib/mlxcel-core/src/weights.rs:310`), `mlx/backend/cuda/load.cpp` (`Load::eval_gpu`).
- `src/models/sanitize.rs` (`load_gemma4_family_weights_with_backing`), the eager precedent already in tree.
- `MLXCEL_PROFILE_TTFT`, the pre-first-token instrumentation #1545 added, documented in `docs/environment-variables.md`.
Contributor guide
Research direction
Start with load_weights_from_dir in src/lib/mlxcel-core/src/weights.rs, the eager precedent in src/models/sanitize.rs, and Load::eval_gpu under src/lib/mlx-cpp/patches/mlx/backend/cuda/. Run the documented warm-cache MLXCEL_PROFILE_TTFT and nsys commands on qwen3.8-27B-4bit, then compare the targeted legs, decode slope, and peak memory. Done means a uniform load-versus-first-token decision is integrated and all acceptance measurements are reported.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, rust
- Domain
- backend, machine-learning, performance
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 38/100