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

heyuhhh · 2026-08-25T10:24:04Z

Description

This is the focused VisualGen VSA layer stacked on the generic PrimTS block-sparse PR #2. Core sparse runtime transport and FMHA dispatch remain in #2; SOL remains in PR #9.

Stack

NVIDIA/TensorRT-LLM#17399 → generic block-sparse #2 → this VSA layer → SOL #9

VSA commit: 6e348388d3 (refactor: align VisualGen sparse attention workflow).

Summary

Add VisualGen VSA configuration, metadata, prediction, and post-processing.
Share one VSA predictor and post-processing path between TRTLLM and CuTeDSL; only fine-attention execution differs.
Publish BlockSparseForwardInputs through unified SparseRuntimeParams and reuse normal core FMHA dispatch.
Preserve compact Q/K/V for dense fallback so ragged padding does not enter softmax.
Keep route buffers and shape-dependent metadata stable for CUDA Graph replay.
Preserve VSA gates across packed QKV, TP, Ulysses, and asynchronous Ulysses execution.
Keep VisualGen SkipSoftmax as a thin concrete TRTLLM backend using the shared core lifecycle.

Scope

This single-commit layer contains only VisualGen code, public VisualGen configuration, documentation, and tests. It adds no core FMHA implementation, vendor pin, SOL algorithm, model-owned timestep scope, or benchmark artifacts.

Validation

Focused metadata/VSA/SkipSoftmax/config/Wan suite: 187 passed.
VSA integration: 4 passed (TRTLLM/CuTe equivalence, ragged parity, CUDA Graph live routes).
New TP/Ulysses/asynchronous-Ulysses nodes: 4 passed.
Final latest-base manager/core-handoff smoke: 9 passed.
Pre-commit, DCO, diff, and secret checks passed.

Known validation gap: the 8-GPU real Wan VSA parity node remains outside tolerance; it reproduces with identical max/mean error on the original VSA source commit, so it is not introduced by this restack.

lfr-0531 · 2026-09-03T09:12:46Z

+    return (x_sum / valid_counts).to(x_tiled.dtype)
+
+
+class VSAPreprocessor:


It'd be better to move VSAPreprocessor to the backend.py.

lfr-0531 · 2026-09-03T09:13:15Z

+        return torch.index_select(x, 1, untile_idx)
+
+
+def _normalize_qkv_inputs(


Same as above.

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

This reverts commit dad7a6b. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

Restore the standalone decode and MLA trace dispatchers now provided by FlashInfer 0.6.18, while retaining compatibility patches for unavailable block-sparse traces and stale wrapper trace schemas. Repin the vendored source to the published paged-context V-tail fix and refresh the security-scanning lock for the rebased FlashInfer dependency. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

Use native CSR page tables and max_kv_len for reusable context plans, scope trace fallbacks to the vendored package, and sanitize invalid paged V rows before PV MMA. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

Re-vendor the PrimTS source after restoring its upstream trace imports and preserve TensorRT-LLM package compatibility in the generated vendor patch. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

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:18:57Z

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 Aug 25, 2026

heyuhhh mentioned this pull request Aug 25, 2026

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

Draft

heyuhhh force-pushed the user/yuhangh/prims-ts-block-sparse-fmha branch 2 times, most recently from 38a4f33 to 8d5c9c6 Compare September 1, 2026 10:19

heyuhhh force-pushed the user/yuhangh/vsa-block-sparse-integration branch 2 times, most recently from 22b6caa to 02c8016 Compare September 3, 2026 04:19

heyuhhh force-pushed the user/yuhangh/prims-ts-block-sparse-fmha branch from 8d5c9c6 to 4e08643 Compare September 3, 2026 04:19

heyuhhh changed the title ~~[None][feat] integrate VisualGen VSA with block-sparse FMHA~~ [None][feat] align VisualGen VSA with block-sparse FMHA Sep 3, 2026

heyuhhh mentioned this pull request Sep 3, 2026

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

Closed

heyuhhh force-pushed the user/yuhangh/vsa-block-sparse-integration branch from 02c8016 to c9cd717 Compare September 3, 2026 04:54

heyuhhh force-pushed the user/yuhangh/prims-ts-block-sparse-fmha branch from 4e08643 to ea50060 Compare September 3, 2026 04:54

heyuhhh force-pushed the user/yuhangh/vsa-block-sparse-integration branch from c9cd717 to 97727ad Compare September 3, 2026 07:37

lfr-0531 reviewed Sep 3, 2026

View reviewed changes

heyuhhh force-pushed the user/yuhangh/vsa-block-sparse-integration branch 3 times, most recently from 90ce012 to 6a5b339 Compare September 3, 2026 19:56

heyuhhh force-pushed the user/yuhangh/prims-ts-block-sparse-fmha branch from ea50060 to 3797309 Compare September 3, 2026 19:56

This was referenced Sep 3, 2026

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

Closed

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

Merged

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

heyuhhh force-pushed the user/yuhangh/prims-ts-block-sparse-fmha branch from 4f46724 to f41044c Compare September 4, 2026 04:54

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

yuxianq added 7 commits September 4, 2026 11:06

[None][feat] integrate PrimTS FMHA kernels

8c87b57

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][chore] address PrimTS review feedback

