[5/7] vLLM GDN/KDA state-only fake quantization

[kaix-nv]

Linear-attention series — 7 PRs

Order	PR	Depends on
1/7	#2497 GDN state/W QAT foundation	main
2/7	#2519 GDN/KDA state QAT + INT8	#2497
3/7	#2657 Megatron Bridge linear attention QAT/QAD example	#2519
4/7	#2788 ReplaySSM state QAT and two-precision design	#2657
5/7	#2541 vLLM GDN/KDA state-only fake quantization	#2788
6/7	#2503 GDN/KDA prefill GEMM quantization	#2541
7/7	#2507 Experimental GDN/KDA approximate inverse	#2503

The six open PRs form native GitHub stack #2789 in the order shown. #2497 has landed, so #2519 targets main. #2788 follows #2657 and #2541 is rebased onto it. #2503 and #2507 retain their existing commits; their rebases remain deferred until their immediate parents merge.

#2788 owns the optional INT8/Hadamard ReplaySSM implementation and the two-precision design plan; two-precision support is not implemented. #2541 applies TensorQuantizer before native vLLM prefill/decode calls. Serving-time prefill-GEMM quantization remains deferred until an optimized fused kernel is available.

What does this PR do?

Type of change: New feature.

Stacked on #2788, which follows #2657. Add state-only fake quantization for vLLM 0.30 Qwen GDN and Kimi KDA layers. Before native layer execution, gather initialized active recurrent-cache slots, apply ModelOpt TensorQuantizer, and write them back. Fresh prefill and inactive slots are unchanged. No CUDA/Triton kernels are added or modified.

The default recipe uses dynamic INT8 blocks of 32 value channels per key row, shared with training. Native value-first cache storage is transposed for QDQ so value grouping remains consistent. FP8 E4M3 is also configurable; the cache remains FP32. QDQ runs before each native state read, including scheduler prompt continuations, not at internal prefill chunk boundaries.

The existing FakeQuantWorker skips calibration data for dynamic recipes and restores state-quantizer names and policies through the HF-to-vLLM mapper. A small TensorQuantizer fix lets dynamic blocks follow changing active-batch shapes. Megatron HF weight exports can be served with the state recipe; their Megatron metadata sidecar is not an HF fakequant export.

Usage

VLLM_GDN_DECODE_KERNEL=triton VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE=1 \
python examples/vllm_serve/vllm_serve_fakequant.py /path/to/model \
  --modelopt-recipe-path examples/vllm_serve/linear_attention_state_int8.yaml \
  --dtype bfloat16 --enforce-eager --no-async-scheduling --no-enable-prefix-caching \
  --mamba-cache-mode none --mamba-cache-dtype float32 \
  --additional-config '{"gdn_prefill_backend":"triton","kda_prefill_backend":"triton","kda_decode_backend":"triton"}'

Training alignment additionally requires matching prefill scheduling, BF16 projections, unbounded Kimi gates and packed GDN decode. Mixed batches retain native vLLM dispatch and are outside the training profile's arithmetic guarantee. The serving adapter supports eager synchronous execution with PP=DP=CP=1; speculative decoding, prefix caching, state transfer and ReplaySSM are excluded.

Testing

Validated source, rebased onto #2657 0b90b4fea6; published head ef9aa9bec0:

• Pre-commit passed for the changed files.
• 79 local CPU checks passed: TensorQuantizer, quantizer reload, linear-attention plugins, and saved-policy name mapping.
• DFW H100, vLLM 0.30.0, Torch 2.13.0+cu130: 44 TensorQuantizer CPU tests passed. All three focused native-worker/mapping tests passed (job 20128119, 371.44 seconds including startup and compilation), using the compiled ModelOpt CUDA extension. The fixture explicitly selects the same Triton packed-decode settings observed in this run.

The two tiny offline worker cases cover GDN and Kimi-Linear KDA: active-cache QDQ before native execution, untouched inactive/fresh slots, disabled behavior, and restored generation. These are functional tests, not pretrained-quality, cache-compression or performance measurements. TP=1 is tested on this revision.

Earlier attempts are retained in local receipts: the first stopped before collection because the launcher lacked cuDNN discovery settings; the next passed KDA with the Torch QDQ fallback but hit a native GDN compiler error for unequal K/V dimensions. The final fixture uses K=V=32 and a corrected CUDA-header search path.

Before your PR is "Ready for review"

• Backward compatible: opt-in; native behavior is retained when disabled.
• Copied code or new dependencies: none.
• Necessary tests: focused CPU checks and native worker integration.
• Changelog: updated.
• Claude approval: not requested for this revision; keep draft.

Additional Information

Prefill GEMM quantization remains deferred to an optimized-kernel follow-up. This PR adds only state fake quantization.

Stack update: rebased onto ReplaySSM head 291b2c1124; this head is 7c6dc6f1b3. Only documentation conflicts required resolution. 90 focused CPU tests passed across state QDQ, ReplaySSM policy, checkpoint lifecycle, recipe docs and TensorQuantizer; changed-file pre-commit and git diff --check passed. No GPU test was rerun for this rebase.

[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.}

[codecov]

Codecov Report

❌ Patch coverage is 4.93827% with 77 lines in your changes missing coverage. Please review.
✅ Project coverage is 71.00%. Comparing base (942115a) to head (d532b10).

Files with missing lines	Patch %	Lines
...orch/quantization/plugins/vllm_linear_attention.py	0.00%	76 Missing ⚠️
modelopt/torch/quantization/plugins/__init__.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@                          Coverage Diff                          @@
##           kaix/linear-attention-qat-example    #2541      +/-   ##
=====================================================================
- Coverage                              71.08%   71.00%   -0.09%     
=====================================================================
  Files                                    658      659       +1     
  Lines                                  72115    72195      +80     
=====================================================================
- Hits                                   51264    51261       -3     
- Misses                                 20851    20934      +83     

Flag	Coverage Δ
unit	59.97% <4.93%> (-0.08%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.