Repository navigation
Conversation
`prepare_chunk_indices` derives per-sequence chunk counts with `.tolist()`, so passing it a device `cu_seqlens` blocks the caller. The count also fixes the Triton grid, so it has to reach the host either way; compute it from the host-side cu_seqlens we already have and copy the result across, as `gdn_attn` already does for its own prefill slice. Both models served from `vllm/models/kimi_k3` hit this on prefill batches, on NVIDIA and ROCm. The existing `chunk_indices` describes the decode-peeled `prefill_query_start_loc`, so add `non_spec_chunk_indices` alongside it and thread an optional `chunk_indices` through `chunk_kda_with_fused_gate`. It defaults to None and falls back to deriving on device, so other callers are unaffected. Building it costs 200-650us per prefill step, so it is opt-in via a builder flag: the Kimi paths opt in, Qwen3-Next pays nothing. Verified host- and device-derived indices match in value, dtype and length. Not run on Kimi hardware. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Nick Hill <nickhill123@gmail.com>
njhill
force-pushed
the
kda-chunk-indices-nosync
branch
from
August 10, 2026 20:03
5406779 to
0eedbb7
Compare
njhill
marked this pull request as ready for review
August 10, 2026 20:09
njhill
requested review from
ZJY0516,
dllehr-amd,
hongxiayang,
tdoublep,
tjtanaa,
tomeras91 and
vadiklyutiy
as code owners
August 10, 2026 20:09
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Member
Author
|
/ci run |
|
✅ Triggered Buildkite CI #83206 for commit |
njhill
requested review from
AndreasKaratzas,
DarkLight1337 and
ywang96
as code owners
August 10, 2026 22:46
Member
Author
|
/ci run |
|
✅ Triggered Buildkite CI #83229 for commit |
4 tasks
Collaborator
|
hi @njhill thanks for your PR! I saw you've also changed the path for ROCm side; I wonder if we have tested this PR on any AMD hardware? |
Member
|
Test-wise, full CI eval on ROCm is in parity with nightly: But as @Fangzhou-Ai mentioned, let's ensure perf-wise that we are not introducing any regressions. |
3 of 13 tasks
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
justtestingthingsx
pushed a commit
to meandmyboiclaude/vllm
that referenced
this pull request
Aug 16, 2026
…sted GDN metadata vllm-project#52297 (picked as the two commits below this one) hoists the per-batch GDN spec metadata out of per-group build() into compute_common_gdn_attn_metadata. Two carried club mechanisms lived in the deleted region and are relocated into the hoisted shape: 1. ReplaySSM spec-sequence mask (use_cache_spec_kernel): the mask must come from ~is_prefilling & (qlen > 0), not from num_decode_draft_tokens_cpu (stale on draft-less steps; -1 for dropped-draft rows). Now an opt-in branch of compute_common_gdn_attn_metadata, driven from prepare_attn via cache_config.use_replayssm_spec + the is_prefilling tensor, and from build_for_cudagraph_capture (q=1 capture batches must still route spec). 2. Stale zero-accept guard (vllm-project#51508-subset + vllm-project#40756 clamp, which upstream's hoist dropped): the batch-level half (stale mask + two-sided clamp) lives in the hoisted function; the per-group half (NULL_BLOCK_ID masked_fill on spec_state_indices_tensor) stays in GDN build(), fed by a new stale_spec_reqs element threaded through MambaHybridAttnMetadata (GDN-only kwarg; spec-row-shaped, unsliced). The context_lens_tensor unconditional compute (072646c) auto-merged into the hoisted build() and needed no relocation. Verified: py_compile all three; test_gdn_rollback_identity 49/49; gather repro unchanged (documents upstream vllm-project#51540, fixed by chain patch at boot). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
33 of 46 tasks
Member
|
Fresh exact H200 evidence for this fix: the All three die on the first request in |
6 tasks done
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
prepare_chunk_indicesderives per-sequence chunk counts with.tolist(), so passing it a devicecu_seqlensblocks the caller. The count also fixes the Triton grid, so it has to reach the host either way; compute it from the host-side cu_seqlens we already have and copy the result across, asgdn_attnalready does for its own prefill slice.Models from
vllm/models/kimi_k3using FLA hit this on prefill batches, on NVIDIA and ROCm. The existingchunk_indicesdescribes the decode-peeledprefill_query_start_loc, so addnon_spec_chunk_indicesalongside it and thread an optionalchunk_indicesthroughchunk_kda_with_fused_gate. It defaults to None and falls back to deriving on device, so other callers are unaffected.Verified host- and device-derived indices match in value, dtype and length. Not run on Kimi hardware.
Note this applies to Kimi Linear but not Kimi K3 with cuda since the latter uses FlashKDA rather than FLA.
Claude was used for this.