Repository navigation
[Perf][GDN] Cut host op calls in the eager prefill path of mixed steps - #61034
kimseunghyun-kr wants to merge 10 commits into
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
|
This pull request has merge conflicts that must be resolved before it can be |
…xing Before the chunk kernel, each GDN layer gathered the prefill states with ssm_state[idx], zeroed the rows without initial state with a boolean __setitem__, and scattered the final states back with ssm_state[idx] = x. The two Python advanced-indexing calls and the per-layer mask inversion are host work in every layer of a mixed step. Use index_select, masked_fill_ and index_copy_ instead. The metadata builder now computes the inverted mask once per step, shaped to broadcast over the [num_prefills, HV, V, K] state. The values are bit-identical: the state indices are non-negative and unique per prefill. Test: tests/kernels/mamba/test_gdn_forward_core_mixed.py compares the state that the chunk kernel receives and the pool after the step with the indexing reference, bit for bit (CUDA). Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: kimseunghyun <incongnito12@gmail.com>
In a mixed step with spec decodes, each GDN layer split the tokens with four index_select calls (mixed_qkv, a, b for spec; mixed_qkv for non-spec), two more for the non-spec a and b, and merged the outputs with two index_copy_ calls. The batch reorder puts decode rows first, so the spec-decode rows usually lead the batch. The metadata builder now records this on the CPU (spec_tokens_first). Then the spec tokens are [0, num_spec_decode_tokens) and the non-spec tokens follow, so the layer uses slices and two copy_ calls. Other batches keep the gather path. The slices have the same strides as the non-spec path, which passes mixed_qkv to the same kernels directly. The values are bit-identical. Tests: - tests/v1/attention/test_gdn_metadata_builder.py: the flag is True exactly when the gather indices are the two ranges (CPU). - tests/kernels/mamba/test_gdn_forward_core_mixed.py: the slice path and the gather path give bit-identical outputs, conv states and SSM states (CUDA). Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: kimseunghyun <incongnito12@gmail.com>
rearrange_mixed_qkv copies q, k and v into one contiguous buffer: one split, three reshape copies, one cat, three slices and three views. In _forward_core, every caller passes the result to fused_sigmoid_gating_delta_rule_update, which calls .contiguous() on q, k and v itself. Add a contiguous=False mode that returns the three strided views, and use it at the three recurrent-update call sites (spec decodes, peeled decodes, decode-only batches). This removes about seven host calls and one copy kernel per call site. The other callers (CPU ops, OLMo) keep the contiguous buffer. The values are bit-identical. Test: tests/kernels/mamba/test_gdn_forward_core_mixed.py compares the strided views with the contiguous copies for spec, mixed spec, mixed non-spec and decode-only batches, bit for bit (CUDA). Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: kimseunghyun <incongnito12@gmail.com>
index_copy_ only accepts int64 indices, but the GDN state indices are int32, so every prefill step raised "index_copy_(): Expected a long tensor for index, but got Int" (found by the CUDA tests on A100). The metadata builder (build and update_block_table) now also stores prefill_state_indices_int64, cast once per step, and the layer uses it for index_select and index_copy_. prefill_state_indices keeps its dtype for its other users. Test: tests/v1/attention/test_gdn_metadata_builder.py checks the int64 copy against prefill_state_indices and after update_block_table (CPU). Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: kimseunghyun <incongnito12@gmail.com>
9c4f4fa to
8286a88
Compare
|
Rebased onto Checked on 1 × A100-40GB (job image vLLM 5555a2d, torch 2.13, with the GDN files of
|
|
Additional data point, different config: DCP 4, concurrency 32. No measurable gain there. The results in the PR body (C20, no DCP) stand. This comment adds a second config, where this PR shows no measurable effect. Setup:
Serving A/B (1 × 60 min per run, alternating B/D on one node, 15 Oct; SLO as in the PR body):
The two complete pairs differ in opposite directions (−1.05 and +2.81 points). The spread within B is about 2 points. So the gain, if any, is about 1 point or less, and these runs cannot resolve it. nsys (one capture per arm, mixed prefill + decode steps, NVTX step ranges, host time per step and rank):
Limits:
Reading: in this config, the mixed step has about 1,350 eager launches. This PR changes how some of them are issued, but not how many there are, so the host time per step does not move. At C20 without DCP, the PR body measured 11 ms shorter mixed steps; I have no explanation yet for the difference between the two configs. A larger cut for this model would need the GDN prefill path to leave eager mode (static-grid metadata + capture); I am preparing that as a separate PR. |
|
I close this PR. Reasons:
If the eager path matters again (for example with internal prefill checkpoints, where the capture is off), I will measure it there first and open a new PR only with a clear end-to-end gain. |
In mixed (prefill + decode) steps the Qwen GDN core ran as an eager break of the breakable CUDA graph: its metadata was new in every step, its Python path depended on the batch, and its kernel grids came from the batch. Measured on Qwen3.8-Flash-Next-FP8, 8x A100, TP8, DCP 4, C32 (nsys): a mixed step takes 99.6 ms of host time against about 65-70 ms of GPU work, and the 36 GDN breaks take 55.2 ms of that host time and 720 of 1,348 launches (20 per layer). Removing Python-level indexing (vllm-project#61034) did not change it (99.3 vs 99.6 ms), so the launches themselves must go into the graph. - _forward_core_static: one branch-free path for all batches of a captured token count. It reads only the capture buffers at full length and sizes every grid by the token count and the buffer capacity max_num_seqs: spec decodes (padded to max_num_seqs, as the FULL decode path), the decodes peeled off the non-spec rows, and the prefills (FLA use_graph=True, causal_conv1d_fn num_programs). A part that the batch lacks has only zero-length sequences or pad programs, which exit. Each part runs the kernels that _forward_core runs for a batch with prefills. Pad rows write only the null state block and a scratch output row; the output merge uses fixed-length index copies. - The capturable predicate of both GDN core ops (plain and fused-norm) allows the capture only inside a breakable PIECEWISE capture with staged capture buffers that hold the batch; otherwise the op breaks as before (other backends, align mode, ROCm AITER, V1 runner, token counts above max_num_batched_tokens). The fused-norm op also applies its output RMS norm to every row there; its fused CUDA MTP decode kernel stays for eager and FULL decode steps. - The first capture of each shape runs the static core once outside the capture (run_uncaptured), so its Triton kernels compile, and autotune if no config is cached, before the capture records them. Not changed: FULL decode graphs, the eager forward, and every other GDN backend. Tests: GPU (Triton/FLA, not SM10x): capture _forward_core (plain and fused-norm, bf16 and fp32 state) once for 64 tokens, replay it for 11 batch kinds (spec only, prefill only, spec+prefill, decode+prefill, 1-token prefills with and without spec decodes, 16 sequences, all 64 tokens, zero-length rows with and without spec decodes, decode only) and compare with the eager _forward_core: core output and every non-null conv and SSM state block bitwise; decode only (eager uses causal_conv1d_update there) with assert_close. CPU: the predicate captures only inside a PIECEWISE capture with buffers that hold the batch, and never the AITER core. Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: kimseunghyun <hellonypcteam@gmail.com>
In mixed (prefill + decode) steps the Qwen GDN core ran as an eager break of the breakable CUDA graph: its metadata was new in every step, its Python path depended on the batch, and its kernel grids came from the batch. Measured on Qwen3.8-Flash-Next-FP8, 8x A100, TP8, DCP 4, C32 (nsys): a mixed step takes 99.6 ms of host time against about 65-70 ms of GPU work, and the 36 GDN breaks take 55.2 ms of that host time and 720 of 1,348 launches (20 per layer). Removing Python-level indexing (vllm-project#61034) did not change it (99.3 vs 99.6 ms), so the launches themselves must go into the graph. - _forward_core_static: one branch-free path for all batches of a captured token count. It reads only the capture buffers at full length and sizes every grid by the token count and the buffer capacity max_num_seqs: spec decodes (padded to max_num_seqs, as the FULL decode path), the decodes peeled off the non-spec rows, and the prefills (FLA use_graph=True, causal_conv1d_fn num_programs). A part that the batch lacks has only zero-length sequences or pad programs, which exit. Each part runs the kernels that _forward_core runs for a batch with prefills. Pad rows write only the null state block and a scratch output row; the output merge uses fixed-length index copies. - The capturable predicate of both GDN core ops (plain and fused-norm) allows the capture only inside a breakable PIECEWISE capture with staged capture buffers that hold the batch; otherwise the op breaks as before (other backends, align mode, ROCm AITER, V1 runner, token counts above max_num_batched_tokens). The fused-norm op also applies its output RMS norm to every row there; its fused CUDA MTP decode kernel stays for eager and FULL decode steps. - The first capture of each shape runs the static core once outside the capture (run_uncaptured), so its Triton kernels compile, and autotune if no config is cached, before the capture records them. Not changed: FULL decode graphs, the eager forward, and every other GDN backend. Tests: GPU (Triton/FLA, not SM10x): capture _forward_core (plain and fused-norm, bf16 and fp32 state) once for 64 tokens, replay it for 11 batch kinds (spec only, prefill only, spec+prefill, decode+prefill, 1-token prefills with and without spec decodes, 16 sequences, all 64 tokens, zero-length rows with and without spec decodes, decode only) and compare with the eager _forward_core: core output and every non-null conv and SSM state block bitwise; decode only (eager uses causal_conv1d_update there) with assert_close. CPU: the predicate captures only inside a PIECEWISE capture with buffers that hold the batch, and never the AITER core. Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: kimseunghyun <hellonypcteam@gmail.com>
In mixed (prefill + decode) steps the Qwen GDN core ran as an eager break of the breakable CUDA graph: its metadata was new in every step, its Python path depended on the batch, and its kernel grids came from the batch. Measured on Qwen3.8-Flash-Next-FP8, 8x A100, TP8, DCP 4, C32 (nsys): a mixed step takes 99.6 ms of host time against about 65-70 ms of GPU work, and the 36 GDN breaks take 55.2 ms of that host time and 720 of 1,348 launches (20 per layer). Removing Python-level indexing (vllm-project#61034) did not change it (99.3 vs 99.6 ms), so the launches themselves must go into the graph. - _forward_core_static: one branch-free path for all batches of a captured token count. It reads only the capture buffers at full length and sizes every grid by the token count and the buffer capacity max_num_seqs: spec decodes (padded to max_num_seqs, as the FULL decode path), the decodes peeled off the non-spec rows, and the prefills (FLA use_graph=True, causal_conv1d_fn num_programs). A part that the batch lacks has only zero-length sequences or pad programs, which exit. Each part runs the kernels that _forward_core runs for a batch with prefills. Pad rows write only the null state block and a scratch output row; the output merge uses fixed-length index copies. - The capturable predicate of both GDN core ops (plain and fused-norm) allows the capture only inside a breakable PIECEWISE capture with staged capture buffers that hold the batch; otherwise the op breaks as before (other backends, align mode, ROCm AITER, V1 runner, token counts above max_num_batched_tokens). The fused-norm op also applies its output RMS norm to every row there; its fused CUDA MTP decode kernel stays for eager and FULL decode steps. - The first capture of each shape runs the static core once outside the capture (run_uncaptured), so its Triton kernels compile, and autotune if no config is cached, before the capture records them. Not changed: FULL decode graphs, the eager forward, and every other GDN backend. Tests: GPU (Triton/FLA, not SM10x): capture _forward_core (plain and fused-norm, bf16 and fp32 state) once for 64 tokens, replay it for 11 batch kinds (spec only, prefill only, spec+prefill, decode+prefill, 1-token prefills with and without spec decodes, 16 sequences, all 64 tokens, zero-length rows with and without spec decodes, decode only) and compare with the eager _forward_core: core output and every non-null conv and SSM state block bitwise; decode only (eager uses causal_conv1d_update there) with assert_close. CPU: the predicate captures only inside a PIECEWISE capture with buffers that hold the batch, and never the AITER core. Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: kimseunghyun <hellonypcteam@gmail.com>
In mixed (prefill + decode) steps the Qwen GDN core ran as an eager break of the breakable CUDA graph: its metadata was new in every step, its Python path depended on the batch, and its kernel grids came from the batch. Measured on Qwen3.8-Flash-Next-FP8, 8x A100, TP8, DCP 4, C32 (nsys): a mixed step takes 99.6 ms of host time against about 65-70 ms of GPU work, and the 36 GDN breaks take 55.2 ms of that host time and 720 of 1,348 launches (20 per layer). Removing Python-level indexing (vllm-project#61034) did not change it (99.3 vs 99.6 ms), so the launches themselves must go into the graph. - _forward_core_static: one branch-free path for all batches of a captured token count. It reads only the capture buffers at full length and sizes every grid by the token count and the buffer capacity max_num_seqs: spec decodes (padded to max_num_seqs, as the FULL decode path), the decodes peeled off the non-spec rows, and the prefills (FLA use_graph=True, causal_conv1d_fn num_programs). A part that the batch lacks has only zero-length sequences or pad programs, which exit. Each part runs the kernels that _forward_core runs for a batch with prefills. Pad rows write only the null state block and a scratch output row; the output merge uses fixed-length index copies. - The capturable predicate of both GDN core ops (plain and fused-norm) allows the capture only inside a breakable PIECEWISE capture with staged capture buffers that hold the batch; otherwise the op breaks as before (other backends, align mode, ROCm AITER, V1 runner, token counts above max_num_batched_tokens). The fused-norm op also applies its output RMS norm to every row there; its fused CUDA MTP decode kernel stays for eager and FULL decode steps. - The first capture of each shape runs the static core once outside the capture (run_uncaptured), so its Triton kernels compile, and autotune if no config is cached, before the capture records them. Not changed: FULL decode graphs, the eager forward, and every other GDN backend. Tests: GPU (Triton/FLA, not SM10x): capture _forward_core (plain and fused-norm, bf16 and fp32 state) once for 64 tokens, replay it for 11 batch kinds (spec only, prefill only, spec+prefill, decode+prefill, 1-token prefills with and without spec decodes, 16 sequences, all 64 tokens, zero-length rows with and without spec decodes, decode only) and compare with the eager _forward_core: core output and every non-null conv and SSM state block bitwise; decode only (eager uses causal_conv1d_update there) with assert_close. CPU: the predicate captures only inside a PIECEWISE capture with buffers that hold the batch, and never the AITER core. Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: kimseunghyun <hellonypcteam@gmail.com>
In mixed (prefill + decode) steps the Qwen GDN core ran as an eager break of the breakable CUDA graph: its metadata was new in every step, its Python path depended on the batch, and its kernel grids came from the batch. Measured on Qwen3.8-Flash-Next-FP8, 8x A100, TP8, DCP 4, C32 (nsys): a mixed step takes 99.6 ms of host time against about 65-70 ms of GPU work, and the 36 GDN breaks take 55.2 ms of that host time and 720 of 1,348 launches (20 per layer). Removing Python-level indexing (vllm-project#61034) did not change it (99.3 vs 99.6 ms), so the launches themselves must go into the graph. - _forward_core_static: one branch-free path for all batches of a captured token count. It reads only the capture buffers at full length and sizes every grid by the token count and the buffer capacity max_num_seqs: spec decodes (padded to max_num_seqs, as the FULL decode path), the decodes peeled off the non-spec rows, and the prefills (FLA use_graph=True, causal_conv1d_fn num_programs). A part that the batch lacks has only zero-length sequences or pad programs, which exit. Each part runs the kernels that _forward_core runs for a batch with prefills. Pad rows write only the null state block and a scratch output row; the output merge uses fixed-length index copies. - The capturable predicate of both GDN core ops (plain and fused-norm) allows the capture only inside a breakable PIECEWISE capture with staged capture buffers that hold the batch; otherwise the op breaks as before (other backends, align mode, ROCm AITER, V1 runner, token counts above max_num_batched_tokens). The fused-norm op also applies its output RMS norm to every row there; its fused CUDA MTP decode kernel stays for eager and FULL decode steps. - The first capture of each shape runs the static core once outside the capture (run_uncaptured), so its Triton kernels compile, and autotune if no config is cached, before the capture records them. Not changed: FULL decode graphs, the eager forward, and every other GDN backend. Tests: GPU (Triton/FLA, not SM10x): capture _forward_core (plain and fused-norm, bf16 and fp32 state) once for 64 tokens, replay it for 11 batch kinds (spec only, prefill only, spec+prefill, decode+prefill, 1-token prefills with and without spec decodes, 16 sequences, all 64 tokens, zero-length rows with and without spec decodes, decode only) and compare with the eager _forward_core: core output and every non-null conv and SSM state block bitwise; decode only (eager uses causal_conv1d_update there) with assert_close. CPU: the predicate captures only inside a PIECEWISE capture with buffers that hold the batch, and never the AITER core. Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: kimseunghyun <hellonypcteam@gmail.com>
Overview
Cut host-side PyTorch op calls in the GDN (gated delta net) layer's eager prefill path. That path runs in every mixed prefill + decode step, between the piecewise CUDA graph pieces. This PR changes no kernel and no numerics; outputs and states are bit-identical.
CUDA parity tests pass on A100 (below). The end-to-end serving A/B gives +1.0 point SLO pass and +2.4 % goodput at concurrency 20, and lets concurrency 18 pass in all three repeats. nsys shows mixed steps 11 ms shorter; py-spy shows the host-time drop in the GDN path.
No new CPU/GPU sync: the new
spec_tokens_firstflag is computed from a CPU tensor in the metadata builder, and py-spy shows 0 % of the worker's main thread in host syncs with the change. Validated outside CI (no A100 in CI). The new kernel tests need SM80+ and run on any such CI GPU. Written with AI assistance (Claude Code); the submitter reviewed every line.Claims
index_select, onemasked_fill_andindex_copy_. The metadata builder computes the mask and an int64 copy of the prefill state indices once per step.spec_tokens_first, when all spec-decode rows lead the batch. MRv2sort_batch_req_idsnormally gives that order. With the flag set, the layer uses slices and 2copy_instead of 6index_selectand 2index_copy_. Other orders keep the gather path.rearrange_mixed_qkv(..., contiguous=False)returns strided views at the three call sites that feed the recurrent update kernel. That kernel already calls.contiguous(). Other callers are unchanged.Validation
Why (measured, py-spy at 100 Hz on all TP workers, MRv2, Qwen3.8-Flash-Next-FP8, TP8 + EP, 8 × A100-SXM4-40GB, MLPerf agentic set at 16 concurrent sessions):
CUDA parity tests (A100-SXM4-40GB, vLLM 81198e9 image with this change backported, unchanged otherwise):
tests/kernels/mamba/test_gdn_forward_core_mixed.py(SM80+). It checks outputs, conv states and SSM states bit-exactly (atol = 0) against the previous code path:test_prefill_state_gather_and_scatter_match_indexing[prefills, spec-decodes-then-prefills]; fresh-prefill state slots hold NaN, to prove they are not read;test_spec_token_slices_match_gathers;test_strided_qkv_views_match_contiguous_copies[4 batch types].tests/v1/attention/test_gdn_metadata_builder.py(CPU): newtest_spec_tokens_first_matches_token_indices, newtest_prefill_state_indices_int64_matches_prefill_state_indices, and the new fields added totest_update_block_table_matches_build.test_gdn_metadata_builder.py,test_gdn_fused_mtp.py,test_gdn_forward_core_split.pyandtests/models/qwen4_exp/test_qsa_reference.py, without and with the change:ruff check/ruff format --check: clean.End to end (measured, 11 Oct 2026): 8 × A100-SXM4-40GB, one node (the same node for both arms, runs 8 h apart, not interleaved). Qwen3.8-Flash-Next-FP8, TP8, MTP 3, FP8 QSA KV, bf16 SSM state, native CPU KV offload 256 GiB. ArtificialAnalysis-style multi-turn agentic replay at concurrency 20, 3 × 60 min per arm. SLO (Artificial Analysis percentile convention): P95 TTFT ≤ 2.93 s and P95 TPOT ≤ 21.74 ms (P5 output speed ≥ 46 tokens/s). "SLO pass" is the share of requests that meet both limits; the SLO holds when it is ≥ 95 %.
Higher load, concurrency 18 (same setup, measured): without this PR 96.49 / 94.72 % (one repeat under the 95 % target); with this PR 97.81 / 96.44 / 96.55 % (all three repeats above 95 %), goodput 2.06 / 1.95 / 1.92 req/s. TPOT p95 per 10-min bin stays at or under 19.9 ms for the whole hour of repeat 1.
All three repeats with the change are above all three without it. The gain is small: the workload is limited by KV cache capacity, and this PR changes only host time in mixed prefill + decode steps. No request errors in either arm. Host profile (py-spy, measured shares): C16, the same replay window, 100 Hz for 300 s on every TP worker, without (6 workers) and with (7 workers) this PR. Share of main-thread samples:
Converted with the same factor as the baseline analysis (about 1,100 mixed steps in the window, an estimate), the GDN host time per mixed step falls from about 47 ms to about 37 ms. The two traces ran on different nodes of the same type. No host syncs in either trace.
GPU timeline (Nsight Systems, device 0, 256 steps per trace, measured):
Decode steps do not change. Mixed steps are about 11 ms shorter, although the window with this PR had larger prefill chunks (which costs more GPU time, not less).
Reproduce the serving runs. The serving setup, client and dataset are the same as in #60602 (see its "Reproduce the serving runs" section). Both arms use the same server command (below). The two arms differ only in this PR's two files. A second node of the same type, with this PR, gave 94.70 / 93.37 / 94.33 % at concurrency 20 (no control run on that node).
Server command (both arms)
Details
NULL_BLOCK_ID0.causal_conv1d_updatewrites in place into its input. With (b), that input is a slice ofmixed_qkv, so the spec rows ofmixed_qkvare overwritten. The spec-only and decode-only paths already do this, and no later code reads those rows.fused_post_conv_prep. Both need kernel changes. With them the saving would be about 15–18 ms (estimate).qwen_gdn_linear_attn.py+66 / −13 (with the int64 fix),gdn_attn.py+24; tests +371.index_copy_. The A100 run caught that bug: the GDN state indices are int32. I am happy to squash.Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.
AI assistance: drafted with Claude Code. The author reviewed every change. The profile and test results come from runs on our cluster.
🤖 Generated with Claude Code