[None][feat] add VisualGen VSA and SOL sparse attention on PrimTS block-sparse FMHA by heyuhhh · Pull Request #18079 · NVIDIA/TensorRT-LLM

heyuhhh · 2026-08-21T16:16:30Z

Description

VisualGen sparse attention algorithms on top of the generic PrimTS block-sparse FMHA merged in #18815. This PR
contains the VisualGen (AIGV) integration plus the two small core additions it needs: an optional fmha_state dict on
TrtllmAttention that lets layers share FMHA plan caches, and attention/backends/sparse/timestep_phase.py, the shared
reduction from a denoising timestep to a dense-or-sparse phase. The attention-core hooks it relies on
(SparseRuntimeParams carrier, block_sparse_attn_predict, prepare_sparse_runtime_params,
Fmha.supports_block_sparse_inputs) are already on main. The stack is rebased past #18329, whose fused CuTeDSL
Sol-Attn backend is kept and now sits next to the PrimTS-based backend behind one SolAttentionConfig.

VisualGen TRTLLM wrapper: forward takes sparse_backend_args and passes separate Q/K/V into the core whenever
block-sparse routes are present, a quant_attention_config is set, or the backend rejects fused QKV; both algorithms
use the core block_sparse_attn_predict hook and the core prepare_sparse_runtime_params aggregation (SkipSoftmax
scheduling stays in that core helper). Every layer of a model component shares its PrimTS block-sparse plan caches
through the fmha_state dict the wrapper hands to the core constructor, so one static profile is planned once per
component instead of once per layer, and libraries rebuilt by update_quant_config read the same dict. The
wrapper's metadata adapter also keeps the host value of the denoising timestep in the same component state.
Timestep schedule: one core reduction (timestep_phase.py) is used by the SkipSoftmax scheduler, both SOL backends
and a single sparse_attn_phase CUDA Graph key that models register for every sparse config with a cutoff
(BaseSparseAttentionConfig.resolve_disabled_until_timestep). Per-token timesteps reduce to their largest live
value, so a Wan 2.2 5B I2V conditioning frame at timestep zero no longer disables the configured dense prefix. The
TRTLLM metadata adapter prepares the timestep as a host value during graph warmup and reuses it under capture, so
SkipSoftmax with a cutoff captures with a device timestep tensor; the CuTeDSL backends read the phase the CUDA Graph
runner publishes.
VSA: configuration, metadata, prediction and post-processing shared between the TRTLLM and CuTeDSL backends in
attention_backend/sparse/vsa, replacing the CuTeDSL-only module. The TRTLLM backend runs the coarse stage first,
hands the complete BlockSparseForwardInputs (including the tile-padding kv_valid_bits only the VSA predictor
knows) to the core through SparseBackendForwardArgs.block_sparse_inputs, and blends the fine and coarse outputs
afterward; compact Q/K/V dense fallback; CUDA Graph-stable route buffers; packed QKV, TP, Ulysses and async Ulysses.
The predictor's tiling, cube-mean, route sorting and blending run as Triton kernels with unchanged numerics: on B200
with Wan 14B shapes the predictor drops from 34.6 ms to 5.1 ms per call for 81 frames and the post-process from
4.1 ms to 0.9 ms.
SOL: SolAttentionConfig (tau, thresh_type, disabled_until_timestep, dense_layers) lowers into one
SolParams for the two backends in attention_backend/sparse/sol. SOLTrtllmAttention runs SOL in two stages: a
functional predictor operator, trtllm::visual_gen_sol_predictor (Triton block pooling, key statistics, thresholds
and exact-block selection, with the diag and exact threshold policies of the fused kernel's preprocess), produces
the exact-block bitmask and K/V proxy summaries inside the core hook, and the shared PrimTS block-sparse FMHA executes
them (proxy-compensated attention from feat(prims-ts): support proxy-compensated block-sparse attention flashinfer-ai/flashinfer#4872). The operator allocates its outputs per call,
so it composes with CUDA Graph capture and torch.compile without plan state. SOLCuTeDSLAttention is [TRTLLM-15917][feat] Integrate Sol-Attn sparse attention into VisualGen #18329's fused
kernel backend moved next to it, consuming the same SolParams. Dense layers and dense timestep phases return no
routes; cross-attention, context parallelism and attention quantization are rejected.

Commits

refactor: align VisualGen sparse attention workflow (VSA)
perf: fuse VSA predictor preprocessing into Triton kernels
fix: tolerate backend-specific kwargs in VisualGen TRTLLM attention - the wrapper keeps the shared VisualGen
backend contract (the attention module forwards a superset of keyword arguments, every backend ignores what it does
not consume), so HunyuanVideo 1.5 and GLM-Image, which always pass key_padding_mask, keep working on TRTLLM.
refactor: share the sparse timestep schedule across VisualGen algorithms
feat: add the TRTLLM SOL backend and unify VisualGen SOL under one config
refactor: inject shared FMHA state through the TRTLLM attention constructor - replaces the post-construction
plan-cache binding with the fmha_state constructor argument (core TrtllmAttention, PrimsTSBlockSparseFmha).
refactor: keep the prepared sparse timestep in the component attention state - the metadata adapter owns the host
timestep value used under CUDA Graph capture.
refactor: make the SOL predictor a functional operator and add the exact threshold - removes the predictor plan
classes, the capture guard and the SOL-specific fullgraph rejection; adds the exact threshold policy to the TRTLLM
backend with Triton key-statistics and threshold kernels.
refactor: unify sparse algorithm naming in the VisualGen attention dispatch
fix: resolve review feedback on VisualGen VSA and SOL sparse attention - drops the 16-shape cap on the VSA
metadata and route caches (they grow with the CUDA Graph set and are released with it), makes the SOL block pooling
kernel mask the block boundary for any block size, registers the VSA and SOL kernel suites in the test lists, moves
the TRTLLM skip-softmax CUDA Graph capture test to the GPU stage under an SM100 gate, and tightens the review-flagged
assertions and annotations.

Behavior notes: SkipSoftmax reduced per-token timesteps by their first element before, which on Wan 2.2 5B I2V read
the conditioning frame's zero and enabled the sparse phase from the first step; the configured cutoff now takes effect
there. SkipSoftmax dense and sparse CUDA graphs are keyed per phase on every model with a cutoff, including LTX-2.

Test Coverage

Run on B200 with an SM100 build (CUTLASS DSL 4.8.0.dev0; the changes are Python-only) after rebasing onto main
a1c6c2b (past #18329):

tests/unittest/_torch/visual_gen/sparse_attention/ (SkipSoftmax, SOL attention including the real B200 CUDA Graph
cases, SOL predictor and kernels), attention/sparse/test_timestep_phase.py, test_attention_vsa.py,
test_trtllm_attention_metadata.py, test_attention_cute_dsl_sol_attn.py, test_visual_gen_args.py: 305 passed.
test_attention_integration.py, test_fa4_cutlass_compatibility.py, test_attention_vsa_kernels.py: 129 passed.
Core test_attention_mla.py, test_fmha_manager.py, test_fmha_registry.py, sparse/test_flashinfer_utils.py:
205 passed, 1 skipped (the block-sparse adapter tests are part of the 305 above).
test_ltx2_pipeline.py and test_ltx2_transformer.py: 67 passed. HunyuanVideo 1.5 and GLM-Image
test_fp8_trtllm_attention: 2 passed.
CUDA hidden (the CPU stages): sparse_attention/, test_visual_gen_args.py, test_trtllm_attention_metadata.py,
test_timestep_phase.py: 188 passed, 29 skipped.
Multi-GPU Wan VSA TP/Ulysses parity was verified on the previous stack; the 8-GPU Wan VSA parity case has a
pre-existing numerical tolerance gap that reproduces on the original VSA source commit.
Commit 10 on top of the above: the touched CPU suites 218 passed, 41 skipped; test_sol_predictor_kernels.py,
test_trtllm_attention_metadata.py, test_attention_vsa_kernels.py and test_attention_vsa.py on B200 136 passed;
test_attention_integration.py -k vsa 4 passed; test_wan_vsa_pipeline.py::TestWanVsa14B_PipelineCorrectness::test_cosine_similarity
1 passed.
l0_b200.yml registers test_attention_vsa.py, test_attention_vsa_kernels.py, test_trtllm_attention_metadata.py
(including the TRTLLM skip-softmax CUDA Graph capture test), test_sol_predictor.py, test_sol_predictor_kernels.py,
the two real-B200 SOL CUDA Graph cases and [TRTLLM-15917][feat] Integrate Sol-Attn sparse attention into VisualGen #18329's test_attention_cute_dsl_sol_attn.py; l0_cpu.yml registers the
SOL (attention, predictor, kernels) and SkipSoftmax CPU cases.

PR Checklist

PR title follows the template
Commits are signed off (DCO)
pre-commit hooks pass

Dev Engineer Review

Adds VSA and SOL support through generic PrimTS block-sparse FMHA.
Adds TRT-LLM and CuTeDSL VSA backends, dense fallback, packed/separate QKV, TP, Ulysses, and CUDA Graph support.
Adds SOL configuration, routing, Triton kernels, and shape-specialized predictor plans.
Centralizes timestep reduction and sparse graph-phase handling.
Removes the legacy CuTeDSL VSA implementation and exports. Verify downstream imports.
VSA rejects quantized attention. VSA and SOL reject context parallelism. SOL rejects Torch fullgraph compilation.
LTX-2 now passes raw denoising timesteps to attention.
Verify fallback paths, custom operators, plan preparation, graph capture, tensor validation, and multi-GPU execution.
CI follow-up is required because reported merge pipelines include failures and an unstable result.

QA Engineer Review

Adds VSA, SOL, multi-GPU, LTX-2, TRT-LLM metadata, SkipSoftmax, and configuration tests.
Covers dispatch, validation, fallback, metadata, routing, kernels, numerical parity, compilation, custom operators, CUDA Graph replay, TP, and Ulysses behavior.
Removes the legacy CuTeDSL VSA test module.
Updates l0_b200.yml and l0_cpu.yml with VSA, SOL, CUDA Graph, metadata, backend, and predictor entries.
No manual-QA list changes are provided.
Coverage verdict: needs follow-up because CI reported failures and multi-GPU execution requires approval.

Per-File QA Perspective

VSA and SOL source files: Verify backend dispatch, validation, dense fallback, metadata caching, route prediction, post-processing, Triton/PyTorch parity, and CUDA Graph behavior.
TRT-LLM and attention files: Verify timestep scheduling, sparse phase selection, QKV routing, cache reuse, and sparse argument forwarding.
VisualGen configuration files: Verify SOL fields, dense-layer parsing, compatibility rules, quantization restrictions, and fullgraph validation.
Model and pipeline files: Verify raw timestep propagation, graph keys, VSA builder lifecycle, TP sharding, and async Ulysses forwarding.
Removed CuTeDSL VSA files and exports: Verify no stale imports remain.
Documentation files: Verify documented backend support, fallback behavior, validation rules, and CUDA Graph behavior.
l0_b200.yml: Adds B200 VSA, SOL, CUDA Graph, and metadata coverage.
l0_cpu.yml: Adds CPU SOL backend and predictor coverage.
VSA tests: Cover backend contracts, metadata, fallback, routing, kernels, parity, packed QKV, and graph replay. Full CI mapping is not established.
SOL tests: Cover backend behavior, predictor plans, kernels, validation, compilation, custom operators, and graph replay. CPU and B200 entries are listed.
Multi-GPU tests: Cover async gate redistribution and TP-local dimensions. CI execution requires the reported approval label.
LTX-2, TRT-LLM metadata, SkipSoftmax, argument, and WAN tests: Cover sparse phase keys, raw timesteps, metadata lifecycle, dispatch, configuration validation, TP-aligned gates, and Ulysses forwarding.

Restore the shared VisualGen backend contract in the TRTLLM wrapper: the attention module forwards a superset of keyword arguments and every backend ignores the ones it does not consume. Rejecting unknown names made TRTLLM the only backend that raised, and HunyuanVideo 1.5 and GLM-Image always pass key_padding_mask, so their TRTLLM paths failed. timestep is now an explicit parameter; other keyword arguments are accepted and ignored as before. Guard the SOL graph-phase resolution with torch.cuda.is_available() so the cpu_only SOL tests run on hosts without a GPU, and pin the SM version in the SkipSoftmax and SAGE combination test like the neighbouring int8 SAGE tests. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>

Reduce a denoising timestep to a dense-or-sparse phase in one place, tensorrt_llm/_torch/attention/backends/sparse/timestep_phase.py, and use it from the SkipSoftmax scheduler, the CuTeDSL Sol-Attn backend and the CUDA Graph key. Per-token timesteps reduce to their largest live value, so a Wan I2V conditioning frame at timestep zero no longer disables the configured dense prefix. The VisualGen TRTLLM attention wrapper owns the timestep schedule: it prepares the timestep as a host value once per eager call and reuses it during CUDA Graph capture, so SkipSoftmax with a timestep cutoff can be captured with a device timestep tensor, and it answers whether a layer runs sparse from the cutoff and dense_layers of its sparse parameters. Models register a single sparse_attn_phase CUDA Graph key for every sparse attention config with a timestep cutoff, resolved through the new BaseSparseAttentionConfig.resolve_disabled_until_timestep. The CuTeDSL Skip Softmax and Sol-Attn backends read that key from the CUDA Graph runner during capture, and the Sol-Attn dense-prefix decision uses the shared reduction instead of the scheduler classmethod, which is removed. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>

…nfig SOL is served by two backends behind SolAttentionConfig (algorithm "sol_attn"): - SOLTrtllmAttention runs SOL in two stages through the generic TRTLLM sparse lifecycle. A TRT-LLM-owned predictor derives the exact block bitmask and K/V proxy summaries inside the core block_sparse_attn_predict hook, from the flattened Q/K/V, the batch layout in the attention metadata and the prepared timestep, and the shared PrimTS block-sparse FMHA executes that route. Each layer owns its predictor; the predictor keeps one shape-specialized plan with graph-stable route and summary buffers per shape. - SOLCuTeDSLAttention is the fused CuTeDSL kernel backend, moved from attention_backend/cute_dsl/sol_attn.py into attention_backend/sparse/sol next to the TRTLLM backend, the way the VSA backends share attention_backend/sparse/vsa. It consumes the same lowered SolParams and keeps its per-call dense delegation for cross-attention, masks and inputs the kernel cannot serve. SolAttentionConfig lowers into SolParams for both backends: tau, disabled_until_timestep and dense_layers (a list of layer indices) drive both, while thresh_type stays a CUTEDSL kernel knob; the TRTLLM predictor implements the diag policy only, so AttentionConfig rejects "exact" with the TRTLLM backend. Quantized attention and torch.compile fullgraph stay rejected for SOL on either backend, and TRTLLM SEPARATE_QKV cross-attention falls back to VANILLA like dense TRTLLM. The SOL predictor kernels, the TRTLLM SOL tests and their test-list entries come with the backend; the CuTeDSL SOL tests follow the moved module and the shared sparse_attn_phase key. The feature guide describes both backends in one SOL section. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>

…ructor TrtllmAttention takes an optional fmha_state dict, defaulting to a private dict per instance, and PrimsTSBlockSparseFmha keeps its wrapper-plan caches under its own key in that dict. Layers constructed with one shared dict plan each static block-sparse profile once and reuse its route workspace; the sharing granularity is the lifetime of the dict the constructor caller hands in. Libraries rebuilt by update_quant_config read the same dict, so the caches survive a rebuild without any rebinding. The VisualGen wrapper passes the fmha_caches entry of its component-scoped attention_metadata_state to the core, which removes bind_plan_cache, the wrapper's update_quant_config override and the metadata adapter's get_fmha_cache_state, and lets the wrapper build its metadata adapter after the core constructor again. The developer guide and the state docstring describe the constructor contract. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>

…n state The host value of the denoising timestep is per-step state of a whole model component, not of one attention layer, so the VisualGen TRTLLM metadata adapter now owns it: prepare_timestep reduces the timestep tensor during eager calls, including the warmup that precedes CUDA Graph capture, stores the value in the component-scoped attention_metadata_state next to the metadata cache and the FMHA state, and returns the stored value under capture, where the tensor cannot be read. The wrapper keeps only the schedule policy: it asks the adapter when a cutoff is configured and answers should_use_sparse from the cutoff and dense_layers of its sparse parameters. The per-layer _prepared_timestep and _timestep_prepared fields are gone; the state is the value itself, and its presence marks the timestep as prepared. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>

…act threshold The TRTLLM SOL predictor is now one functional custom operator, trtllm::visual_gen_sol_predictor(q, k, v, block_size, tau, sm_scale, thresh_type) -> (exact_block_bits, k_summary, v_summary), that allocates its outputs per call. Every output is produced and consumed inside the same transformer forward, so under CUDA Graph capture the tensors come from the graph pool and are replayed with the graph, and the runner's eager warmup already compiles the Triton kernels before capture. That removes the plan and geometry classes, the per-instance plan cache, the capture guard and the host-side graph break, so SOL no longer rejects torch.compile fullgraph and the corresponding VisualGenArgs validator is gone. The routing threshold of every query block comes from two Triton kernels that mirror the fused kernel's preprocess: a key-statistics kernel reduces the key block summaries to their per-channel mean and variance and, for the exact policy, their full second moment with tensor-core dots on the 16-bit summaries, and a threshold kernel projects the query block centroids onto those statistics with an IEEE fp32 dot. diag models each key channel independently and exact uses the full key covariance. The selection kernel consumes the threshold, and the PyTorch implementation of the same rule stays as the CPU path and the reference, computed in float64 because the TF32 matmuls TensorRT-LLM enables would perturb the fp32 second moment. SolParams and SolAttentionConfig lower thresh_type to both backends, so the rule that rejected exact with the TRTLLM backend is removed. Tests: the predictor tests cover both policies against an fp64 oracle, a correlated-key case where the policies differ, fresh outputs under CUDA Graph replay, a fullgraph compile, threshold parity with the fused kernel's preprocess and Triton-versus-reference thresholds; the SOL attention tests stub the predictor functions instead of a predictor object. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>

…spatch The attention module and the backend factory use the same names for the sparse dispatch: sparse_config for the user config, sparse_algorithm for its discriminator, is_vsa and is_sol for the two algorithms, and SOL as the display name in messages and comments, including the CuTeDSL SOL backend moved from the Sol-Attn module. The attention module also drops a duplicate context-parallelism check that the shared one already covered. The Wan pipeline registry no longer lists the FastVideo Wan2.1 VSA checkpoint: the checkpoint resolves through its pipeline class like before, and this change set is about the sparse attention backends, not about adding models. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>

Drop the 16-shape cap on the VSA metadata and route caches. The CUDA Graph runner that captures the cached tensor addresses keeps one graph per shape without a cap, so the limit only turned a new resolution/frame profile into a runtime error. Both caches now keep one entry per distinct shape for as long as the graphs that reference them and are released together with the graphs in cleanup. Make the SOL block pooling kernel mask the block boundary so block sizes that are not a multiple of the load width match the PyTorch fallback, and reject non-positive block sizes. Add the VSA and SOL predictor kernel suites to the CPU and B200 test lists. Move the TRTLLM skip-softmax CUDA Graph capture test out of the cpu_only module into the TRTLLM metadata suite with an SM100 gate so it runs on the GPU stage. Assert that TRTLLM VSA executes the block-sparse FMHA rather than only predicting routes, check the forwarded VSA gate values, verify that the Wan VSA pipeline reference takes the SDPA fallback, and reject non-finite tensor timesteps. Annotate the helpers and tests flagged during review and read the sparse algorithm discriminator directly instead of through getattr. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>

heyuhhh · 2026-09-20T08:19:49Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-20T08:27:14Z

PR_Github #74620 [ run ] triggered by Bot. Commit: 85b6b4d Link to invocation

tensorrt-cicd · 2026-09-20T11:48:14Z

PR_Github #74620 [ run ] completed with state SUCCESS. Commit: 85b6b4d
/LLM/main/L0_MergeRequest_PR pipeline #61407 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

heyuhhh · 2026-09-20T12:43:55Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-20T12:50:32Z

PR_Github #74635 [ run ] triggered by Bot. Commit: 4fe7022 Link to invocation

tensorrt-cicd · 2026-09-20T17:37:52Z

PR_Github #74635 [ run ] completed with state SUCCESS. Commit: 4fe7022
/LLM/main/L0_MergeRequest_PR pipeline #61421 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

heyuhhh · 2026-09-21T01:38:22Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-21T01:45:53Z

PR_Github #74700 [ run ] triggered by Bot. Commit: 4fe7022 Link to invocation

tensorrt-cicd · 2026-09-21T03:03:25Z

PR_Github #74700 [ run ] completed with state SUCCESS. Commit: 4fe7022
/LLM/main/L0_MergeRequest_PR pipeline #61485 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:

Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, or
Wait for the PR to be fully approved — the label is added automatically once approval is complete.
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

heyuhhh · 2026-09-21T03:19:25Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-21T03:25:08Z

PR_Github #74724 [ run ] triggered by Bot. Commit: 4fe7022 Link to invocation

tensorrt-cicd · 2026-09-21T03:50:56Z

PR_Github #74724 [ run ] completed with state FAILURE. Commit: 4fe7022
/LLM/main/L0_MergeRequest_PR pipeline #61508 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:

Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, or
Wait for the PR to be fully approved — the label is added automatically once approval is complete.
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

heyuhhh · 2026-09-21T06:22:03Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-21T06:29:04Z

PR_Github #74745 [ run ] triggered by Bot. Commit: 4fe7022 Link to invocation

tensorrt-cicd · 2026-09-21T06:57:07Z

PR_Github #74745 [ run ] completed with state FAILURE. Commit: 4fe7022
/LLM/main/L0_MergeRequest_PR pipeline #61526 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:

Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, or
Wait for the PR to be fully approved — the label is added automatically once approval is complete.
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

heyuhhh · 2026-09-21T07:20:29Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-09-21T07:27:52Z

PR_Github #74760 [ run ] triggered by Bot. Commit: 0f0ac14 Link to invocation

karljang

Thanks for the rebase and for improving the sparse-attention framework. Overall looks good. Two inline comments are worth fixing before merge.

karljang · 2026-09-21T07:29:30Z

+
+        if self.quant_attention_config is not None:
+            raise ValueError("SOL sparse execution does not support quant_attention_config")
+        if not any(


This only checks the PrimTS lib exists. If it rejects the request, the manager falls through to a dense lib and SOL silently does nothing. Can we raise here instead?

Thanks for the comments! Do you mean we can fallback to dense fmha if there is no support for block sparse attention?

karljang · 2026-09-21T07:29:30Z

+                prediction hooks. A ``block_sparse_inputs`` payload selects the
+                generic block-sparse FMHA.
+            timestep: Denoising timestep consumed by the sparse prediction hooks.
+            **kwargs: Backend-specific keyword arguments forwarded by the attention


HunyuanVideo 1.5 and GLM-Image always pass key_padding_mask; ignoring it here attends to padded tokens. Please raise or at least warn.

Seems that i don't have changes about key_padding_mask, could you explain your points in detail? Thanks!

github-actions Bot assigned heyuhhh Aug 21, 2026

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

This was referenced Aug 25, 2026

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

Closed

[None][feat] align VisualGen VSA with block-sparse FMHA heyuhhh/TensorRT-LLM#7

Closed

heyuhhh force-pushed the user/yuhangh/update_vsa branch from 10fbec9 to 7446583 Compare August 25, 2026 14:26

heyuhhh changed the title ~~[None][feat] enable PrimTS block-sparse attention for VisualGen VSA~~ [None][feat] integrate VisualGen VSA with block-sparse FMHA Aug 25, 2026

heyuhhh force-pushed the user/yuhangh/update_vsa branch 11 times, most recently from ad0165b to 6e34838 Compare September 4, 2026 14:55

heyuhhh mentioned this pull request Sep 7, 2026

[None][feat] add generic PrimTS block-sparse FMHA and unify sparse attention runtime inputs #18815

Merged

4 tasks

heyuhhh force-pushed the user/yuhangh/update_vsa branch from 6e34838 to 39c3ebb Compare September 7, 2026 09:18

heyuhhh changed the title ~~[None][feat] integrate VisualGen VSA with block-sparse FMHA~~ [None][feat] VisualGen VSA and SOL sparse attention on PrimTS block-sparse FMHA (overview) Sep 7, 2026

heyuhhh mentioned this pull request Sep 7, 2026

[None][feat] add VisualGen SOL sparse attention heyuhhh/TensorRT-LLM#9

Closed

heyuhhh force-pushed the user/yuhangh/update_vsa branch 9 times, most recently from a912915 to e4d992d Compare September 8, 2026 03:32

heyuhhh added 7 commits September 19, 2026 06:05

heyuhhh force-pushed the user/yuhangh/update_vsa branch from a62bc6d to c63a310 Compare September 20, 2026 06:39

heyuhhh mentioned this pull request Sep 20, 2026

[None][fix] Resolve review feedback on the VisualGen VSA and SOL sparse attention PR heyuhhh/TensorRT-LLM#11

Closed

Merge branch 'main' into user/yuhangh/update_vsa

4fe7022

Merge branch 'main' into user/yuhangh/update_vsa

0f0ac14

karljang reviewed Sep 21, 2026

View reviewed changes

Conversation

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

Uh oh!

Description

Commits

Test Coverage

PR Checklist

Dev Engineer Review

QA Engineer Review

Per-File QA Perspective

Uh oh!

heyuhhh commented Sep 20, 2026

Uh oh!

tensorrt-cicd commented Sep 20, 2026

Uh oh!

tensorrt-cicd commented Sep 20, 2026

Uh oh!

heyuhhh commented Sep 20, 2026

Uh oh!

tensorrt-cicd commented Sep 20, 2026

Uh oh!

tensorrt-cicd commented Sep 20, 2026

Uh oh!

heyuhhh commented Sep 21, 2026

Uh oh!

tensorrt-cicd commented Sep 21, 2026

Uh oh!

tensorrt-cicd commented Sep 21, 2026

Uh oh!

heyuhhh commented Sep 21, 2026

Uh oh!

tensorrt-cicd commented Sep 21, 2026

Uh oh!

tensorrt-cicd commented Sep 21, 2026

Uh oh!

heyuhhh commented Sep 21, 2026

Uh oh!

tensorrt-cicd commented Sep 21, 2026

Uh oh!

tensorrt-cicd commented Sep 21, 2026

Uh oh!

heyuhhh commented Sep 21, 2026

Uh oh!

tensorrt-cicd commented Sep 21, 2026

Uh oh!

karljang left a comment

Choose a reason for hiding this comment

Uh oh!

karljang Sep 21, 2026

Choose a reason for hiding this comment

Uh oh!

heyuhhh Sep 21, 2026

Choose a reason for hiding this comment

Uh oh!

karljang Sep 21, 2026

Choose a reason for hiding this comment

Uh oh!

heyuhhh Sep 21, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

heyuhhh commented Aug 21, 2026 •

edited

Loading