Skip to content

[Perf][GDN] Cut host op calls in the eager prefill path of mixed steps - #61034

Closed
kimseunghyun-kr wants to merge 10 commits into
vllm-project:mainfrom
kimseunghyun-kr:gdn-mixed-step-host-diet
Closed

kimseunghyun-kr wants to merge 10 commits into
vllm-project:mainfrom
kimseunghyun-kr:gdn-mixed-step-host-diet

Conversation

@kimseunghyun-kr

@kimseunghyun-kr kimseunghyun-kr commented Oct 11, 2026 •

Copy link
Copy Markdown
Contributor

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_first flag 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

  • (a) Prefill SSM state gather / zero / scatter: Python indexing becomes index_select, one masked_fill_ and index_copy_. The metadata builder computes the mask and an int64 copy of the prefill state indices once per step.
  • (b) Spec / non-spec token split and merge: the builder sets a CPU flag, spec_tokens_first, when all spec-decode rows lead the batch. MRv2 sort_batch_req_ids normally gives that order. With the flag set, the layer uses slices and 2 copy_ instead of 6 index_select and 2 index_copy_. Other orders keep the gather path.
  • (c) q/k/v: 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.
  • Expected effect (estimate): about 9–14 ms less host time per mixed step on Qwen3.8-Flash-Next (36 GDN layers), A100, TP8. The estimate comes from a py-spy profile, summarised below.

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):

  • Mixed steps are host-paced. The worker's main thread spends about 80 ms of Python per mixed step, and 0 % of it is in host syncs.
  • 71 % of that piecewise-forward host time is the eager GDN core.
  • The three blocks changed here take about 8.0 ms (a), 8.1 ms (b) and 6.1 ms (c) per mixed step. These ms values are estimated from the sample shares.

CUDA parity tests (A100-SXM4-40GB, vLLM 81198e9 image with this change backported, unchanged otherwise):

  • New 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): new test_spec_tokens_first_matches_token_indices, new test_prefill_state_indices_int64_matches_prefill_state_indices, and the new fields added to test_update_block_table_matches_build.
  • Regression run of test_gdn_metadata_builder.py, test_gdn_fused_mtp.py, test_gdn_forward_core_split.py and tests/models/qwen4_exp/test_qsa_reference.py, without and with the change:
    • Old tests: the same 11 failures in both runs. They are not caused by this change: the fused MTP kernel is not built in the cu12.9 image, and the FP8 QSA selection tests need sm_89+.
    • New tests: they fail without the change and pass with it.
  • 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 %.

Arm SLO pass r1 / r2 / r3 (%) Mean Goodput r1 / r2 / r3 (req/s) Mean
Without this PR 92.01 / 90.92 / 92.18 91.70 1.851 / 1.875 / 1.872 1.866
With this PR 92.58 / 92.62 / 93.06 92.75 1.885 / 1.885 / 1.960 1.910 (+2.4 %)

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:

Part of the mixed-step (piecewise) forward Without With Change
GDN eager path (36 layers) 24.4 % 19.1 % −22 % relative
QSA + indexer eager path (12 layers) 8.5 % 8.0 % —
Whole piecewise forward 34.5 % 28.6 % −17 % relative
Worker idle (waiting for the next step) 33.7 % 36.5 % +2.8 points

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):

Without With
Decode step, GPU period (median) 24.67 ms 24.40 ms
Mixed prefill + decode step, GPU period (median) 103.7 ms 92.6 ms
Mixed step, host NVTX range (median) 96.0 ms 85.4 ms
Prefill tokens per mixed step (median) 296 672

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)
vllm serve Qwen/Qwen3.8-Flash-Next-FP8 --max-num-seqs 32 --enable-prefix-caching --no-enable-flashinfer-autotune --enable-expert-parallel --tensor-parallel-size 8 --gpu-memory-utilization 0.9 --enable-auto-tool-choice --tool-call-parser qwen3_xml --reasoning-parser qwen3 --max-model-len 262144 --distributed-executor-backend mp --speculative-config '{"method": "mtp", "num_speculative_tokens": 3}' --cudagraph-capture-sizes 1 2 4 8 16 24 32 40 48 56 64 96 128 192 256 384 512 768 1024 1536 2048 --max-num-batched-tokens 1024 --prefix-match-unit 16 --engram-config '{"cpu_offload": true}' --attention-config '{"indexer_kv_dtype": "bf16"}' --mamba-cache-dtype bfloat16 --mamba-ssm-cache-dtype bfloat16 --kv-cache-dtype fp8 --kv-offloading-backend native --kv-offloading-size 256

