Skip to content

[None][feat] VisualGen VSA and SOL sparse attention on PrimTS block-sparse FMHA (overview) - #18079

Draft
heyuhhh wants to merge 6 commits into
NVIDIA:mainfrom
heyuhhh:user/yuhangh/update_vsa
Draft

[None][feat] VisualGen VSA and SOL sparse attention on PrimTS block-sparse FMHA (overview)#18079
heyuhhh wants to merge 6 commits into
NVIDIA:mainfrom
heyuhhh:user/yuhangh/update_vsa

Conversation

@heyuhhh

@heyuhhh heyuhhh commented Aug 21, 2026

Copy link
Copy Markdown
Collaborator

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

  1. Generic PrimTS block-sparse FMHA and sparse-attention framework refactor (4 commits, targets main), including
    the core TrtllmAttention.block_sparse_attn_predict hook, the SparseBackendForwardArgs.block_sparse_inputs
    carrier that algorithms use to hand routes to the core, prepare_sparse_runtime_params as the single aggregation
    point (all prediction hooks plus SkipSoftmax scheduling), and the FA4 WorkTileInfo compatibility fix:
    NVIDIA/TensorRT-LLM#18815.
  2. VisualGen VSA (predicts before the core forward and hands routes through sparse_backend_args) and SOL (overrides
    the 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.
  3. FlashInfer kernels: proxy-compensated block-sparse attention feat(prims-ts): support proxy-compensated block-sparse attention flashinfer-ai/flashinfer#4872 (merged) and PrimTS decode
    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 is yuxianq/flashinfer:trtllm-prims-ts-dev
    (the tree main pins) 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 2D block_tables for both paged
    block-sparse entry points, a dense-style validate switch on the wrappers, and the PrimTS block-sparse trace
    templates rewritten as literal schemas).

Commits

  1. chore: update FlashInfer PrimTS pin
  2. feat: add generic PrimTS block-sparse support
  3. refactor: unify sparse attention runtime inputs
  4. fix: keep FA4 WorkTileInfo independent of CUTLASS task scheduling
  5. refactor: align VisualGen sparse attention workflow
  6. feat: add VisualGen SOL sparse attention

Test Coverage

See the split PRs for the per-layer test runs.

@heyuhhh
heyuhhh force-pushed the user/yuhangh/update_vsa branch from 10fbec9 to 7446583 Compare August 25, 2026 14:26
@heyuhhh 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
heyuhhh force-pushed the user/yuhangh/update_vsa branch 11 times, most recently from ad0165b to 6e34838 Compare September 4, 2026 14:55
@heyuhhh
heyuhhh force-pushed the user/yuhangh/update_vsa branch from 6e34838 to 39c3ebb Compare September 7, 2026 09:18
@heyuhhh 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
heyuhhh force-pushed the user/yuhangh/update_vsa branch 9 times, most recently from a912915 to e4d992d Compare September 8, 2026 03:32
@heyuhhh
heyuhhh force-pushed the user/yuhangh/update_vsa branch 4 times, most recently from b7e8c8a to b230bf9 Compare September 9, 2026 02:00
Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
@heyuhhh
heyuhhh force-pushed the user/yuhangh/update_vsa branch from b230bf9 to d4a1527 Compare September 9, 2026 02:20
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
heyuhhh force-pushed the user/yuhangh/update_vsa branch from d4a1527 to b8d9b03 Compare September 9, 2026 03:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant