[TRTLLM-15917][feat] Integrate Sol-Attn sparse attention into VisualGen - #18329
Merged
Merged
Conversation
4 tasks
karljang
force-pushed
the
feat/sol-attn-visualgen-reduced
branch
6 times, most recently
from
August 31, 2026 22:14
34d16d7 to
60ef12a
Compare
karljang
marked this pull request as ready for review
September 3, 2026 15:20
karljang
requested review from
QiJune,
StanleySun639,
ZhanruiSunCh,
arysef,
asfiyab-nvidia,
luyiyun1021,
niukuo,
nv-guomingz,
pengbowang-nv,
schetlur-nv,
tongyuantongyu,
yibinl-nvidia and
yunruis
September 3, 2026 15:20
Seven findings from automated review of b25ca70, each verified against the source before being applied. Correctness: - The MHA invariant was an `assert`, which `python -O` strips. GQA/MQA would then reach the kernel wrapper, which sees unequal Q/K shapes and takes its dense fallback -- degrading silently instead of rejecting an unsupported configuration. Now a `ValueError`; the test asserts the new type. - `SolAttentionConfig.dense_layers` accepted malformed specs. A non-numeric token raised from `_parse_dense_layers` during attention construction, far from the config that caused it; worse, a descending range such as `4-2` raised nothing at all -- `range(4, 3)` is empty, so the layers the user asked to force dense quietly stayed sparse. A `field_validator` now rejects both at config time. Test quality: - `test_sol_attn_self_attention_is_served_under_separate_qkv` allocated a CUDA tensor with no skip guard, so it errored rather than skipped on a CPU-only host. `_can_serve` compares shapes and a layer index and never touches the device, so the test now builds CPU tensors and runs everywhere -- strictly more coverage than adding the skip marker its two CUDA neighbours carry. Housekeeping: - `cute_dsl_kernels/blackwell/sol_attn_backend.py`, added by this PR, was missing the NVIDIA SPDX header every sibling file carries. - Complete the type annotations in `sol_attn.py`: `frozenset[int]`, `set[int]`, and the parameters of `_delegate`/`_dense_by_step`. - `sparse-attention.md`: add the missing `Sol-Attn` table-of-contents entry (nested, since the section is an h3 under Overview like `Algorithms`), and fix "dense dense attention", a duplicated word spanning a line break that a flat grep missed. Tests: 95 passed, 1 skipped across the Sol-Attn, VSA and dense CuTeDSL suites, up from 86 -- nine new cases covering the `dense_layers` validator on both the accept and reject paths. `pre-commit run` clean on the changed files. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
The vendored `sm100/` kernel already runs on both datacenter Blackwell steppings; only this port's gates said otherwise. Why it works. `cute.compile()` takes no architecture argument -- the CuTe DSL JIT targets whatever device it compiles on -- and the kernel body uses no SM100-exclusive construct. The dense `cute_dsl_fmha_fwd` in this repository already serves `sm_100a` and `sm_103a` from a single class for exactly that reason, differing only by an exp2-emulation flag that Sol-Attn does not use. Measured, not assumed. The vendored package was driven directly on a B300 SXM6 AC and on a B200 as a control -- same seed, shape, and pinned toolchain (nvidia-cutlass-dsl 4.6.2, flash-attn-4 4.0.0b19). Driving the package rather than the TensorRT-LLM wrapper is deliberate: the wrapper turns an ineligible architecture into a silent dense fallback, which would read as success. Against a dense SDPA reference at tau 0.0/1.0/2.0, cosine was 0.826578/0.686830/0.629194 on SM103 versus 0.826579/0.686839/0.629211 on SM100, with mean absolute error equal to printed precision. The residual appears only in the maximum element, consistent with reduction order across a different SM count. Three gates are widened. `SUPPORTED_ARCHS` and `_CUTE_BACKENDS` gain `(10, 3)`. The third, in `_sol_attn_cute`, was a hardcoded `arch != (10, 0)` and is now keyed off `_CUTE_BACKENDS`: it sits deeper than `_backend_for_arch` and no test reached it, so widening the other two would have left it as the only thing still rejecting B300 -- with the suite green. `test_no_arch_literal_outside_the_dispatch_map` closes that hole and was verified by reintroducing the literal, which made it the sole failure. This satisfies the invariant SM120 was dropped for: Sol-Attn's architecture set stays a subset of the dense CuTe DSL FMHA kernel's, so its dense fallback always matches its own backend. THIRD_PARTY_NOTICES records the divergence from upstream's SM100-only packaging next to that note, with the measurements above. Corrects two claims that are now false: `SolAttentionConfig` and the `SolAttention` docstring said sm100 "only", which read as a hardware limit rather than what had been validated. "Only the sm100 kernels are carried" is untouched -- that is about vendoring scope and remains true. Registers the suite in `l0_b300.yml`, which previously carried no VisualGen attention tests at all. Not covered: end-to-end accuracy and performance on B300. `tau=2.0`, the 0.9090 dense prefix and the 0.25 LPIPS gate were calibrated on B200, and the 1.43x speedup is a B200 number; neither transfers without measurement. Tests: 97 passed, 1 skipped across the Sol-Attn, VSA and dense CuTeDSL suites. `pre-commit run` clean on the changed files. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Correctness. The kernel-launch `except Exception` spanned backend loading, argument resolution, the launch itself, and bookkeeping, so a `TypeError` or `NameError` silently became a dense run -- an integration bug reading as "Sol-Attn just didn't speed anything up". `CODING_GUIDELINES.md` asks for the smallest exception set. The guarded region is now the launch alone, with `_resolve_kv_splits` moved out (a bad value is a configuration error and must surface) and the counter moved after it. Narrowing this needs care: CuTe DSL derives `DSLBaseError` from `Exception`, not `RuntimeError`, so the obvious `(ImportError, OSError, RuntimeError)` tuple would have stopped catching real JIT and codegen failures and turned a previously degrading path into a hard crash. `_degradable_kernel_errors()` names it explicitly and resolves it lazily, since this module defers every CuTe DSL import to first use. `SolAttentionConfig.dense_layers` silently accepted empty entries, so `","`, `"0,,2"` and `" "` forced fewer layers dense than written. They now raise. Kernel-level accuracy test. There was no enabled numerical test -- the only one was a skipped placeholder, because Sol-Attn's routing is score-derived and no tau provably forces full dense routing. A single KV block sidesteps that: with `tokens == BLOCK_SIZE` there is nothing to route away, so sparsity is structurally impossible and the kernel must reproduce dense attention whatever tau says. `test_cute_kernel_matches_dense_on_a_single_block` asserts that at rtol/atol=2e-2, and asserts `kernel_calls` advanced so a dense fallback cannot make it pass by comparing dense against itself. The two architecture tests added in the previous commit were also weak: one checked only that the keys exist rather than that `_backend_for_arch` resolves them, and the literal check rejected only `(10, 0)`, so a hardcoded `(10, 3)` would have passed. The latter is now an AST check that rejects any architecture tuple compared against `arch`. Structure and docs. Merges the separate VSA and Sol-Attn `cp_size > 1` guards -- every sparse algorithm here routes over the whole sequence, so none can be split across context-parallel ranks. Moves the `SEPARATE_QKV` rationale out of `modules/attention.py` into `SolAttention._can_serve`, where the behaviour lives. `THIRD_PARTY_NOTICES.md` drops from 119 to 49 lines, keeping attribution and licensing (the FlashAttention BSD-3 retention, the `_vendor` exclusion, the Triton/CUTLASS/cuda-python note) and losing the engineering narrative. The divergence list and the `torch.library.custom_op` comparison move into `sol_attn_backend.py`'s module docstring and the guard comment, so a re-sync still finds them next to the code they constrain. Marks VSA as supported in the algorithms table and applies the reviewer's wording for the backend-compatibility paragraph. Tests: 101 passed, 0 skipped across the Sol-Attn, VSA and dense CuTeDSL suites, up from 97 passed and 1 skipped. `pre-commit run` clean. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
… phase `sol_attn_graph_phase` duplicated `SkipSoftmaxScheduler.get_graph_phase_for_timestep` line for line, `_as_float` included, and `models/modeling.py` called both: the shared classmethod for skip-softmax and the copy for Sol-Attn, twenty lines apart in the same function. The copy's own docstring said it had the same contract, which is reason to call the original rather than restate it. Both branches now call the classmethod. That deletes the duplicate and its `_as_float` helper, and drops `sol_attn_graph_phase` from the `cute_dsl` exports; it was introduced by this PR and has no external users. Behaviour is unchanged -- the bodies were identical -- and the phase tests now exercise the shared implementation against Sol-Attn's cutoff. Also fixes an unrelated guard in `test_dense_paths_use_cutedsl_backend`. It skipped on `torch.cuda.is_available()` alone, but its premise is that the dense paths reach `cute_dsl_fmha_fwd`, which exists only on sm100/sm103, so on any other CUDA device it failed rather than skipping -- observed on H200. It now checks `_cute_dense_available()`. Tests: 56 passed, 45 skipped on H200; the skips are the sm100-only paths, which have no kernel on that device. The full 101-test Blackwell run predates this commit, so the sm100 paths should be re-confirmed on B200 before merge. `pre-commit run` clean. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
… they name `test_ineligible_reason_is_reported` parametrized three cases with ids `cpu-ok-shape`, `cpu-wrong-head-dim` and `cpu-wrong-rank`, but every case built a CPU tensor and `sol_attn_ineligible_reason` checks `is_cuda` first, so all three returned "not a CUDA tensor". The rank and head_dim branches were never exercised despite the ids claiming otherwise. Each case now uses a tensor-like stub that reports `is_cuda=True`, so it passes that first gate and fails exactly one later check: rank, head_dim, or dtype (newly covered). No GPU is needed; the architecture check comes after these and is not reached. Also brings the module docstring up to date: it still described kernel-level numerical equivalence as deferred and pointed at a placeholder that `53150fc` replaced with `test_cute_kernel_matches_dense_on_a_single_block`. Tests: 50 passed, 2 skipped on H200 (the skips are the sm100-only paths). `pre-commit run` clean. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
…2 timesteps Four review findings on the Sol-Attn integration, three of them P1. Masks. `_can_serve` treated equal Q/K lengths as unmasked self-attention, but HunyuanVideo1.5 and GLM-Image pass a `[B, S]` `key_padding_mask` and Cosmos3 passes `CAUSAL`. The sparse kernel is noncausal and takes no mask, and the in-family dense fallback (`CuTeDSLAttention._fwd(**kwargs)`) swallows `key_padding_mask` silently, so padded tokens took part in attention with no error. Masked calls are now disqualified from the kernel and routed to a backend that honors the mask: `CAUSAL` to dense CuTeDSL, which supports it; `key_padding_mask` to a `VanillaAttention` instance (HND layout, transposed in and out), the only backend that consumes it. `_sdpa` honors both instead of dropping them. Tests check the output against a masked reference, so a mask that is routed but then dropped still fails. CUDA-graph capture. `_dense_by_step` resolved the dense-prefix phase from a CUDA timestep tensor with `.item()`, a device-to-host sync that stream capture forbids; `torch.compiler.disable` does not help, it only excludes Dynamo. The runner already resolves every extra key host-side to build the graph key, so it now republishes them through a contextvar for the duration of warmup and capture, and both consumers prefer that: Sol-Attn reads `sol_attn_phase`, and the CuTeDSL skip-softmax path passes `skip_softmax_phase` into a new optional `graph_phase=` on `SkipSoftmaxScheduler.get_runtime_params`. Skip-softmax had the identical exposure through the shared scheduler. Verified on B200 by capturing one graph on each side of the cutoff with a CUDA timestep, no sync error, `kernel_calls` advancing only in the sparse phase. LTX2. Three defects, two of them pre-existing. The base pipeline passed `step_index / num_steps` -- ascending -- as the graph-key timestep, inverting the dense prefix; it now passes the scheduler sigma, already in [0, 1] with the contract's sense. Blocks passed `video.timesteps` / `audio.timesteps` -- AdaLN modulation output -- to every attention call as if it were the scheduler time, so any `disabled_until_timestep` compared against a learned activation; this affected skip-softmax on LTX2 too. `BasicAVTransformerBlock.forward` gains `timestep`, all nine sites use it, and `LTXModel.forward` threads it. The two-stages pipeline passed no timestep and never registered the phase hook, so both phases would have shared one captured graph; it now does both. The config docstring's claim that every pipeline supplies the timestep is corrected. There is no LTX2 end-to-end test in these suites; this is verified by unit tests and review. Kernel-vs-reference test with approximation. Adds a PyTorch reference of the vendored kernel's routing and block-mean approximation -- `kc` = block mean of K, `vc` = block sum of V, `mean + tau*std` threshold per query block in the log2 domain, exact when the column-mean score clears it or the block is within one of the diagonal, approximated blocks contributing `exp*vc` and `exp*block_len`. With Gaussian inputs at 256 tokens, 37.5% of (q_block, kv_block, head) pairs are approximated -- exactly the |dblock| >= 2 pairs -- kernel vs reference max 1.7e-3, reference vs dense max 0.45, and kernel vs dense equal to it, so the kernel's deviation from dense is fully accounted for by the modeled approximation. Tolerance is 5e-3 from that calibration. The test asserts the reference's own mask shows approximated blocks and that `kernel_calls` advanced, so neither an all-exact run nor a dense fallback can pass; clustered inputs were rejected because they route mass-free blocks to the approximation and pass vacuously. Tests: B200 106 passed / 0 skipped before the mask routing, H200 62 passed / 47 skipped after it (the skips are sm100-only paths); B200 re-run pending. `pre-commit run` clean on the changed files. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
… test `nb` was assigned and then `del`-ed to quiet the linter; the assignment itself was the leftover. No behavior change. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
…aph key contract `LTX2TwoStagesPipeline._setup_cuda_graphs` now registers the transformer's extra CUDA-graph key functions on the runner, matching the single-stage pipeline. The `TinyTransformer` stand-ins in `test_ltx2_pipeline.py` bypass `BaseModel`, so they lacked that method and the two CUDA-graph setup tests failed with `AttributeError` in CI. Give the stubs a recording no-op and assert the pipeline registers on the runner it creates. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
…e override `_LTX2CUDAGraphRunner.capture()` re-implements capture for `Modality` inputs and did not enter `resolved_extra_keys_scope`, so with a Sol-Attn dense prefix configured the callee fell back to reading the CUDA timestep with `.item()` inside capture, which CUDA rejects. Wrap LTX-2's warmup and capture in the same scope the base runner uses. The two-stage runner inherits the override, so it is covered as well. The capture/replay tests are parametrized over the base runner and both LTX-2 runners; without this change the two LTX-2 cases fail with "operation failed due to a previous error during capture". Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
…onConfig
Both accepted values ("auto" and "1") resolved to a single KV split, the
only count the shipped sm100/sm103 kernels support, so the field offered
a choice with no effect. Remove it from the public config and pass
`kv_splits=1` internally; the kernel interface still rejects any other
value. The option can return if a second useful strategy ships.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Register the vendored Sol-Attn tree as the `sana-sol-attn` entry in the vendor lock, pinned to NVlabs/Sana `sol-engine` @ 5fe5febd. The destination is restored to upstream bytes for the selected files, and the TensorRT-LLM adaptations live in `3rdparty/vendor_patches/sana-sol-attn.patch`: imports point at the pip `flash_attn` CuTe helpers instead of the removed `_vendor` copy and use relative `..common` paths, the interface dispatches to the SM100 kernel for SM100/SM103 only (the SM89/SM90/SM120 and Triton paths are not shipped), and the notices file records the pin and scope. The tree is excluded from ruff and pre-commit formatting so the patch stays free of formatting noise, and `vendor_sources.py check` passes offline. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
…t flag, no kernel-error recovery
* `SolAttentionConfig.dense_layers` is a `list[int]` of layer indices instead
of a parsed "0,2-4" string; the validator rejects negative indices and
normalizes to a sorted, deduplicated list.
* The strict flag is `TRTLLM_SOL_ATTN_STRICT`, following the repository's
environment-variable convention.
* The eligibility check uses the shared `get_sm_version()` helper; the
supported set is `{100, 103}`.
* Kernel exceptions are no longer caught and answered with dense attention.
A failed CuTe launch can leave the device in a bad state, so only inputs
known up front to be unservable (shape, dtype, architecture) are routed to
dense; everything else propagates.
* The construction-time dense-backend comment states plainly that the
measured compile/eager mismatch was a graph-structure effect of
`torch.compiler.disable`, not a kernel or torch.compile bug, and the
cross-attention length heuristic is documented as such with the explicit
signal tracked in TRTLLM-16475.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Explain that construction-time availability checks keep capability queries and attribute initialization out of dense-path tracing. Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
…ss-attention flag Rebase follow-up: main now gates the fallback on separate_qkv_cross_attention (the separate_qkv_is_self_attention flag from NVIDIA#18147); keep that condition and the Sol-Attn note together. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
karljang
force-pushed
the
feat/sol-attn-visualgen-reduced
branch
from
September 18, 2026 03:36
748b4a8 to
b901f29
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #74307 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #74242 [ run ] completed with state |
juney-nvidia
approved these changes
Sep 18, 2026
juney-nvidia
left a comment
juney-nvidia
left a comment
Collaborator
There was a problem hiding this comment.
Approved from OSS compliance perspective
chang-l
enabled auto-merge (squash)
September 18, 2026 05:16
Collaborator
|
PR_Github #74307 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #74390 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #74390 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #74485 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #74485 [ run ] completed with state |
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
11 participants
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
Adds Sol-Attn (arXiv:2607.24027) as a third
sparse-attention algorithm for VisualGen, alongside
skip_softmaxand VSA. It foldsdynamic block routing, sparse computation, and an approximation-correction term into a
single online-softmax pass.
Configured through
SolAttentionConfig, dispatched viacreate_attentionthe sameway
skip_softmaxandvsaare:disabled_until_timestepfollows skip-softmax's field of the same name and the samesense: dense while the normalized denoising timestep is at or above the cutoff, sparse
below it. The value arrives as a forward kwarg that every VisualGen pipeline already
supplies, so no per-pipeline wiring is needed.
Scope and behaviour
sm100 (B200/GB200) only,
head_dim=128, bf16, MHA. Context parallelism andquantized attention are rejected explicitly, mirroring VSA's guards. An unsupported
shape, dtype or architecture degrades to dense with a
warning_onceand adense_fallback_callscounter;TRTLLM_SOL_ATTN_STRICT=1raises instead.Non-sparse work stays on the configured backend. Sol-Attn is self-attention only and
does not run its kernel on every step, so three paths do dense attention: the
dense_layersguard, thedisabled_until_timestepprefix, and kernel-ineligibilityfallback. All three use
cute_dsl_fmha_fwd— the dense kernel of the selected backend —not
torch.nn.functional.scaled_dot_product_attention. Cross-attention (SEPARATE_QKV)likewise falls back within the backend family rather than to VANILLA.
This matters beyond tidiness: with
disabled_until_timestep=0.0001, so the sparse kernelnever fires, Sol-Attn is byte-identical to a plain
backend: CUTEDSLrun (LPIPS0.0000). Any measured difference is therefore sparsity and nothing else. The pipeline is
bit-deterministic on this workload — a repeated identical config also scores 0.0000 — so
that is an exact statement, not an approximate one.
Performance
Wan2.2-T2V-A14B, 720x1280x81f, 40 steps (the model default,
models/wan/defaults.py), B200, seed 42,torch.compileenabled (the productiondefault). Baseline is dense CuTeDSL under the same compile setting.
Eager, for reference: 474.1 s -> 341.9 s, 1.386x, 27.9 % (single repetition).
Protocol: one warmup generation then two timed repetitions; the figures above are
their mean. Within-run spread is at most 0.13 %. Each prompt's baseline and
candidate were measured in the same allocation, which is what makes the
ratios comparable -- absolute times drift by ~2 % between allocations (different
node, different clock state), while the speedups do not.
Speedup is
T_base / T_new; time saved is(1 - T_new / T_base) x 100.Accuracy
LPIPS against the dense CuTeDSL baseline at the same compile setting, gate 0.25,
worst-prompt governs.
torch.compileKEEP in both modes. Enabling
torch.compiledoes not cost quality at this operatingpoint -- it is marginally better on every prompt.
Test coverage
tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py, registered inl0_b200.yml. Covers backend-factory dispatch, cross-attention staying in-family,context-parallel and quantized-attention rejection, GQA/MQA rejection, the
dense_layersguard, dense-prefix phase semantics either side of the cutoff, fail-open on a missing
timestep, both CUDA-graph key cases, kernel-eligibility reasons,
TRTLLM_SOL_ATTN_STRICT,dense-fallback numerics and counters, arch-list drift between
SUPPORTED_ARCHSand_CUTE_BACKENDS, and that both Dynamo-opacity boundaries stay decorated.34 passed, 1 documented skip in this PR's own suite on B200.
Run together with the VSA and dense CuTeDSL suites -- both touched by the
cross-attention change in
modules/attention.py-- the three total 84 passed,1 skipped, so neither neighbouring backend regresses.
Divergence from upstream
Not a byte-faithful vendoring.
sol_attn/THIRD_PARTY_NOTICES.mdrecords the pin and everydeliberate difference; start there for a currency check.
triton_ref/not carried_vendor/flash_attn/not carriedflash-attn-4; verified bit-identical on B200@torch.compiler.disableon the launch boundarysol_attn()unguarded; without it Dynamo traces into the CuTe DSL JIT builderscripts/vendor_sources.py(lock entrysana-sol-attn, 4-file patch)cute_dsl_fmha_fwdlogger.warning_oncereplacesprint()dense_fallback_calls,sol_attn_ineligible_reason(),TRTLLM_SOL_ATTN_STRICTon the eligibility pathsol_attn_backend.pyis itself adapted from upstream's file of the same name, which sitsoutside the vendored package; only the kernel-wrapper subset is carried. Both projects are
Apache-2.0. Upstream guards the
torch.compilepath withtorch.library.custom_op+register_fake, which keeps the kernel in the graph rather than breaking at it — areasonable follow-up, not adopted here because this PR's measurements were taken with the
disableform.Dev Engineer Review
SolAttentionConfigand CUTEDSL dispatch for Sol-Attn.head_dim=128.SkipSoftmaxScheduler.get_graph_phase_for_timestepfor CUDA-graph phase selection.sol_attn_graph_phasehelper.#71298failed before the rebase. A post-rebase run and B200 reconfirmation remain necessary.QA Engineer Review
tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py.tests/integration/test_lists/test-db/l0_b200.ymlandl0_b300.yml.Per-File QA Perspective
tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py: Covers Sol-Attn dispatch, validation, fallback, graph phases, kernel behavior, and dense delegation. It is listed in both B200 and B300 CI lists.tests/integration/test_lists/test-db/l0_b200.yml: Adds the Sol-Attn test to B200 pre-merge CI.tests/integration/test_lists/test-db/l0_b300.yml: Adds the Sol-Attn test to B300 pre-merge CI.tensorrt_llm/visual_gen/sparse_attention.py: AddsSolAttentionConfigand validates thresholds, KV splits, dense prefixes, and forced-dense layers.tensorrt_llm/visual_gen/args.py: Adds Sol-Attn to the sparse-attention union and quantization compatibility rules.tensorrt_llm/visual_gen/__init__.py: ExposesSolAttentionConfig.tensorrt_llm/_torch/visual_gen/attention_backend/utils.py: Dispatches CUTEDSL Sol-Attn and forwards its configuration.tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py: ExposesSolAttention.tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py: RetainsSolAttentionand removes thesol_attn_graph_phaseexport.tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py: Uses the shared scheduler for timestep gating and graph phases.tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py: Adds eligibility checks, fallback and strict behavior, counters, and dense delegation.tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/interface.py: Validates BF16 contiguous inputs, supported architectures, sink parameters, and kernel runtime availability.tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/preprocess.py: Adds threshold preprocessing and routing summaries.tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/layout_utils.py: Adds CuTe layout transformations used by the kernel.tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/runtime.py: Converts tensors to aligned CuTe tensors through DLPack.tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/selector.py: Adds routing-mask bit operations and exact-route selection.tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/mainloop.py: Adds the SM100 forward kernel. QA should verify routing, online softmax, synchronization, and BF16 output.tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/math.py: Adds the tensor-core GEMM helper.tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/softmax.py: Adds online-softmax and accumulator-rescaling helpers.tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/tmem.py: Adds TMEM load, store, wait, and output-copy helpers.tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/__init__.py,common/__init__.py,sm100/__init__.py, andsm100/kernel.py: Add package exports and vendoring metadata.tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.mdandsm100/LICENSE.flash-attention: Document licensing and provenance. QA should verify packaging and notice compliance.tensorrt_llm/_torch/visual_gen/models/modeling.py: Adds Sol-Attn CUDA-graph phase-key handling.tensorrt_llm/_torch/visual_gen/modules/attention.py: Adds Sol-Attn selection and rejects context parallelism.docs/source/visual-gen/features/sparse-attention.md: Documents hardware support, inputs, fallback behavior, strict mode, timestep gating, and quantization constraints.