[TSUMUGI-XE]

Purpose

vLLM's RecoverSSM (--use-replayssm, vllm-project/vllm#58863) verifies drafts from a read-only checkpoint: each request owns one state block, the verify step records a float32 replay entry per token, and the accepted prefix is committed after sampling. gated_delta_rule_spec instead writes one state per draft position, so on XPU RecoverSSM currently has to fall back to the Triton/FLA path (vllm-project/vllm#59982), which is slower on our hardware.

This adds gated_delta_rule_spec_recoverssm, a variant of the spec-decode delta stage with the RecoverSSM contract. The math is the same as gated_delta_rule_spec_kernel; the differences are:

• the initial state is read from spec_state_indices_tensor[n] (1D, one block per request) and never written;
• per token it writes replay_state[block, hv, t, :] = [corr (V), normalised k (K), decay (1)] in float32, matching the layout of vLLM's gdn_recoverssm_verify;
• a request whose block index is <= null_block_id gets zero outputs and its replay block is not touched;
• the gate input a + dt_bias is added in float32 so that the recorded decay matches vLLM's float32 verify/commit.

The conv stage needs no kernel change: we call causal_conv1d_spec with the request's block repeated across the index row and num_accepted_tokens = 1. End to end in our vLLM tree, per-token logprobs stayed within run-to-run noise versus the Triton path; we did not compare conv states directly. Existing ops are unchanged.

vllm-project/vllm#58863 is still open. The record layout matches its head (a77830311b), and we will follow it if it changes; this can wait to merge until #58863 lands.

Test Plan

Built GDN-only for bmg-g31-a0 (oneAPI 2026.1, torch 2.14.0+xpu) and run on one Intel Arc Pro B70:

GDN_KERNELS_ENABLED=ON XPU_SPECIFIC_KERNELS_ENABLED=ON VLLM_XPU_ENABLE_XE2=ON \
BASIC_KERNELS_ENABLED=OFF FA2_KERNELS_ENABLED=OFF MOE_KERNELS_ENABLED=OFF \
MQA_LOGITS_KERNELS_ENABLED=OFF MHC_KERNELS_ENABLED=OFF XPUMEM_ALLOCATOR_ENABLED=OFF \
VLLM_XPU_ENABLE_ONEDNN=OFF VLLM_XPU_ENABLE_XE3P=OFF VLLM_XPU_ENABLE_XE_DEFAULT=OFF \
VLLM_XPU_AOT_DEVICES=bmg-g31-a0 VLLM_XPU_XE2_AOT_DEVICES=bmg-g31-a0 \
  python setup.py build_ext --inplace
pytest tests/gdn_attn/test_gated_delta_rule_recoverssm.py \
       tests/gdn_attn/test_gated_delta_rule.py -k "recoverssm or mtp"

The comparison against vLLM's Triton verify used a local script (not part of this PR) that feeds the same causal_conv1d_spec outputs to both.

Test Result

• tests/gdn_attn/test_gated_delta_rule_recoverssm.py (new, 384 cases): float32 reference for outputs and replay records, checkpoint unchanged, folding the records reproduces the recurrence state, ragged batches, a null-block row, spec_token_indx given or None, fp16/bf16, fp32/activation-dtype state, HV 32/48. With tests/gdn_attn/test_gated_delta_rule.py -k mtp: 448 passed.

  • The record tolerance (1e-4) is tight on purpose: a build where a + dt_bias was rounded to bf16 failed 160 of these cases (decay off by ~1e-3).

• Against vLLM's Triton gdn_recoverssm_verify ([Spec Decode][Qwen4Exp] RecoverSSM for Qwen GDN and the PLE short conv vllm#58863) on the same causal_conv1d_spec outputs, Qwen3.8-Flash-Next dims (H=16, HV=48, K=V=128, 2 tokens per request; 8 batches of 1–13 requests incl. ragged, null block, no token_indx): max output difference 3.1e-5 (outputs ~3–5e-2), replay records within 2.4e-7 (corr ~2.7), 6e-8 (k) and 1.8e-7 (decay), checkpoint unchanged.

• Speed. Verify stage only, one GDN layer, XPU graph replay, Flash-Next dims, 2 tokens per request; µs, median of 3 runs in one process. The last column is for reference only and is not like-for-like: without RecoverSSM the kernel writes one state per draft position, while RecoverSSM writes small replay records here and commits one state per request after sampling (a separate step in vLLM for both RecoverSSM columns, not included).

  requests	Triton RecoverSSM path	causal_conv1d_spec + this op	(for reference) gdn_attention, no RecoverSSM
  1	28.8	13.7	17.8
  8	123.1	37.2	79.3
  13	203.2	54.2	170.7

  End to end the difference is small (measured with this kernel wired into vLLM locally); the main benefit is that RecoverSSM on XPU no longer has to fall back to the Triton/FLA path.

• Not tested: other devices; the XE2 chunk path is not touched.

AI assistance (Claude Code) was used; every changed line was reviewed by the submitter.

[jikunshang]

cc @mayuyuace @yma11