[4/6] vLLM GDN/KDA state-only fake quantization

[kaix-nv]

Linear-attention series — 6 PRs

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

The five open PRs form one native GitHub stack in the order shown. #2497 has landed, so #2519 targets main. #2541 now targets #2657. Rebase each remaining descendant after its immediate parent merges.

#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 #2657 (kaix/linear-attention-qat-example). 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.

The serving example README includes the state-quantization launch command, numerical boundaries, and runtime requirements.

Prefill/decode quantization boundary

State QDQ applies to both prefill and decode, specifically at native vLLM invocation boundaries:

• Prefill: quantize the incoming initialized recurrent state once before each native prefill invocation. With scheduler-level chunked prefill, each invocation creates one QDQ boundary; internal kernel chunks do not.
• Decode: gather and quantize only the active recurrent-cache slots immediately before each native decode invocation. Inactive slots are unchanged.
• State write: native kernels update and write recurrent state in FP32 without additional final-write quantization. The next prefill or decode invocation quantizes that state before reading it.

This is floating-point fake QDQ for numerical qualification. It does not store the persistent recurrent cache in packed INT8 or FP8, execute state updates with low-precision MMA, or establish memory-capacity, bandwidth, or latency savings.

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 examples/vllm_serve/README.md#linear-attention-state-quantization for the exact rounding cadence and runtime limits.

Testing

Stack metadata has been updated to remove #2562. The branch history still needs to be replayed onto current #2657; the runtime results below belong to the earlier revision.

Compilation-fixture update: this branch is restacked on #2497's separate follow-up commit 95766de709ac. The changed test modules passed at #2497 (26 passed, 20 hardware skips in each cold/warm run) and at the #2507 stack tip (64 passed, 20 hardware skips on two RTX A6000 GPUs); intermediate PRs were not separately rerun. Functional calls created no new tracked kernel binaries and kept the default 120-second cap. All six source trees match the validated trees, and commit hooks passed. Native FP8-state/Hopper cases remain hardware-gated.

Earlier scope-specific validation follows.

For the earlier documentation/example amendment, pre-commit, Markdown links and anchors, command/Python syntax, source-manifest readability, and stale-path checks passed. Runtime kernels were not changed; model training, distributed integration, and quality studies were not rerun for this amendment.

Prior runtime validation 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

• Autopilot · 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.