Xqa decode kernels by DanBlanaru · Pull Request #43232 · vllm-project/vllm

DanBlanaru · 2026-05-20T17:46:54Z

Purpose

Enable the FlashInfer TRTLLM/XQA decode path on Hopper (SM90) for vLLM V1 FlashInfer attention.

This makes TRTLLM attention support phase-aware, so decode can use FlashInfer's trtllm_batch_decode_with_kv_cache XQA path while prefill can remain on the existing FlashInfer native path when TRTLLM prefill is not supported. The branch also tracks prefill and decode query dtypes separately, since SM90 XQA decode requires BF16/FP16 query tensors with FP8 KV cache

Test Plan

git clone -b xqa_decode_kernels https://github.com/DanBlanaru/vllm.git
cd vllm
uv venv --python 3.12 --seed
source .venv/bin/activate
VLLM_USE_PRECOMPILED=1 uv pip install -e . --torch-backend=auto
uv pip install --pre flashinfer-python --index-url https://flashinfer.ai/whl/nightly/ --no-deps

You can run the xqa backend with:

MODEL=Qwen/Qwen3-30B-A3B-FP8
VLLM_KV_CACHE_LAYOUT=HND FLASHINFER_DISABLE_VERSION_CHECK=1 \
vllm serve "$MODEL" --attention-backend FLASHINFER --kv-cache-dtype fp8 \
  --max-model-len 32768 --max-num-seqs 16 --gpu-memory-utilization 0.85 \
  --enable-prefix-caching --no-enable-log-requests &
vllm bench serve --backend vllm --model "$MODEL" --dataset-name random \
  --random-input-len 32000 --random-output-len 500 --num-prompts 16 \
  --max-concurrency 8 --ignore-eos --ready-check-timeout-sec 1200

and the fav3 with:

vllm serve "$MODEL" --attention-backend FLASH_ATTN --kv-cache-dtype fp8 \
  --max-model-len 32768 --max-num-seqs 16 --gpu-memory-utilization 0.85 \
  --enable-prefix-caching --no-enable-log-requests &
vllm bench serve --backend vllm --model "$MODEL" --dataset-name random \
  --random-input-len 32000 --random-output-len 500 --num-prompts 16 \
  --max-concurrency 8 --ignore-eos --ready-check-timeout-sec 1200

all of these tests pass on SM90:

pytest tests/kernels/attention/test_use_trtllm_attention.py
pytest tests/kernels/attention/test_trtllm_kvfp8_dequant.py
pytest tests/kernels/attention/test_flashinfer_trtllm_attention.py
pytest tests/v1/attention/test_trtllm_attention_integration.py
pytest tests/engine/test_arg_utils.py -k attention_config

gemini-code-assist

Code Review

This pull request refactors the FlashInfer attention backend to support TRTLLM attention independently for prefill and decode phases, introducing granular capability checks and support for distinct query data types across phases. Feedback identifies critical issues in the Triton dequantization kernel regarding tl.arange constraints and memory hazards from per-layer buffer allocations. Additionally, reviewers noted a logic error in KV cache layout assertions that could cause crashes on SM100 and raised concerns about overriding user-specified flags when TRTLLM attention is explicitly disabled.

DanBlanaru · 2026-05-20T18:03:13Z

FI + XQA is better in FP8 in both ttft and tpot

FP8

fav3:

============ Serving Benchmark Result ============
Successful requests:                     16
Failed requests:                         0
Maximum request concurrency:             8
Benchmark duration (s):                  27.71
Total input tokens:                      512000
Total generated tokens:                  8000
Request throughput (req/s):              0.58
Output token throughput (tok/s):         288.69
Peak output token throughput (tok/s):    768.00
Peak concurrent requests:                11.00
Total token throughput (tok/s):          18764.71
---------------Time to First Token----------------
Mean TTFT (ms):                          3464.68
Median TTFT (ms):                        3119.35
P99 TTFT (ms):                           8039.12
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          20.54
Median TPOT (ms):                        21.25
P99 TPOT (ms):                           25.53
---------------Inter-token Latency----------------
Mean ITL (ms):                           20.57
Median ITL (ms):                         11.07
P99 ITL (ms):                            331.14
==================================================

native FlashInfer

