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

heyuhhh · 2026-09-03T19:56:31Z

Description

This is the focused VisualGen SOL layer stacked on VisualGen VSA #7. Generic sparse runtime transport and PrimTS block-sparse execution remain in PR #2.

Dependencies

PrimTS base: NVIDIA/TensorRT-LLM#17399.
Generic TensorRT-LLM integration: yuxianq/TensorRT-LLM#2.
VisualGen VSA base: heyuhhh/TensorRT-LLM#7.
Proxy-compensated core attention: flashinfer-ai/flashinfer#4872.

SOL commit: b5e6d2d354 (feat: add VisualGen SOL sparse attention).

Summary

Add SolAttentionConfig and the VisualGen SOLTrtllmAttention backend.
Predict exact-block bitmasks and K/V proxy summaries, then pass them through SparseRuntimeParams.block_sparse_inputs.
Reuse generic PrimTS block-sparse FMHA rather than adding an SOL-specific core FMHA library.
Keep predictor and plan state model-scoped and route buffers live across CUDA Graph replay.
Resolve dense/sparse phases before capture and key graphs using live modality timesteps.
Preserve raw LTX-2 video/audio denoising timesteps for SOL cutoff decisions instead of AdaLN-transformed embeddings.
Reject unsupported masks, quantization, cross-attention, layouts, and tensor envelopes explicitly.

Scope

This single-commit layer contains no benchmark-results/ and makes no separate performance claim. Kernel implementation and performance evidence remain in FlashInfer NVIDIA#4872; this PR covers TensorRT-LLM VisualGen integration and correctness.

Validation

SOL predictor: 10 passed.
SOL backend: 30 passed, including real B200 eager/CUDA Graph dense-reference and mixed-proxy 256/257-token cases.
VSA preservation: 34 passed.
Metadata/config/pipeline: 97 passed.
Fresh final B200/graph/timestep gate: 5 passed.
Pre-commit, DCO, diff, scope, and secret checks passed.

Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>

## 📌 Description This PR extends PrimTS contiguous block-sparse attention with **proxy compensation** and a packed **bitmask** sparse format. The proxy path keeps selected blocks exact while approximating omitted KV blocks with one arithmetic-mean K summary and one summed-V summary per semantic KV block. Route preparation emits score-keep metadata so a block already covered by an exact route is not counted again by its proxy. The softmax denominator is corrected by the represented token mass, including the true length of the final partial KV block. ### Sol-Attn motivation and scope This PR is intended to provide the PrimTS **core-attention** substrate for [Sol-Attn's sparse-attention flow](https://github.com/NVlabs/Sana/tree/sol-engine/techniques/sparse_backends), described in the [Sol-Attn paper](https://arxiv.org/abs/2607.24027). Sol-Attn identifies critical KV blocks for exact computation and reuses block-level proxy information from omitted blocks to approximate their contribution during online softmax. At the API boundary, the intended flow is: `prediction / routing (out of scope) → exact-block BSR or bitmask + K/V summaries → PrimTS route preparation → proxy-compensated core attention` This PR implements only the PrimTS route-preparation and core-attention stages. It does **not** include the Sol-Attn predictor, thresholding and exact-block selection, summary generation, or the end-to-end Sana integration; callers provide the selected-block metadata and K/V summaries. ### Bitmask format In addition to canonical BSR `block_indptr` / `block_indices`, contiguous block-sparse attention can now consume packed `UInt32` exact-block bitmaps owned per `(batch, KV head, Q block)`. Each bit selects one semantic KV block and out-of-range padding bits in the final word are ignored. BSR and bitmask inputs are prepared into the same internal route stream, so the attention core remains format-independent. Bitmask and proxy modes can be combined: set bits use raw K/V exact routes, while unset blocks are represented by proxy summaries. Both reusable plan/run and one-shot contiguous APIs are supported. Proxy execution currently requires a dense mask; paged-KV proxy execution is outside this PR. The PR also teaches Trace Apply to honor a runtime-selected trace template, which keeps BSR/bitmask and exact/proxy schemas isolated instead of binding every call to template zero. ## 📊 Proxy overhead Controlled A/B on one NVIDIA B200 with identical inputs, exact-block pattern, capacity, and static scheduler; only `use_proxy_routes` and the required summaries change. | Path | Proxy disabled | Proxy enabled | Additional overhead | | --- | ---: | ---: | ---: | | Public `BlockSparseTSWrapper.run` | 1.3172 ms | 1.4298 ms | **+8.55%** | | Prevalidated adapter | 1.3162 ms | 1.4318 ms | **+8.78%** | Workload: BF16 Q/K/V `[1, 32760, 12, 128]`, Q-block 64, KV-block 64, physical KV route 256, dense bitmask, and 84 of 512 exact blocks per row (16.4%). Proxy adds two prepared routes per row (21 → 23). Results use five balanced rounds with 80 CUDA-event samples per arm (400 samples/arm). Timings include route preparation and attention; summary construction is excluded. Each mode was checked against its own independent reference because proxy compensation intentionally changes the output semantics. Kernel split diagnostic: attention p50 increased by 7.64% (1.1261 → 1.2121 ms); prepare increased by 0.0030 ms (0.0145 → 0.0176 ms). ## 🧪 Tests - [x] `tests/trace_apply/test_trace_apply.py`: 26 passed - [x] `tests/trace/test_fi_trace_template_consistency.py`: 774 passed - [x] `tests/attention/test_attention_ts_block_sparse.py`: 140 passed, 88 skipped - [x] BK8 SWAPS and BK64 Keeps GPU proxy validation passed for both BSR and bitmask - [x] BSR/bitmask results are bitwise identical in the focused GPU cases; final-word padding bits are ignored - [x] Repository pre-commit hooks passed, including mypy, Ruff check, and Ruff format ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have run the hooks manually and fixed all reported issues. ### ✅ Tests - [x] Tests have been added or updated as needed. - [x] All targeted tests are passing. ## Reviewer Notes The commits are separated into generic runtime trace dispatch, block-sparse functionality, focused correctness/trace tests, and review-driven cleanup. `experimental/sol_attention` is not part of this PR.  ## Summary by CodeRabbit * **New Features** * Added BSR and bitmask routing options for block-sparse attention. * Added proxy-route support using caller-provided K/V summaries. * Added route-specific validation and metadata handling for wrapper and one-shot APIs. * Expanded tracing to support all sparse format and route-mode combinations. * Improved trace application to select and execute only the runtime-matched template. * **Documentation** * Documented contiguous bitmask routing, proxy routes, capacity semantics, and summary behavior. * **Tests** * Added GPU coverage for BSR and bitmask proxy routes and multi-template trace dispatch.  ## Current integration stack (2026-09-04) - This branch is rebased onto FlashInfer `main` `60b49158ab4fb81718aef486c2d3c89aec4c1901`; current PR head: `6b5b2d31f3314be77ccf6655c899b7b971863a1a`. - The TRT-LLM composite pin is [heyuhhh/flashinfer@`71bf7842`](https://github.com/heyuhhh/flashinfer/tree/yuhangh/tmp-sol-attn-trtllm-pr17399), which reapplies all eight commits from Yuxian's newest `trtllm-prims-ts` head (`edddf6f5`) on top of this PR. - TensorRT-LLM stack: [#17399 PrimTS base](NVIDIA/TensorRT-LLM#17399) → [general block-sparse FMHA](yuxianq/TensorRT-LLM#2) → [VisualGen VSA/shared workflow](heyuhhh/TensorRT-LLM#7) → [VisualGen SOL integration](heyuhhh/TensorRT-LLM#9). - Final B200 integration verification includes raw/proxy/paged generic routes, VSA CUDA Graph live routes, SOL all-exact parity, and mixed proxy parity at sequence lengths 256 and 257. --------- Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>

