[None][feat] integrate PrimTS FMHA kernels by yuxianq · Pull Request #17399 · NVIDIA/TensorRT-LLM

yuxianq · 2026-08-07T04:16:29Z

Description

Integrate FlashInfer's experimental task-scheduled PrimTS attention kernels into the PyTorch TensorRT-LLM attention backend for supported Blackwell workloads. The integration is opt-in and depends on flashinfer-ai/flashinfer#4829.

Integration and supported scope

Add PrimsTSFmha for SM100 and SM103. Enable it with TLLM_FMHA_LIBS=+prims_ts; it precedes flashinfer_trtllm_gen and fallback in canonical order but is excluded from the default FMHA set. Unsupported requests continue through the remaining enabled FMHA libraries.
Support BF16/FP16 paged context and single-token decode for standard MHA/GQA, plus BF16 DeepSeek-style MLA decode.
Supported standard-attention geometry currently includes page sizes 16/32/64/128, context head dimensions 128/256, decode head dimensions 64/128/256, and decode GQA ratios up to 32. MLA decode requires BF16, one logical KV head, at most 128 local query heads, KV-LoRA rank 512, and RoPE head dimension 64.
Fail closed for unsupported features, including quantized KV cache, speculative decoding, beam search, sparse attention, chunked context, and cyclic sliding-window page tables.
Keep FMHA selection caching limited to values that can vary from request to request. Layer/model invariants do not belong in _FmhaCacheKey; this invariant is now documented next to the key. TensorRT-LLM does not enable CUDA Graph capture for context-only or mixed batches, and fused-QKV/cache-update mode is fixed for a layer within an LLM instance.

Unified page-table metadata

Context, decode, and MLA now use the same native PrimTS contract: an int32 fixed row-strided page table [B, C] plus live int32 sequence lengths [B].

After selecting the layer's KV pool and active batch, TensorRT-LLM exposes a graph-stable [B, 2, C] view. Plane 0 contains the K-page IDs used by PrimTS; plane 1 is the V-page field. The adapter passes block_tables[:B, 0, :] directly, producing a zero-copy [B, C] view with stride (2*C, 1), while the normalized K/V cache views retain their respective pool displacement.
Logical page p for request b resolves as block_tables[b, p]. seq_lens[b] determines the active page count, so unused capacity columns are never semantically consumed.
Context now consumes params.sequence_lengths[:batch_size] directly instead of recovering the same lengths by differencing cu_kv_seqlens.
The old CSR staging buffers and per-launch CSR transformation are removed from the TensorRT-LLM adapter for context, decode, and MLA.

Planning, preprocessing, and workspace ownership

Cache one context, decode, and MLA wrapper/plan per active batch size; other geometry and capacity are invariant for the layer/model instance. Bind current request metadata at run() time.
For standard attention, reuse TensorRT-LLM's fused context/generation preprocessing for QKV extraction, RoPE, KV-cache append, and page metadata. MLA decode reuses the KV metadata builder.
Add the backward-compatible internal skip_fmha_workspace=False THOP option. PrimTS uses it to omit the unused fixed 32 MiB TRTLLM-Gen FMHA slice while retaining compact preprocessing buffers.
Context needs no PrimTS scratch. Standard decode receives a separate aligned workspace tail, and MLA sizes only its own scratch.
Reset only plan_state.workspace.split_kv_counter when the fused global split-reduction path needs it, rather than clearing unrelated workspace bytes.
Make PhasedFmha populate both phase batch fields and compute the active layer's page-pool bound. V1 and V2 KV-cache managers use explicit type-specific bounds; V2 queries the local layer's Role.KEY upper bound.

Vendoring and dependencies

Vendor flashinfer/attention/prims_ts from yuxianq/flashinfer@20d8afb4 through the locked-source workflow. The lock records the source revision, persistent TensorRT-LLM compatibility patch, patch digest, and materialized-tree digest and supports offline and remote reconstruction checks.
Keep TensorRT-LLM-specific trace compatibility downstream: the vendor patch removes trace imports/decorators that are absent or schema-incompatible in TensorRT-LLM's installed FlashInfer package. The upstream PrimTS source has no TensorRT-LLM-specific trace loader.
Require CUTLASS DSL 4.7.0, QuACK 0.5.0, FlashAttention 4 4.0.0b19, and TVM-FFI 0.1.13.post2. The NGC 26.05 base image and system CUDA toolkit are unchanged.
Leave security_scanning/poetry.lock unchanged; its refresh is intentionally handled separately.

