[GDN] Accept graph-padded spec metadata for DFlash decode#391
Open
ehartford wants to merge 3 commits into
Open
[GDN] Accept graph-padded spec metadata for DFlash decode#391ehartford wants to merge 3 commits into
ehartford wants to merge 3 commits into
Conversation
added 2 commits
June 3, 2026 14:50
Signed-off-by: Eric Hartford <eric@quixi.ai>
Signed-off-by: Eric Hartford <eric@quixi.ai>
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.
Summary
Fixes #389.
The XPU GDN attention op already accepts graph-padded leading dimensions for the active token tensors, but the speculative-decoding metadata path still requires exact metadata sizes:
spec_query_start_loc.size(0) == num_spec_decodes + 1spec_state_indices_tensor.size(0) == num_spec_decodesnum_accepted_tokens.size(0) == num_spec_decodesUnder vLLM graph capture / compiled execution, DFlash/spec-decode metadata can be allocated for the captured graph shape while only the first
num_spec_decodesentries are active. In that case the op rejects otherwise valid inputs with:This PR relaxes those spec metadata checks to accept graph-padded leading dimensions and narrows the active prefix before launching the native kernels. It also guards the XE2 chunk GDN prefill path so it is used only for head dimensions supported by the chunk kernels.
Changes
spec_query_start_loc,spec_state_indices_tensor, andnum_accepted_tokensto have leading dimensions at least as large as the active spec decode count.Why This Is Not #344 Again
#344 fixed graph-padded active token tensors (
core_attn_out,z,projected_states_qkvz,projected_states_ba). This PR addresses the remaining speculative metadata tensors used by the DFlash/spec-decode path.#368 added native Qwen MTP support. This PR makes that path tolerate graph-padded metadata shapes produced by captured execution.
Test Plan
Built this branch independently from
origin/mainin a clean worktree with:Ran focused GDN test coverage:
Test Result