Repository navigation
Conversation
A prefill pass allocated every intermediate tensor per layer: a 1024-token pass issued about 460 cudaMallocAsync calls plus the matching frees, and the host time between layers showed up as a gap at every layer boundary. The pass now owns one buffer per role (normed rows, q/k/v, the linear-attention projections, the MLP gate/up and down, the attention output) and reuses it for every layer of that kind, so layers write into mapped pages instead of fresh allocations. The second residual of a layer accumulates into the block output while that accumulator is also the source, which the existing add_batch_into cannot express because it takes source and destination separately. The elementwise module gains add_batch_inplace for that shape; the device kernel reads both inputs before writing, so the alias is safe. Measured on one H20 with a vLLM 0.29.0 client, two alternating rounds of the reference and the buffered binary, 1024-token prompts unless noted: prefill-only, 64 prompts at concurrency 32: 5.45 -> 5.28 s, 11.74 -> 12.13 prompts/s prefill-only, 4096-token prompts: 3.33 -> 3.26 s c16, 1024 in / 256 out: 11.56 -> 11.39 s, 1418 -> 1438 tok/s, TPOT 10.02 -> 9.90 ms qps16, 1024 in / 128 out: 7.37 -> 7.16 s, 1112 -> 1145 tok/s, TPOT 21.53 -> 20.98 ms Single-request and 8192-token units do not regress. Greedy output is byte-identical to the reference binary on 4 x 500 and 2 x 2000-token prompts, and the unified_forward tests pass. Signed-off-by: CAICAIIs <3360776475@qq.com>
Codex Review Summary: ✅ CompletedThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3847cb2b12
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| HiddenStates::zeros(&self.ctx, self.config.hidden_size, chunk_len)?; | ||
| let mut chunk_scratch = | ||
| GdrChunkwiseScratch35::new(&self.ctx, &self.config, self.geometry, chunk_len)?; | ||
| let mut pass_buffers = self.prefill_pass_buffers(chunk_len)?; |
There was a problem hiding this comment.
For a prompt near the 20,000-token PREFILL_CHUNK_LEN on a memory-tight GPU, this call keeps the full-attention, linear-attention, and MLP buffers allocated simultaneously for the whole chunk. GdrChunkwiseScratch35::estimate_bytes, which weights.rs and tp_executor/worker.rs use to reserve memory before spending the rest on KV/decode capacity, still budgets shared_layer + max(full_attn_temps, mlp_temps) and does not include the linear set; for Qwen3.5-4B TP1 the omission is about 1.76 GiB at 20k tokens, so startup can succeed and the first long prefill can OOM despite the advertised reserve. Update that estimator, or alias mutually exclusive roles, to match this lifetime.
Useful? React with 👍 / 👎.
| acc_ptr as *const ffi::Half, | ||
| b_ptr as *const ffi::Half, | ||
| acc_ptr as *mut ffi::Half, |
There was a problem hiding this comment.
Every Qwen3.5 prefill layer now reaches this call for its second residual, but add_cuda launches add_kernel with both a and out declared __restrict__ in csrc/shared/elementwise.cu:12-15. Passing the same acc_ptr in those positions violates the kernel's non-aliasing contract; optimized NVCC code is allowed to assume they never overlap, so residual values can be miscomputed even though the source happens to load before its store. Use a dedicated in-place accumulator kernel, or remove the conflicting restriction, instead of aliasing this API.
Useful? React with 👍 / 👎.
|
This pull request has merge conflicts with main. Please rebase it so it can be reviewed and merged. |
What changes
A prefill pass allocated its intermediate tensors per layer. At 1024 tokens the pass issued roughly 460
cudaMallocAsynccalls and the matching frees, and the host time between layers showed up as a gap at every layer boundary. The pass now owns one buffer per role and reuses it for every layer of that kind:PrefillPassBuffersholds the normed rows, q/k/v, the linear-attention projections (qkv,z,b,a), the conv and GDN outputs, the gated-norm output, the MLP gate/up and down, and the attention output.prefill_chunk_forwardwalks the layers with two hidden blocks ping-ponged between them, so a layer writes its residual into a mapped page instead of a fresh allocation.add_batch_inplace: the post-MLP residual accumulates into the block output while that output is also the source, whichadd_batch_intocannot express with two live borrows. The device kernel reads both inputs before writing, so the alias is safe.Evidence
One H20, single card, vLLM 0.29.0 client, two alternating rounds of the reference and the buffered binary, 1024-token prompts unless noted. Each cell averages both rounds.
Single-request and long-prompt units do not regress: bs1 246.3 → 247.1 tok/s, 8192-token prefill 196.7 → 197.5 tok/s (41.64 → 41.48 s).
Correctness
cargo test --release -p pegainfer-qwen35 --features qwen35 --lib unified_forward: 2 passed.cargo fmt --all -- --checkandcargo clippy -p pegainfer-qwen35 --features qwen35 --all-targets -- -D warningsclean.Notes
main. The numbers above were taken with both changes present, since that is the tree that was built.