[8/10] Add two-precision ReplaySSM state quantization serving

[kaix-nv]

Linear-attention series — 10 PRs

Order	PR	Depends on
1/10	#2497 GDN state/W QAT foundation	main
2/10	#2519 GDN/KDA state QAT + INT8/FP8	#2497
3/10	#2657 Megatron Bridge QAT/QAD example	#2519
4/10	#2805 NVFP4 recurrent-state QAT	#2657
5/10	#2829 TensorQuantizer stochastic rounding	#2805
6/10	#2788 ReplaySSM state QAT and two-precision GDN	#2829
7/10	#2541 vLLM state-only fake quantization	#2788
8/10	#2878 ReplaySSM state quantization serving	#2541
9/10	#2503 GDN/KDA prefill GEMM quantization	#2878
10/10	#2507 Experimental approximate inverse	#2503

The nine open PRs form native GitHub stack #2879 in this order. #2497 has landed, so #2519 targets main.

#2829 supplies the shared rounding policy before ReplaySSM training and serving. #2788 owns the ReplaySSM training foundation; #2541 owns ordinary state-QDQ serving; #2878 adds the experimental two-precision ReplaySSM serving backend and native stochastic NVFP4 read integration. The serving branch inherits rounding from #2829 without a temporary backport.

#2503 and #2507 retain their existing code and commit heads. Their functional ports and rebases remain deferred until their immediate parents land; their bases and table positions express the intended merge order. Serving-time prefill-GEMM QDQ remains deferred until an optimized fused kernel is available.

What does this PR do?

Type of change: New feature.

Add an opt-in vLLM worker for two-precision GDN ReplaySSM state quantization. It installs the recipe before cache allocation so vLLM owns the FP32 recurrence anchor, FP8/NVFP4 read checkpoint and 16-token replay buffers. Chunked prompt prefill, decode cache refresh, request-slot reuse and decode CUDA graphs share the same cache policy as training.

NVFP4 read checkpoints optionally use gdn_read_quantizer.rounding="stochastic", supplied by #2829. E4M3 block scales retain nearest rounding and the recurrence anchor stays FP32. Training saves the sampled reads for its identity-STE backward; serving uses per-request device RNG seeds that work with CUDA graphs. The temporary rounding backport has been removed; the API comes from the parent stack.

This is a new backend with a coupled worker, cache layout and refresh kernel. The worker and cache must land together to have a runnable intermediate state. It reuses the optional upstream ReplaySSM emulation kernels rather than copying their implementation. This is floating-storage numerical emulation, not packed-cache compression.

Usage

Install the optional ReplaySSM dependency pinned in examples/llm_qat/linear_attention/REPLAYSSM.md, alongside vLLM 0.30.x. Use an exported Hugging Face checkpoint and the same explicit recipe as training:

export PYTHONPATH="$PWD/examples/vllm_serve:$PYTHONPATH"
export RECIPE_PATH=general/ptq/gdn_replayssm_nvfp4_read
VLLM_GDN_DECODE_KERNEL=triton vllm serve /path/to/hf-checkpoint \
  --worker-cls replayssm_worker.ReplaySSMWorker \
  --dtype bfloat16 --max-num-seqs 1 \
  --compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE"}' \
  --no-enable-prefix-caching --no-async-scheduling \
  --mamba-cache-mode none --mamba-cache-dtype float32 \
  --additional-config '{"gdn_prefill_backend":"triton"}'

The example README includes the complete CUDA-graph launch configuration and the YAML override for stochastic NVFP4 reads. Supports GDN with K=128, one GPU and one active request; speculative decoding and cache transfer are unsupported. Standard inference export still rejects stochastic quantizers: this worker installs the supported policy from its explicit recipe.

Testing

• Restack: 166 focused CPU tests passed across quantizers, policy/checkpoint lifecycle, state decode/replay/reference tests and attention dispatch. Focused pre-commit and git diff --check passed.
• The restacked ReplaySSM kernels, native training adapter, serving plugin and GPU test files match the previously validated db9053881f source byte-for-byte.
• Prior JHB GB300 validation: 18 ReplaySSM GPU cases passed. After correcting the engine fixture to use a YAML recipe, the final rerun passed three stochastic training cases (including packed gradients) and the stochastic vLLM engine test with full decode CUDA graphs, repeated requests and seed replay.
• JHB runtime: vLLM 0.30.0, PyTorch 2.13.0+cu130, Triton 3.7.1 and fla-core 0.5.1. The engine test uses an offline tiny Qwen3-Next model with dummy weights; no model-accuracy, throughput or compression result is claimed. GPU tests have not been rerun for the stack reorder.

Before your PR is "Ready for review"

• Backward compatible: Yes; the worker and recipe are opt-in.
• Copied code or new dependency: The optional upstream ReplaySSM emulation dependency is shared with [6/10] ReplaySSM state QAT and two-precision GDN #2788; kernels are imported, not copied.
• Necessary tests: Cache/reset/CUDA-graph tests, native training/gradient cases, recipe round-trip and a tiny offline engine test.
• Changelog: Updated the existing experimental ReplaySSM feature entry.
• Claude review: Pending; keep this PR draft.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}