Details

  • Assumptions:
    • Prefill state indices are non-negative and unique. Padded rows use NULL_BLOCK_ID 0.
    • causal_conv1d_update writes in place into its input. With (b), that input is a slice of mixed_qkv, so the spec rows of mixed_qkv are overwritten. The spec-only and decode-only paths already do this, and no later code reads those rows.
  • Not done: the FLA chunk kernel reading and writing the state in place through indices, and folding (c) into fused_post_conv_prep. Both need kernel changes. With them the saving would be about 15–18 ms (estimate).
  • Size: qwen_gdn_linear_attn.py +66 / −13 (with the int64 fix), gdn_attn.py +24; tests +371.
  • Commits: (a), (b), (c), plus a fix that uses int64 indices for 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-checklist skill. (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

@github-actions

Copy link
Copy Markdown

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: 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.

🚀

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@mergify

mergify Bot commented Oct 14, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @kimseunghyun-kr.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify Bot added the needs-rebase label Oct 14, 2026
kimseunghyun-kr and others added 4 commits October 14, 2026 16:34
…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>
@kimseunghyun-kr
kimseunghyun-kr force-pushed the gdn-mixed-step-host-diet branch from 9c4f4fa to 8286a88 Compare October 14, 2026 08:01
@kimseunghyun-kr

Copy link
Copy Markdown
Contributor Author

Rebased onto main ec8a372 (resolves the conflict). Only the first commit changed: main's #60050 (FlashInfer prefill checkpoints) now wraps the prefill call, and this PR's index_select + masked_fill_ gather feeds both paths. The checkpoint export still runs before the cache write, as on main. The other 3 commits are unchanged.

Checked on 1 × A100-40GB (job image vLLM 5555a2d, torch 2.13, with the GDN files of main and of this PR bound in). The torch 2.14 / Triton 3.8 bump on main was not covered by this run.

Test file main ec8a372 This PR
tests/kernels/mamba/test_gdn_forward_core_mixed.py (not on main) 7 passed
tests/v1/attention/test_gdn_metadata_builder.py 62 passed 75 passed
tests/kernels/mamba/test_gdn_prefill_flashinfer.py 1 passed, 4 skipped 1 passed, 4 skipped
test_gdn_forward_core_split.py, test_gdn_prefill_cutedsl.py skipped (SM10x only) skipped (SM10x only)

@mergify mergify Bot removed the needs-rebase label Oct 14, 2026
@kimseunghyun-kr

Copy link
Copy Markdown
Contributor Author

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:

  • Qwen3.8-Flash-Next-FP8 on 8 × A100-SXM4-40GB.
  • TP8 + EP, DCP 4 ([Model][DCP] Support Qwen4Exp QSA #59279 with its follow-ups), MTP 3, FP8 KV, --max-num-batched-tokens 2048, MRv2.
  • The same agentic multi-turn replay at concurrency 32.
  • Arms: B = without this PR. D = B + this PR's commits, from before the rebase (917e03bb3..f35e92558).

Serving A/B (1 × 60 min per run, alternating B/D on one node, 15 Oct; SLO as in the PR body):

Arm SLO pass r1 / r2 / r3 Mean Goodput (req/s) mean
B 80.53 / 78.54 / 79.03 % 79.37 % 2.29
D 79.48 / (client metrics lost) / 81.84 % 80.66 % 2.35

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):

B D
Mixed-step host time, P50 99.6 ms 99.3 ms
GDN eager segments, P50 55.2 ms, 720 launches 54.8 ms, 720 launches
One GDN eager segment, P50 1.52 ms, 20 launches 1.52 ms, 20 launches

Limits:

  • I could not confirm from the trace that D ran the new code path. Both paths launch kernels with the same names on this PyTorch build.
  • One capture per arm, and the prefill tokens per mixed step differ (P50 500 in B vs 296 in D).

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.

@kimseunghyun-kr

Copy link
Copy Markdown
Contributor Author

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.

kimseunghyun-kr pushed a commit to kimseunghyun-kr/vllm that referenced this pull request Oct 20, 2026
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>
kimseunghyun-kr pushed a commit to kimseunghyun-kr/vllm that referenced this pull request Oct 20, 2026
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>
kimseunghyun-kr pushed a commit to kimseunghyun-kr/vllm that referenced this pull request Oct 25, 2026
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>
kimseunghyun-kr pushed a commit to kimseunghyun-kr/vllm that referenced this pull request Oct 27, 2026
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>
kimseunghyun-kr pushed a commit to kimseunghyun-kr/vllm that referenced this pull request Oct 27, 2026
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant