[GDN] Spec-decoding-aware attention kernel

[jasonboukheir]

Purpose

Adds a spec-decoding code path to the SYCL gdn_attention kernel that mirrors
the FLA Triton IS_SPEC_DECODING semantics, so vLLM's hybrid-linear-attention
models (Qwen3-Next / Qwen3.5 / Qwen3.6) can run MTP / EAGLE-style draft +
verify on XPU through the native kernel instead of falling back to the FLA
Triton path on every spec step.

Problem

The current SYCL gdn_attention schema doesn't carry the two tensors the
spec-decode verify step needs (spec_state_indices_tensor,
num_accepted_tokens), so vLLM's dispatcher routes the spec path to the FLA
Triton fallback in _gdn_xpu_spec_python_path. That fallback works but is
~10× slower than the native SYCL kernel on the same captures (0.29 ms vs
3.0 ms median on K=1 and K=4 layer captures — measured via the replay
harness in the next section).

The deeper issue is the conv-state layout contract. Under spec decoding, FLA
treats the per-sequence conv ring as a rolled history at slot 0 of
cache_indices — only that slot is updated, slots 1..K-1 are untouched. A
prior native spec attempt that scattered per-token snapshots across K slots
diverged from FLA (max conv-state diff ≈ 47 on K=4) and propagated through
the SSM, making the kernel unusable for the verify step.

Changes

• csrc/xpu/gdn_attn/causal_conv1d.hpp — two-pass spec-aware conv1d.

  • Pass 1 (in-kernel): stages each candidate token's per-token conv
    state into the existing conv_states_tmp buffer (resized to
    (batch, state_len, conv_elems) under spec). Per-work-item, this
    is one write of the current input at
    tmp[batch, state_len - K_call + t_in_seq, dim] plus, for
    t_in_seq < state_len - K_call, a shift copy from the slot's
    pre-state.
  • Pass 2 (spec_state_roll_kernel): consolidates tmp into the
    rolled history at slot[cache_indices[batch], :, :]. Pass 2 is
    the sole writer of the slot — race-free with pass 1, which only
    wrote tmp.
  • The IS_SPEC load now reads from cache_indices[batch_id] with a
    load_offset_shift = n_acc - 1, matching FLA's column-0 contract.
  • state_len is plumbed through kernel_launcher /
    KERNEL_LAUNCHER_IMPL, computed at the launcher as
    conv_states_stride_0 / conv_elems.

• csrc/xpu/gdn_attn/gated_delta_rule.hpp — IS_SPEC template
  parameter. When IS_SPEC, the per-sequence load slot comes from
  spec_state_indices[i_n, num_accepted_tokens[i_n]-1], the per-token
  store slot comes from spec_state_indices[i_n, t-seq_start], and a
  raw slot value ≤ 0 (NULL_BLOCK_ID) is treated as "no valid prior
  state". The non-spec path is template-specialized and behaviorally
  unchanged.

• csrc/xpu/gdn_attn/gdn_attn_interface.cpp — dispatch. Both
  spec_state_indices_tensor and num_accepted_tokens must be provided
  together; mismatched pairs are rejected at the interface boundary.
  When both are set, routes through native kernels with IS_SPEC=true.

• csrc/xpu/ops.h, csrc/xpu/torch_bindings.cpp. Extends the
  gdn_attention schema with the two new tensors as
  Optional[Tensor] = None. Existing non-spec callers see no schema
  change; passing None for both preserves bit-exact pre-PR behavior.

The SSM kernel is unchanged — FLA's fused_sigmoid_gating already uses
per-token spec slots for the SSM ring, and the existing SYCL
gated_delta_rule.hpp SSM path was already correct under that contract.

Test Plan

Paired vLLM PR: vllm-project/vllm#42382
(draft). It adds the dispatcher that calls into the new schema and
the pytest replay harness. The kernel changes in this PR are usable
today via the SYCL op directly; the paired PR wires them into vLLM's
GDN dispatch behind VLLM_XPU_USE_SYCL_SPEC_GDN={unset|0|auto|1} so
existing users see no behavior change until they opt in.

The kernel is exercised by a replay harness (in the paired PR) at
tests/kernels/xpu/test_spec_gdn_replay.py. It captures
(metadata, inputs, outputs) tuples from a live Qwen3.6 MTP run via
VLLM_XPU_DUMP_SPEC_GDN=<dir> (default cap 200 tuples), then replays
each capture through the native kernel and diffs against an inline FLA
oracle within bf16 tolerance (atol=rtol=2e-2).

The harness covers:

• Single-batch match vs. FLA — every captured tuple, both non_spec
  and spec flavors.
• Synthetic mixed batch (test_sycl_mixed_batch_matches_per_subset) —
  pairs each non_spec capture with the same-layer
  spec_K4_min1_max1 capture, builds a unified pool with disjoint slot
  ranges, replicates the production split / two-kernel-call /
  index_copy_ scatter, and diffs each subset against its own
  reference.
• reorder_input=False equivalence
  (test_sycl_reorder_input_false_equivalence) — repacks captured
  projections into the per-k_head interleaved layout used in
  production Qwen3-Next traffic and verifies the kernel produces
  equivalent outputs under both reorder_input=True and
  reorder_input=False on fresh pools.