============ Serving Benchmark Result ============
Successful requests:                     16        
Failed requests:                         0         
Maximum request concurrency:             8         
Benchmark duration (s):                  56.00     
Total input tokens:                      512000    
Total generated tokens:                  8000      
Request throughput (req/s):              0.29      
Output token throughput (tok/s):         142.86    
Peak output token throughput (tok/s):    200.00    
Peak concurrent requests:                11.00     
Total token throughput (tok/s):          9285.66   
---------------Time to First Token----------------
Mean TTFT (ms):                          3475.30   
Median TTFT (ms):                        3489.27   
P99 TTFT (ms):                           7712.39   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          48.78     
Median TPOT (ms):                        50.37     
P99 TPOT (ms):                           53.18     
---------------Inter-token Latency----------------
Mean ITL (ms):                           48.78     
Median ITL (ms):                         40.66     
P99 ITL (ms):                            309.08    
==================================================

xqa:

============ Serving Benchmark Result ============
Successful requests:                     16
Failed requests:                         0
Maximum request concurrency:             8
Benchmark duration (s):                  25.76
Total input tokens:                      512000
Total generated tokens:                  8000
Request throughput (req/s):              0.62
Output token throughput (tok/s):         310.59
Peak output token throughput (tok/s):    760.00
Peak concurrent requests:                11.00
Total token throughput (tok/s):          20188.65
---------------Time to First Token----------------
Mean TTFT (ms):                          3167.72
Median TTFT (ms):                        2893.08
P99 TTFT (ms):                           7247.78
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          19.20
Median TPOT (ms):                        19.79
P99 TPOT (ms):                           23.92
---------------Inter-token Latency----------------
Mean ITL (ms):                           19.21
Median ITL (ms):                         11.03
P99 ITL (ms):                            287.28
==================================================

BF16

It seems vLLM FAv3 is better in BF16 in both ttft and tpot

fav3:

============ Serving Benchmark Result ============
Successful requests:                     16
Failed requests:                         0
Maximum request concurrency:             8
Benchmark duration (s):                  28.55
Total input tokens:                      512000
Total generated tokens:                  8000
Request throughput (req/s):              0.56
Output token throughput (tok/s):         280.22
Peak output token throughput (tok/s):    688.00
Peak concurrent requests:                12.00
Total token throughput (tok/s):          18214.55
---------------Time to First Token----------------
Mean TTFT (ms):                          4503.61
Median TTFT (ms):                        4759.79
P99 TTFT (ms):                           7676.50
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          19.47
Median TPOT (ms):                        18.28
P99 TPOT (ms):                           26.96
---------------Inter-token Latency----------------
Mean ITL (ms):                           19.48
Median ITL (ms):                         12.10
P99 ITL (ms):                            306.79
==================================================

xqa:

============ Serving Benchmark Result ============
Successful requests:                     16
Failed requests:                         0
Maximum request concurrency:             8
Benchmark duration (s):                  30.82
Total input tokens:                      512000
Total generated tokens:                  8000
Request throughput (req/s):              0.52
Output token throughput (tok/s):         259.58
Peak output token throughput (tok/s):    623.00
Peak concurrent requests:                13.00
Total token throughput (tok/s):          16872.59
---------------Time to First Token----------------
Mean TTFT (ms):                          4790.86
Median TTFT (ms):                        4899.46
P99 TTFT (ms):                           8255.20
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          21.16
Median TPOT (ms):                        19.85
P99 TPOT (ms):                           28.90
---------------Inter-token Latency----------------
Mean ITL (ms):                           21.17
Median ITL (ms):                         13.12
P99 ITL (ms):                            337.08
==================================================

Qwen3-30B FP8 MMLU

xqa fp8:

Groups	Version	Filter	n-shot	Metric	Value	Stderr
mmlu	2	none		acc	0.7752	± 0.0033
humanities	2	none	0	acc	0.6752	± 0.0064
other	2	none	0	acc	0.8136	± 0.0067
social sciences	2	none	0	acc	0.8739	± 0.0059
stem	2	none	0	acc	0.7900	± 0.0070

fav3 fp8:

Groups	Version	Filter	n-shot	Metric	Value	Stderr
mmlu	2	none		acc	0.7765	± 0.0033
humanities	2	none	0	acc	0.6812	± 0.0064
other	2	none	0	acc	0.8146	± 0.0067
social sciences	2	none	0	acc	0.8707	± 0.0059
stem	2	none	0	acc	0.7894	± 0.0070