b116b80

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][infra] add locked source vendoring

4c7dac3

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][infra] simplify vendored source updates

4c5de1d

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][perf] reuse PrimTS plans by batch size

94e0087

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] make PrimTS FMHA opt-in

8f78011

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] align PrimTS with latest main

0f6dfac

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

yuxianq and others added 22 commits September 4, 2026 11:06

[None][chore] update vendored FlashInfer PrimTS

a510a3c

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][refactor] compile PrimTS kernels during plan

3ea9bed

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][chore] refresh vendored FlashInfer PrimTS

35e2d43

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] avoid async compiler thread in PrimTS test

a9e3fea

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] fall back for PrimTS context on SM103

4ac6292

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] revert temporary PrimTS SM103 fallback

08849c3

This reverts commit dad7a6b. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] zero-fill PrimTS paged context V tails

6b2a7c2

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][test] isolate FlashInfer support fixtures

02df73a

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][perf] skip unused PrimTS decode control resets

2d1ae62

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][refactor] unify PrimTS page-table metadata

8da8d58

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] Resolve rebase fallout

b0e540f

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][chore] defer Poetry lock refresh

6f951c3

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] distinguish graph FMHA cache entries

f9e01c0

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] key FMHA cache by input layout

5e23907

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] limit FMHA cache key to request state

45862a0

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

chore: update FlashInfer PrimTS pin

3f3c119

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

feat: add generic PrimTS block-sparse support

4f4a9f5

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

refactor: unify sparse attention runtime inputs

4130a1e

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

refactor: align VisualGen sparse attention workflow

6e34838

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

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

heyuhhh force-pushed the user/yuhangh/prims-ts-block-sparse-fmha branch from f41044c to 4130a1e Compare September 4, 2026 14:56

heyuhhh mentioned this pull request Sep 4, 2026

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

Closed

heyuhhh force-pushed the user/yuhangh/prims-ts-block-sparse-fmha branch 2 times, most recently from 2ef79ea to 5c51360 Compare September 7, 2026 09:10

heyuhhh closed this Sep 7, 2026

Navigation Menu

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

[None][feat] align VisualGen VSA with block-sparse FMHA#7
heyuhhh wants to merge 42 commits into
user/yuhangh/prims-ts-block-sparse-fmhafrom
user/yuhangh/vsa-block-sparse-integration

heyuhhh commented Aug 25, 2026 •

edited

Loading

Uh oh!

Uh oh!

lfr-0531 Sep 3, 2026

Uh oh!

lfr-0531 Sep 3, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

heyuhhh commented Sep 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		return (x_sum / valid_counts).to(x_tiled.dtype)


		class VSAPreprocessor:

		return torch.index_select(x, 1, untile_idx)


		def _normalize_qkv_inputs(

Conversation

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

Uh oh!

Description

Stack

Summary

Scope

Validation

Uh oh!

Uh oh!

lfr-0531 Sep 3, 2026

Choose a reason for hiding this comment

Uh oh!

lfr-0531 Sep 3, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

heyuhhh commented Sep 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

heyuhhh commented Aug 25, 2026 •

edited

Loading