• Determinism — 10× byte-equal core_attn_out + conv_state +
  ssm_state on K=1 and K=4 captures.

E2E benchmark via vllm bench serve on Intel Arc B70:

Model:   palmfuture/Qwen3.6-35B-A3B-GPTQ-Int4
Quant:   INC (GPTQ-sym-int4 MoE → xpu_fused_moe(is_int4=True))
KV:      turboquant_k3v4_nc (3-bit MSE-Lloyd-Max K + 4-bit V)
Spec:    MTP, num_speculative_tokens=2
Dataset: random 256-in / 512-out, 20 prompts, --max-concurrency 1 --ignore-eos

Test Result

Replay harness (kernel-level correctness vs. FLA):

Rung	Test	Result
Fully accepted K=1	single-batch vs. FLA	90/90
Fully accepted K=4	single-batch vs. FLA	169/170 spec (1 borderline bf16, 2 cells over atol=2e-2)
num_accepted=1, K=3	single-batch vs. FLA	180/180
Mixed batch (synthetic)	spec + non-spec split	29/30 (1 pre-existing layer-0 non-spec non-determinism, also fails the single-batch test on the same capture, not introduced by this PR)
reorder_input=False equivalence	200 captures	200/200
Determinism (10×)	bytes equal on core_attn_out + conv_state + ssm_state	PASS

Perf (replay harness, median over captures on Intel Arc B70):

	FLA Triton fallback	This PR (native SYCL)
K=1 capture	~3.0 ms	~0.29 ms (~10.3×)
K=4 capture	~3.0 ms	~0.29 ms (~10.4×)

E2E (single-stream vllm bench serve, baseline = no spec decoding):

metric	baseline	spec (MTP K=2)	delta
output throughput (tok/s)	58.85	75.80	+28.8%
wall-clock (s)	173.99	135.09	−22.4%
median TPOT (ms)	16.92	13.08	−22.7%
median TTFT (ms)	94.61	115.70	+21 ms (one-time MTP draft graph dispatch)
acceptance rate	—	64.5%	—
mean accepted length	—	2.29 / 3	—
per-position acceptance	—	P0 74.7%, P1 54.2%	—

20 requests, no per-step fallback to FLA, no failures.
Raw JSON available on request.

Related issues

• Orthogonal to please expose gdn::causal_conv1d and gdn::gated_delta_rule #189 (request to expose gdn::causal_conv1d and
  gdn::gated_delta_rule as standalone torch ops). This PR keeps the
  existing single-op surface and extends its signature; it does not
  split the kernel.
• Tangentially related to GDN kernel rejects padded inputs from torch.compile / cudagraph capture (incompatible with vLLM PIECEWISE) #320 (GDN kernel rejects padded inputs
  from torch.compile / cudagraph capture). Same file
  (gdn_attn_interface.cpp) but a separate bug — that one is about the
  size(0) == num_actual_tokens precondition under PIECEWISE capture
  and is unaffected by this PR.

[mayuyuace]

Hello, we are planning to support this feature.
Can you provide UT for this PR, and mtp results of the model Qwen Next running?

[jasonboukheir]

@mayuyuace — thanks!

UT: added tests/gdn_attn/test_gdn_attn_spec.py on the branch — 14 tests, all green on Arc Pro B70. Covers K=N aligned-ring ≡ non-spec native decode (bit-exact), NULL_BLOCK_ID suppresses pool I/O, and three arg-validation negatives. Full FLA-Triton oracle coverage stays in the paired vLLM replay harness (vllm-project/vllm#42382).

MTP results: Arc Pro B70 is 32 GiB VRAM and the public Qwen3-Next-80B-A3B quants don't fit (Intel AutoRound-int4-mixed ≈ 47 GiB, AWQ-4bit / MLX-4bit ≈ 40 GiB, NVFP4 ≈ 38 GiB; no public weights-only INT3). Ran on Lorbus/Qwen3.6-27B-int4-AutoRound instead — same hybrid GDN architecture as Qwen3-Next, just a smaller dense backbone. Settings match the PR description (random 256-in / 512-out, 20 prompts, max-concurrency 1, --ignore-eos, KV turboquant_k3v4_nc):

metric	baseline	spec MTP=2	delta
output throughput (tok/s)	11.45	17.81	+55.5%
wall (s)	894.34	575.10	−35.7%
median TPOT (ms)	86.04	49.04	−43.0%
acceptance rate	—	66.96% (P0 77.5% / P1 56.4%)	—
mean accepted length	—	2.34 / 3	—

20/20 requests successful, no FLA fallback. Acceptance lines up with the 64.5% the PR describes on the 35B-A3B run. (The 35B MoE flavor still needs vllm-project/vllm#41426 to start up; once that lands I'll rerun on the same model from the PR description.)

If Intel has a 32-GiB-fitting Qwen3-Next variant I should be testing against, happy to grab it.

[mayuyuace]

Can you provide accuracy data of the models also?

[jasonboukheir]

@mayuyuace I see you have a PR here #368 . Looking through it, I think it's best to close this out and take that!