mergify · 2026-05-23T10:41:52Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @DanBlanaru.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

DanBlanaru · 2026-05-29T00:00:09Z

I also ran benchmarks on B200 to make sure we don't affect perf on the sm100 path
current:

============ Serving Benchmark Result ============
Successful requests:                     128
Failed requests:                         0
Maximum request concurrency:             64
Benchmark duration (s):                  149.95
Total input tokens:                      3840000
Total generated tokens:                  128000
Request throughput (req/s):              0.85
Output token throughput (tok/s):         853.59
Peak output token throughput (tok/s):    3072.00
Peak concurrent requests:                67.00
Total token throughput (tok/s):          26461.39
---------------Time to First Token----------------
Mean TTFT (ms):                          15855.81
Median TTFT (ms):                        4176.47
P99 TTFT (ms):                           55192.74
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          58.31
Median TPOT (ms):                        69.38
P99 TPOT (ms):                           71.77
---------------Inter-token Latency----------------
Mean ITL (ms):                           58.40
Median ITL (ms):                         21.49
P99 ITL (ms):                            278.14
==================================================

git reset before my commits:

============ Serving Benchmark Result ============
Successful requests:                     128
Failed requests:                         0
Maximum request concurrency:             64
Benchmark duration (s):                  149.85
Total input tokens:                      3840000
Total generated tokens:                  128000
Request throughput (req/s):              0.85
Output token throughput (tok/s):         854.18
Peak output token throughput (tok/s):    3072.00
Peak concurrent requests:                67.00
Total token throughput (tok/s):          26479.49
---------------Time to First Token----------------
Mean TTFT (ms):                          16027.95
Median TTFT (ms):                        4251.61
P99 TTFT (ms):                           55530.76
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          58.08
Median TPOT (ms):                        68.92
P99 TPOT (ms):                           71.38
---------------Inter-token Latency----------------
Mean ITL (ms):                           58.17
Median ITL (ms):                         21.28
P99 ITL (ms):                            278.92
==================================================

mergify · 2026-06-01T16:15:55Z

Hi @DanBlanaru, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

pavanimajety

First round of comments.

Restore NVFP4 and speculative-decode guards for the XQA path, and shorten the SM90 query dtype comments. Signed-off-by: Dan Blanaru <[email protected]>

Keep SM90 XQA advertised as single-token decode until vLLM wires the speculative decode mask path. Signed-off-by: Dan Blanaru <[email protected]>

Avoid applying q_scale when SM90 XQA decode uses model-dtype queries with FP8 KV cache. Signed-off-by: Dan Blanaru <[email protected]>

MatthewBonanni

LGTM, thanks for iterating on this!

#43232 moved Q-dtype selection into a classmethod that re-derives its inputs from the global config, dropping two pieces of resolved builder state: the arch capability gate (FP8-Q crashed engine init on fa2-only archs like SM89/SM120) and the per-KV-group cache dtype (unquantized groups from --kv-cache-dtype-skip-layers wrongly got FP8-Q). Make it an instance method over self.cache_dtype and restore the capability gate (FI native fa3 on SM90, trtllm-gen/XQA on SM100). Co-authored-by: Claude Co-Authored-By: Claude Fable 5 <[email protected]> Claude-Session: https://claude.ai/code/session_012jb3BF13uN6cBp8TtGDQYe Signed-off-by: mgoin <[email protected]>

