[None][feat] VC-Attention-QK(16,8) recipes for PrimsTS FMHA context by harrisonzhy · Pull Request #20068 · NVIDIA/TensorRT-LLM

harrisonzhy · 2026-10-09T21:30:38Z

Description

Adds the VC-Attention-QK16 and VC-Attention-QK8 recipes to the PrimTS FMHA in the VisualGen TRTLLM
attention backend, on top of the PrimTS context FMHA already on main.

Vendored prims_ts pinned to harrisonzhy/trtllm-prims-ts-dev_vc-attn-20261029 @ 7087b06c2: the canonical
trtllm-prims-ts-dev-20261012 tip (aa709061) plus one squashed commit per FlashInfer PR:
flashinfer PR 6296 (VC-Attention-QK16),
flashinfer PR 6306 (VC-Attention-QK8) and
flashinfer PR 6721 (preparation fused into the
pre-pass kernels), cherry-picked. Lock and compatibility patch regenerated by scripts/vendor/manage.py; the
patch content is unchanged.
QuantAttentionConfig.vc_attention runs a recipe as VC-Attention: (bf16, fp8, (0, 0, 1)) is
VC-Attention-QK16 and (fp8, fp8, (128, 128, 1)) VC-Attention-QK8, on SM100/SM103 with
TLLM_FMHA_LIBS=+prims_ts. vc_repair_budget appends that fraction of tokens as V repair rows instead
of restoring the tile means (0 keeps the means) and requires vc_attention.
The VisualGen TRTLLM attention layer prepares the VC operands per layer (K/V, and Q under QK8) with
the repair budget or on the V-Smooth schedule, and passes them to the attention op as
AttentionForwardArgs.vc_operands. The PrimTS adapter only runs the matching plan, keyed on the
demean phase and repair tile count, which are compile-time plan options of the kernel. Requires equal Q
and K/V head counts, head dimension 128 and a dense mask; VC requests with a KV cache are rejected.
The V-Smooth schedule follows the normalized timestep the attention already receives, like the
timestep-scheduled sparse backends (sparse/timestep_phase.py): the top quarter of the timestep range
smooths V with four regroupings, the rest runs plain E4M3 V. With CUDA graphs that phase is a graph key
(vc_smooth_phase), so each phase replays its own graph.

Vendor Promotion

schema_version: 1
vendors:
  flashinfer-prims-ts:
    upstream_prs:
      - https://github.com/flashinfer-ai/flashinfer/pull/6296
      - https://github.com/flashinfer-ai/flashinfer/pull/6306
      - https://github.com/flashinfer-ai/flashinfer/pull/6721
    resolutions:
      6e9018fc753a1c648a6504bba8a8df42339ef50b:
        upstream_pr: https://github.com/flashinfer-ai/flashinfer/pull/6296
        reason: Squash of the ten PR 6296 commits cherry-picked onto the canonical pin aa709061. The upstream PR diff is taken against a newer flashinfer main whose prims_ts files differ from the pin in neighbouring lines, so the aggregate similarity is 79%.
      77faea27329cb6bbe4698843aa058ed67ea27bb0:
        upstream_pr: https://github.com/flashinfer-ai/flashinfer/pull/6306
        reason: Squash of the two PR 6306 commits cherry-picked onto the canonical pin; the PR branch was stacked on PR 6296, so its diff includes that PR and scores 51% against this commit.

Test Coverage

tests/unittest/_torch/visual_gen/test_visual_gen_args.py: vc_attention on both recipes,
vc_repair_budget acceptance, bounds and rejection, non-TRTLLM and non-Blackwell rejection.
tests/unittest/_torch/visual_gen/test_vc_timestep_phase.py: the timestep schedule and the graph key.
tests/unittest/_torch/attention/test_prims_ts_attention_backend.py (GPU): both recipes with tile means
and with a repair budget, VC operands passed, SDPA match; VC rejected on paged-KV requests.
End to end on B200: Wan2.2-T2V-A14B through trtllm-serve, FLUX.3 klein-base text-to-image through the
VisualGen API, both recipes with and without vc_repair_budget: 0.005.

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The 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, please comment /bot help.