Performance Analysis

Current-head paired B200 performance

These measurements were taken at current PR head 895c2c2 with FlashInfer pin e500966b.

One exclusive B200 ran BF16, TP/PP/EP=1, page size 32, and fixed 1024-token input/output. C1 used 10 measured requests plus one warmup; C64 used 64 measured requests plus 64 warmups at concurrency 64. Every arm used a fresh process in N/P/P/N order. N used TLLM_FMHA_LIBS=-prims_ts; P used TLLM_FMHA_LIBS=+prims_ts. Values are geometric means of the forward and reverse passes. Positive speedup means that enabling PrimTS is better.

Model and workload	N latency (ms)	P latency (ms)	Latency speedup	N output tok/s	P output tok/s	Throughput speedup
Qwen2-7B-Instruct, C1 latency	3,285.937	3,306.617	-0.625%	311.627	309.677	-0.626%
DeepSeek-V3-Lite MLA, C1 latency	3,486.057	3,511.996	-0.739%	293.735	291.566	-0.738%
Qwen2-7B-Instruct, C64 offline throughput	4,829.636	4,841.853	-0.252%	13,557.248	13,523.167	-0.251%
DeepSeek-V3-Lite MLA, C64 offline throughput	5,604.950	5,200.513	+7.777%	11,681.900	12,588.366	+7.760%

The forward/reverse primary-metric ratio spans were 0.055 percentage points for Qwen C1 latency, 0.683 for DeepSeek C1 latency, 0.170 for Qwen C64 throughput, and 0.084 for DeepSeek C64 throughput. DeepSeek C64 therefore shows a clear repeatable gain; the other three cells retain small regressions.

All 16 counted arms passed import, provenance, dispatch, requested-length, and repeatability checks. DeepSeek output tokens matched exactly between N and P. Qwen output tokens were exact across repeats within each condition but differed between the two FMHA backends; cross-backend equality was informational rather than gating. Dispatch probes confirmed Qwen P used PrimTS context and decode, while DeepSeek P used fallback context and PrimTS MLA decode.

One Qwen C64 P-reverse process hit a KV-cache-manager-v2 introspection segfault before timing. That partial attempt was excluded; the single retry with the sealed, unchanged configuration passed and supplied the counted sample.

Historical paired B200 baseline

The following measurements were taken on PR revision 3063375c with FlashInfer pin 5a6a73a2, before the unified fixed-row metadata path. They remain a historical baseline, not a performance sign-off for the current head.

One exclusive B200 ran BF16, TP/PP/EP=1, page size 32, fixed 1024-token input/output, and fresh processes in ABBA order. ON used TLLM_FMHA_LIBS=+prims_ts; OFF used TLLM_FMHA_LIBS=-prims_ts.

Model and workload	PrimTS disabled output tok/s	PrimTS enabled output tok/s	Enabled vs disabled
Qwen2-7B-Instruct, C1 latency	313.159	305.174	-2.550%
DeepSeek-V3-Lite MLA, C1 latency	300.097	291.768	-2.776%
Qwen2-7B-Instruct, C64 offline throughput	13,573.93	13,322.88	-1.850%
DeepSeek-V3-Lite MLA, C64 offline throughput	11,904.13	11,539.71	-3.061%

All 16 runs completed and dispatch instrumentation proved the expected ON/OFF backend selection. PrimTS used less average GPU power in all four cells; DeepSeek C64 reduced measurement-window energy by 4.426% and improved output-token energy efficiency by 4.624% despite its throughput gap.

Test Coverage