Sync 2026-07-03: picks up Model Runner V2 improvements (motivated by V2 viability reopening via PR#26). Notable: MRV2 scheduler req-slot bugfix (vllm-project#46974), MRV2 spec-decode int32 block-verify overflow fix (vllm-project#47383), MRV2 DSA-indexer prefill warmup for GLM-5.2/DSA (vllm-project#47285), V2-default for dense models (vllm-project#44443 — DSv4 is MoE so stays V1), DeepGEMM tag -> nv-dev a6b593d2 for SM120 (vllm-project#47304), TLI heterogeneous-vocab spec decode (vllm-project#38174), Delete PagedAttention (vllm-project#47361 — no impact, our sparse-MLA runner is a flashinfer symbol), Xqa decode kernels (vllm-project#43232). Conflict: vllm/v1/spec_decode/llm_base_proposer.py (our DSv4 MTP spec_step_idx routing vs upstream TLI) — resolved keeping BOTH: TLI's heterogeneous-vocab branch routed through our _compute_logits(spec_step_idx). Our SM12x stack preserved (persistent_topk, topk.cu, PR#26 padded-Q kernel, dspark, attention, R1 hardening all UNCHANGED); config/vllm.py DSv4->V1 default + DSpark env-control intact.

Signed-off-by: Dan Blanaru <[email protected]> Co-authored-by: Matthew Bonanni <[email protected]>

ZhanqiuHu · 2026-07-07T17:41:42Z

Seems like this is break Nemotron + FlashInfer w/ FP8 kv cache, on Hopper

Signed-off-by: Dan Blanaru <[email protected]> Co-authored-by: Matthew Bonanni <[email protected]>

Signed-off-by: Dan Blanaru <[email protected]> Co-authored-by: Matthew Bonanni <[email protected]> Signed-off-by: mayuyuace <[email protected]>

Signed-off-by: Dan Blanaru <[email protected]> Co-authored-by: Matthew Bonanni <[email protected]>

DanBlanaru requested review from mgoin, pavanimajety and vadiklyutiy as code owners May 20, 2026 17:46

DanBlanaru changed the title ~~DraftXqa decode kernels~~ Draft: Xqa decode kernels May 20, 2026

DanBlanaru marked this pull request as draft May 20, 2026 17:47

gemini-code-assist Bot reviewed May 20, 2026

View reviewed changes

Comment thread vllm/v1/attention/backends/flashinfer.py Outdated

Comment thread vllm/v1/attention/backends/flashinfer.py

Comment thread vllm/v1/attention/backends/flashinfer.py Outdated

Comment thread vllm/utils/flashinfer.py Outdated

mergify Bot added nvidia v1 labels May 20, 2026

github-project-automation Bot added this to NVIDIA May 20, 2026

mergify Bot added the needs-rebase label May 23, 2026

DanBlanaru force-pushed the xqa_decode_kernels branch from 4c0975c to 006b817 Compare May 27, 2026 14:50

mergify Bot removed the needs-rebase label May 27, 2026

DanBlanaru force-pushed the xqa_decode_kernels branch from 87f2371 to 3a3b951 Compare May 29, 2026 00:07

DanBlanaru marked this pull request as ready for review May 29, 2026 00:11

DanBlanaru requested review from AndreasKaratzas, LucasWilkinson, MatthewBonanni, WoosukKwon, tlrmchlsmth, yewentao256 and zyongye as code owners May 29, 2026 00:11

DanBlanaru changed the title ~~Draft: Xqa decode kernels~~ Xqa decode kernels May 29, 2026

pavanimajety added the verified Run pre-commit for new contributors without triggering other tests label Jun 1, 2026

pavanimajety reviewed Jun 1, 2026

View reviewed changes

Comment thread vllm/v1/attention/backends/flash_attn.py Outdated

Comment thread vllm/v1/attention/backends/flashinfer.py

Comment thread vllm/v1/attention/backends/flashinfer.py

pavanimajety added the ready ONLY add when PR is ready to merge/full CI is needed label Jun 1, 2026

pavanimajety reviewed Jun 1, 2026

View reviewed changes

Comment thread vllm/v1/attention/backends/flashinfer.py

MatthewBonanni reviewed Jun 24, 2026

View reviewed changes

Comment thread vllm/v1/attention/backends/flashinfer.py Outdated

Comment thread vllm/v1/attention/backends/flashinfer.py Outdated

DanBlanaru added 2 commits June 24, 2026 09:18

Address FlashInfer XQA review comments

0643f83

Restore NVFP4 and speculative-decode guards for the XQA path, and shorten the SM90 query dtype comments. Signed-off-by: Dan Blanaru <[email protected]>

fix: limit FlashInfer XQA spec decode graphs

9556657

Keep SM90 XQA advertised as single-token decode until vLLM wires the speculative decode mask path. Signed-off-by: Dan Blanaru <[email protected]>

MatthewBonanni reviewed Jun 29, 2026

View reviewed changes

Comment thread vllm/v1/attention/backends/flashinfer.py

Comment thread vllm/v1/attention/backends/flashinfer.py

MatthewBonanni reviewed Jun 29, 2026

View reviewed changes

Comment thread vllm/v1/attention/backends/flashinfer.py

MatthewBonanni reviewed Jun 29, 2026

View reviewed changes

Comment thread vllm/v1/attention/backends/flashinfer.py

fix: scale FlashInfer XQA decode queries by dtype

90f6f85

Avoid applying q_scale when SM90 XQA decode uses model-dtype queries with FP8 KV cache. Signed-off-by: Dan Blanaru <[email protected]>

MatthewBonanni approved these changes Jun 29, 2026

View reviewed changes

github-project-automation Bot moved this to Ready in NVIDIA Jun 29, 2026

MatthewBonanni enabled auto-merge (squash) June 29, 2026 19:21

MatthewBonanni added 3 commits June 30, 2026 11:18

Merge branch 'main' into xqa_decode_kernels

19e0b66

Merge branch 'main' into xqa_decode_kernels

38716db

Merge branch 'main' into xqa_decode_kernels

f8b1d94

mgoin approved these changes Jul 2, 2026

View reviewed changes

vllm-bot merged commit d29125c into vllm-project:main Jul 2, 2026
97 of 99 checks passed

github-project-automation Bot moved this from Ready to Done in NVIDIA Jul 2, 2026

mgoin mentioned this pull request Jul 2, 2026

[BugFix] Derive FlashInfer Q dtype from resolved per-group builder state #47485

Merged

vllm-agent mentioned this pull request Jul 3, 2026

Revert "Xqa decode kernels" (#43232) #47496

Draft

jakki-amd pushed a commit to jakki-amd/vllm that referenced this pull request Jul 6, 2026

Xqa decode kernels (vllm-project#43232)

baa2a96

Signed-off-by: Dan Blanaru <[email protected]> Co-authored-by: Matthew Bonanni <[email protected]>

ZhanqiuHu mentioned this pull request Jul 7, 2026

[Bug] Nemotron + FlashInfer + FP8 KV cache crashes on Hopper: assert query.is_contiguous() in maybe_quant_query #47905

Closed

MatthewBonanni mentioned this pull request Jul 7, 2026

[Bugfix] Allow non-contiguous query in FlashInfer FP8 query quantization #47908

Merged

2 tasks

lkk12014402 pushed a commit to lkk12014402/vllm that referenced this pull request Jul 8, 2026

Xqa decode kernels (vllm-project#43232)

18c6851

Signed-off-by: Dan Blanaru <[email protected]> Co-authored-by: Matthew Bonanni <[email protected]>

mayuyuace pushed a commit to mayuyuace/vllm that referenced this pull request Jul 9, 2026

Xqa decode kernels (vllm-project#43232)

d3a4f84

Signed-off-by: Dan Blanaru <[email protected]> Co-authored-by: Matthew Bonanni <[email protected]> Signed-off-by: mayuyuace <[email protected]>

jethac mentioned this pull request Jul 11, 2026

[Attention][Quantization] NVFP4 KV cache on consumer/SoC Blackwell (sm120/sm121) for Gemma 3/4 via FlashInfer FA2 #46329

Open

NickLucche pushed a commit to NickLucche/vllm that referenced this pull request Jul 15, 2026

Xqa decode kernels (vllm-project#43232)

4528e49

Signed-off-by: Dan Blanaru <[email protected]> Co-authored-by: Matthew Bonanni <[email protected]>

philippesic pushed a commit to philippesic/vllm-semantic-cache that referenced this pull request Jul 19, 2026

Xqa decode kernels (vllm-project#43232)

1ba62c3

Signed-off-by: Dan Blanaru <[email protected]> Co-authored-by: Matthew Bonanni <[email protected]>

plasticchris pushed a commit to plasticchris/vllm that referenced this pull request Jul 20, 2026

Xqa decode kernels (vllm-project#43232)

f6962e6

Signed-off-by: Dan Blanaru <[email protected]> Co-authored-by: Matthew Bonanni <[email protected]>

Uh oh!

Uh oh!

Conversation

DanBlanaru commented May 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

DanBlanaru commented May 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mergify Bot commented May 23, 2026

Uh oh!

DanBlanaru commented May 29, 2026

Uh oh!

mergify Bot commented Jun 1, 2026

Uh oh!

pavanimajety left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

MatthewBonanni left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ZhanqiuHu commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

DanBlanaru commented May 20, 2026 •

edited

Loading

DanBlanaru commented May 20, 2026 •

edited

Loading

ZhanqiuHu commented Jul 7, 2026 •

edited

Loading