yuxianq · 2026-10-09T21:33:37Z

@harrisonzhy — vendor promotion needs author action:

Upstream attribution is unresolved for these source commits:

6b00497398f491136be91fb0e60ac3f314da53ba: no exact match; best similarities: upstream PR #6721: 1.37%; no PR meets the 90.00% similarity threshold
ae3851ab5f958ba22f154059df5b38321f6369d2: no exact match; best similarities: upstream PR #6721: 2.42%; no PR meets the 90.00% similarity threshold
c12d2aaed9e8bd1ea44f1fdc051ad501e99ba70f: no exact match; best similarities: upstream PR #6721: 53.21%; no PR meets the 90.00% similarity threshold
6af6d33e49ea9eb31d7decfec2b043fb0f78a1f3: no exact match; best similarities: upstream PR #6721: 35.20%; no PR meets the 90.00% similarity threshold
70cee28e32912ff5fe52c297b2c97e0900e84ef0: no exact match; best similarities: upstream PR #6721: 15.10%; no PR meets the 90.00% similarity threshold
4948e810217540bc9ce1966651626b7a3764f511: no exact match; best similarities: upstream PR #6721: 23.22%; no PR meets the 90.00% similarity threshold
37d9543933d27a21049b73b603eaf43917fde224: no exact match; best similarities: upstream PR #6721: 2.10%; no PR meets the 90.00% similarity threshold
3b4959843e694a731c0ddf678dda88da98e15670: no exact match; best similarities: upstream PR #6721: 5.90%; no PR meets the 90.00% similarity threshold
64c2e14937c572d0e15727595c9b0ba62605b3e1: no exact match; best similarities: upstream PR #6721: 2.27%; no PR meets the 90.00% similarity threshold
76d34ef9269d6fd263921d76b2683221f8bab0e3: no exact match; best similarities: upstream PR #6721: 15.42%; no PR meets the 90.00% similarity threshold
62a4002b306b5c16631684c41cf97052f6a7c3c0: no exact match; best similarities: upstream PR #6721: 59.29%; no PR meets the 90.00% similarity threshold
47dd9da34c120de5a6cc69b9b553be2b03806453: no exact match; best similarities: upstream PR #6721: 6.25%; no PR meets the 90.00% similarity threshold
16341a21a1dc3f86105803aee46fe1632721eac1: no exact match; best similarities: upstream PR #6721: 1.65%; no PR meets the 90.00% similarity threshold
219fda5bddf3c483d24df065f2e6dfcd705e6929: no exact match; best similarities: upstream PR #6721: 0.51%; no PR meets the 90.00% similarity threshold
c7271b1ede6cb337b654d058531ec3ae01eba502: no exact match; best similarities: upstream PR #6721: 8.34%; no PR meets the 90.00% similarity threshold
24dc58ca9ec4590dab6df43f0b8a02e4ee728057: no exact match; best similarities: upstream PR #6721: 7.58%; no PR meets the 90.00% similarity threshold
dc38cf843e2c2ea01bc79c2115fb319b8cf8fd31: no exact match; best similarities: upstream PR #6721: 19.03%; no PR meets the 90.00% similarity threshold
31f1199d783887a016361cb892e957567febbfb9: no exact match; best similarities: upstream PR #6721: 2.48%; no PR meets the 90.00% similarity threshold
6fa6865e9a96c0d12b50d3540d2ee2a4cf0355b3: no exact match; best similarities: upstream PR #6721: 43.99%; no PR meets the 90.00% similarity threshold

Fix: In the PR description, add missing PRs to vendors.flashinfer-prims-ts.upstream_prs, remove incorrectly listed PRs, or correct/rebase the source pin and update the vendored snapshot. For an intentional adaptation or ambiguous ownership, merge a targeted resolutions entry into the existing vendor metadata as shown below. Replace the SHA and URL and give a specific reason; preserve other entries. No complete commit_map is required. A resolution records an author assertion, not verified equivalence; it is retained during refresh until separately verified.

