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

[kaix-nv]

Linear-attention PR stack — 5 drafts

Order	PR	Depends on
1/5	#2497 GDN state/W QAT foundation	main
2/5	#2519 GDN/KDA decode QAT + INT8	#2497
3/5	#2541 vLLM GDN/KDA state-only fake quantization	#2519
4/5	#2503 GDN/KDA prefill GEMM quantization	#2541
5/5	#2507 Experimental GDN/KDA approximate inverse	#2503

All five drafts form a linear GitHub stack in the review order shown above. #2541 applies TensorQuantizer before native vLLM prefill/decode calls.

A separate vLLM prefill-GEMM PR will wait for an optimized fused kernel. #2506 and #2509 are superseded and closed.

What does this PR do?

Type of change: New feature.

Stacked on #2519 (kaix/linear-attention-decode-first). Add a state-only ModelOpt fake-quant plugin for vLLM GDN and KDA. It applies TensorQuantizer to the incoming recurrent state immediately before each native prefill or decode call. Native kernels and cache management remain in use; this PR changes no CUDA or Triton kernels.

Prefill quantizes the initialized state tensor passed to the original chunk kernel. Decode gathers only active cache slots, applies QDQ, and writes them back before the original recurrent kernel. Each slot/head has its own dynamic scale over [Dk,Dv]. FP8 E4M3 and signed symmetric INT8 are supported, with FP32 dequantized state.

This replaces the earlier draft's custom attention execution and separate request cache. There is no additional persistent state allocation or worker memory reservation. Quantization happens once per native invocation, including scheduler-level prompt continuations; it does not round every internal prefill chunk or the final-state write. Quantizer configuration and state names survive save/restore through the HF-to-vLLM mapper.

The worker validates adapter policy before calibration or warmup and after loading quantizer state. It discovers adapters from the model instead of retaining a second list; shared runtime capability checks run once per binding.

The execution-config definition comes from #2519. This serving adapter continues to accept only the default execution policy and state quantizers; it does not enable the training prefill-GEMM or replay paths. The native integration test now propagates its import paths to spawned workers as well as through PYTHONPATH.

Usage

PYTHONPATH=.:examples/vllm_serve \
RECIPE_PATH=examples/vllm_serve/linear_attention_state_int8.yaml \
python examples/vllm_serve/vllm_serve_fakequant.py /path/to/model \
  --tensor-parallel-size 2 --enforce-eager --no-async-scheduling \
  --no-enable-prefix-caching --mamba-cache-dtype float32

The recipe sets algorithm: null: dynamic state scales require no calibration dataset. Set state num_bits: [4, 3] for FP8. Existing worker weight/activation calibration remains available. Use MODELOPT_STATE_PATH instead of RECIPE_PATH to restore saved quantizer configuration.

The state-only adapter rejects nondefault execution policies, including the previous draft's replay policy. Use the new state-only recipe. See docs/linear_attention_vllm.md for the exact rounding cadence and runtime limits.

Testing

Current restacked source on two RTX A6000 GPUs, Torch 2.9.1+cu128, and Triton 3.5.1:

• 13 CPU export/reload tests passed.
• All 5 native vLLM integration tests passed (578.19 seconds): GDN and KDA at TP=1/2 plus saved-state name mapping. Tests exercise INT8 and FP8 state QDQ, native-kernel controls, inactive cache slots, prompt continuations, repeated generation, calibration, checkpoint reload, and disabled-path agreement.
• Pre-commit, diff checks, and commit-signature verification passed.

The integration run used the clean pinned vLLM checkout 930288170c31e8568290fff407dca8caf17d16ad, whose recurrent-state ABI is key-first, with existing local compiled artifacts. The initially selected editable vLLM checkout used value-first state and correctly failed the runtime guard; that run is not included in the passing result.

The passing run used the existing local pytest harness to set gpu_memory_utilization=0.04, with the same 128 MiB cache budget and unchanged numerical assertions. This avoids requiring nearly all GPU memory for tiny synthetic models. NCCL_P2P_DISABLE=1 is required on this host. These are validation settings, not changes to the serving worker or kernels.

PYTHONPATH=. python -m pytest -q \
  tests/unit/torch/export/test_vllm_fakequant_hf.py \
  tests/unit/torch/export/test_vllm_quantizer_reload.py

NCCL_P2P_DISABLE=1 \
PYTHONPATH=.:tmp/vllm-runtime:examples/vllm_serve:tmp/design-review \
python -m pytest -p low_memory_vllm -q \
  tests/gpu_vllm/torch/quantization/test_vllm_linear_attention.py

The local low_memory_vllm fixture wraps this test module's LLM constructor with functools.partial(LLM, gpu_memory_utilization=0.04). It is not part of the PR. Models use tiny offline Qwen3-Next/Kimi Linear configurations and synthetic weights; this is functional qualification, with no pretrained-quality or performance claim.

Before your PR is "Ready for review"

• Is this change backward compatible?: Opt-in; disabled-path agreement is tested. The unmerged draft's replay-serving recipe and execution policies are intentionally replaced by state-only QDQ.
• If you copied code or added a dependency, did you follow CONTRIBUTING.md?: N/A; no copied third-party kernel or new dependency. Wrappers call native vLLM functions.
• Did you write necessary tests?: Yes; native-kernel and worker integration coverage.
• Did you update Changelog?: Yes.
• Did you get Claude approval?: No; keep draft.

Additional Information

The supported runtime is vLLM 0.15.x with key-first FP32 state, eager synchronous execution, TP=1/2 and PP=DP=CP=1. Speculative decoding, prefix caching, state transfer, and CUDA graphs remain unsupported.

A separate prefill-GEMM PR is deferred until an optimized fused kernel is available. It will reuse #2503's eight numerical sites; this adapter does not expose the materialized PyTorch prefill backend. Megatron training/backward qualification and vLLM forward/cache qualification remain separate.

[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

---

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

[codecov]

Codecov Report

❌ Patch coverage is 0% with 73 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (kaix/linear-attention-decode-first@757f337). Learn more about missing BASE report.

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

Additional details and impacted files

@@                          Coverage Diff                          @@
##             kaix/linear-attention-decode-first    #2541   +/-   ##
=====================================================================
  Coverage                                      ?   70.20%           
=====================================================================
  Files                                         ?      619           
  Lines                                         ?    68475           
  Branches                                      ?        0           
=====================================================================
  Hits                                          ?    48070           
  Misses                                        ?    20405           
  Partials                                      ?        0           

Flag	Coverage Δ
unit	57.57% <0.00%> (?)

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.