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

[ehartford]

Summary

The XPU GDN attention op currently requires several speculative-decoding metadata tensors to have exact leading dimensions derived from the active decode count. Under vLLM graph capture / compiled execution, the data tensors may be padded to a captured graph size while the active metadata still describes the real request count. This can make DFlash/speculative decode fail at runtime even though the active prefix of the metadata is valid.

The observed failure is:

RuntimeError: spec_query_start_loc must have size [num_spec_decodes + 1]

This is the same class of graph-padding contract issue as padded token tensors, but on the speculative metadata tensors rather than core_attn_out / projected state tensors.

Relevant code

At current main (dae50e2aa58301d3161c26d15bfd867dc5954c12), gdn_attention validates exact sizes for spec metadata:

• spec_query_start_loc->size(0) == num_spec_decodes + 1
  

  vllm-xpu-kernels/csrc/xpu/gdn_attn/gdn_attn_interface.cpp

  Lines 152 to 164 in dae50e2

  	if (num_spec_decodes > 0) {
  	TORCH_CHECK(
  	spec_query_start_loc->is_contiguous(),
  	"spec_query_start_loc must be contiguous");
  	TORCH_CHECK(
  	spec_query_start_loc->dtype() == torch::kInt32,
  	"spec_query_start_loc must be of int32 dtype");
  	TORCH_CHECK(
  	spec_query_start_loc->dim() == 1,
  	"spec_query_start_loc must be 1D of shape [num_spec_decodes + 1]");
  	TORCH_CHECK(
  	spec_query_start_loc->size(0) == num_spec_decodes + 1,
  	"spec_query_start_loc must have size [num_spec_decodes + 1]");

• spec_state_indices_tensor->size(0) == num_spec_decodes
  

  vllm-xpu-kernels/csrc/xpu/gdn_attn/gdn_attn_interface.cpp

  Lines 176 to 190 in dae50e2

  	TORCH_CHECK(
  	spec_state_indices_tensor->is_contiguous(),
  	"spec_state_indices_tensor must be contiguous");
  	TORCH_CHECK(
  	spec_state_indices_tensor->dtype() == torch::kInt32,
  	"spec_state_indices_tensor must be of int32 dtype");
  	TORCH_CHECK(
  	spec_state_indices_tensor->dim() == 2,
  	"spec_state_indices_tensor must be 2D of shape [num_spec_decodes, "
  	"num_speculative_tokens + 1]");
  	TORCH_CHECK(
  	num_spec_decodes > 0 &&
  	spec_state_indices_tensor->size(0) == num_spec_decodes,
  	"spec_state_indices_tensor must have size [num_spec_decodes, "
  	"num_speculative_tokens + 1]");

• num_accepted_tokens->size(0) == num_spec_decodes
  

  vllm-xpu-kernels/csrc/xpu/gdn_attn/gdn_attn_interface.cpp

  Lines 193 to 204 in dae50e2

  	TORCH_CHECK(
  	num_accepted_tokens->is_contiguous(),
  	"num_accepted_tokens must be contiguous");
  	TORCH_CHECK(
  	num_accepted_tokens->dtype() == torch::kInt32,
  	"num_accepted_tokens must be of int32 dtype");
  	TORCH_CHECK(
  	num_accepted_tokens->dim() == 1,
  	"num_accepted_tokens must be 1D of shape [num_spec_decodes]");
  	TORCH_CHECK(
  	num_accepted_tokens->size(0) == num_spec_decodes,
  	"num_accepted_tokens size must be num_spec_decodes");

There is also an unconditional XE2 chunk path for prefill:

vllm-xpu-kernels/csrc/xpu/gdn_attn/gdn_attn_interface.cpp

Lines 403 to 410 in dae50e2

	#ifdef VLLM_XPU_ENABLE_XE2
	// XE2 chunk path handles all non-spec tokens whenever there are prefills,
	// even when spec_decodes are also present. The XE2 kernels accept an
	// optional token_indx so they can read mixed_qkvz/mixed_ba and write z /
	// core_attn_out directly at the interleaved global slots indicated by
	// non_spec_token_indx, avoiding host-side gather/scatter.
	if (num_prefills > 0) {
	int batch_size = non_spec_query_start_loc->size(0) - 1;

The exact-size checks are brittle when the caller allocates graph-padded metadata buffers but only the first num_spec_decodes rows are active.

Repro

Hardware/software example where this reproduces:

• Intel Arc Pro B60/B70 class XPU
• vLLM XPU backend with graph capture enabled
• Hybrid GDN Qwen model using DFlash/speculative decode
• Model: Intel/Qwen3.6-35B-A3B-int4-mixed-AutoRound
• DFlash draft model: z-lab/Qwen3.6-35B-A3B-DFlash
• tensor_parallel_size=4
• long-context run, e.g. max_model_len=32768, max_num_batched_tokens=32768
• speculative config with DFlash and num_speculative_tokens=15

Representative launch shape:

export VLLM_TARGET_DEVICE=xpu
export ONEAPI_DEVICE_SELECTOR=level_zero:0,1,2,3
export ZE_AFFINITY_MASK=0,1,2,3
export VLLM_XPU_ENABLE_XPU_GRAPH=1

vllm serve Intel/Qwen3.6-35B-A3B-int4-mixed-AutoRound \
  --tensor-parallel-size 4 \
  --dtype auto \
  --max-model-len 32768 \
  --max-num-batched-tokens 32768 \
  --gpu-memory-utilization 0.85 \
  --speculative-config '{"method":"dflash","model":"z-lab/Qwen3.6-35B-A3B-DFlash","num_speculative_tokens":15}'

Then issue a request whose prompt/decode path enters the captured graph/spec-decode path. The engine can fail during GDN attention before model execution completes with the exact metadata assertion above.

A minimal unit-level repro should be possible by calling torch.ops._xpu_C.gdn_attention with:

• num_spec_decodes = N
• spec_query_start_loc allocated with leading size larger than N + 1
• spec_state_indices_tensor allocated with leading size larger than N
• num_accepted_tokens allocated with leading size larger than N
• valid active data in the prefix only

Expected behavior is that the op should use the active prefix described by num_spec_decodes, not reject the padded allocation.

Expected behavior

The GDN op should tolerate graph-padded speculative metadata tensors when the active prefix is valid. In particular:

• spec_query_start_loc.size(0) should be allowed to be at least num_spec_decodes + 1.
• spec_state_indices_tensor.size(0) should be allowed to be at least num_spec_decodes.
• num_accepted_tokens.size(0) should be allowed to be at least num_spec_decodes.
• Kernel launches should consume only the active prefix.

This would align the metadata contract with graph-padded tensor execution and avoid DFlash/spec-decode startup/runtime failures on valid captured shapes.

Suggested tests

Add coverage for graph-padded spec metadata, including:

1. exact-size metadata tensors still pass;
2. padded metadata tensors with valid active prefix pass and match exact-size output;
3. padded tail values are ignored;
4. undersized metadata tensors still fail;
5. mixed prefill + decode + spec-decode shapes continue to route correctly;
6. unsupported XE2 chunk shapes fall back safely instead of entering a chunk path that assumes dimensions aligned to the chunk size.