vendors:
  flashinfer-prims-ts:
    upstream_prs:
      - https://github.com/flashinfer-ai/flashinfer/pull/123
    resolutions:
      FULL_40_CHARACTER_SOURCE_SHA:
        upstream_pr: https://github.com/flashinfer-ai/flashinfer/pull/123
        reason: Explain the downstream adaptation or why this PR owns the change.

The bot will recheck automatically on a later poll.

yuxianq

Vendor validation result: no source-validation objection.

Latest vendor validation and step-by-step fix.

Automated vendor metadata issue resolved; code review remains required.

yuxianq

Vendor validation result: no source-validation objection.

Latest vendor validation and step-by-step fix.

Automated vendor metadata issue resolved; code review remains required.

yuxianq

Vendor validation result: no source-validation objection.

Latest vendor validation and step-by-step fix.

yuxianq · 2026-10-13T08:23:52Z

@harrisonzhy — the PrimTS upstream refresh #20155 has merged. Please update this source-update PR before merging it.

The new canonical source branch is yuxianq/flashinfer:trtllm-prims-ts-dev-20261012, initially pinned at ae312f7cc1723b06fda5168105ade0f328ee47d8.

Rebase your temporary FlashInfer branch onto the new canonical branch; retain only changes not already included upstream, resolving any API differences.
Rebase this TRT-LLM PR onto latest main, including the refreshed compatibility patch and integration changes.
Keep vendors.flashinfer-prims-ts in 3rdparty/vendor_sources.lock.yaml pointing to your temporary source branch, update its immutable commit, and regenerate the selected vendored files and digests. Do not restore the old canonical branch/pin while resolving conflicts.
Update the PR description's vendor-promotion block if the paired upstream PRs or explicit commit resolutions change after rebasing.
Verify the vendored snapshot and rerun the relevant tests and CI.

The publishing bot now targets the new dated canonical branch. After this source-update PR merges, its lock-only promotion will target that branch. The old trtllm-prims-ts-dev branch is intentionally unchanged.

Automated vendor metadata issue resolved; code review remains required.

yuxianq

Vendor validation result: no source-validation objection.

Latest vendor validation and step-by-step fix.

harrisonzhy · 2026-11-01T07:26:48Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-11-01T07:34:26Z

PR_Github #81830 [ run ] triggered by Bot. Commit: 8952905 Link to invocation

tensorrt-cicd · 2026-11-01T11:24:44Z

PR_Github #81830 [ run ] completed with state SUCCESS. Commit: 8952905
/LLM/main/L0_MergeRequest_PR pipeline #67483 completed with status: 'UNSTABLE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Either:

Wait for the PR to be fully approved — the label is added automatically once approval is complete. Having unresolved open comments is fine, or
If needed, ask a member of NVIDIA/trt-llm-ci-approvers to add the label manually.
Then re-trigger CI with the same bot command (no rebase needed).

⚠️ 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

Link to invocation

harrisonzhy · 2026-11-02T09:05:37Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-11-02T09:13:37Z

PR_Github #81964 [ run ] triggered by Bot. Commit: 5cbc0d7 Link to invocation

tensorrt-cicd · 2026-11-02T10:36:01Z

PR_Github #81964 [ run ] completed with state FAILURE. Commit: 5cbc0d7
/LLM/main/L0_MergeRequest_PR pipeline #67607 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

harrisonzhy · 2026-11-02T17:53:07Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-11-02T18:01:10Z

PR_Github #82051 [ run ] triggered by Bot. Commit: 0ab2727 Link to invocation

chang-l

Approved w/ minor comments on doc changes

github-actions · 2026-11-02T21:00:30Z

Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements.

tensorrt-cicd · 2026-11-03T03:20:30Z

PR_Github #82051 [ run ] completed with state SUCCESS. Commit: 0ab2727
/LLM/main/L0_MergeRequest_PR pipeline #67685 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

harrisonzhy · 2026-11-03T18:39:12Z

/bot run

tensorrt-cicd · 2026-11-03T18:56:15Z

PR_Github #82337 [ run ] triggered by Bot. Commit: db30816 Link to invocation

