Repository navigation
[None][feat] Support PrimsTS SageAttention for VisualGen VSA and SOL - #20224
Conversation
|
@heyuhhh — vendor promotion needs author action: The pinned temporary source commit Fix: Fetch The bot will recheck automatically on a later poll. |
heyuhhh
changed the title
Automated vendor metadata issue resolved; code review remains required.
da333b2 to
e77b710
Compare
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
WalkthroughThis change adds SageAttention Q/K/V quantization with optional K/V mean smoothing. It connects eligible quantized inputs to contiguous PrimTS FMHA and enables supported TRTLLM visual-generation sparse paths to use quantized attention. It also updates recipe validation, documentation, and tests. ChangesSageAttention quantization and sparse FMHA
Priority: ➖ Normal Estimated code review effort: 4 (Complex) | ~60 minutes Change: Feature Sequence Diagram(s)sequenceDiagram
participant TrtllmAttention
participant SageQuantize
participant TorchOperator
participant SageQuantKernel
participant PrimTSFmha
TrtllmAttention->>SageQuantize: Quantize eligible Q, K, and V
SageQuantize->>TorchOperator: Submit tensors and options
TorchOperator->>SageQuantKernel: Launch quantization kernels
SageQuantKernel-->>SageQuantize: Write quantized tensors and scales
SageQuantize-->>TrtllmAttention: Return quantized inputs and metadata
TrtllmAttention->>PrimTSFmha: Pass quantized inputs and metadata
Suggested reviewers: Merge Risk: 🔵 Low · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 30.97% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 155 functions across 24 files. (1 skipped… | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Title check The title follows the required format and clearly identifies the main change: adding PrimTS SageAttention support for VisualGen VSA and SOL. Description check The description includes the required Description, Test Coverage, and PR Checklist sections. It explains the motivation, implementation, API impact, performance, accuracy, vendored source, and relevan… Full details: Description check
-
Explanation
The description includes the required Description, Test Coverage, and PR Checklist sections. It explains the motivation, implementation, API impact, performance, accuracy, vendored source, and relevant tests in sufficient detail.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 30.97% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 155 functions across 24 files. (1 skipped: 1 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
✨ Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
|
/bot run --disable-fail-fast |
|
PR_Github #79735 [ run ] triggered by Bot. Commit: |
Vendor validation:
|
|
PR_Github #79735 [ run ] completed with state |
…ts-dev Update tensorrt_llm/_torch/attention/backends/prims_ts to FlashInfer commit 0e0761ef5 of the heyuhhh/flashinfer trtllm-prims-ts-dev branch: the tip of yuxianq/flashinfer trtllm-prims-ts-dev-20261012 (ae3851ab5) plus flashinfer-ai/flashinfer#6014. The two commits added to trtllm-prims-ts-dev-20261012 since the current pin work around a Blackwell driver TMA descriptor issue in the context kernel: on SM100 and SM103 devices, ragged Q tensor maps use conservative out-of-bounds addressing unless the Q layout proves a contiguous 128 KiB prefix. flashinfer-ai/flashinfer#6014 tunes the block-sparse decode kernels for SM103 and adds a Sage plan that prepares its scales ahead of attention: K scales are read in place from the prepared buffer, V scales and the optional V mean are per batch entry ([B, Hkv, D]), and the route preparation launches as a dependent of the scale preparation. The vendor patch is unchanged apart from refreshed line offsets. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
….compile - Run PrimsTSBlockSparseFmha.forward outside torch.compile: planning records an external CUDA event that Dynamo rejects. - Resolve the dense-layer flag at construction in the VisualGen TRTLLM and CuTeDSL SOL backends. Testing layer_idx per call made torch.compile specialize the attention frames per layer; with SOL sparse on every layer the core forward was then traced 40 times per call on real tensors, and the traces kept each layer's Q, K, V and output alive until a full garbage collection (OOM on B200 at 720p, 81 frames). - Test that compiled callers share one graph across SOL layers. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
The contiguous PrimsTSBlockSparseFmha path now consumes the SageAttention recipe from AttentionForwardArgs: - After route prediction, Q/K/V are quantized with FlashInfer's trtllm_sage_attention_quantize, honoring smooth_k. - Proxy K/V means are quantized by a new Triton op, trtllm::quantize_sage_summaries: one scale per K mean, while V means reuse the V channel scales so exact and proxy routes share one PV scale. - Plans carry explicit K and V dtypes and are keyed by the recipe. Sage needs BF16 input; paged block-sparse attention rejects it, and PrimTS plans check which tile profiles take it. The dense PrimsTSFmha still rejects Sage, now in its own support check instead of the shared feature gate. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
- AttentionConfig accepts a Sage quant_attention_config together with VSA or SOL when the backend is TRTLLM; other backends still reject the combination. - The TRTLLM VisualGen backend, VSA and SOL pass the recipe to the block-sparse path instead of rejecting it, so PrimsTSBlockSparseFmha quantizes after prediction. VSA falls back to the dense fine stage when Sage meets non-BF16 inputs. - Document the composition and test Sage block-sparse attention through the VisualGen TRTLLM backend. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
…nd SM103 The trtllm-gen FMHA cubins already include dense FP8 Sage context kernels with 16-token K scale blocks for SM100 and SM103, the attention op accepts Q/K blocks of 1, 4 and 16, and the PrimTS block-sparse path supports them. Only the VisualGen recipe table left FP8 (1, 16, 1) out. Add it for SM100 and SM103, list it in the quantized-attention guide, and cover it in the recipe validation and the dense numerical test. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
…e kernel The PrimTS block-sparse path quantized Q/K/V with FlashInfer's copy of the TensorRT-LLM Sage quantization kernel, which only offers 1, 4, or 16 tokens per Q/K scale and reduces the V scales and the K-smoothing mean over the whole batch. PrimTS now takes V scales per batch entry and accepts more K block sizes, so the path moves to the TensorRT-LLM kernel and extends it: - Per-sequence statistics: given KV cumulative lengths, the kernel reduces the V scales and the K mean over each sequence, [batch, heads_kv, head_dim], as SageAttention does. Without them it keeps the whole-batch reduction the trtllm-gen kernels consume. - Block sizes: contiguous Q/K scale blocks of 1, 4, 16, and now 32 tokens. Blocks larger than 8 tokens are read in 8-token chunks twice (maximum, then quantize), with the chunk loop kept rolled, so registers and code size stay fixed; smaller blocks keep the single pass. Other powers of two up to 256 only need a dispatch case. With 16 tokens per K scale, quantizing a real Wan VSA layer on GB300 takes 2.02 ms instead of 2.75 ms with FlashInfer's copy. - Rounding: Q/K are now scaled and converted in FP32 for every block size. The single pass converted the BF16/FP16 product to INT8, which truncates, so INT8 Q/K carried up to 1.2 grid steps of error (mean 0.44) instead of half a step (mean 0.25). This also applies to the dense INT8 Sage path. - Quantization outside the FMHA libraries: after sparse prediction and before FMHA selection, the attention core quantizes the Q/K/V of block-sparse calls and the K/V means of their proxy routes, and passes the result in AttentionForwardArgs.sage_quantized. It does so when a library of the layer declares supports_sage_quantized_inputs; PrimsTSBlockSparseFmha does, and serves Sage only with these inputs. Dense SageAttention still quantizes inside the attention op, which ignores them. AttentionForwardArgs.sage_attn_qk_dtype gives the core and the library the recipe's Q/K dtype. - Op: trtllm::sage_quantize quantizes Q/K/V, or, with reuse_stats, only K/V with the V scales and K mean of an earlier call; the Python sage_quantize takes that call's result as stats. The proxy summaries are quantized this way, one key per mean with a scale of its own, which replaces their Triton quantizer. The op has a fake implementation for torch.compile. The adapter no longer checks the quantizer's block sizes and input dtypes; the op validates its own inputs. VSA no longer limits Sage to BF16 inputs either: the scales are computed in FP32, so zero or tiny FP16 Q/K blocks no longer overflow them. - Tests compare against a PyTorch reference over all block sizes, INT8 and FP8, K smoothing, GQA, ragged sequences, and proxy summaries, check which calls the attention core quantizes, and cover zero Q/K with FP16 inputs. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
SageAttention can subtract the per-channel V mean before quantizing V and add it back to the output: every row's attention weights sum to one, so the output is unchanged while the range V has to cover in FP8 narrows. PrimTS takes this mean (v_mean), and the Sage path now supplies it: - Kernel: with a V-mean buffer, the statistics launch of sageQuant.cu collects each channel's V sum and signed extrema, and the quantizing launch derives scales that cover V minus its mean and subtracts the mean. V smoothing is a template parameter of the statistics launch, so the launch without it keeps its code and registers; one dispatch instantiates only the smoothing combinations that exist. trtllm::sage_quantize fills v_mean, and subtracts a given one with reuse_stats, so the V means of proxy routes subtract the mean the kernel adds back. - The attention op cannot add the mean back, so the attention core also quantizes dense contiguous calls with V smoothing, and PrimsTSBlockSparseFmha serves dense calls with core-quantized inputs through a PrimTS dense plan with 128-token Q/KV blocks, which select the faster Q128/KV128 Sage profile. With V smoothing, its Sage plans enable v_mean and take the mean. The attention op serves the remaining calls with V quantized unsmoothed and logs a warning once. - QuantAttentionConfig.smooth_v (default False) enables it for VisualGen on the TRTLLM backend; AttentionForwardArgs.sage_attn_smooth_v carries it. - Tests: quantizer references with V smoothing for Q/K/V and summaries, including -0.0 extrema; core quantization of dense calls with V smoothing; dense-plan selection; and VisualGen parity for mixed SOL routes, SOL dense timesteps, and a dense model. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
The Sage quantization kernel reduced the K-smoothing and V-smoothing means with float atomics: every warp adds its partial mean, and the order the atomics land in changes from run to run, so the last bits of the means did too. Those bits flip some quantized K/V values, and a diffusion model amplifies the flips over its steps: two runs of the same request with smoothing produced visibly different videos. The statistics launch now sums the means in fixed point (32 fraction bits) in int64. Every warp's partial mean is rounded once and integer additions are exact, so the sum no longer depends on the order. Partial means are clamped to 2^22 so that the sums cannot overflow. Any mean keeps smoothing exact as long as every consumer uses the same one. - A small kernel after the statistics launch converts the sums into the float means that the quantizing launch and the consumers read. With V smoothing it also turns the V maximum, now collected in the V scale buffer itself, and the V minimum into the scales of V minus its mean, so the quantizing launch reads the V scales as without smoothing. - Another small kernel before the statistics launch starts every reduction from its identity, so callers no longer clear buffers: the attention op drops its memsets, and trtllm::sage_quantize allocates its sums and V minima uninitialized. - The attention op's K-smoothing workspace gains the int64 sum, which also makes dense SageAttention with smooth_k reproducible; the workspace test checks the order of the Sage slices. - Test that repeated sage_quantize calls with K and V smoothing give bitwise identical means and quantized K and V. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
The quantizer and VisualGen Sage tests multiplied parameters whose combinations exercise no further code: every smoothing pair with every head dimension for each block size, 96 summary cases that include V scales no real call produces, and both Q/K dtypes where the path under test does not depend on them. Every kernel instantiation, block size, Q/K dtype, and SOL route kind stays covered, as do the torch.compile and CUDA Graph paths, the reproducibility check, and the edge cases; 192 cases become 71. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
… flag TrtllmAttention resolves at construction whether its layer is dense, and should_use_sparse reads that flag. The helper that builds the attention with object.__new__ did not set it, so test_forward_hands_separate_qkv_to_core_when_backend_rejects_fused_qkv failed with an AttributeError. The B200 test list also drops its separate test_sage_quant.py entry, which the unittest/_torch/attention entry already runs. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
…the V-scale constants - The mean conversion after a statistics launch writes the mean of every sum it is given and, for V, reads the minimum. invokeSageQuant now rejects a K or V mean sum without them, as it already rejected smoothing without the sums. - The K mean of a sequence is read at row seqIdx * numHeadsV + headIdx, the row the statistics launch writes, instead of relying on the K launch having as many heads as V. - The FP8 E4M3 maximum and the V-scale floor are named constants shared by the smoothed and unsmoothed V scales; the FP8 Q/K scales use the same maximum. The quantized outputs are bitwise unchanged. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
The checked-in schema for trtllm-serve --visual_gen_args now includes QuantAttentionConfig.smooth_v. It is regenerated with scripts/generate_trtllm_serve_schemas.py; the other two schemas are left unchanged, since their differences come from LLM arguments this change does not touch. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
…ooth_v on SM90 The attention core quantizes the Q/K/V of block-sparse and V-smoothing calls without asking which FMHA libraries of the layer take quantized inputs; each library decides in its support check whether it uses them, and the attention op ignores them. Fmha.supports_sage_quantized_inputs had no other reader and is removed. V smoothing needs the PrimTS kernels, which add the V mean back to the attention output and run only on SM100 and SM103, and the core quantizer has no layout for the SM90 recipe. AttentionConfig now rejects smooth_v on other architectures instead of quantizing inputs that no kernel there can use. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
8102192 to
08d8cfa
Compare
Vendor validation:
|
|
/bot run --disable-fail-fast |
Vendor validation:
|
|
PR_Github #79904 [ run ] triggered by Bot. Commit: |
|
PR_Github #79904 [ run ] completed with state Resource maintenance in LLM/main/L0_Test-x86_64-Multi-GPU #2981: B200 multi-GPU stages are paused to relieve the CI queue backlog (109 running / 459 pending on 2026-10-22). Owner: @sharabeshj-nvda. The CI team decides when to lift it. |
----- BEGIN VENDOR PROMOTION V1 -----
{
"attribution": {
"0e0761ef5ef10cb0762cc1b890a5c1a329bb49bf": {
"method": "upstream-aggregate",
"upstream_base": "ae312f7cc1723b06fda5168105ade0f328ee47d8",
"upstream_head": "1bec389b38a1534b7da21c135dde001a6abfd123"
},
"6b00497398f491136be91fb0e60ac3f314da53ba": {
"method": "same-edits",
"upstream_base": "2895577720741ddfedff118202c89584996b8875",
"upstream_head": "14770b419ec3261c2afc55e401051c92bb7f3fd9"
},
"ae3851ab5f958ba22f154059df5b38321f6369d2": {
"comparison": {
"source_base": "6b00497398f491136be91fb0e60ac3f314da53ba",
"source_head": "ae3851ab5f958ba22f154059df5b38321f6369d2",
"upstream_base": "873d21fd565adfa98748c53f90fd4f8aa4d5ed60",
"upstream_head": "a9883a5e8ada0d25754ecb31ae7eb83109c403cc"
},
"method": "similar-edits",
"refresh_policy": "retain-until-verified",
"similarity": 0.9859154929577465,
"similarity_metric": "sequence-matcher-lines-v1",
"similarity_threshold": 0.9,
"upstream_base": "2895577720741ddfedff118202c89584996b8875",
"upstream_head": "14770b419ec3261c2afc55e401051c92bb7f3fd9",
"upstream_pr": 6482
}
},
"attribution_scope": {
"include": [
"**/*.py"
],
"source": "flashinfer/attention/prims_ts"
},
"canonical_branch": "trtllm-prims-ts-dev-20261012",
"canonical_url": "https://github.com/yuxianq/flashinfer.git",
"changes": [
{
"commits": [
"0e0761ef5ef10cb0762cc1b890a5c1a329bb49bf"
],
"upstream_base": "ae312f7cc1723b06fda5168105ade0f328ee47d8",
"upstream_head": "1bec389b38a1534b7da21c135dde001a6abfd123",
"upstream_merge": "388e7a311e9f79941134dcc298cbc942b5c607a1",
"upstream_pr": "flashinfer-ai/flashinfer#6014"
},
{
"commits": [
"6b00497398f491136be91fb0e60ac3f314da53ba",
"ae3851ab5f958ba22f154059df5b38321f6369d2"
],
"upstream_base": "2895577720741ddfedff118202c89584996b8875",
"upstream_head": "14770b419ec3261c2afc55e401051c92bb7f3fd9",
"upstream_merge": "f57c58b898cc0f96ddea03066a3fd66ad201337a",
"upstream_pr": "flashinfer-ai/flashinfer#6482"
}
],
"consumer_base": "main",
"consumer_merge": "791a402fc8d514ba73ece0fe37af06036b08a18e",
"consumer_pr": "#20224",
"contributor_metadata": {
"metadata_digest": "ca3752e66016a893957f16ade23d2cf69bf0b75d087817f241402f29d47a5d6c",
"pr_author": "heyuhhh",
"source_head": "08d8cfa731414f27a20c31bbe9967f254a711cd3"
},
"kind": "promotion",
"previous": "ae312f7cc1723b06fda5168105ade0f328ee47d8",
"promoted": "0e0761ef5ef10cb0762cc1b890a5c1a329bb49bf",
"source_branch": "trtllm-prims-ts-dev",
"source_url": "https://github.com/heyuhhh/flashinfer.git",
"upstream_base": "ae312f7cc1723b06fda5168105ade0f328ee47d8",
"upstream_branch": "main",
"upstream_observed_main": "13d7fcf4041635f5b7f573bb0dbab23570ae1034",
"upstream_repository": "flashinfer-ai/flashinfer",
"vendor": "flashinfer-prims-ts",
"version": 1
}
----- END VENDOR PROMOTION V1 -----
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
…ted canonical tip TensorRT-LLM main (NVIDIA#20224, promoted by NVIDIA#20870) moved the flashinfer-prims-ts canonical pin from ae3851ab to 0e0761ef (the NVIDIA#6014 SM103 block-sparse/sage tuning). The linear source branch is rebased onto that tip as yunruis/flashinfer trtllm-prims-ts-dev-20261023 @ aa709061: the duplicate NVIDIA#6014 cherry-pick drops out and the remaining five flashinfer main PrimTS commits (NVIDIA#6089, NVIDIA#5351, NVIDIA#6482, NVIDIA#6529, NVIDIA#5434) replay without conflicts. The vendored tree, patch and digests are byte-identical to the previous pin (f924fd45); only the lock's branch and commit change. AI-assisted. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com>
…ted canonical tip TensorRT-LLM main (NVIDIA#20224, promoted by NVIDIA#20870) moved the flashinfer-prims-ts canonical pin from ae3851ab to 0e0761ef (the NVIDIA#6014 SM103 block-sparse/sage tuning). The linear source branch is rebased onto that tip as yunruis/flashinfer trtllm-prims-ts-dev-20261023 @ aa709061: the duplicate NVIDIA#6014 cherry-pick drops out and the remaining five flashinfer main PrimTS commits (NVIDIA#6089, NVIDIA#5351, NVIDIA#6482, NVIDIA#6529, NVIDIA#5434) replay without conflicts. The vendored tree, patch and digests are byte-identical to the previous pin (f924fd45); only the lock's branch and commit change. AI-assisted. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com>
Dev Engineer Review
trtllm::sage_quantizeoperator, PrimTS support for dense and block-sparse calls, and optional V smoothing throughQuantAttentionConfig.smooth_v, which defaults toFalse.QA Engineer Review
test_sage_block_sparse.pytotests/integration/test_lists/test-db/l0_b200.yml. The supplied change summary does not establish whether the other changed tests appear in CI or manual-QA lists.Per-File QA Perspective
3rdparty/vendor_patches/flashinfer-prims-ts.patch: Verify the patched imports, decorators, and tracing behavior against the pinned PrimTS source.3rdparty/vendor_sources.lock.yaml: Verify that the source and patch digests match the intended vendor tree.cpp/tensorrt_llm/common/attentionOp.cpp: Verify workspace sizing, K mean-sum wiring, and supported block-size errors.cpp/tensorrt_llm/common/attentionWorkspace.h: Verify alignment, ordering, and sizing of the new workspace slice.cpp/tensorrt_llm/common/sageQuant.cu: Verify smoothing, per-sequence statistics, mean accumulation, dtypes, and block-size handling.cpp/tensorrt_llm/common/sageQuant.h: Verify the new statistics pointers and sequence-length inputs.cpp/tensorrt_llm/thop/CMakeLists.txt: Verify that visual-generation targets compile and link the Sage quantization operator.cpp/tensorrt_llm/thop/sageQuantizeOp.cpp: Verify validation and output behavior for normal quantization andreuse_stats.cpp/tests/unit_tests/common/attentionWorkspaceTest.cpp: Verify workspace ordering and sizes.docs/source/features/visualgen-quantized-attention.md: Verify the recipes and smoothing behavior documentation.docs/source/features/visualgen-sparse-attention.md: Verify backend compatibility and restrictions.tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md: Verify documented quantized-input and smoothing behavior.tensorrt_llm/_torch/attention/backends/fmha/fallback.py: Verify when the V-smoothing warning is emitted.tensorrt_llm/_torch/attention/backends/fmha/interface.py: Verify the default capability flag for Sage-quantized inputs.tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py: Verify unsupported Sage configurations still fail.tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py: Verify plan selection, cache keys, and paged-request rejection.tensorrt_llm/_torch/attention/backends/interface.py: Verify Sage metadata and dtype selection.tensorrt_llm/_torch/attention/backends/trtllm.py: Verify quantization eligibility and proxy-statistics reuse.tensorrt_llm/_torch/attention/kernels/sage_quant.py: Verify allocation, smoothing, and statistics reuse.tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py: Verify the fake operator signature matches the registered operator.tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/backend.py: Verify dense-layer resolution and quantized SOL routing.tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/backend.py: Verify VSA eligibility and fallback behavior.tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py: Verify dense-layer selection and V-smoothing propagation.tensorrt_llm/visual_gen/args.py: Verify defaults, recipe validation, and backend restrictions.tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json: Verify thatsmooth_vis a boolean with afalsedefault and that its description matches the supported kernel behavior.tests/integration/test_lists/test-db/l0_b200.yml: Verify that the Sage block-sparse test entry resolves to the intended test.tests/unittest/_torch/attention/kernels/test_sage_quant.py: Covers quantization, smoothing, validation, graph replay, and reproducibility. Its test-list status is not established.tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py: Covers quantization eligibility, proxy statistics, dense behavior, and rejection paths. Its test-list status is not established.tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py: Covers dense routing and backend composition. Its test-list status is not established.tests/unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py: Covers VSA/SOL outputs, plans, smoothing, routing, replay, and edge cases. It is listed in the B200 pre-merge database.tests/unittest/_torch/visual_gen/test_attention_trtllm_sage.py: Adds an SM100 FP8 case. Its test-list status is not established.tests/unittest/_torch/visual_gen/test_visual_gen_args.py: Covers attention composition, recipe validation, and smoothing restrictions. Its test-list status is not established.tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py: Initializes dense-layer state in the test helper. Its test-list status is not established.Description
VisualGen VSA and SOL on the TRTLLM backend run their block-sparse attention on the vendored PrimTS kernels, but rejected a SageAttention
quant_attention_config. This PR enables SageAttention for them: the attention core quantizes Q/K/V with TensorRT-LLM's Sage quantization kernel before FMHA selection, and PrimTS runs Sage attention on the quantized inputs. It also adds V smoothing (QuantAttentionConfig.smooth_v) and makes the smoothing means bitwise reproducible. The commit messages describe the individual changes.API:
QuantAttentionConfig.smooth_vis a new prototype field (defaultFalse); existing configurations are unchanged (api-compatible).Setup
Wan2.1-VSA-T2V-14B, 720p, 81 frames, 50 steps, CFG 5.0, torch.compile. The Sage recipe quantizes Q per token, K per 16 tokens, and V per channel (FP8) with K and V smoothing; Q/K are FP8 on GB300 (SM103) and INT8 on B200 (SM100). "No Sage" is BF16 attention of the same kind (dense, SOL, or VSA with the same sparsity).
Performance
Kernel level (one self-attention call of a 720p layer, Sage vs BF16, quantization excluded). Isolated: one real layer called back to back; in the pipeline: median of the calls in denoise steps 9-11.
End to end (50-step denoise time; median over 4 GB300 / 6 B200 GPUs, speedups measured on the same GPU):
Accuracy
LPIPS (AlexNet, averaged over the 81 frames) between videos of the same prompt and seed, mean over 16 videos (8 prompts x 2 seeds):
Vendored source: the PrimTS pin is the latest tip of the dated canonical FlashInfer branch trtllm-prims-ts-dev-20261012, which adds the merged flashinfer-ai/flashinfer#6482 over the current pin, plus the merged flashinfer-ai/flashinfer#6014, cherry-picked as an exact upstream commit.
Test Coverage
tests/unittest/_torch/attention/kernels/test_sage_quant.py: quantizer against a PyTorch reference (all block sizes, INT8/FP8, K and V smoothing, GQA, ragged sequences, proxy summaries, -0.0 extrema), and bitwise identical results of repeated calls with K and V smoothingtests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py: core-side quantization, Sage plan rules, dense plan with V smoothingcpp/tests/unit_tests/common/attentionWorkspaceTest.cpp: order of the Sage workspace slicestests/unittest/_torch/attention/test_attention_op_sync.py,test_fmha_manager.py,test_fmha_interface.py,test_attention_failure_logging.pytests/unittest/_torch/visual_gen/sparse_attention(SOL and VSA),test_attention_trtllm_sage.py,test_visual_gen_args.pytests/unittest/_torch/attention/test_prims_ts_fmha.py,test_prims_ts_attention_backend.pyThe current head passes
test_sage_quant.py,test_prims_ts_block_sparse.py,sparse_attention/test_sage_block_sparse.py, the SOL and VSA tests,test_attention_trtllm_sage.py, andattentionWorkspaceTeston B200 (SM100); the full list passed on GB300 (SM103) when the PR was opened. Compiled Wan2.1-VSA-T2V-14B generations with VSA and SOL plus FP8 Sage, with and without V smoothing, run end to end on GB300.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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.