An earlier branch revision passed a full clean Release build. The unified fixed-row revision passed an incremental Release build/install plus exact-worktree native import checks on B200 with CUTLASS DSL 4.7.0.
On that fixed-row revision, the full natural-order tests/unittest/_torch/attention/test_prims_ts_attention_backend.py passed 18/18 on both B200 SM100 and B300 SM103, with threadleak enabled. Coverage includes context/generation/mixed Qwen2 GQA, V1/V2 KV-cache managers, FP16 context, V-tail sanitization, DeepSeek MLA, fixed-row metadata mutation, workspace reset/sharing, generation CUDA Graph replay, and fallback behavior.
Latest-head CPU suites passed: test_fmha_manager.py (29), test_prims_ts_fmha.py (77), and test_combined_fmha.py (152).
Vendor verification passed locally and offline at 20d8afb4; the vendor test suite passed 25/25. Ruff, formatting, Python compilation, diff checks, DCO, and applicable pre-commit hooks passed.
Integration accuracy cases for DeepSeek-V3-Lite BF16 and Qwen3-8B BF16 are registered in l0_b200.yml and instrument real PrimTS phase calls.
CI was re-triggered with /bot run --disable-fail-fast. The latest full pipeline is in progress; see the PR checks for live status.

The B200/B300 18/18 runs were made on the fixed-row revision before the final rebase and subsequent vendor/manager follow-ups. The later CPU/vendor checks cover those follow-ups; the historical performance section is not a current-head performance claim.

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why.
PR follows the TRT-LLM coding guidelines to the best of my knowledge.
Test cases are provided for new code paths.
No public TensorRT-LLM API change is introduced; skip_fmha_workspace is an internal THOP parameter and defaults to False.
New and updated dependencies were checked through the security-scanning declarations and locked-source attribution workflow.
CODEOWNERS and attention documentation are updated.
The base container and system CUDA toolkit remain unchanged.
Reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, comment /bot help.

yuxianq · 2026-08-28T21:15:18Z

/bot run

tensorrt-cicd · 2026-08-28T21:21:35Z

PR_Github #70067 [ run ] triggered by Bot. Commit: 69753c5 Link to invocation

tensorrt-cicd · 2026-08-29T00:29:37Z

PR_Github #70067 [ run ] completed with state SUCCESS. Commit: 69753c5
/LLM/main/L0_MergeRequest_PR pipeline #57339 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

yuxianq · 2026-08-29T04:08:27Z

/bot run

yuxianq · 2026-08-29T06:23:47Z

/bot run

yuxianq · 2026-08-29T09:40:52Z

/bot run

tensorrt-cicd · 2026-08-29T09:46:34Z

PR_Github #70140 [ run ] triggered by Bot. Commit: 586531d Link to invocation

tensorrt-cicd · 2026-08-29T12:10:22Z

PR_Github #70140 [ run ] completed with state SUCCESS. Commit: 586531d
/LLM/main/L0_MergeRequest_PR pipeline #57403 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

Re-vendor the PrimTS source after restoring its upstream trace imports and preserve TensorRT-LLM package compatibility in the generated vendor patch. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

yuxianq · 2026-09-05T11:30:01Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-05T11:37:33Z

PR_Github #71688 [ run ] triggered by Bot. Commit: 895c2c2 Link to invocation

tensorrt-cicd · 2026-09-05T11:39:53Z

PR_Github #71663 [ run ] completed with state ABORTED. Commit: 63aaadc

Link to invocation

yuxianq · 2026-09-05T16:28:40Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-05T16:36:11Z

PR_Github #71714 [ run ] triggered by Bot. Commit: 895c2c2 Link to invocation

tensorrt-cicd · 2026-09-05T16:38:34Z

PR_Github #71688 [ run ] completed with state ABORTED. Commit: 895c2c2

Link to invocation

tensorrt-cicd · 2026-09-05T22:19:16Z

PR_Github #71714 [ run ] completed with state SUCCESS. Commit: 895c2c2
/LLM/main/L0_MergeRequest_PR pipeline #58795 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

yuxianq · 2026-09-05T22:31:59Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-05T22:38:59Z

PR_Github #71724 [ run ] triggered by Bot. Commit: 895c2c2 Link to invocation

tensorrt-cicd · 2026-09-06T02:16:23Z

