b70-optimization-lab

Codex GDN Spec Parity Fix

Date: 2026-06-20

Root Cause

The XPU GDN verifier path did not leave speculative recurrent state columns equal to a one-token-at-a-time target-model decode. The original native _xpu_C.gdn_attention interface only accepts one state index per sequence and writes one final recurrent state. It has no way to receive the full spec_state_indices_tensor table, so it cannot satisfy the verifier contract for every target/draft/bonus column.

The first Python fix made the state columns exact by looping over speculative positions and calling the heavy FLA chunk kernel once per position per layer. That passed the synthetic recurrent-state check but failed the real endpoint canary and was far too slow: the canary diverged at repeat 20/96 and throughput fell to about 6.46 tok/s.

The definitive fix requires the spec column table to be handled inside the native captured recurrent path, not by a Python per-position loop.

Iteration 2 Changes

Implemented a native spec-decode path in vllm-xpu-kernels:

Updated vllm/_xpu_ops.py so XPU spec rows call the new native spec op instead of the Python _xpu_gdn_exact_spec_recurrent per-position loop. No-spec native decode was left on its existing path.

Also tested two follow-up acceleration attempts:

Neither follow-up was canary-safe.

Validation

Direct-source native spec path, short endpoint probe:

cd /home/steve/llm-optimizations
MODEL_PATH=/mnt/fast-ai/qwen36-quark-int8-fp8-mtp-hybrid \
SERVER_LAUNCHER=scripts/launch-qwen36-quark-int8-accepted.sh \
VLLM_QWEN35_MTP_FORCE_FP8_BLOCK=1 \
VLLM_EXTRA_ARGS='--speculative-config {"method":"mtp","num_speculative_tokens":1}' \
COMPILATION_CONFIG='{"cudagraph_mode":"PIECEWISE","max_cudagraph_capture_size":128}' \
XPU_GRAPH=1 VLLM_XPU_ENABLE_XPU_GRAPH=1 VLLM_XPU_FORCE_GRAPH_WITH_COMM=1 \
VLLM_XPU_GRAPH_NOOP_COMM_CAPTURE=1 VLLM_XPU_GDN_NATIVE_FALLBACK=prefill \
VLLM_XPU_GDN_PREFILL_RECURRENT_FALLBACK=1 VLLM_XPU_DISABLE_PREFILL_CUDAGRAPH_REPLAY=1 \
VLLM_XPU_GREEDY_SAMPLE_TOPK_FALLBACK=1 VLLM_XPU_INT8_MOE_MIXED_WORKSPACE=1 \
GPU_MEMORY_UTILIZATION=0.95 ABLATION_FAST_GRAPH_AUTOCONFIG=0 \
VLLM_XPU_GDN_PROMOTE_ACCEPTED_SPEC_STATE=1 \
VLLM_XPU_GDN_NONSPEC_POSTPROCESS_ACCEPTED_STATE=0 \
JSON_REPEATS=7 COLOR_REPEATS=0 \
bash scripts/run-qwen36-ablation-candidate.sh tp4-mtp-k1-direct-source-json7

Result:

Earlier native spec path with postprocessing enabled:

Native indexed no-spec decode acceleration attempt:

Fused spec kernel acceleration attempt:

Synthetic direct-source partial-acceptance check:

Current Blocker

The native spec-table expansion is parity-capable on a short endpoint JSON probe, but it is still far below the hard throughput target. The best canary-safe native result was 42.3866 tok/s, versus the requested target of roughly within 20% of the 93.55 tok/s no-spec baseline.

The remaining cost is still inside the captured native spec path: each GDN layer/spec row materializes intermediate q/k/v/b/a tensors and launches the spec conv plus spec recurrent kernels. Removing the Python per-position loop fixed the catastrophic 6 tok/s path, but this two-kernel-plus-intermediate native path is still about 2x too slow.

The attempted one-kernel fused native path is the likely required performance direction, but the first implementation was not canary-safe. The most likely remaining issue is an ordering/math mismatch in the fused path around conv-state sourcing/writes, BA/Z gating, or recurrent-state publication for partially accepted requests.

No commit was made because the required endpoint canary target was not met.

Iteration 3 Changes

Implemented the untried sequence-path direction:

The XPU kernels were rebuilt with scripts/build-vllm-xpu-kernels-xpu-c-only.sh, then _xpu_C.abi3.so, libgrouped_gemm_xe_2.so, and libgdn_attn_kernels_xe_2.so were copied into vllm_xpu_kernels/.

Iteration 3 Validation

Required endpoint canary:

Static PIECEWISE metadata short probe:

Partial-accept diagnostic:

Earlier diagnostic contrast:

Iteration 3 Blocker

The sequence path still does not pass the required endpoint canary and remains far below the requested speed band near the 93.55 tok/s no-spec baseline. No commit was made.

The narrowest non-masking failure is partial-accept state publication for the GDN recurrent cache. With normal partial acceptance, the endpoint diverges after metrics on the JSON canary; with all-or-nothing acceptance, the same short after-metrics JSON probe passes. The earlier synthetic direct-source check also had exact z/conv output while core/recurrent state still drifted, so conv history threading is not the first proven divergent state.

