Skip to content

[GDN] Accept graph-padded spec metadata for DFlash decode#391

Open
ehartford wants to merge 3 commits into
vllm-project:mainfrom
QuixiAI:pr/gdn-dflash
Open

[GDN] Accept graph-padded spec metadata for DFlash decode#391
ehartford wants to merge 3 commits into
vllm-project:mainfrom
QuixiAI:pr/gdn-dflash

Conversation

@ehartford

Copy link
Copy Markdown

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 + 1
  • spec_state_indices_tensor.size(0) == num_spec_decodes
  • num_accepted_tokens.size(0) == num_spec_decodes

Under vLLM graph capture / compiled execution, DFlash/spec-decode metadata can be allocated for the captured graph shape while only the first num_spec_decodes entries are active. In that case the op rejects otherwise valid inputs with:

spec_query_start_loc must have size [num_spec_decodes + 1]

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

  • Allow spec_query_start_loc, spec_state_indices_tensor, and num_accepted_tokens to have leading dimensions at least as large as the active spec decode count.
  • Narrow those metadata tensors to the active prefix before passing them to the native spec conv / gated-delta kernels.
  • Add a supported-shape guard before using the XE2 chunk GDN path.
  • Extend GDN tests to cover graph-padded spec metadata and unsupported chunk-shape fallback behavior.
  • Update default attention kernel config comments/entries used by the reduced build/test profile.

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/main in a clean worktree with:

source ~/.venv/bin/activate
PATH=/opt/intel/oneapi/compiler/2026.0/bin:$PATH \
CMPLR_ROOT=/opt/intel/oneapi/compiler/2026.0 \
VLLM_PAGED_DECODE_CONFIG=paged_decode_default.conf \
VLLM_CHUNK_PREFILL_CONFIG=chunk_prefill_default.conf \
python setup.py build_ext --inplace

Ran focused GDN test coverage:

source ~/.venv/bin/activate
XPU_KERNEL_TEST_SCOPE=mini pytest -q tests/gdn_attn/test_gdn_attn.py --tb=short

Test Result

20 passed in 3.41s

Eric Hartford added 2 commits June 3, 2026 14:50
Signed-off-by: Eric Hartford <eric@quixi.ai>
Signed-off-by: Eric Hartford <eric@quixi.ai>
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.

GDN spec metadata shape checks reject graph-padded DFlash decode batches

1 participant