[Core] Use FlashInfer for pre-SM100 NVFP4 KV cache updates

[lesj0610]

Purpose

I wanted to use --kv-cache-dtype nvfp4 on my RTX 3090 setup, but it was only usable on the SM100 native path. This PR makes the pre-SM100 path use FlashInfer's slot-mapping API for NVFP4 paged KV appends, so Ampere and Hopper systems can use the smaller NVFP4 KV cache without relying on the Blackwell-only path.

The main benefit is KV cache capacity. On the tested RTX 3090 setup, NVFP4 gives about 3x more KV cache tokens, with quality on par with the auto dtype in the MRCR runs below. The tables compare against auto and TurboQuant 4bit_nc; TurboQuant did not start for the Gemma4 models on this setup.

AI assistance: Codex and Claude Fable 5 were used during implementation and PR preparation; the submitter reviewed the changes.

Changes

• Use FlashInfer's nvfp4_quantize_append_paged_kv_cache_with_slot_mapping for pre-SM100 NVFP4 KV cache updates.
• Add an FA2 prefill fast path that dequants FP4 KV pages to BF16 and runs flash_attn_varlen_func. If the FlashInfer paged-dequant helper is unavailable, this path falls back.
• Keep native FA2 query tensors in model dtype. FP8-Q is only used when the trtllm-gen path is selected. This extends the Q dtype gating from [BugFix] Derive FlashInfer Q dtype from resolved per-group builder state vllm-project/vllm#47485 to the nvfp4 branch, which that fix did not cover.
• Keep TQ cache layout selection on the attention spec. After Support nvfp4 kv with kv-cache-dtype-skip-layers sliding_window vllm-project/vllm#42890, specs with KVQuantMode.NONE default to the unquantized auto layout for shape calculation; TQ specs still use a packed layout, so they preserve the configured turboquant_* cache dtype.
• Handle mixed K/V head dimensions by splitting packed NVFP4 data/scale views with explicit layout checks.
• Cache the derived FlashInfer KV views per bound KV cache tensor. The update path and forward path now share the same canonicalized split views.
• Centralize KV cache shape/stride handling for attention backends without changing backend method signatures.
• Log the NVFP4 FA2 prefill scratch workspace reservation at startup. Sliding-window layers currently reserve full-context scratch; a future window clamp can reduce this to the active window.

Dependencies

No FlashInfer version bump is needed. Every API this PR uses is present in the
release vLLM already pins (flashinfer-python==0.6.17,
flashinfer-cubin==0.6.17 in requirements/cuda.txt):

• nvfp4_quantize_append_paged_kv_cache_with_slot_mapping — required, drives the pre-SM100 NVFP4 paged KV append.
• nvfp4_kv_dequantize_paged — optional; enables the FA2 prefill fast path. That path falls back when the helper is unavailable.
• Large-head FA2 support, needed by the same prefill path.

Verified against the pinned release:

python -c "
import flashinfer
from flashinfer.quantization import nvfp4_kv_dequantize_paged
print(flashinfer.__version__,
      hasattr(flashinfer, 'nvfp4_quantize_append_paged_kv_cache_with_slot_mapping'),
      callable(nvfp4_kv_dequantize_paged))
"
# 0.6.17 True True

Gemma4 selector support: "[Model] Allow Gemma4 to use FlashInfer when FA4 is unavailable" allows Gemma4 to choose FlashInfer when FA4 is
unavailable.

Test Plan

Unit tests:

pytest tests/kernels/attention/test_flashinfer.py -q -k nvfp4
pytest tests/kernels/attention/test_flashinfer.py -q -k "q_dtype or q_quantization_disable"
pytest tests/kernels/attention/test_flashinfer.py -q -k fa2_prefill_reservation
pytest tests/v1/worker/test_attn_utils.py -q
pytest tests/quantization/test_turboquant.py -q -k metadata_builder

Static checks:

python -m py_compile <modified Python files>
ruff check <modified Python files>
ruff format --check <modified Python files>
pre-commit run --files <modified files>

Runtime smoke on RTX 3090 (SM86): Qwen3.6-27B GPTQ int4, TP=1, --kv-cache-dtype nvfp4, without disable_flashinfer_q_quantization.

Benchmarks on RTX 3090 x 2, TP=2: KV cache capacity, MRCR quality (30 samples, 32K context), and vllm bench serve (random 8K prompts, output length 64, 16 requests, max_num_seqs=64).

Test Result

Unit tests on the current head after merging latest main: 23 passed (nvfp4), 5 passed (q dtype), 3 passed (reservation), 6 passed (attn_utils), 2 passed (TurboQuant metadata builder). Static checks passed.

Runtime smoke: resolved query dtypes logged as prefill=torch.bfloat16, decode=torch.bfloat16, decode_backend=flashinfer-native; server ready; /v1/completions returned 200 OK.

The TQ cache-dtype regression from the vllm-project#42890 merge is covered by test_turboquant_spec_preserves_configured_cache_dtype_for_shape, which routes through TurboQuantAttentionBackend.get_kv_cache_shape() and verifies the configured turboquant_4bit_nc dtype is not converted to auto during cache-shape calculation.