The remaining bad state is therefore the SSM/recurrent state column consumed after a partial accept/reject boundary. For num_speculative_tokens=1, this is the transition that chooses column 0 for accepted length 1 versus column 1 for accepted length 2. The first endpoint-visible wrong token in the static-metadata probe is completion token index 8, where the unit key is replaced by 历史悠久 Pioneer after several prior tokens were correct.

I could not capture a trustworthy per-layer first-difference line: every state-tensor trace that reads conv/SSM values back to CPU changes the timing and makes the short after-metrics JSON probe pass. Existing non-invasive evidence points to SSM/recurrent publication after partial acceptance, not conv state, but the exact GDN layer cannot be named from a non-masking trace.

Iteration 4 Changes

Implemented the ReplaySSM spec-verify recurrent reconstruction as native _xpu_C SYCL kernels:

Updated vllm/model_executor/layers/mamba/gdn_linear_attn.py so the ReplaySSM path calls the native recurrent kernel for max_spec_len <= 2 and max_cache_len in (2, 4), and calls the native commit kernel when available. The old torch implementations remain as env-gated fallbacks.

The XPU kernels were rebuilt with /home/steve/llm-optimizations/scripts/build-vllm-xpu-kernels-xpu-c-only.sh, then _xpu_C.abi3.so and libgrouped_gemm_xe_2.so were copied into /home/steve/src/vllm-xpu-kernels/vllm_xpu_kernels/.

Iteration 4 Validation

Focused native-vs-torch parity:

Micro-timing for _xpu_C.gdn_replayssm_spec_decode with endpoint dimensions:

Endpoint attempts:

The required json/color canaries were not allowed to complete on the slow Iteration 4 runs; the live throughput was already orders of magnitude outside the requested speed band and the runs were stopped to avoid burning the full 512-token probes.

Iteration 4 Blocker

The fused recurrent ReplaySSM kernel itself is not the speed blocker: direct timing is around 0.020 ms per call for the endpoint shape. The remaining endpoint speed bottleneck is the conv/spec scaffolding that is still executed through torch before the recurrent kernel:

This is exactly the unfused “conv state thread as a sequence, no slot-copy” part of the requested work. The recurrent reconstruction and cursor/flush metadata are now native and parity-matched, but the current path still builds a temporary conv-state table and writes pending conv history through torch for every GDN layer and verify step. That keeps the endpoint at 0.1-0.4 tok/s instead of near the 93.55 tok/s baseline.

The next required kernel is a native sequence-threaded causal-conv/spec staging path that consumes the base conv state plus raw spec sequence directly and produces q/k/v/a/b and pending conv history without Python torch materialization or slot-copy state. No commit was made because the endpoint gate did not pass.

Iteration 5 Changes

Implemented the ReplaySSM conv/spec staging fusion requested by Iteration 4:

Rebuilt with /home/steve/llm-optimizations/scripts/build-vllm-xpu-kernels-xpu-c-only.sh and copied _xpu_C.abi3.so plus libgrouped_gemm_xe_2.so into /home/steve/src/vllm-xpu-kernels/vllm_xpu_kernels/.

Iteration 5 Validation

Focused native-vs-torch parity:

Speed gate did not pass, so the required 96/96 json and 96/96 color canary was not run to completion:

Profiling result:

The old torch conv/spec staging sequence is no longer the steady-state dominant op. The speed gate is still missed because the endpoint has large first-spec-step warmup/recapture plus draft/proposer/runtime overhead, and the client/vLLM decode histogram still reports multi-second decode time that is not explained by steady native GDN kernel timings. No commit was made because the required speed/canary gate did not pass.

Iteration 6 Changes

Implemented the graph-capture lifecycle change in vLLM only:

Rebuilt with /home/steve/llm-optimizations/scripts/build-vllm-xpu-kernels-xpu-c-only.sh and copied _xpu_C.abi3.so plus libgrouped_gemm_xe_2.so into /home/steve/src/vllm-xpu-kernels/vllm_xpu_kernels/.

Iteration 6 Validation

Lightweight code checks:

Required endpoint canary attempt:

Targeted all-rank timing profile:

The first captured spec step is the concrete remaining outlier:

Later spec steps show the graph-captured target forward is mostly steady:

There is still a smaller later outlier:

Conclusion: ReplaySSM ensure_state and conv staging are no longer the visible steady bottleneck, and the target forward is graph-captured enough to reach ~15-17 ms on steady spec steps. The endpoint speed gate is still missed because the first post-prefill spec step blocks for multi-second draft/rejection-sampler runtime across all TP ranks, with occasional smaller draft/proposer outliers later. No commit was made because the required speed/canary gate did not pass.

Iteration 7 Warmup Isolation Attempt

Finding from the fast-but-incorrect full-sampler warmup:

Implemented isolation changes in vLLM:

Lightweight validation:

Endpoint validation status:

Iteration 8 Serving Crash Fix and Stop at Third Layer

Fixed the first real-request crash:

Validation of the crash fix:

Native stage-conv correctness issue found and fixed:

