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
…gured backend Enabling `sol_attn` silently swapped the attention kernel in two places that have nothing to do with sparsity, so an A/B against a `backend: CUTEDSL` dense baseline was measuring a backend difference, not the algorithm. Self-attention: the `dense_layers` guard, the `disabled_until_timestep` prefix and every kernel-ineligibility fallback called `torch.nn.functional.scaled_dot_product_attention`, while the baseline ran `cute_dsl_fmha_fwd`. The prefix alone covers ~24 % of the work at the certified operating point, so this was not a rare edge case. All three paths now route through `CuTeDSLAttention`, using upstream's existing `dense_fn` hook for the third. SDPA is retained only where the CuTe kernel cannot serve the device, and says so once. Cross-attention: `modules/attention.py` routes `SEPARATE_QKV` to VANILLA when the sparse algorithm is vsa/sol_attn, but plain `CUTEDSL` does not match that condition and keeps CuTeDSL. WAN's `attn2` is `SEPARATE_QKV` in every block, so merely enabling the feature moved cross-attention to torch SDPA everywhere, regardless of `tau`, `disabled_until_timestep`, or whether the sparse kernel ever ran. Sol-Attn now falls back within its own backend family; `create_attention` re-selects the sparse class from `attention_config`, so the cross-attention module is built with `sparse_attention_config=None`. TRTLLM keeps VANILLA, since `TrtllmAttention` genuinely cannot serve `SEPARATE_QKV`. Verification. With sparsity disabled entirely (`disabled_until_timestep=0.0001`, so the sparse kernel never fires) Sol-Attn is now **byte-identical** to a plain `backend: CUTEDSL` run: LPIPS 0.0000, against 0.1279 before. That is an exact result, not an approximate one -- a repeated identical config also scores 0.0000, so the pipeline is bit-deterministic on this workload and any nonzero value is signal. At the certified operating point (`tau=2.0`, `disabled_until_timestep=0.9090`) on Wan2.2-T2V-A14B, 720x1280x81f, 50 steps, B200, p01, against the now-valid baseline: | | denoise | S | delta | LPIPS | previously | |---|---|---|---|---|---| | eager | 427.54 s | 1.373x | 27.2 % | 0.1936 | 0.2477 | | torch.compile | 364.11 s | 1.418x | 29.5 % | 0.2337 | 0.4159 | Both inside the 0.25 gate. The compiled figure moved from 166 % of gate to 93 %: the apparent collapse of quality under `torch.compile` was entirely the reference mismatch, amplified because `cute_dsl_fmha_fwd` is `@torch.compiler.disable`'d and bit-identical either way while the SDPA path is not. VSA has the identical cross-attention defect. It is deliberately not changed here, since that alters a separate feature; tracked as TRTLLM-16105. Tests: 84 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>
…ention Cold review found that routing Sol-Attn's `SEPARATE_QKV` fallback to CUTEDSL also caught *self*-attention that merely uses that qkv mode, which is a regression rather than a fix. `QwenImageAttention` is `SEPARATE_QKV` with `separate_qkv_is_self_attention=True` (`models/qwen_image/transformer_qwen_image.py`). Redirecting it flipped `attn_backend` from VANILLA to CUTEDSL, and `_supports_qwen_key_padding_mask` tests for the literal string "VANILLA", so with `ulysses_size > 1` the model raised `NotImplementedError` on a configuration that worked before. WAN's `attn1` is likewise `SEPARATE_QKV` under async Ulysses. The fallback is now gated on `not separate_qkv_is_self_attention`, so only genuine cross-attention moves in-family and those paths keep VANILLA. Adds `test_dense_paths_use_cutedsl_backend`, a CUDA test asserting that all three dense paths -- the `dense_layers` guard, the `disabled_until_timestep` prefix, and the `dense_fn` ineligibility fallback -- reach the configured backend's dense kernel. The existing dense tests build CPU tensors, so `_dense` takes its SDPA branch by construction and cannot observe this; the two are renamed so they no longer read as asserting the old behaviour. Reverts `l0_gb202.yml` to base: dropping sm120 left a "Visual Gen tests" header with no test under it, mislabelling unrelated BERT and Qwen3 entries. Corrects stale "dense SDPA" wording in the module and config docstrings, the two runtime fallback messages, and the user-facing sparse-attention doc, all of which became false when the dense paths moved in-family. That doc's example cutoff also moves to the validated 0.9090. Records the dense-path routing in THIRD_PARTY_NOTICES.md, which claimed to list every deliberate divergence and omitted this one -- exactly what a re-sync would overwrite. Softens the `torch.compile` latency citation from a flat "69x (2496.9 s vs 36.2 s)" to "near two orders of magnitude (2496.9 s without it)". The 2496.9 s is archived; the post-fix figure was measured while another job shared the GPU and its result file was later overwritten, so the precise ratio is not reproducible from artifacts. Verification. The byte-identity control now runs in the mode the PR reports: `disabled_until_timestep=0.0001` with `torch.compile` enabled at 40 steps scores LPIPS 0.0000 against the same dense anchor as the headline numbers (denoise 414.24 s vs 414.36 s). Previously that control had only been run eager at 50 steps, while every reported number was compiled at 40 -- and at an earlier fix stage compile tripled the residual, so the extrapolation was unsafe. Tests: 85 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>
`SolAttnAttention` stutters: the algorithm is already named "Sol-Attn", so the class read as Attn-Attention. Upstream has no equivalent name to preserve -- it exposes dispatch functions and a `_SolContext` dataclass, not an attention backend class, so this follows only this repository's own `<Name>Attention(AttentionBackend)` convention alongside `CuTeDSLAttention`, `VSAAttention`, `TrtllmAttention` and `VanillaAttention`. `SolAttnAttentionConfig` renames to `SolAttentionConfig` for the same reason and to match `SkipSoftmaxAttentionConfig` / `VideoSparseAttentionConfig`. Neither name has shipped, so this costs no compatibility. Mechanical: 43 references across 11 files, no behaviour change. One incidental reformat -- the shorter name lets an import in `models/modeling.py` fit on one line. Tests: 85 passed, 1 skipped across the Sol-Attn, VSA and dense CuTeDSL suites. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
…r module Sol-Attn is a policy over a dense backend, not a peer of one. It answers "should the sparse kernel run for *this* call", and delegates everything else to the dense CuTe DSL backend it wraps. This commit makes the code say that. `SolAttention` gains `_can_serve` -- cross-attention, the `dense_layers` guard, and the `disabled_until_timestep` prefix all become one predicate -- and `_delegate`, the single exit to the inner backend. `forward` reduces to "serve it, or hand it over", and the `dense_fn` ineligibility hook routes to the same place, so all four dense paths now leave through one function. Removes Sol-Attn from the `SEPARATE_QKV` rule in `modules/attention.py`, and with it the `model_copy(sparse_attention_config=None)` special case at the `create_attention` call. That rule had to infer cross-attention from `qkv_mode`, which describes how Q/K/V are *projected*, not whether K/V come from another sequence. The inference is wrong wherever SEPARATE_QKV is chosen for other reasons -- Qwen-Image always, WAN's `attn1` under async Ulysses -- and each wrong guess silently cost that module its configured backend. The predicate compares `k.shape[1]` against `q.shape[1]` instead, which is the thing actually being asked. VSA keeps the old rule; it has the same defect, tracked separately as TRTLLM-16105. Behaviour preservation. Wan2.2-T2V-A14B, 720x1280x81f, 40 steps, B200, seed 42, `torch.compile` on, at the operating point this PR reports: the output tensor digest is `d43f9af3...` before and after, bit-identical. That value reproduces across five executions in four processes -- committed HEAD, both refactor variants, and two repetitions of the prior measurement. Denoise 290.34 s vs 290.45 s (0.04 %, within run-to-run spread). Tests: 86 passed, 1 skipped across the Sol-Attn, VSA and dense CuTeDSL suites. Adds `test_sol_attn_self_attention_is_served_under_separate_qkv`, which pins the async-Ulysses case the old rule got wrong, and reworks the cross-attention test to assert that `SolAttention` remains the backend and delegates, rather than being replaced at construction. `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>
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>
karljang
force-pushed
the
feat/sol-attn-visualgen-reduced
branch
from
September 16, 2026 13:29
bc69b7c to
ef1caf5
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #73850 [ run ] triggered by Bot. Commit: |
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>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #73888 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #73850 [ run ] completed with state |
Collaborator
|
PR_Github #73888 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #73915 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #73915 [ run ] completed with state
|
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>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #74034 [ run ] triggered by Bot. Commit: |
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
9 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.