The benchmarks below were measured on this branch before the latest main merge; the attention code paths are semantically unchanged by the merge. Gemma4 rows were measured with the Gemma4 FlashInfer selector change (vllm-project#47547) stacked locally.

KV Cache Capacity

Model	KV dtype	GPU KV cache size	vs auto	Backend
Qwen3.6-27B	auto	291,999 tokens	1.00x	Flash-attn
Qwen3.6-27B	nvfp4	884,736 tokens	3.03x	FlashInfer
Qwen3.6-27B	TQ 4bit_nc	943,250 tokens	3.23x	TurboQuant
Qwen3.6-35B-A3B	auto	724,640 tokens	1.00x	Flash-attn
Qwen3.6-35B-A3B	nvfp4	2,129,920 tokens	2.94x	FlashInfer
Qwen3.6-35B-A3B	TQ 4bit_nc	2,380,148 tokens	3.29x	TurboQuant
Gemma4-31B	auto	122,360 tokens	1.00x	Triton
Gemma4-31B	nvfp4	392,369 tokens	3.21x	FlashInfer
Gemma4-31B	TQ 4bit_nc	startup failed*	-	TurboQuant
Gemma4-26B-A4B	auto	590,793 tokens	1.00x	Triton
Gemma4-26B-A4B	nvfp4	1,802,224 tokens	3.05x	FlashInfer
Gemma4-26B-A4B	TQ 4bit_nc	startup failed*	-	TurboQuant

* Gemma4 TQ 4bit_nc did not start on this setup.

Quality — MRCR

30 samples, 32K context.

Model	KV dtype	Match ratio	n8 match ratio	tok/s	Backend
Qwen3.6-27B	auto	0.9388	0.8245	29.39	Flash-attn
Qwen3.6-27B	nvfp4	0.9388	0.8245	29.47	FlashInfer
Qwen3.6-27B	TQ 4bit_nc	0.9433	0.8378	27.86	TurboQuant
Qwen3.6-35B-A3B	auto	0.8973	0.6999	118.94	Flash-attn
Qwen3.6-35B-A3B	nvfp4	0.9006	0.7097	120.74	FlashInfer
Qwen3.6-35B-A3B	TQ 4bit_nc	0.8983	0.7028	64.61	TurboQuant
Gemma4-31B	auto	0.7109	0.4714	14.59	Triton
Gemma4-31B	nvfp4	0.7134	0.4767	14.14	FlashInfer
Gemma4-26B-A4B	auto	0.4337	0.3584	50.59	Triton
Gemma4-26B-A4B	nvfp4	0.4950	0.3262	49.86	FlashInfer

NVFP4 was on par with auto in the Qwen runs. The Gemma4 results are from 30-sample runs with some metric movement: Gemma4-26B-A4B shows a higher overall match ratio (+0.06) and a lower n8 match ratio (-0.03) in this run.

Serving Benchmark

vllm bench serve, random 8K prompts, output length 64, 16 requests, max_num_seqs=64.

Model	KV dtype	Mean TTFT	Mean TPOT	Output tok/s	Status	Backend
Qwen3.6-27B	auto	42.07 s	574.22 ms	12.93	completed	Flash-attn
Qwen3.6-27B	nvfp4	41.31 s	567.70 ms	13.16	completed	FlashInfer
Qwen3.6-27B	TQ 4bit_nc	75.03 s	1038.18 ms	7.24	completed	TurboQuant
Qwen3.6-35B-A3B	auto	9.57 s	131.23 ms	55.40	completed	Flash-attn
Qwen3.6-35B-A3B	nvfp4	9.50 s	129.67 ms	55.85	completed	FlashInfer
Qwen3.6-35B-A3B	TQ 4bit_nc	18.12 s	260.88 ms	29.20	completed	TurboQuant
Gemma4-31B	auto	83.60 s	1040.14 ms	6.47	completed	Triton
Gemma4-31B	nvfp4	65.97 s	937.40 ms	8.13	completed	FlashInfer
Gemma4-26B-A4B	auto	25.45 s	362.40 ms	21.00	completed	Triton
Gemma4-26B-A4B	nvfp4	15.58 s	224.26 ms	34.08	completed	FlashInfer

Notes

• SM100 is not locally verified. The trtllm-gen direct path is unchanged. On SM100, wrapper-based fallback cases where trtllm-gen is not selected now use the auto backend with model-dtype output instead of forcing trtllm-gen with FP8 output.
• OpenPanGU and MIMO V2 mixed-head runtime is not verified with real model weights. The mixed-layout split view behavior is covered by unit tests.
• Sliding-window FA2 prefill currently reserves full-context scratch workspace. A window clamp can reduce this later; this PR only logs the current reservation size.
• The FA2 fast path is optional. If the FlashInfer paged-dequant helper is not present, the code falls back instead of failing startup.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results.
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[chatgpt-codex-connector]

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Advanced

Run ID: 6949579c-425a-4ef3-b7e9-914f623cd776

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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