tensorrt-cicd · 2026-11-03T22:10:23Z

PR_Github #82337 [ run ] completed with state FAILURE. Commit: db30816
/LLM/main/L0_MergeRequest_PR pipeline #67956 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

…S FMHA QuantAttentionConfig.vc_attention runs a recipe as VC-Attention on the TRTLLM backend: (bf16, fp8, (0, 0, 1)) keeps Q/K in BF16 with V as E4M3 residuals around 128-token tile means that the kernel restores (VC-Attention-QK16), and (fp8, fp8, (128, 128, 1)) adds Hadamard-rotated E4M3 Q/K scaled per 128 tokens (VC-Attention-QK8). vc_repair_budget replaces the tile means by that fraction of tokens as V repair rows and requires vc_attention. The VisualGen backend sends the recipe down the separate-QKV path as forward args, the cache-free adapter owns the per-layer V-Smooth preparation over the denoising steps and keys its plan on the demean phase and repair tile count, and FMHA libraries without the supports_vc_attention capability reject the request. Signed-off-by: Harrison Zhang <harrisonz@nvidia.com>

… 7087b06c2 Pins flashinfer-prims-ts to the fork lineage branch harrisonzhy/trtllm-prims-ts-dev_vc-attn-20261029: the canonical trtllm-prims-ts-dev-20261012 tip aa709061 plus the VC-Attention commits of flashinfer PRs 6296 (QK16), 6306 (QK8) and 6721 (fused preparation), cherry-picked. The vendored tree is a byte copy of the fork plus the compatibility patch, regenerated at the new revision with unchanged content. Signed-off-by: Harrison Zhang <harrisonz@nvidia.com>

harrisonzhy · 2026-11-04T06:14:18Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-11-04T06:22:10Z

PR_Github #82418 [ run ] triggered by Bot. Commit: 5b5c3a2 Link to invocation

github-actions Bot assigned harrisonzhy Oct 9, 2026

github-actions Bot added the VisualGen label Oct 9, 2026

harrisonzhy changed the title ~~[None][feat] primsTS dense bf16, qk-bf16/pv-fp8, fp8 fmha context perf bringup~~ [None][feat] VC-Attention-QK16 recipe for PrimTS dense context FMHA Oct 9, 2026

yuxianq previously requested changes Oct 9, 2026

View reviewed changes

harrisonzhy closed this Oct 10, 2026

harrisonzhy deleted the harrisonzhy/primsts_fmha_vc-attn-qk16v8 branch October 10, 2026 02:13

harrisonzhy restored the harrisonzhy/primsts_fmha_vc-attn-qk16v8 branch October 10, 2026 02:14

harrisonzhy reopened this Oct 10, 2026

harrisonzhy mentioned this pull request Oct 10, 2026

feat(prims-ts): VC-Attention FP8 recipe for fmha context flashinfer-ai/flashinfer#6306

Merged

5 of 11 tasks

yuxianq previously requested changes Oct 10, 2026

View reviewed changes

yuxianq mentioned this pull request Oct 12, 2026

[None][chore] Refresh PrimTS vendor source onto FlashInfer main #20155

Merged

1 task done

harrisonzhy changed the title ~~[None][feat] VC-Attention-QK16 recipe for PrimTS dense context FMHA~~ [None][feat] VC-Attention-QK(16,8) recipe for PrimTS dense context FMHA Oct 12, 2026

yuxianq previously requested changes Oct 13, 2026

View reviewed changes

harrisonzhy force-pushed the harrisonzhy/primsts_fmha_vc-attn-qk16v8 branch from 920e6e4 to d8bd052 Compare October 13, 2026 17:23

harrisonzhy requested a review from yuxianq October 13, 2026 17:48

chang-l mentioned this pull request Oct 13, 2026

[None][perf] primsTS dense bf16, qk-bf16/pv-fp8, fp8 fmha context bringup #20006

Merged

1 task done

harrisonzhy force-pushed the harrisonzhy/primsts_fmha_vc-attn-qk16v8 branch 4 times, most recently from 58eb3de to 4b29c74 Compare October 19, 2026 21:36