PR_Github #71724 [ run ] completed with state FAILURE. Commit: 895c2c2
/LLM/main/L0_MergeRequest_PR pipeline #58805 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

yuxianq · 2026-09-06T13:24:01Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-06T13:31:13Z

PR_Github #71761 [ run ] triggered by Bot. Commit: 895c2c2 Link to invocation

tensorrt-cicd · 2026-09-06T14:20:20Z

PR_Github #71761 [ run ] completed with state SUCCESS. Commit: 895c2c2
/LLM/main/L0_MergeRequest_PR pipeline #58842 completed with status: 'SUCCESS'

CI Report

Link to invocation

… backends NVIDIA#17899 moved tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py under the new kv_cache/ subpackage. NVIDIA#17399, which added the phased and PrimTS FMHA libraries, was in flight at the same time and still imports the old path, so main now fails at: ModuleNotFoundError: No module named 'tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2' fmha/__init__.py imports phased and prims_ts, and attention/backends/__init__.py reaches fmha through trtllm.py, so the whole attention.backends package is unimportable and every unit test that touches an attention backend fails collection. Point the three fmha modules and the three FMHA unit tests at pyexecutor.kv_cache.kv_cache_manager_v2, matching the sibling backends (interface.py, sparse/dsa/cache_manager.py). Adds a CPU-only import guard so a future move of this module is caught by the unit tests. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>

…ak main Two module moves landed while dependent PRs were in flight, and main now fails to import the PyTorch runtime: 1. NVIDIA#17899 moved tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py under the new kv_cache/ subpackage. NVIDIA#17399, which added the phased and PrimTS FMHA libraries, still imports the old path: ModuleNotFoundError: No module named 'tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2' fmha/__init__.py imports phased and prims_ts, and attention/backends/__init__.py reaches fmha through trtllm.py, so the whole attention.backends package is unimportable. 2. NVIDIA#17968 moved tensorrt_llm/_torch/attention_backend/ to tensorrt_llm/_torch/attention/backends/. NVIDIA#18652 extracted pyexecutor/engine/lora.py with the old path, and model_engine.py imports engine.lora, so the executor is unimportable even with (1) fixed. Point the three fmha modules, engine/lora.py and the affected unit tests at the current paths, matching the sibling modules. Adds a CPU-only import guard covering both chains so a future move is caught by the unit tests. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>

…ak main Two module moves landed while dependent PRs were in flight, and main now fails to import the PyTorch runtime: 1. NVIDIA#17899 moved tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py under the new kv_cache/ subpackage. NVIDIA#17399, which added the phased and PrimTS FMHA libraries, still imports the old path: ModuleNotFoundError: No module named 'tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2' fmha/__init__.py imports phased and prims_ts, and attention/backends/__init__.py reaches fmha through trtllm.py, so the whole attention.backends package is unimportable. 2. NVIDIA#17968 moved tensorrt_llm/_torch/attention_backend/ to tensorrt_llm/_torch/attention/backends/. NVIDIA#18652 extracted pyexecutor/engine/lora.py with the old path, and model_engine.py imports engine.lora, so the executor is unimportable even with (1) fixed. Point the three fmha modules, engine/lora.py, the visual-gen FlashInfer backend and the affected unit tests at the current paths, matching the sibling modules. Adds a CPU-only import guard covering both chains so a future move is caught by the unit tests. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>

…ak main Two module moves landed while dependent PRs were in flight, and main now fails to import the PyTorch runtime: 1. NVIDIA#17899 moved tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py under the new kv_cache/ subpackage. NVIDIA#17399, which added the phased and PrimTS FMHA libraries, still imports the old path: ModuleNotFoundError: No module named 'tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2' fmha/__init__.py imports phased and prims_ts, and attention/backends/__init__.py reaches fmha through trtllm.py, so the whole attention.backends package is unimportable. 2. NVIDIA#17968 moved tensorrt_llm/_torch/attention_backend/ to tensorrt_llm/_torch/attention/backends/. NVIDIA#18652 extracted pyexecutor/engine/lora.py with the old path, and model_engine.py imports engine.lora, so the executor is unimportable even with (1) fixed. Point the three fmha modules, engine/lora.py, the visual-gen FlashInfer backend and the affected unit tests at the current paths, matching the sibling modules. Adds a CPU-only import guard covering both chains so a future move is caught by the unit tests. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com> (cherry picked from commit 0c221e3) Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>