heyuhhh · 2026-09-07T09:19:00Z

Superseded by the re-stacked pull requests now that NVIDIA#17399 has merged into main: the generic block-sparse layer is NVIDIA#18815 and the VisualGen VSA + SOL layer is #10. NVIDIA#18079 stays as the overview of the full stack. Closing this stacked draft.

github-actions Bot assigned heyuhhh Sep 3, 2026

This was referenced Sep 3, 2026

[None][feat] align VisualGen VSA with block-sparse FMHA #7

Closed

feat(prims-ts): support proxy-compensated block-sparse attention flashinfer-ai/flashinfer#4872

Merged

[None][feat] add VisualGen SOL sparse attention backend #8

Closed

heyuhhh force-pushed the user/yuhangh/vsa-block-sparse-integration branch from 6a5b339 to 81d4fb7 Compare September 3, 2026 20:04

heyuhhh force-pushed the user/yuhangh/sol-attention-integration branch from e86650d to 318bf98 Compare September 3, 2026 20:04

heyuhhh force-pushed the user/yuhangh/vsa-block-sparse-integration branch from 81d4fb7 to 1f4e3ad Compare September 4, 2026 04:11

heyuhhh force-pushed the user/yuhangh/sol-attention-integration branch from 318bf98 to 52a3d5c Compare September 4, 2026 04:20

heyuhhh force-pushed the user/yuhangh/vsa-block-sparse-integration branch from 1f4e3ad to ad0165b Compare September 4, 2026 04:54

heyuhhh force-pushed the user/yuhangh/sol-attention-integration branch from 52a3d5c to 48769dc Compare September 4, 2026 04:54

feat: add VisualGen SOL sparse attention

b5e6d2d

Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>

heyuhhh force-pushed the user/yuhangh/sol-attention-integration branch from 48769dc to b5e6d2d Compare September 4, 2026 14:55

heyuhhh force-pushed the user/yuhangh/vsa-block-sparse-integration branch from ad0165b to 6e34838 Compare September 4, 2026 14:55

This was referenced Sep 4, 2026

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

Closed

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

Draft

heyuhhh closed this Sep 7, 2026

Navigation Menu

[None][feat] add VisualGen SOL sparse attention - #9

[None][feat] add VisualGen SOL sparse attention#9
heyuhhh wants to merge 1 commit into
user/yuhangh/vsa-block-sparse-integrationfrom
user/yuhangh/sol-attention-integration

heyuhhh commented Sep 3, 2026 •

edited

Loading

Uh oh!

heyuhhh commented Sep 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

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

Uh oh!

Description

Dependencies

Summary

Scope

Validation

Uh oh!

heyuhhh commented Sep 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

heyuhhh commented Sep 3, 2026 •

edited

Loading