[None][feat] VisualGen VSA and SOL sparse attention on PrimTS block-sparse FMHA (overview) - #18079
Draft
heyuhhh wants to merge 6 commits into
Draft
[None][feat] VisualGen VSA and SOL sparse attention on PrimTS block-sparse FMHA (overview)#18079heyuhhh wants to merge 6 commits into
heyuhhh wants to merge 6 commits into
Conversation
This was referenced Aug 25, 2026
heyuhhh
force-pushed
the
user/yuhangh/update_vsa
branch
from
August 25, 2026 14:26
10fbec9 to
7446583
Compare
heyuhhh
force-pushed
the
user/yuhangh/update_vsa
branch
11 times, most recently
from
September 4, 2026 14:55
ad0165b to
6e34838
Compare
Open
4 tasks
heyuhhh
force-pushed
the
user/yuhangh/update_vsa
branch
from
September 7, 2026 09:18
6e34838 to
39c3ebb
Compare
heyuhhh
force-pushed
the
user/yuhangh/update_vsa
branch
9 times, most recently
from
September 8, 2026 03:32
a912915 to
e4d992d
Compare
heyuhhh
force-pushed
the
user/yuhangh/update_vsa
branch
4 times, most recently
from
September 9, 2026 02:00
b7e8c8a to
b230bf9
Compare
Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
heyuhhh
force-pushed
the
user/yuhangh/update_vsa
branch
from
September 9, 2026 02:20
b230bf9 to
d4a1527
Compare
Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
Importing cutlass.experimental.task_scheduling rewrites the shared cutlass.utils.WorkTileInfo class in place so that its constructor unpacks tile_idx into exactly three scalars. FlashAttention 4 subclasses that class with a four-axis coordinate and inherits the constructor, so any process that probes or plans the vendored PrimTS kernels turns every later FA4 kernel trace into a ValueError. Install the upstream tuple semantics on the FA4 subclass from the existing CuTe DSL compatibility layer so the parent rewrite cannot reach it, and cover the worst-case import order in the FA4 compatibility tests. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
heyuhhh
force-pushed
the
user/yuhangh/update_vsa
branch
from
September 9, 2026 03:08
d4a1527 to
b8d9b03
Compare
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
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.
Description
Overview PR for VisualGen VSA and SOL sparse attention on generic PrimTS block-sparse FMHA, kept for end-to-end
review of the full stack. The mergeable units are the split PRs below; this branch carries all six commits so the
complete change can be read and tested in one place.
Review stack
main), includingthe core
TrtllmAttention.block_sparse_attn_predicthook, theSparseBackendForwardArgs.block_sparse_inputscarrier that algorithms use to hand routes to the core,
prepare_sparse_runtime_paramsas the single aggregationpoint (all prediction hooks plus SkipSoftmax scheduling), and the FA4
WorkTileInfocompatibility fix:NVIDIA/TensorRT-LLM#18815.
sparse_backend_args) and SOL (overridesthe core hook) algorithms (2 commits, stacked on the branch above): [None][feat] add VisualGen VSA and SOL sparse attention on PrimTS block-sparse FMHA heyuhhh/TensorRT-LLM#10.
optimizations perf(prims-ts): Optimize&refine PrimsTS block sparse attention flashinfer-ai/flashinfer#5002 (pending); the vendored PrimTS tree is pinned to
heyuhhh/flashinfer@61454c5c(yuhangh/tmp-sol-attn-trtllm-dev), which isyuxianq/flashinfer:trtllm-prims-ts-dev(the tree
mainpins) plus [TRTLLM-5589] feat: Integrate TRT-LLM Gen FP8 Batched GEMM with Pytorch workflow kernel autotuner #4872, Feature request: pin versions of python packages in CI #5002 and one follow-up commit (fixed 2Dblock_tablesfor both pagedblock-sparse entry points, a dense-style
validateswitch on the wrappers, and the PrimTS block-sparse tracetemplates rewritten as literal schemas).
Commits
chore: update FlashInfer PrimTS pinfeat: add generic PrimTS block-sparse supportrefactor: unify sparse attention runtime inputsfix: keep FA4 WorkTileInfo independent of CUTLASS task schedulingrefactor: align VisualGen sparse attention workflowfeat: add VisualGen SOL sparse attentionTest Coverage
See the split PRs for the per-layer test runs.