[Bug]: Spec-decode draft with head_dim=256 on a hybrid (GDN) target crashes FA2 paged-KV shape check at warmup

[arthurgao2003]

Environment

• vLLM 0.25.0 (pip wheel), torch 2.11.0+cu130, CUDA 13.0
• 8× NVIDIA RTX PRO 5000 Blackwell (sm_120), TP=2
• Target: Qwen3.5-family 27B hybrid-attention model (GDN linear attention + full attention every 4th layer), bf16
• Drafts: DeepSpec-format DSpark/DFlash block-diffusion drafts (5 layers, hidden 5120)

Description

Serving a hybrid GDN target with a DSpark speculative draft whose attention geometry is
24 heads / 4 KV heads / head_dim 256 crashes during engine warmup
(kernel_warmup → _dummy_run → speculator.propose) inside the draft's forward:

RuntimeError: k must have shape (num_blocks, page_block_size, num_heads_k, head_size)

raised by torch.ops._vllm_fa2_C.varlen_fwd via
vllm/v1/attention/backends/flash_attn.py:952 → vllm_flash_attn/flash_attn_interface.py:315.

The same command with a draft of identical per-token KV bytes but
32 heads / 8 KV heads / head_dim 128 starts and serves normally.
head_dim looks like the discriminating variable, not the method or the padded block size:

draft geometry	spec method	"Setting attention block size to N tokens"	result
32 q / 8 kv / head_dim 128	dflash	N = 816	✅ serves
32 q / 8 kv / head_dim 128	dspark	N = 832	✅ serves
24 q / 4 kv / head_dim 256	dspark	N = 816	💥 warmup crash

(N from the hybrid page-alignment path: interface.py:890 "Setting attention block size to N tokens to ensure that attention page size is >= mamba page size", followed by
interface.py:914 padding mamba page size to match. Both 816 and 832 are multiples of 16.
The target itself also uses head_dim 256 for its full-attention layers and serves fine
without a draft, and with the head_dim-128 drafts.)

VLLM_ATTENTION_BACKEND=FLASHINFER and VLLM_ATTENTION_BACKEND=TRITON_ATTN do not
change the failing path — the traceback still goes through _vllm_fa2_C.varlen_fwd,
so there appears to be no backend-selection workaround for the draft.

For what it's worth, the same draft checkpoint served on an L20 (sm_89) box running a
2026-07-09 nightly (0.23.1rc1.dev968+g2c17d33f4); we have not isolated whether the
regression is version- or GPU-arch-related.

Reproduce

CUDA_VISIBLE_DEVICES=6,7 python -m vllm.entrypoints.cli.main serve <27B-hybrid-target> \
  --served-model-name grpo27b -tp 2 --max-num-seqs 64 --max-model-len 16384 \
  --gpu-memory-utilization 0.90 --port 40021 \
  --speculative-config '{"method":"dspark","model":"<draft-headdim-256>","num_speculative_tokens":6}'

Draft config (DeepSpec save_pretrained layout, model_type patched to qwen3):
5 layers, hidden 5120, 24 heads / 4 kv / head_dim 256, block_size 7,
dflash_config={"mask_token_id":248063,"target_layer_ids":[3,15,31,43,59]}.

Traceback

Full worker traceback

WorkerProc hit an exception.
Traceback (most recent call last):
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/vllm/v1/executor/multiproc_executor.py", line 996, in worker_busy_loop
    output = func(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/vllm/tracing/otel.py", line 178, in sync_wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/vllm/v1/worker/gpu_worker.py", line 758, in compile_or_warm_up_model
    kernel_warmup(self)
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/vllm/model_executor/warmup/kernel_warmup.py", line 95, in kernel_warmup
    flashinfer_autotune(worker.model_runner)
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/vllm/model_executor/warmup/kernel_warmup.py", line 157, in flashinfer_autotune
    runner._dummy_run(
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/vllm/v1/worker/gpu/eplb_utils.py", line 38, in wrapper
    result = fn(self, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/vllm/v1/worker/gpu/model_runner.py", line 602, in _dummy_run
    self.speculator.propose(
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py", line 404, in propose
    self._generate_draft(
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py", line 163, in _generate_draft
    head_hidden = self._run_model(
                  ^^^^^^^^^^^^^^^^
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py", line 185, in _run_model
    last_hidden_states = self.model(
                         ^^^^^^^^^^^
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/vllm/model_executor/models/qwen3_dflash.py", line 710, in forward
    return self.model(input_ids, positions, inputs_embeds)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/vllm/compilation/decorators.py", line 520, in __call__
    return self.aot_compiled_fn(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/torch/_dynamo/aot_compile.py", line 224, in __call__
    return self.fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/vllm/model_executor/models/qwen3_dflash.py", line 594, in forward
    def forward(
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/vllm/compilation/caching.py", line 217, in __call__
    return self.optimized_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<string>", line 38, in execution_fn
  File "<string>", line 6, in __vllm_inlined_submods__1
  File "/data2/dspark/env_clone/lib/python3.11/site-packages/torch/_ops.py", line 1269, in __call__
    return self._op(*args, **kwargs)

Related

• [Hybrid][KVCache] Adjust alignment block size according attn supported kernel sizes #36317 (hybrid KV cache alignment block size vs attention kernel supported sizes — open)
• [HMA]Move hybrid blksize to update_block_size_for_backend to fix attn supported block size is not 16 issue #37467 (hybrid block size moved to update_block_size_for_backend)
• [Bugfix] Split attention groups by num_heads_q for spec-decode drafts #43543 (split attention groups by num_heads_q for spec-decode drafts)