Repository navigation
Conversation
…ts-dev Update tensorrt_llm/_torch/attention/backends/prims_ts to FlashInfer commit dae2097db of the heyuhhh/flashinfer trtllm-prims-ts-dev branch: the current pin (yuxianq/flashinfer trtllm-prims-ts-dev-20261012, FlashInfer main ae312f7cc) plus flashinfer-ai/flashinfer#6014. That change 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. - Ops: trtllm::sage_quantize quantizes Q/K/V, and trtllm::sage_quantize_kv quantizes K/V with the V scales and K mean of an earlier call. The proxy summaries use the latter with one scale per K mean, which replaces their Triton quantizer. Both ops have fake implementations for torch.compile. The adapter no longer checks the quantizer's block sizes and input dtypes; the ops validate their 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. trtllm::sage_quantize fills v_mean, and trtllm::sage_quantize_kv subtracts a given one, 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 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>
…ction VSA chooses between the PrimTS block-sparse fine stage and the compact dense fine stage before it predicts routes, and made that choice from its own copy of the PrimTS rules. The copy did not know the plan rules of a Sage recipe, so VSA could build routes that no FMHA library accepts. PrimsTSBlockSparseFmha now answers the question with bsr_unsupported_reason, which validates the same plan key as a real call, and VSA keeps only its own geometry checks. - VSAMetadata.topk gives the predictor and the check one routing width. - The VisualGen TRTLLM backend builds the Sage forward arguments in one place, _sage_forward_kwargs. - Tests: the check accepts the VSA route shape and returns the plan validator's reason otherwise; VSA falls back to the dense fine stage when PrimTS cannot plan its routes. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
Collaborator
|
Vendor promotion metadata and source attribution are valid. This is not a code-review approval. Promotion still requires the source PR to merge. |
yuxianq
previously requested changes
Oct 13, 2026
yuxianq
left a comment
yuxianq
left a comment
Collaborator
There was a problem hiding this comment.
Missing vendor-promotion markers in the PR description: <!-- vendor-promotion:start -->, <!-- vendor-promotion:end -->.
See the vendor-bot comment for the specific fix and the latest validation result.
heyuhhh
changed the title
yuxianq
dismissed
their stale review
October 13, 2026 13:21
Automated vendor metadata issue resolved; code review remains required.
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. A small kernel then converts the sums into the float means that the quantizing launch and the consumers read; the quantizing launch is unchanged. Any mean keeps smoothing exact as long as every consumer uses the same one. Partial means are clamped to 2^22 so that the sums cannot overflow. - sageQuant: ptrKMeanSum and ptrVMeanSum carry the zeroed sums of the statistics launch. - trtllm::sage_quantize allocates the sums; k_mean and v_mean are written for every row, empty sequences included. - The attention op's K-smoothing workspace gains an int64 sum, which also makes dense SageAttention with smooth_k reproducible. - 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>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
@coderabbitai summary
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 the combination and quantizes the PrimTS Sage path with TensorRT-LLM's own Sage quantization kernel, extended for it.Commits:
trtllm-prims-ts-dev(dae2097db): the current pin from [None][chore] Refresh PrimTS vendor source onto FlashInfer main #20155 (FlashInfer main ae312f7cc) plus perf(prims-ts): tune the block-sparse/sage kernel for SM103 flashinfer-ai/flashinfer#6014, which tunes the block-sparse decode kernels for SM103 and adds a Sage plan with per-batch V scales and an optional V mean.PrimsTSBlockSparseFmha.forwardruns outside Dynamo, and the dense-layer flag is resolved at construction (testinglayer_idxper call re-traced the core forward per layer with SOL on every layer and ran out of memory on B200 at 720p, 81 frames).PrimsTSBlockSparseFmha.AttentionForwardArgs.sage_quantized,supports_sage_quantized_inputs). New opstrtllm::sage_quantizeandtrtllm::sage_quantize_kvreplace FlashInfer's quantizer and the Triton summary quantizer.QuantAttentionConfig.smooth_v): the quantization kernel subtracts the per-channel V mean and PrimTS adds it back; dense Sage calls with V smoothing run on a PrimTS dense plan.PrimsTSBlockSparseFmhawhether its routes can run before it predicts them, instead of keeping its own copy of the PrimTS rules.API:
QuantAttentionConfig.smooth_vis a new prototype field (defaultFalse); existing configurations are unchanged (api-compatible).Quantization cost on GB300 (a real Wan2.1 14B 720p layer, 16 tokens per K scale): VSA 2.03 ms (FlashInfer's copy of the kernel: 2.75 ms), SOL 1.69 ms (2.28 ms).
End-to-end results (Wan2.1-VSA-T2V-14B, 720p, 81 frames, 50 steps; B200 with INT8 Q/K, GB300 with FP8 Q/K, 16 tokens per K scale, K and V smoothing; 16 videos per configuration):
Vendored source: the PrimTS pin is the dated canonical FlashInfer branch trtllm-prims-ts-dev-20261012 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 smoothing, VSA route checktests/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.pyAll of the above pass on GB300 (SM103);
test_sage_quant.pyalso passes on B200 (SM100). 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.