harrisonzhy changed the title ~~[None][feat] VC-Attention-QK(16,8) recipe for PrimTS dense context FMHA~~ [None][feat] VC-Attention-QK(16,8) recipes for PrimTS FMHA context Oct 19, 2026

harrisonzhy changed the title ~~[None][feat] VC-Attention-QK(16,8) recipes for PrimTS FMHA context~~ [None][feat] VC-Attention-QK(16,8) recipes for PrimsTS FMHA context Oct 19, 2026

harrisonzhy force-pushed the harrisonzhy/primsts_fmha_vc-attn-qk16v8 branch from 4b29c74 to ee4d7d5 Compare October 19, 2026 23:25

yuxianq previously requested changes Oct 21, 2026

View reviewed changes

harrisonzhy force-pushed the harrisonzhy/primsts_fmha_vc-attn-qk16v8 branch from b3b8856 to adb2aa8 Compare October 21, 2026 22:20

farazkh80 reviewed Nov 1, 2026

View reviewed changes

harrisonzhy force-pushed the harrisonzhy/primsts_fmha_vc-attn-qk16v8 branch from 8952905 to 5cbc0d7 Compare November 2, 2026 00:47

harrisonzhy requested a review from farazkh80 November 2, 2026 09:05

farazkh80 approved these changes Nov 2, 2026

View reviewed changes

harrisonzhy force-pushed the harrisonzhy/primsts_fmha_vc-attn-qk16v8 branch from 5cbc0d7 to 0ab2727 Compare November 2, 2026 17:52

chang-l approved these changes Nov 2, 2026

View reviewed changes

Comment thread docs/source/features/visualgen-quantized-attention.md Resolved

Comment thread docs/source/features/visualgen-quantized-attention.md Resolved

trtllm-agent added the ci: full pre-merge approved label Nov 2, 2026

harrisonzhy force-pushed the harrisonzhy/primsts_fmha_vc-attn-qk16v8 branch from 0ab2727 to db30816 Compare November 2, 2026 21:39

harrisonzhy assigned xrq-phys and unassigned xrq-phys Nov 2, 2026

harrisonzhy enabled auto-merge (squash) November 2, 2026 21:53

xrq-phys approved these changes Nov 2, 2026

View reviewed changes

harrisonzhy added 2 commits November 4, 2026 02:00

harrisonzhy force-pushed the harrisonzhy/primsts_fmha_vc-attn-qk16v8 branch from db30816 to 5b5c3a2 Compare November 4, 2026 02:03

Conversation

harrisonzhy commented Oct 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Vendor Promotion

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

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

Uh oh!

Uh oh!

yuxianq left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yuxianq left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yuxianq left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yuxianq commented Oct 13, 2026

Uh oh!

yuxianq left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

harrisonzhy commented Nov 1, 2026

Uh oh!

tensorrt-cicd commented Nov 1, 2026

Uh oh!

tensorrt-cicd commented Nov 1, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

harrisonzhy commented Nov 2, 2026

Uh oh!

tensorrt-cicd commented Nov 2, 2026

Uh oh!

tensorrt-cicd commented Nov 2, 2026

Uh oh!

harrisonzhy commented Nov 2, 2026

Uh oh!

tensorrt-cicd commented Nov 2, 2026

Uh oh!

chang-l left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Nov 2, 2026

Uh oh!

tensorrt-cicd commented Nov 3, 2026

Uh oh!

harrisonzhy commented Nov 3, 2026

Uh oh!

tensorrt-cicd commented Nov 3, 2026

Uh oh!

tensorrt-cicd commented Nov 3, 2026

Uh oh!

harrisonzhy commented Nov 4, 2026

Uh oh!

tensorrt-cicd commented Nov 4, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

11 participants

harrisonzhy commented Oct 9, 2026 •

edited

Loading

yuxianq commented Oct 9, 2026 •

edited

Loading

yuxianq left a comment •

edited

Loading

yuxianq left a comment •

edited

Loading

yuxianq left a comment •

edited

Loading

yuxianq left a comment •

edited

Loading