Stage-conv parity after the kernel rebuild:

Endpoint result after the stage-conv fix:

Stop reason:

Remaining issue set:

No commit was made because the required endpoint gate did not pass.

Follow-up (2026-06-22): ReplaySSM + XPU graph was fast but not valid

The ReplaySSM correctness path (snapshot/restore on partial-reject and full-accept, --mamba-cache-mode align) passed BOTH canaries but ran fully eager at ~13 tok/s. Two findings explained the gap, but did not produce a valid >150 tok/s result:

  1. Snapshot/restore copies were NOT the bottleneck. Batching the per-layer state[block_id].clone() / .copy_(saved) loops into single torch._foreach_copy_ calls (new helper _xpu_batched_restore_snapshots + batched snapshot in _xpu_snapshot_mamba_running_states_for_spec, gpu_model_runner.py) is correctness-neutral but moved throughput only 13.08 -> 13.26 tok/s (noise). It was later reverted because it did not help the endpoint and was not the root cause.
  2. Eager execution was the bottleneck. The MTP run never captured an XPU graph. Enabling it required overriding the launcher default VLLM_XPU_DISABLE_PREFILL_CUDAGRAPH_REPLAY=1 to =0, because the align/ReplaySSM spec-decode batches register as non-uniform, and PREFILL_REPLAY_OFF=1 skips every non-uniform PIECEWISE capture (gpu_model_runner.py:12318). With the flag at 1, all 19 captures were skipped (finished in 1 secs, took 0.00 GiB), so decode hit an uncaptured shape and tried to capture at runtime -> RuntimeError: CUDA graph capturing detected at an inappropriate time (compilation/cuda_graph.py:1595). With =0, capture actually ran (finished in 14 secs) and decode replayed the graph.

Smoke result tp4-mtp-k1-replayssm-iter14-graph-capnonuni (8+8 canaries, 2 metric repeats):

CORRECTION: the max_capture=128 config above fails color AT SCALE

The 8+8 smoke passed, but the full 96/96 run of the max_cudagraph_capture_size=128 config FAILED color: json 96/96, but color produced 1 garbage output ~1 in 30 requests (e.g. 伪 伪 伪 ... or a spurious prefix token like _kategori/institution/aff followed by the correct blue, green, orange, red). Throughput held at ~73-76 tok/s.

Isolation runs (all original code, graph on):

PARTIAL: cap capture size so prefill stays eager (fixes garbage, NOT robust)

Set max_cudagraph_capture_size=4. Decode-spec batches are tiny (1 seq x 2 tokens = size 2), so sizes 1/2/4 are still captured (fast decode preserved), but prefill (30-512 tokens) exceeds 4 and runs eager -> no captured prefill. This DID eliminate the first-token garbage corruption (signature changed away from 伪/garbage).

But it is NOT a robust pass. Initial run iter19-graph-capsmall (32 json + 96 color) passed 96/96 color — but that was luck at the lower bug rate. The full lock iter19-graph-capsmall-full (96 json + 96 color), SAME config, FAILED:

So there are two graph-vs-eager races:

  1. captured-prefill garbage -> fixed by max_cudagraph_capture_size=4;
  2. eager full-accept restore racing the captured decode replay (~1/50-1/100) -> NOT fixed. The restore is an eager state[block_id].copy_(saved) between captured decode replays; under graph the rollback intermittently fails to take, so the next step double-processes the bonus token -> plausible-but-wrong output.

VLLM_XPU_SYNC_CUDAGRAPH_REPLAY=1 (tried at max=128) only reduced severity and halved throughput (75 -> 43, below the 74.84 gate), so a global sync is not viable.

Bottom line / status

Repro for the (partial) graph config:

cd /home/steve/llm-optimizations
VLLM_XPU_GDN_REPLAYSSM_SPEC=1 \
VLLM_XPU_SPEC_DECODE_RESTORE_DRAFT_PARTIAL_REJECT_GDN_STATE=1 \
VLLM_XPU_SPEC_DECODE_RESTORE_DRAFT_FULL_ACCEPT_GDN_STATE=1 \
VLLM_EXTRA_ARGS="--no-async-scheduling --mamba-cache-mode align" \
COMPILATION_CONFIG='{"cudagraph_mode":"PIECEWISE","max_cudagraph_capture_size":4}' \
XPU_GRAPH=1 VLLM_XPU_ENABLE_XPU_GRAPH=1 VLLM_XPU_FORCE_GRAPH_WITH_COMM=1 \
VLLM_XPU_GRAPH_NOOP_COMM_CAPTURE=1 \
VLLM_XPU_DISABLE_PREFILL_CUDAGRAPH_REPLAY=0 \
SERVER_LAUNCHER=/home/steve/llm-optimizations/scripts/launch-qwen36-quark-int8-hybrid-mtp.sh \
FI_TCP_IFACE=eno1 CCL_KVS_IFACE=eno1 GPU_MEMORY_UTILIZATION=0.95 \
JSON_REPEATS=96 COLOR_REPEATS=96 \
bash scripts/run-qwen36-ablation-candidate.sh <label>
This site is open source. Improve this page.