github-actions Bot assigned yuxianq Aug 7, 2026

github-actions Bot added the VisualGen label Aug 7, 2026

This was referenced Aug 21, 2026

[None][feat] integrate VisualGen VSA with block-sparse FMHA #18079

Draft

[None][feat] add generic PrimTS block-sparse FMHA yuxianq/TensorRT-LLM#2

Draft

yuxianq force-pushed the feat/prims-ts-fmha-lib branch 2 times, most recently from 74f5bf0 to 69753c5 Compare August 28, 2026 21:11

yuxianq force-pushed the feat/prims-ts-fmha-lib branch from 9969311 to 586531d Compare August 29, 2026 09:39

yuxianq marked this pull request as ready for review August 29, 2026 12:21

yuxianq requested review from a team as code owners August 29, 2026 12:21

yuxianq requested review from EmmaQiaoCh, kris1025, luyiyun1021, mlefeb01, pengbowang-nv and yibinl-nvidia August 29, 2026 12:21

yuxianq added 8 commits September 5, 2026 11:20

[None][chore] defer Poetry lock refresh

6cdeb5f

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] distinguish graph FMHA cache entries

a787939

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] key FMHA cache by input layout

770c239

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] keep FlashInfer trace compatibility downstream

86c5a19

Re-vendor the PrimTS source after restoring its upstream trace imports and preserve TensorRT-LLM package compatibility in the generated vendor patch. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] limit FMHA cache key to request state

1ead82d

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][chore] update vendored FlashInfer PrimTS

f3e8291

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][chore] update vendored FlashInfer PrimTS

aca8fac

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] align PrimTS with attention package

895c2c2

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

yuxianq force-pushed the feat/prims-ts-fmha-lib branch from 63aaadc to 895c2c2 Compare September 5, 2026 11:29

yuxianq merged commit b6dd8af into NVIDIA:main Sep 6, 2026
13 checks passed

brnguyen2 mentioned this pull request Sep 6, 2026

[None][fix] Update stale attention and KV-cache import paths that break main #18771

Merged

1 task

chang-l mentioned this pull request Sep 6, 2026

[https://nvbugs/6676844][test] Unwaive test_wan_t2v_example after CI checkpoint storage fix #18732

Open

1 task

Uh oh!

Conversation

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

Uh oh!

Description

Integration and supported scope

Unified page-table metadata

Planning, preprocessing, and workspace ownership

Vendoring and dependencies

Performance Analysis

Current-head paired B200 performance

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

yuxianq commented Aug 28, 2026

Uh oh!

tensorrt-cicd commented Aug 28, 2026

Uh oh!

tensorrt-cicd commented Aug 29, 2026

Uh oh!

yuxianq commented Aug 29, 2026

Uh oh!

yuxianq commented Aug 29, 2026

Uh oh!

yuxianq commented Aug 29, 2026

Uh oh!

tensorrt-cicd commented Aug 29, 2026

Uh oh!

tensorrt-cicd commented Aug 29, 2026

Uh oh!

yuxianq commented Sep 5, 2026

Uh oh!

tensorrt-cicd commented Sep 5, 2026

Uh oh!

tensorrt-cicd commented Sep 5, 2026

Uh oh!

yuxianq commented Sep 5, 2026

Uh oh!

tensorrt-cicd commented Sep 5, 2026

Uh oh!

tensorrt-cicd commented Sep 5, 2026

Uh oh!

tensorrt-cicd commented Sep 5, 2026

Uh oh!

yuxianq commented Sep 5, 2026

Uh oh!

tensorrt-cicd commented Sep 5, 2026

Uh oh!

tensorrt-cicd commented Sep 6, 2026

Uh oh!

yuxianq commented Sep 6, 2026

Uh oh!

tensorrt-cicd commented Sep 6, 2026

Uh oh!

tensorrt-cicd commented Sep 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

yuxianq commented Aug 7, 2026 •

edited

Loading