From 73a0b940cfccd61d50909f1904fe9700bf5bea41 Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Thu, 27 Aug 2026 17:40:20 -0700 Subject: [PATCH 01/22] [TRTLLM-15917][feat] Integrate Sol-Attn sparse attention into VisualGen Adds Sol-Attn (arXiv:2607.24027) as a third sparse-attention algorithm for VisualGen, alongside `skip_softmax` and VSA. It folds dynamic block routing, sparse computation, and an approximation-correction term into a single online-softmax pass. Config surface: `SolAttnAttentionConfig` in `visual_gen/args.py` / `sparse_attention.py` -- `tau` (routing threshold), `thresh_type` (`diag`/`exact`), `kv_splits`, `disabled_until_timestep` (dense-prefix cutoff), and `dense_layers` (comma/range layer-skip spec). Dispatch goes through `create_attention` the same way `skip_softmax` and `vsa` do. Cross-attention (`SEPARATE_QKV`) falls back to VANILLA, and context-parallel (`cp_size > 1`) and quantized attention are both rejected, mirroring VSA's existing guards. Dense prefix ------------ `disabled_until_timestep` follows skip-softmax's field of the same name and the same sense: the layer runs dense while the normalized denoising timestep is at or above the cutoff, and switches to the sparse kernel below it. The value arrives as a forward kwarg, which `modules/attention.py` already threads to every backend and every VisualGen pipeline normalizes by `num_train_timesteps`, so no per-pipeline wiring is needed and there is no process-wide state. `models/wan/pipeline_wan.py` is untouched. Because the prefix swaps kernels without changing tensor shapes, the two phases must not share a captured CUDA graph; `register_cuda_graph_extra_key_fns` registers `sol_attn_phase` from the same `kwargs["timestep"]` source as `skip_softmax_phase`. `dense_layers` needs no key, being fixed per layer at construction. Kernel scope ------------ The kernel is vendored from its reference implementation (see `cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md` for the upstream pin and its currency check). Only the two architectures with hardware evidence are carried: sm100 (B200/GB200) and sm120 (RTX Blackwell). Upstream's sm89 and sm90 kernels and its Triton reference path are not included; sm90 covers H100/H200/GH200 and should return in a follow-up with measurements behind it rather than ship unvalidated. Every vendored file carries an SPDX Apache-2.0 header naming its NVlabs/Sana origin; the two files that derive from FlashAttention additionally cite BSD-3-Clause and point at `sm100/LICENSE.flash-attention`, and the cuDNN Frontend license the SM120 kernel adapts is vendored at `sm120/LICENSE.cudnn-frontend` at the commit the notices cite. Upstream also vendors a copy of FlashAttention's CuTe DSL helpers. That copy is not carried: TensorRT-LLM already depends on flash-attn-4, which provides the same `flash_attn.cute` modules, verified on B200 to give bit-identical output. `preprocess.py` implements the routing/threshold stage in Triton, so Triton is a required runtime dependency on every Sol-Attn path. Failure behaviour ----------------- Inputs the kernel cannot serve -- unsupported architecture, `head_dim` other than 128, non-bf16 dtype, or mismatched k/v -- fall back to dense SDPA with a `warning_once` naming the specific reason, and increment `dense_fallback_calls` alongside `kernel_calls`. Kernel exceptions take the same path. `SOL_ATTN_STRICT=1` raises instead, for both arms. Without this the feature degrades to a silent no-op for a whole run and surfaces only as absent speedup. Docs ---- `docs/source/visual-gen/features/sparse-attention.md` gains a `sol_attn` row and a section covering the YAML surface, the sm100/sm120 + head_dim=128 + bf16 + MHA constraints, the cutoff semantics, and the fallback/`SOL_ATTN_STRICT` behaviour. Its claim that VSA is the only CUTEDSL algorithm mutually exclusive with quantized attention is corrected, since Sol-Attn now is too. Tests ----- New `tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py`, registered in `l0_b200.yml` (sm100) and `l0_gb202.yml` (sm120): backend-factory dispatch, cross-attention VANILLA fallback, context-parallel and quantized-attention rejection, GQA/MQA rejection, the `dense_layers` guard, dense-prefix phase semantics at and either side of the cutoff (including tensor-valued timesteps), fail-open on a missing timestep, both CUDA-graph key cases, kernel-eligibility reasons, `SOL_ATTN_STRICT` on the eligibility path, dense-fallback numerics and counters, arch-list drift between `SUPPORTED_ARCHS` and `_CUTE_BACKENDS`, and `kv_splits` rejection. 32 tests plus one documented skip for GPU kernel-vs-dense equivalence at full routing. Validation ---------- * B200 (sm100): 31/31 pass, and 68 passed alongside `test_attention_cute_dsl.py`, which #17781 extended. Kernel output bit-identical across a 12-point (shape, tau) sweep; `kernel_calls=12`, `dense_fallback_calls=0` under `SOL_ATTN_STRICT=1`. Denoise time on B200, 50 steps, mean of 2 reps after 1 warmup, against a dense CuTeDSL baseline: Wan2.2-TI2V-5B 1.127x without CUDA graphs and 1.200x with them; Wan2.2-T2V-A14B 1.451x without and 1.406x with. Enabling graphs helps the 5B and slightly hurts A14B; the cause is not established, so the best A14B configuration remains graphs-off. Run-to-run spread was under 0.06% throughout. * RTX 5090 (sm120): resolves to `cute_sm120`; 9/9 sweep points ran with no dense fallback. End-to-end generation was not possible on that GPU because 32 GB is insufficient for the models used here, so sm120 has kernel-level evidence only. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../features/visualgen-sparse-attention.md | 33 +- .../visual_gen/attention_backend/__init__.py | 2 + .../attention_backend/cute_dsl/__init__.py | 8 +- .../attention_backend/cute_dsl/sol_attn.py | 205 +++ .../visual_gen/attention_backend/utils.py | 11 +- .../blackwell/sol_attn/THIRD_PARTY_NOTICES.md | 57 + .../blackwell/sol_attn/__init__.py | 10 + .../blackwell/sol_attn/common/__init__.py | 10 + .../blackwell/sol_attn/common/layout_utils.py | 135 ++ .../blackwell/sol_attn/common/runtime.py | 19 + .../blackwell/sol_attn/common/selector.py | 171 ++ .../blackwell/sol_attn/interface.py | 334 ++++ .../blackwell/sol_attn/preprocess.py | 454 +++++ .../sol_attn/sm100/LICENSE.flash-attention | 29 + .../blackwell/sol_attn/sm100/__init__.py | 10 + .../blackwell/sol_attn/sm100/kernel.py | 10 + .../blackwell/sol_attn/sm100/mainloop.py | 1530 +++++++++++++++++ .../blackwell/sol_attn/sm100/math.py | 34 + .../blackwell/sol_attn/sm100/softmax.py | 137 ++ .../blackwell/sol_attn/sm100/tmem.py | 138 ++ .../sol_attn/sm120/LICENSE.cudnn-frontend | 204 +++ .../blackwell/sol_attn/sm120/__init__.py | 10 + .../blackwell/sol_attn/sm120/kernel.py | 24 + .../blackwell/sol_attn/sm120/mainloop.py | 1003 +++++++++++ .../blackwell/sol_attn_backend.py | 212 +++ .../_torch/visual_gen/models/modeling.py | 59 +- .../_torch/visual_gen/modules/attention.py | 30 +- tensorrt_llm/visual_gen/__init__.py | 3 + tensorrt_llm/visual_gen/args.py | 31 +- tensorrt_llm/visual_gen/sparse_attention.py | 72 +- .../test_lists/test-db/l0_b200.yml | 1 + .../test_lists/test-db/l0_gb202.yml | 4 + .../test_attention_cute_dsl_sol_attn.py | 421 +++++ 33 files changed, 5364 insertions(+), 47 deletions(-) create mode 100644 tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/__init__.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/__init__.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/layout_utils.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/runtime.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/selector.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/interface.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/preprocess.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/LICENSE.flash-attention create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/__init__.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/kernel.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/mainloop.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/math.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/softmax.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/tmem.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/LICENSE.cudnn-frontend create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/__init__.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/kernel.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/mainloop.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py create mode 100644 tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py diff --git a/docs/source/features/visualgen-sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md index 38efb717765e..3c069ba6d079 100644 --- a/docs/source/features/visualgen-sparse-attention.md +++ b/docs/source/features/visualgen-sparse-attention.md @@ -21,6 +21,37 @@ Sparse attention in VisualGen is configured through `VisualGenArgs.attention_con |---|---|---| | `skip_softmax` | `SkipSoftmaxAttentionConfig` | Supported | | VSA | TBD | TODO | +| `sol_attn` | `SolAttnAttentionConfig` | Supported (CUTEDSL, sm100/sm120) | + +### Sol-Attn + +Sol-Attn ([arXiv:2607.24027](https://arxiv.org/abs/2607.24027)) folds dynamic block +routing, sparse computation, and an approximation-correction term into one +online-softmax pass. It runs on the **CUTEDSL** backend only, on sm100 +(B200/GB200) and sm120 (RTX Blackwell), and requires `head_dim=128`, bfloat16, +and MHA (`num_kv_heads == num_heads`). + +```yaml +attention_config: + backend: CUTEDSL + sparse_attention_config: + algorithm: sol_attn + tau: 2.0 # routing threshold; higher routes more blocks sparse + thresh_type: diag # or "exact" + disabled_until_timestep: 0.9545 # dense while normalized timestep >= cutoff + dense_layers: '0' # optional: layers forced dense +``` + +`disabled_until_timestep` has the same meaning as it does for Skip Softmax: +attention runs dense while the normalized denoising timestep is at or above the +cutoff, protecting the high-noise prefix, and switches to the sparse kernel +below it. Use `None` rather than `0.0` to disable the prefix. + +On an input the kernel cannot serve — an unsupported architecture, a +`head_dim` other than 128, a non-bfloat16 dtype — Sol-Attn falls back to dense +SDPA, logs the specific reason once, and counts the fallback. Set +`SOL_ATTN_STRICT=1` to raise instead of falling back, which is useful when +benchmarking to confirm the kernel actually ran. ## Skip Softmax Attention @@ -92,7 +123,7 @@ User configuration is supplied through Python or YAML and controls how the check `threshold_scale_factor` and `target_sparsity` are alternatives: if both are present, `threshold_scale_factor` takes precedence and the calibration formula is not used. User-provided `target_sparsity` and `disabled_until_timestep` override checkpoint defaults. Checkpoint `ignore` patterns always disable Skip Softmax Attention for matching layers. -Skip Softmax Attention works with both the **TRTLLM** and **CUTEDSL** attention backends in VisualGen. Set `attention_config.backend` to either when enabling it. On CUTEDSL, Skip Softmax Attention can also be combined with `quant_attention_config`'s block-scaled Q/K recipes (MXFP8, NVFP4); VSA is the only CUTEDSL sparse-attention algorithm that is mutually exclusive with quantized attention. +Skip Softmax Attention works with both the **TRTLLM** and **CUTEDSL** attention backends in VisualGen. Set `attention_config.backend` to either when enabling it. On CUTEDSL, Skip Softmax Attention can also be combined with `quant_attention_config`'s block-scaled Q/K recipes (MXFP8, NVFP4); VSA and Sol-Attn each replace the dense CuTeDSL path and are therefore mutually exclusive with quantized attention. #### Mapping `disabled_until_timestep` to Actual Denoising Steps diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py index 12aa3ec4195c..6acbff9c5b41 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py @@ -24,6 +24,7 @@ from .cute_dsl import ( VSA_TILE_SIZE, CuTeDSLAttention, + SolAttnAttention, VSAAttention, VSAMetadata, VSAMetadataBuilder, @@ -48,6 +49,7 @@ "FlashAttn4Attention", "FlashInferAttention", "RingAttention", + "SolAttnAttention", "TrtllmAttention", "TrtllmAttentionMetadata", "UlyssesAttention", diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py index 9b70421c3b81..20e52c1f0574 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py @@ -15,11 +15,13 @@ """ CuTe DSL attention backend family for visual generation models. - fmha.py — CuTeDSLAttention (dense and blockscaled JIT FMHA) - vsa.py — VSAAttention (Video Sparse Attention, CuTe JIT + SDPA fallback) + fmha.py — CuTeDSLAttention (dense and blockscaled JIT FMHA) + vsa.py — VSAAttention (Video Sparse Attention, CuTe JIT + SDPA fallback) + sol_attn.py — SolAttnAttention (Sol-Attn dynamic block routing, CuTe JIT + SDPA fallback) """ from .fmha import CuTeDSLAttention, _cute_dsl_import_error +from .sol_attn import SolAttnAttention, sol_attn_graph_phase from .vsa import ( VSA_KERNEL_MAX_CUBES, VSA_TILE_SIZE, @@ -42,4 +44,6 @@ "set_vsa_forward_context", "get_vsa_forward_context", "_cute_dsl_import_error", + "SolAttnAttention", + "sol_attn_graph_phase", ] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py new file mode 100644 index 000000000000..cd600f91fcbe --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Sol-Attn backend for visual generation models. + +Sol-Attn (https://arxiv.org/abs/2607.24027) is dynamic block routing + +sparse computation + approximation correction folded into one online-softmax +pass. The kernel is vendored from its reference implementation +(https://github.com/NVlabs/Sana, branch +https://github.com/NVlabs/Sana/tree/sol-engine, pinned at commit +https://github.com/NVlabs/Sana/commit/5fe5feb -- see +``cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md`` for the pin +and its currency-check note) under ``..cute_dsl_kernels.blackwell.sol_attn`` +/ ``sol_attn_backend.py``. Only the sm100 (B200/GB200) and sm120 (RTX +Blackwell) kernels are carried; the upstream sm89/sm90 kernels and the Triton +reference path are not, and the FlashAttention CuTe helpers they needed come +from the ``flash-attn-4`` dependency rather than a vendored copy. + +This file is only the TRT-LLM AttentionBackend adapter around that kernel's +public BTHD entry point, plus the dense_layers layer-skip guard. + +``disabled_until_timestep`` is the dense-prefix control, and mirrors +skip_softmax's field of the same name: sparse attention stays disabled (that +is, the layer runs dense SDPA) while the normalized denoising timestep is at +or above the cutoff, and switches to the sparse kernel once it drops below. + +The timestep arrives as a forward kwarg -- ``modules/attention.py`` already +threads it to every backend, and all VisualGen pipelines normalize it to +``[0, 1]`` by ``num_train_timesteps`` per the ``BaseDiffusionModel.forward`` +contract. Nothing has to be wired per pipeline, and there is no process-wide +state to keep in sync. + +""" + +from typing import Any, Optional + +import torch + +from tensorrt_llm.logger import logger + +from ..interface import AttentionBackend, AttentionTensorLayout + +_sol_attn_import_error = None +try: + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.sol_attn_backend import ( + _run_sol_attn_bthd as _sol_attn_run, + ) +except (ImportError, OSError) as e: + _sol_attn_run = None + _sol_attn_import_error = e + + +def _as_float(timestep: Any) -> Optional[float]: + """Coerce a scalar/0-d/1-element timestep to float, else None.""" + if timestep is None: + return None + if isinstance(timestep, torch.Tensor): + if timestep.numel() == 0: + return None + return float(timestep.reshape(-1)[0].item()) + try: + return float(timestep) + except (TypeError, ValueError): + return None + + +def sol_attn_graph_phase( + timestep: Any, *, disabled_until_timestep: Optional[float] +) -> Optional[int]: + """Return 1 once descending timesteps cross the cutoff, 0 before, else None. + + Same contract and sense as + ``SkipSoftmaxScheduler.get_graph_phase_for_timestep``: phase 0 is the dense + prefix, phase 1 the sparse phase, and ``None`` means there is no phase to + distinguish so the CUDA-graph runner omits the key part. + """ + if disabled_until_timestep is None: + return None + value = _as_float(timestep) + if value is None: + return None + return int(value < disabled_until_timestep) + + +def _parse_dense_layers(spec: Optional[str]) -> frozenset: + layers: set = set() + for item in str(spec or "").split(","): + item = item.strip() + if not item: + continue + if "-" in item: + start, end = item.split("-", 1) + layers.update(range(int(start), int(end) + 1)) + else: + layers.add(int(item)) + return frozenset(layers) + + +class SolAttnAttention(AttentionBackend): + """Sol-Attn dynamic block-routing sparse attention (CuTeDSL, sm100/sm120). + + The kernel wrapper already falls back to dense SDPA on any unsupported + shape/dtype/arch (see ``_run_sol_attn_bthd``); this class only adds the + ``dense_layers`` layer-skip guard (evaluated at construction time, no + external plumbing needed) and forwards the routing knobs from config. + """ + + def __init__( + self, + layer_idx: int = 0, + num_heads: int = 8, + head_dim: int = 128, + num_kv_heads: Optional[int] = None, + dtype: Optional[torch.dtype] = None, + sparse_attention_config=None, + **kwargs, + ): + if _sol_attn_run is None: + raise ImportError( + "SolAttnAttention requires the vendored sol_attn kernel " + f"package; import failed: {_sol_attn_import_error}" + ) + self.layer_idx = layer_idx + self.num_heads = num_heads + self.head_dim = head_dim + self.num_kv_heads = num_kv_heads or num_heads + assert self.num_kv_heads == self.num_heads, ( + f"Sol-Attn is MHA-only (num_kv_heads == num_heads), got " + f"num_kv_heads={self.num_kv_heads}, num_heads={self.num_heads}. " + f"GQA/MQA is not supported." + ) + self.dtype = dtype + cfg = sparse_attention_config + self.tau = getattr(cfg, "tau", 1.0) + self.thresh_type = getattr(cfg, "thresh_type", "diag") + self.kv_splits = getattr(cfg, "kv_splits", "auto") + self.disabled_until_timestep = getattr(cfg, "disabled_until_timestep", None) + self.dense_layers = _parse_dense_layers(getattr(cfg, "dense_layers", None)) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + """q, k, v: [B, S, H, D] (NHD), same original token order in and out.""" + dense_by_layer = self.layer_idx in self.dense_layers + dense_by_step = False + if self.disabled_until_timestep is not None: + phase = sol_attn_graph_phase( + kwargs.get("timestep"), + disabled_until_timestep=self.disabled_until_timestep, + ) + if phase is None: + # Fail open, matching the CuTeDSL skip-softmax path: without a + # timestep we cannot tell which phase we are in, so run the + # sparse kernel rather than silently forcing dense forever. + # This degrades quality rather than raising, so say so once. + logger.warning_once( + "SolAttnAttentionConfig.disabled_until_timestep=" + f"{self.disabled_until_timestep} is set, but no `timestep` reached " + "the Sol-Attn forward call. The dense prefix it requests will not " + "be applied. Ensure the pipeline passes a normalized timestep, or " + "unset disabled_until_timestep.", + key="sol_attn_missing_timestep", + ) + else: + dense_by_step = phase == 0 + if dense_by_layer or dense_by_step: + return torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) + ).transpose(1, 2) + return _sol_attn_run( + q, + k, + v, + tau=self.tau, + thresh_type=self.thresh_type, + kv_splits=self.kv_splits, + ) + + @classmethod + def support_lse(cls) -> bool: + return False + + @property + def preferred_layout(self) -> AttentionTensorLayout: + return AttentionTensorLayout.NHD + + @classmethod + def support_fused_qkv(cls) -> bool: + return False diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py index 4cdb62d28a94..0be1d6299a0d 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py @@ -145,14 +145,17 @@ def create_attention( ) kwargs["attention_metadata_state"] = attention_metadata_state if backend.upper() == "CUTEDSL" and attention_config is not None: - if ( - attention_config.sparse_attention_config is not None - and getattr(attention_config.sparse_attention_config, "algorithm", None) == "vsa" - ): + sparse_algo = getattr(attention_config.sparse_attention_config, "algorithm", None) + if sparse_algo == "vsa": from .cute_dsl.vsa import VSAAttention attn_cls = VSAAttention kwargs["sparse_attention_config"] = attention_config.sparse_attention_config + elif sparse_algo == "sol_attn": + from .cute_dsl.sol_attn import SolAttnAttention + + attn_cls = SolAttnAttention + kwargs["sparse_attention_config"] = attention_config.sparse_attention_config return attn_cls( layer_idx=layer_idx, diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md new file mode 100644 index 000000000000..6d78d739802a --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md @@ -0,0 +1,57 @@ +# Third-party notices + +This package is vendored from +[`github.com/NVlabs/Sana`](https://github.com/NVlabs/Sana), branch +[`sol-engine`](https://github.com/NVlabs/Sana/tree/sol-engine), at commit +[`5fe5feb`](https://github.com/NVlabs/Sana/commit/5fe5feb) (2026-08-17; +best-effort reconstruction from vendoring-date file timestamps and upstream +commit history, not an exact recorded pin from the original port -- see the +pull request for how this was verified). Checked against the current branch +tip ([`83e54df`](https://github.com/NVlabs/Sana/commit/83e54df), 2026-08-20) +on 2026-08-27; the only upstream change since that touches this subset is a +merge whose SM89 work is out of scope here. The rest of that merge (MPS/Metal +Apple Silicon backend, RTX 4090/5090 configs) is likewise out of scope for +this CUDA/Blackwell-only subset. + +**Note for future currency checks.** These files are linted and formatted to +this repository's style (`ruff check` and `ruff format`, line length 100) +rather than kept byte-identical to upstream, so a direct `diff` against +upstream shows formatting noise as well as real changes. Upstream wraps at +roughly 80 columns; most of the difference is expressions joined onto one +line. To compare semantics, run `ruff format` over the upstream copy first and +diff the normalized results -- that is how the currency check above was done. + +## Scope of the vendored subset + +Only the pieces needed for the architectures TensorRT-LLM ships are carried: + +| Carried | Not carried | +|---|---| +| `interface.py`, `preprocess.py`, `common/` | `sm89/`, `sm90/` (incl. `sm90/_compat/`) | +| `sm100/` (B200 / GB200) | `triton_ref/` Triton reference attention | +| `sm120/` (RTX Blackwell) | `_vendor/flash_attn/` (see below) | + +The upstream package vendored a copy of FlashAttention's CuTe DSL helpers +under `sol_attn/_vendor/flash_attn/cute/`. That copy is **not** carried here: +TensorRT-LLM already depends on +[`flash-attn-4`](https://github.com/Dao-AILab/flash-attention) (pinned in +`requirements.txt`), which provides the same `flash_attn.cute` modules, and +the SM100/SM120 kernels import them from that dependency directly. This was +verified on B200 to produce bit-identical output to the vendored copy across a +shape/tau sweep. FlashAttention's BSD-3-Clause license is retained at +`sol_attn/sm100/LICENSE.flash-attention` because portions of the SM100 design +scaffold still derive from that project. + +`preprocess.py` implements the routing/threshold stage in Triton, so Triton is +a required runtime dependency on every Sol-Attn path, not only a fallback. + +The runtime also depends on NVIDIA CUTLASS / CuTe DSL, cuda-python, and +PyTorch. Those dependencies are not redistributed by this repository and +remain subject to their respective licenses. + +The SM120 warp-MMA/TMA execution skeleton and online-softmax helpers are +adapted from +[NVIDIA cuDNN Frontend's block-sparse-attention reference](https://github.com/NVIDIA/cudnn-frontend/tree/74785165de2da954a2c879a5e3e6f95411c2292d) +at commit `74785165de2da954a2c879a5e3e6f95411c2292d`. That source is +licensed under the Apache License 2.0; adapted files retain the +corresponding SPDX header. diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/__init__.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/__init__.py new file mode 100644 index 000000000000..078ee8402512 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/__init__.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# +# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see +# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. +"""Sol-Attn.""" + +from .interface import get_sol_attn_backend, sol_attn + +__all__ = ["get_sol_attn_backend", "sol_attn"] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/__init__.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/__init__.py new file mode 100644 index 000000000000..0c0fb62abe10 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/__init__.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# +# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see +# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. +"""Internal helpers shared by the architecture backends.""" + +from .runtime import to_cute_tensor + +__all__ = ["to_cute_tensor"] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/layout_utils.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/layout_utils.py new file mode 100644 index 000000000000..870efa6835bf --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/layout_utils.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# +# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see +# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. +"""Tensor-layout helpers shared by the two CuTe kernels.""" + +import cutlass.cute as cute +from cutlass import const_expr + + +def transpose_view(tensor: cute.Tensor) -> cute.Tensor: + shape = (tensor.shape[1], tensor.shape[0], *tensor.shape[2:]) + order = (1, 0, *range(2, cute.rank(tensor))) + return cute.composition( + tensor, + cute.make_ordered_layout(shape, order=order), + ) + + +def select(tensor: cute.Tensor, modes: list[int]) -> cute.Tensor: + return cute.make_tensor( + tensor.iterator, + cute.select(tensor.layout, modes), + ) + + +def _accumulator_mn_layout( + layout: cute.Layout, + transpose: bool = False, +) -> cute.Layout: + column_major = cute.make_layout(layout.shape) + shape = ( + (column_major.shape[0][1], column_major.shape[1]), + ( + column_major.shape[0][0], + *column_major.shape[0][2:], + column_major.shape[2], + ), + *column_major.shape[3:], + ) + stride = ( + (column_major.stride[0][1], column_major.stride[1]), + ( + column_major.stride[0][0], + *column_major.stride[0][2:], + column_major.stride[2], + ), + *column_major.stride[3:], + ) + if const_expr(transpose): + shape = (shape[1], shape[0], *shape[2:]) + stride = (stride[1], stride[0], *stride[2:]) + return cute.composition( + layout, + cute.make_layout(shape, stride=stride), + ) + + +def reshape_acc_to_mn( + accumulator: cute.Tensor, + transpose: bool = False, +) -> cute.Tensor: + return cute.make_tensor( + accumulator.iterator, + _accumulator_mn_layout(accumulator.layout, transpose), + ) + + +@cute.jit +def _accumulator_frga_layout(layout: cute.Layout) -> cute.Layout: + if const_expr(cute.rank(layout.shape[0]) == 3): + divisor = 2 if const_expr(layout.shape[0][2] % 2 == 0) else 1 + divided = cute.logical_divide( + layout, + ((None, None, divisor), None, None), + ) + return cute.make_layout( + ( + ( + divided.shape[0][0], + divided.shape[0][1], + divided.shape[0][2][0], + ), + divided.shape[1], + (divided.shape[0][2][1], divided.shape[2]), + ), + stride=( + ( + divided.stride[0][0], + divided.stride[0][1], + divided.stride[0][2][0], + ), + divided.stride[1], + (divided.stride[0][2][1], divided.stride[2]), + ), + ) + + assert layout.shape[2] % 2 == 0 + divided = cute.logical_divide(layout, (None, None, 2)) + return cute.make_layout( + ( + ( + divided.shape[0][0], + divided.shape[0][1], + divided.shape[2][0], + ), + divided.shape[1], + divided.shape[2][1], + ), + stride=( + ( + divided.stride[0][0], + divided.stride[0][1], + divided.stride[2][0], + ), + divided.stride[1], + divided.stride[2][1], + ), + ) + + +def reshape_acc_to_frgA(accumulator: cute.Tensor) -> cute.Tensor: + return cute.make_tensor( + accumulator.iterator, + _accumulator_frga_layout(accumulator.layout), + ) + + +__all__ = [ + "reshape_acc_to_frgA", + "reshape_acc_to_mn", + "select", + "transpose_view", +] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/runtime.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/runtime.py new file mode 100644 index 000000000000..502b6cb6a468 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/runtime.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# +# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see +# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. +"""Small host helpers shared by the architecture backends.""" + +from cutlass.cute.runtime import from_dlpack + + +def to_cute_tensor(tensor): + return from_dlpack( + tensor, + assumed_align=16, + enable_tvm_ffi=True, + ).mark_layout_dynamic(leading_dim=tensor.ndim - 1) + + +__all__ = ["to_cute_tensor"] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/selector.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/selector.py new file mode 100644 index 000000000000..13ab6a23bc73 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/selector.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# +# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see +# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. +"""CTA-local routing-mask helpers shared by the CuTe architecture backends.""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, const_expr +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + + +@dsl_user_op +def sol_attn_bfind_b32( + value: Int32, + *, + loc=None, + ip=None, +) -> Int32: + return Int32( + llvm.inline_asm( + T.i32(), + [Int32(value).ir_value(loc=loc, ip=ip)], + "bfind.u32 $0, $1;", + "=r,r", + has_side_effects=False, + is_align_stack=False, + ) + ) + + +@dsl_user_op +def sol_attn_popc_b32( + value: Int32, + *, + loc=None, + ip=None, +) -> Int32: + return Int32( + llvm.inline_asm( + T.i32(), + [Int32(value).ir_value(loc=loc, ip=ip)], + "popc.b32 $0, $1;", + "=r,r", + has_side_effects=False, + is_align_stack=False, + ) + ) + + +@cute.jit +def _mask_word( + mask0: Int32, + mask1: Int32, + mask2: Int32, + mask3: Int32, + word: Int32, +) -> Int32: + result = mask0 + if word == Int32(1): + result = mask1 + if word == Int32(2): + result = mask2 + if word == Int32(3): + result = mask3 + return result + + +@cute.jit +def _test_exact_bit( + mask0: Int32, + mask1: Int32, + mask2: Int32, + mask3: Int32, + offset: Int32, +) -> cutlass.Boolean: + word = offset // Int32(32) + bit = offset - word * Int32(32) + return (_mask_word(mask0, mask1, mask2, mask3, word) & (Int32(1) << bit)) != Int32(0) + + +@cute.jit +def sol_attn_test_exact_bit_limited_words( + mask0: Int32, + mask1: Int32, + mask2: Int32, + mask3: Int32, + offset: Int32, + group_words: cutlass.Constexpr[int], +) -> cutlass.Boolean: + bit = offset & Int32(31) + if const_expr(group_words == 1): + return (mask0 & (Int32(1) << bit)) != Int32(0) + if const_expr(group_words == 2): + word = mask0 + if offset >= Int32(32): + word = mask1 + return (word & (Int32(1) << bit)) != Int32(0) + if const_expr(group_words == 3): + index = offset // Int32(32) + word = mask0 + if index == Int32(1): + word = mask1 + if index == Int32(2): + word = mask2 + return (word & (Int32(1) << bit)) != Int32(0) + return _test_exact_bit(mask0, mask1, mask2, mask3, offset) + + +@cute.jit +def sol_attn_set_exact_bit( + mask0: Int32, + mask1: Int32, + mask2: Int32, + mask3: Int32, + offset: Int32, +): + word = offset // Int32(32) + bit_value = Int32(1) << (offset - word * Int32(32)) + if word == Int32(0): + mask0 = mask0 | bit_value + if word == Int32(1): + mask1 = mask1 | bit_value + if word == Int32(2): + mask2 = mask2 | bit_value + if word == Int32(3): + mask3 = mask3 | bit_value + return mask0, mask1, mask2, mask3 + + +@cute.jit +def sol_attn_route_is_exact( + q_block: Int32, + kv_block: Int32, + column_mean: Float32, + threshold: Float32, + valid: cutlass.Boolean, +) -> cutlass.Boolean: + distance = q_block - kv_block + if distance < Int32(0): + distance = Int32(0) - distance + return ((column_mean > threshold) or distance <= Int32(1)) and valid + + +@cute.jit +def sol_attn_mask_word_constexpr( + mask0: Int32, + mask1: Int32, + mask2: Int32, + mask3: Int32, + word: cutlass.Constexpr[int], +) -> Int32: + if const_expr(word == 0): + return mask0 + if const_expr(word == 1): + return mask1 + if const_expr(word == 2): + return mask2 + return mask3 + + +__all__ = [ + "sol_attn_bfind_b32", + "sol_attn_mask_word_constexpr", + "sol_attn_popc_b32", + "sol_attn_route_is_exact", + "sol_attn_set_exact_bit", + "sol_attn_test_exact_bit_limited_words", +] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/interface.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/interface.py new file mode 100644 index 000000000000..e136ba93cea3 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/interface.py @@ -0,0 +1,334 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# +# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see +# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. +"""Public Sol-Attn interface.""" + +from __future__ import annotations + +import functools + +import torch + +BLOCK_SIZE = 64 +_CUTE_BACKENDS = { + (10, 0): "cute_sm100", # B200 / GB200 + (12, 0): "cute_sm120", # RTX Pro Blackwell / GeForce Blackwell +} +_compiled = {} + + +def _validate_inputs( + q, + k, + v, + thresh_type, + sink_tokens=0, + sink_start=None, +): + if q.ndim != 4 or q.shape != k.shape or q.shape != v.shape: + raise ValueError("q, k, and v must share shape [B, T, H, 128]") + if q.shape[1] == 0 or q.shape[3] != 128: + raise ValueError("Sol-Attn requires T > 0 and head dimension 128") + if any(x.dtype != torch.bfloat16 for x in (q, k, v)): + raise TypeError("q, k, and v must use torch.bfloat16") + if q.device.type != "cuda" or k.device != q.device or v.device != q.device: + raise ValueError("q, k, and v must be on the same CUDA device") + if not (q.is_contiguous() and k.is_contiguous() and v.is_contiguous()): + raise ValueError("q, k, and v must be contiguous BTHD tensors") + if thresh_type not in ("diag", "exact"): + raise ValueError("thresh_type must be 'diag' or 'exact'") + if not isinstance(sink_tokens, int): + raise TypeError("sink_tokens must be an integer") + if not 0 <= sink_tokens <= q.shape[1]: + raise ValueError("sink_tokens must be in [0, T]") + if sink_start is not None: + if not isinstance(sink_start, int): + raise TypeError("sink_start must be an integer or None") + if not 0 <= sink_start <= q.shape[1]: + raise ValueError("sink_start must be in [0, T]") + if sink_start + sink_tokens > q.shape[1]: + raise ValueError("sink_start + sink_tokens must be <= T") + + return tuple(torch.cuda.get_device_capability(q.device)) + + +@functools.lru_cache(maxsize=1) +def _cute_runtime_available() -> bool: + """Whether the optional CuTe DSL runtime can be imported.""" + + try: + import cuda.bindings.driver # noqa: F401 + import cutlass.cute # noqa: F401 + except ImportError: + return False + return True + + +def _backend_for_arch( + arch: tuple[int, int], + *, + cute_available: bool | None = None, +) -> str: + """Select the CuTe kernel for ``arch``, or raise if there isn't one. + + Unsupported architectures raise rather than silently degrading: the caller + (``_run_sol_attn_bthd``) turns that into an explicit dense-SDPA fallback + with a warning, so a missing kernel is visible instead of showing up only + as absent speedup. + """ + + cute_backend = _CUTE_BACKENDS.get(arch) + if cute_backend is None: + raise RuntimeError( + f"Sol-Attn has no kernel for SM{arch[0]}{arch[1]}; supported " + f"architectures are " + f"{', '.join(f'SM{a}{b}' for a, b in sorted(_CUTE_BACKENDS))}." + ) + available = _cute_runtime_available() if cute_available is None else cute_available + if not available: + raise RuntimeError( + "Sol-Attn requires the CuTe DSL runtime (cutlass.cute and " + "cuda.bindings.driver); neither could be imported." + ) + return cute_backend + + +def get_sol_attn_backend(device: torch.device | str | int | None = None) -> str: + """Return the backend selected for ``device`` without compiling it.""" + + if device is None: + device = torch.cuda.current_device() + return _backend_for_arch(tuple(torch.cuda.get_device_capability(device))) + + +def _validate_cute(arch, tokens, kv_splits): + if kv_splits != 1: + raise ValueError( + "kv_splits=2/4 was an SM90-only path; this build ships SM100/SM120 " + "kernels only, so kv_splits must be 1." + ) + route_groups = ((tokens + 63) // 64 + 63) // 64 + if kv_splits > route_groups: + raise ValueError("each KV split must contain at least one N64 route group") + + +def _stream(device): + import cuda.bindings.driver as cuda + + return cuda.CUstream(torch.cuda.current_stream(device).cuda_stream) + + +def _to_cute_tensors(tensors): + from .common import to_cute_tensor + + return [to_cute_tensor(x) for x in tensors] + + +def _sink_block_range(tokens, sink_start, sink_tokens): + blocks = (tokens + BLOCK_SIZE - 1) // BLOCK_SIZE + if not sink_tokens: + return blocks, blocks + start = tokens - sink_tokens if sink_start is None else sink_start + return ( + start // BLOCK_SIZE, + (start + sink_tokens + BLOCK_SIZE - 1) // BLOCK_SIZE, + ) + + +def _compile_sm100( + key, + tensors, + scale, + sink_start_block, + sink_end_block, + stream, +): + import cutlass.cute as cute + + from .sm100 import forward + + args = _to_cute_tensors(tensors) + compiled = cute.compile( + forward, + *args, + scale, + sink_start_block, + sink_end_block, + stream=stream, + options="--enable-tvm-ffi", + ) + _compiled[key] = compiled + return compiled, args + + +def _compile_sm120( + key, + tensors, + scale, + sink_start_block, + sink_end_block, + stream, +): + import cutlass.cute as cute + + from .sm120 import make_kernel + + operator = make_kernel() + args = _to_cute_tensors(tensors) + compiled = cute.compile( + operator, + *args, + scale, + sink_start_block, + sink_end_block, + stream=stream, + options="--enable-tvm-ffi", + ) + _compiled[key] = compiled + return compiled, args + + +def _sol_attn_cute( + q, + k, + v, + *, + arch, + scale, + tau, + thresh_type, + kv_splits, + sink_tokens, + sink_start, +): + from .preprocess import prepare + + batch, tokens, heads, _ = q.shape + + with torch.cuda.device(q.device): + kc, vc, threshold = prepare( + q, + k, + v, + scale=scale, + tau=tau, + thresh_type=thresh_type, + ) + output = torch.empty_like(v) + lse = torch.empty( + (batch, tokens, heads), + device=q.device, + dtype=torch.float32, + ) + stream = _stream(q.device) + key = (q.device.index, arch, batch, tokens, heads, kv_splits) + + if arch == (10, 0): + sink_start_block, sink_end_block = _sink_block_range( + tokens, + sink_start, + sink_tokens, + ) + tensors = [q, k, v, output, kc, vc, threshold, lse] + compiled = _compiled.get(key) + if compiled is None: + compiled, args = _compile_sm100( + key, + tensors, + scale, + sink_start_block, + sink_end_block, + stream, + ) + else: + args = _to_cute_tensors(tensors) + compiled( + *args, + scale, + sink_start_block, + sink_end_block, + stream=stream, + ) + else: + sink_start_block, sink_end_block = _sink_block_range( + tokens, + sink_start, + sink_tokens, + ) + tensors = [q, k, v, output, kc, vc, threshold, lse] + compiled = _compiled.get(key) + if compiled is None: + compiled, args = _compile_sm120( + key, + tensors, + scale, + sink_start_block, + sink_end_block, + stream, + ) + else: + args = _to_cute_tensors(tensors) + compiled( + *args, + scale, + sink_start_block, + sink_end_block, + stream=stream, + ) + return output + + +def sol_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + scale: float | None = None, + tau: float = 1.0, + thresh_type: str = "diag", + kv_splits: int = 1, + sink_tokens: int = 0, + sink_start: int | None = None, +) -> torch.Tensor: + """Compute noncausal Sol-Attn for contiguous BF16 BTHD tensors. + + ``sink_start`` and ``sink_tokens`` keep every KV block overlapping the + corresponding contiguous token range exact for all queries. Omitting + ``sink_start`` places the range at the token suffix. + """ + + arch = _validate_inputs( + q, + k, + v, + thresh_type, + sink_tokens, + sink_start, + ) + if kv_splits != 1: + raise ValueError( + "kv_splits must be 1; the 2/4 path was SM90-only and this build " + "ships SM100/SM120 kernels only." + ) + _backend_for_arch(arch) # raises on an architecture with no kernel + scale = q.shape[-1] ** -0.5 if scale is None else float(scale) + tau = float(tau) + + _validate_cute(arch, q.shape[1], kv_splits) + return _sol_attn_cute( + q, + k, + v, + arch=arch, + scale=scale, + tau=tau, + thresh_type=thresh_type, + kv_splits=kv_splits, + sink_tokens=sink_tokens, + sink_start=sink_start, + ) + + +__all__ = ["get_sol_attn_backend", "sol_attn"] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/preprocess.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/preprocess.py new file mode 100644 index 000000000000..5c44a3d070cc --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/preprocess.py @@ -0,0 +1,454 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# +# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see +# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. +"""Block summaries and routing thresholds shared by both CuTe kernels.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl +from triton.tools.tensor_descriptor import TensorDescriptor + +BLOCK_SIZE = 64 +HEAD_DIM = 128 +THRESHOLD_GROUP_SIZE = 64 + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=warps, num_stages=stages) + for warps in (4, 8) + for stages in (1, 2, 3, 4) + ], + key=["T"], +) +@triton.jit +def _reduce_kc_kernel( + k_desc, + kc, + T, + H: tl.constexpr, + N: tl.constexpr, + D: tl.constexpr, + BLOCK: tl.constexpr, + TILE_D: tl.constexpr, +): + d_tile, block, batch_head = ( + tl.program_id(0), + tl.program_id(1), + tl.program_id(2), + ) + batch, head = batch_head // H, batch_head % H + block_len = tl.minimum(BLOCK, T - block * BLOCK) + values = k_desc.load([batch, block * BLOCK, head, d_tile * TILE_D]).reshape([BLOCK, TILE_D]) + summary = tl.sum(values, axis=0) / block_len + offsets = d_tile * TILE_D + tl.arange(0, TILE_D) + tl.store( + kc + ((batch * N + block) * H + head) * D + offsets, + summary, + mask=offsets < D, + ) + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=warps, num_stages=stages) + for warps in (4, 8) + for stages in (1, 2, 3, 4) + ], + key=["T"], +) +@triton.jit +def _reduce_vc_kernel( + v_desc, + vc, + T, + H: tl.constexpr, + N: tl.constexpr, + D: tl.constexpr, + BLOCK: tl.constexpr, + TILE_D: tl.constexpr, +): + d_tile, block, batch_head = ( + tl.program_id(0), + tl.program_id(1), + tl.program_id(2), + ) + batch, head = batch_head // H, batch_head % H + values = v_desc.load([batch, block * BLOCK, head, d_tile * TILE_D]).reshape([BLOCK, TILE_D]) + summary = tl.sum(values, axis=0) + offsets = d_tile * TILE_D + tl.arange(0, TILE_D) + tl.store( + vc + ((batch * N + block) * H + head) * D + offsets, + summary, + mask=offsets < D, + ) + + +@triton.autotune( + configs=[triton.Config({}, num_warps=4, num_stages=2)], + key=["N"], +) +@triton.jit +def _reduce_kc_stats_kernel( + kc_desc, + kc_mean, + kc_var_diag, + H: tl.constexpr, + N: tl.constexpr, + D: tl.constexpr, + TILE_D: tl.constexpr, + GROUP: tl.constexpr, +): + d_tile, batch_head = tl.program_id(0), tl.program_id(1) + batch, head = batch_head // H, batch_head % H + block_offsets = tl.arange(0, GROUP) + block_offsets = tl.max_contiguous(block_offsets, GROUP) + d_offsets = d_tile * TILE_D + tl.arange(0, TILE_D) + total = tl.zeros((TILE_D,), dtype=tl.float32) + total_sq = tl.zeros((TILE_D,), dtype=tl.float32) + count = tl.full((), 0.0, dtype=tl.float32) + for start in range(0, N, GROUP): + valid = start + block_offsets < N + values = ( + kc_desc.load([batch, start, head, d_tile * TILE_D]) + .reshape([GROUP, TILE_D]) + .to(tl.float32) + ) + values = tl.where(valid[:, None], values, 0.0) + total += tl.sum(values, axis=0) + total_sq += tl.sum(values * values, axis=0) + count += tl.sum(valid.to(tl.float32), axis=0) + mean = total / count + variance = tl.maximum(total_sq / count - mean * mean, 0.0) + valid_d = d_offsets < D + tl.store( + kc_mean + batch_head * D + d_offsets, + mean, + mask=valid_d, + ) + tl.store( + kc_var_diag + batch_head * D + d_offsets, + variance, + mask=valid_d, + ) + + +@triton.autotune( + configs=[triton.Config({}, num_warps=4, num_stages=2)], + key=["T"], +) +@triton.jit +def _diag_threshold_kernel( + q_desc, + kc_mean, + kc_var_diag, + global_threshold, + softmax_scale, + T, + H: tl.constexpr, + N: tl.constexpr, + D: tl.constexpr, + BLOCK: tl.constexpr, + TILE_D: tl.constexpr, + TAU: tl.constexpr, +): + q_block, batch_head = tl.program_id(0), tl.program_id(1) + batch, head = batch_head // H, batch_head % H + q_start = q_block * BLOCK + q_len = tl.minimum(BLOCK, T - q_start).to(tl.float32) + d_offsets = tl.arange(0, TILE_D) + valid_d = d_offsets < D + q_values = q_desc.load([batch, q_start, head, 0]).reshape([BLOCK, TILE_D]) + q_centroid = tl.sum(q_values.to(tl.float32), axis=0) / q_len + mean_kc = tl.load( + kc_mean + batch_head * D + d_offsets, + mask=valid_d, + other=0.0, + ) + var_kc = tl.load( + kc_var_diag + batch_head * D + d_offsets, + mask=valid_d, + other=0.0, + ) + log2_scale = softmax_scale * 1.4426950408889634 + mean = tl.sum(q_centroid * mean_kc, axis=0) * log2_scale + variance = tl.sum(q_centroid * q_centroid * var_kc, axis=0) * (log2_scale * log2_scale) + std = tl.sqrt(tl.maximum(variance, 0.0) + 1.0e-6) + tl.store( + global_threshold + (batch * N + q_block) * H + head, + mean + TAU * std, + ) + + +@triton.jit +def _pool_query_kernel( + q_desc, + q_bar, + T, + H: tl.constexpr, + N: tl.constexpr, + D: tl.constexpr, + BLOCK: tl.constexpr, + TILE_D: tl.constexpr, +): + q_block, batch_head = tl.program_id(0), tl.program_id(1) + batch, head = batch_head // H, batch_head % H + q_start = q_block * BLOCK + q_len = tl.minimum(BLOCK, T - q_start).to(tl.float32) + offsets = tl.arange(0, TILE_D) + values = q_desc.load([batch, q_start, head, 0]).reshape([BLOCK, TILE_D]) + centroid = tl.sum(values.to(tl.float32), axis=0) / q_len + tl.store( + q_bar + (batch_head * N + q_block) * D + offsets, + centroid, + mask=offsets < D, + ) + + +@triton.jit +def _exact_fused_threshold_kernel( + q_bar, + kc_mean, + kc_second_moment, + global_threshold, + softmax_scale, + H: tl.constexpr, + N: tl.constexpr, + D: tl.constexpr, + BLOCK_M: tl.constexpr, + TILE_D: tl.constexpr, + TAU: tl.constexpr, +): + row_tile, batch_head = tl.program_id(0), tl.program_id(1) + rows = row_tile * BLOCK_M + tl.arange(0, BLOCK_M) + offsets = tl.arange(0, TILE_D) + valid_rows = rows < N + valid_d = offsets < D + + q_centroid = tl.load( + q_bar + (batch_head * N + rows[:, None]) * D + offsets[None, :], + mask=valid_rows[:, None] & valid_d[None, :], + other=0.0, + ) + mean_kc = tl.load( + kc_mean + batch_head * D + offsets, + mask=valid_d, + other=0.0, + ) + second_moment = tl.load( + kc_second_moment + batch_head * D * D + offsets[:, None] * D + offsets[None, :], + mask=valid_d[:, None] & valid_d[None, :], + other=0.0, + ) + + raw_mean = tl.sum(q_centroid.to(tl.float32) * mean_kc[None, :], axis=1) + projected = tl.dot( + q_centroid, + second_moment, + out_dtype=tl.float32, + ) + raw_second_moment = tl.sum( + projected * q_centroid.to(tl.float32), + axis=1, + ) + log2_scale = softmax_scale * 1.4426950408889634 + mean = raw_mean * log2_scale + variance = tl.maximum( + raw_second_moment - raw_mean * raw_mean, + 0.0, + ) * (log2_scale * log2_scale) + threshold = mean + TAU * tl.sqrt(variance + 1.0e-6) + batch, head = batch_head // H, batch_head % H + tl.store( + global_threshold + (batch * N + rows) * H + head, + threshold, + mask=valid_rows, + ) + + +def _reduce_kv( + k: torch.Tensor, + v: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + batch, tokens, heads, head_dim = k.shape + blocks = triton.cdiv(tokens, BLOCK_SIZE) + tile_d = min(128, triton.next_power_of_2(head_dim)) + kc = torch.empty( + (batch, blocks, heads, head_dim), + device=k.device, + dtype=torch.bfloat16, + ) + vc = torch.empty_like(kc) + k_desc = TensorDescriptor.from_tensor( + k, + [1, BLOCK_SIZE, 1, tile_d], + ) + v_desc = TensorDescriptor.from_tensor( + v, + [1, BLOCK_SIZE, 1, tile_d], + ) + grid = (triton.cdiv(head_dim, tile_d), blocks, batch * heads) + _reduce_kc_kernel[grid]( + k_desc, + kc, + tokens, + heads, + blocks, + head_dim, + BLOCK_SIZE, + tile_d, + ) + _reduce_vc_kernel[grid]( + v_desc, + vc, + tokens, + heads, + blocks, + head_dim, + BLOCK_SIZE, + tile_d, + ) + return kc, vc + + +def _compute_diag_threshold( + q: torch.Tensor, + kc: torch.Tensor, + *, + tau: float, + scale: float, +) -> torch.Tensor: + batch, tokens, heads, head_dim = q.shape + blocks = triton.cdiv(tokens, BLOCK_SIZE) + tile_d = min(128, triton.next_power_of_2(head_dim)) + kc_mean = torch.empty( + (batch, heads, head_dim), + device=q.device, + dtype=torch.float32, + ) + kc_var_diag = torch.empty_like(kc_mean) + global_threshold = torch.empty( + (batch, blocks, heads), + device=q.device, + dtype=torch.float32, + ) + q_desc = TensorDescriptor.from_tensor( + q, + [1, BLOCK_SIZE, 1, tile_d], + ) + kc_desc = TensorDescriptor.from_tensor( + kc, + [1, THRESHOLD_GROUP_SIZE, 1, tile_d], + ) + _reduce_kc_stats_kernel[(triton.cdiv(head_dim, tile_d), batch * heads)]( + kc_desc, + kc_mean, + kc_var_diag, + heads, + blocks, + head_dim, + tile_d, + THRESHOLD_GROUP_SIZE, + ) + _diag_threshold_kernel[(blocks, batch * heads)]( + q_desc, + kc_mean, + kc_var_diag, + global_threshold, + scale, + tokens, + heads, + blocks, + head_dim, + BLOCK_SIZE, + tile_d, + tau, + ) + return global_threshold + + +def _compute_exact_threshold( + q: torch.Tensor, + kc: torch.Tensor, + *, + tau: float, + scale: float, +) -> torch.Tensor: + batch, tokens, heads, head_dim = q.shape + blocks = triton.cdiv(tokens, BLOCK_SIZE) + tile_d = min(128, triton.next_power_of_2(head_dim)) + batch_heads = batch * heads + kc_bh = kc.permute(0, 2, 1, 3) + kc_mean = kc_bh.mean(dim=2, dtype=torch.float32) + kc_second_moment = torch.matmul( + kc_bh.transpose(-1, -2), + kc_bh, + ) + kc_second_moment.div_(blocks) + q_bar = torch.empty( + (batch_heads, blocks, head_dim), + device=q.device, + dtype=torch.bfloat16, + ) + global_threshold = torch.empty( + (batch, blocks, heads), + device=q.device, + dtype=torch.float32, + ) + q_desc = TensorDescriptor.from_tensor( + q, + [1, BLOCK_SIZE, 1, tile_d], + ) + _pool_query_kernel[(blocks, batch_heads)]( + q_desc, + q_bar, + tokens, + heads, + blocks, + head_dim, + BLOCK_SIZE, + tile_d, + num_warps=4, + num_stages=1, + ) + block_m = 64 + _exact_fused_threshold_kernel[(triton.cdiv(blocks, block_m), batch_heads)]( + q_bar, + kc_mean, + kc_second_moment, + global_threshold, + scale, + heads, + blocks, + head_dim, + block_m, + tile_d, + tau, + num_warps=4, + num_stages=1, + ) + return global_threshold + + +def prepare( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + tau: float, + scale: float, + thresh_type: str = "diag", +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + kc, vc = _reduce_kv(k, v) + if thresh_type == "exact": + threshold = _compute_exact_threshold(q, kc, tau=tau, scale=scale) + else: + threshold = _compute_diag_threshold(q, kc, tau=tau, scale=scale) + return kc, vc, threshold + + +__all__ = ["prepare"] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/LICENSE.flash-attention b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/LICENSE.flash-attention new file mode 100644 index 000000000000..5860e4b33f3d --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/LICENSE.flash-attention @@ -0,0 +1,29 @@ +BSD 3-Clause License + +Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +* Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/__init__.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/__init__.py new file mode 100644 index 000000000000..81efe5d1ea4d --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/__init__.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# +# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see +# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. +"""Blackwell backend.""" + +from .kernel import forward + +__all__ = ["forward"] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/kernel.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/kernel.py new file mode 100644 index 000000000000..55067dd109c3 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/kernel.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# +# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see +# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. +"""Blackwell kernel entry.""" + +from .mainloop import forward + +__all__ = ["forward"] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/mainloop.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/mainloop.py new file mode 100644 index 000000000000..099e0b375846 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/mainloop.py @@ -0,0 +1,1530 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# +# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see +# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. +# +# Portions derive from the FlashAttention project +# (https://github.com/Dao-AILab/flash-attention), BSD-3-Clause; its license +# text is vendored at sol_attn/sm100/LICENSE.flash-attention. +"""Sol-Attn forward kernel for Blackwell SM100. + +The kernel routes two physical N64 halves at a time and accumulates their exact +indices into one logical G256 stream. Per-column additive masks are built once +in shared memory and reused by the approximate and exact score paths. +""" + +import math + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import flash_attn.cute.pipeline as fa_pipeline +import flash_attn.cute.utils as fa_utils +from cutlass import BFloat16, Float32, Int32 +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cutlass_dsl import T, dsl_user_op +from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned + +from ..common import layout_utils +from ..common.selector import sol_attn_popc_b32, sol_attn_route_is_exact +from . import math as mma_utils +from .softmax import _load_m64_n128_score as _load_pair_score +from .softmax import _online_update_one_half as _online_update_pair +from .softmax import _rescale_m64_partial_o as _rescale_pair_o +from .tmem import ( + _add_physical_tmem_base, + _zero_based_tmem_tensor, + load_m64_o_fp32_256b, + tcgen05_wait_st, +) + +M = 64 +N_MEMBER = 64 +N_PACK_HALF = 128 +D = 128 +DV = 128 +THREADS = 192 +PAIR_STAGES = 1 +TMEM_COLS = 256 +PAIR_SCORE_OFFSET = 0 +PAIR_P_OFFSET = 64 +O_OFFSET = 128 +PACK_QK_INST = (M, N_PACK_HALF, 16) +PACK_QK_TILE = (M, N_PACK_HALF, D) +PACK_PV_INST = (M, DV, 16) +PACK_PV_TILE = (M, DV, N_PACK_HALF) +PACK_QK_QUARTER_INST = (M, N_MEMBER, 16) +PACK_PV_QUARTER_INST = (M, 64, 16) +PACK_QK_GATHER_TILE = (M, N_MEMBER, 64) +PACK_PV_GATHER_TILE = (M, 64, 64) +LOG2E = math.log2(math.e) +LN2 = math.log(2.0) +SEMANTIC_ROW_OFFSET = 16 +LOGICAL_GROUP_SIZE = 256 +ROUTE_TILE_SIZE = 128 +ROUTE_HALVES_PER_GROUP = LOGICAL_GROUP_SIZE // ROUTE_TILE_SIZE +ROUTE_MASK_WORDS = 4 +# masks[0:4], current-half exact count, append base, cumulative exact count, +# logical-terminal-half flag +PACKET_WORDS = 8 +ROUTE_INDEX_CAPACITY = LOGICAL_GROUP_SIZE +PAIR_P_CHUNKS = 4 +PAIR_P_CHUNK_PACKED_COLUMNS = (N_PACK_HALF // 2) // PAIR_P_CHUNKS +PAIR_P_PACKED_REGISTERS_PER_THREAD_PER_CHUNK = 8 +O_PACKED_STORE_VALUES_PER_WORD = 2 +O_PACKED_STORE_ALIGNMENT_BYTES = 4 +O_PACKED_STORE_WRITER_THREADS = 4 * 32 +O_ROWS_PER_OWNER_THREAD = 2 +O_PACKED_WORDS_PER_ROW_PER_THREAD = 16 +O_PACKED_COLUMN_STRIDE = 8 + + +@dsl_user_op +def _cvt_bf16x2_f32( + hi: Float32, + lo: Float32, + *, + loc=None, + ip=None, +) -> Int32: + """Round two FP32 values and pack them as ``{lo, hi}`` BF16 bits.""" + + return Int32( + llvm.inline_asm( + T.i32(), + [ + Float32(hi).ir_value(loc=loc, ip=ip), + Float32(lo).ir_value(loc=loc, ip=ip), + ], + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _store_global_u32_inline( + ptr: cute.Pointer, + value: Int32, + *, + loc=None, + ip=None, +) -> None: + """Store one aligned same-row BF16 pair as a single 32-bit word.""" + + llvm.inline_asm( + None, + [ + ptr.toint().ir_value(), + Int32(value).ir_value(loc=loc, ip=ip), + ], + "st.global.u32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _prmt_b32( + a: Int32, + b: Int32, + sel: Int32, + *, + loc=None, + ip=None, +) -> Int32: + """Select four bytes from packed words ``a`` and ``b``.""" + + return Int32( + llvm.inline_asm( + T.i32(), + [ + Int32(a).ir_value(loc=loc, ip=ip), + Int32(b).ir_value(loc=loc, ip=ip), + Int32(sel).ir_value(loc=loc, ip=ip), + ], + "prmt.b32 $0, $1, $2, $3;", + "=r,r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@cute.jit +def _store_pair_probability_chunked_tmemp( + o_template: cute.Tensor, + probabilities: cute.Tensor, + tmem_base: Int32, + p_offset: Int32, + owner_tidx: Int32, +): + """Store M64xN128 BF16 P as four live-range-bounded x8 chunks. + + Probabilities remain FP32 until each x8 fragment is converted, and every + chunk waits for its St16x64b store before the fragment goes out of scope. + """ + + assert o_template.element_type == Float32 + assert cute.size(o_template) == M * DV + p_chunk_layout = cute.composition( + o_template.layout, + cute.make_layout((M, PAIR_P_CHUNK_PACKED_COLUMNS)), + ) + relative_chunk = _zero_based_tmem_tensor(Float32, p_chunk_layout) + store_atom = cute.make_copy_atom( + tcgen05.copy.St16x64bOp(tcgen05.copy.Repetition(8)), + Float32, + ) + tiled_store = tcgen05.make_tmem_copy(store_atom, relative_chunk) + thread_store = tiled_store.get_slice(owner_tidx) + destination_relative = thread_store.partition_D(relative_chunk) + destination = _add_physical_tmem_base(destination_relative, tmem_base + p_offset) + p_store_coordinates = thread_store.partition_S( + cute.make_identity_tensor((M, PAIR_P_CHUNK_PACKED_COLUMNS)) + ) + lane = owner_tidx % Int32(32) + + for chunk_idx in cutlass.range_constexpr(PAIR_P_CHUNKS): + p_store_registers = cute.make_rmem_tensor(p_store_coordinates.shape, Float32) + assert cute.size(p_store_registers) == PAIR_P_PACKED_REGISTERS_PER_THREAD_PER_CHUNK + assert cute.size(probabilities) == 2 * cute.size(p_store_registers) * PAIR_P_CHUNKS + p_store_words = cute.make_tensor( + cute.recast_ptr(p_store_registers.iterator, dtype=Int32), + p_store_registers.layout, + ) + probability_base = chunk_idx * (2 * cute.size(p_store_registers)) + for i in cutlass.range(cute.size(p_store_registers), unroll_full=True): + low = probability_base + i * 2 + high = low + 1 + own = _cvt_bf16x2_f32( + Float32(probabilities[high]), + Float32(probabilities[low]), + ) + peer = cute.arch.shuffle_sync_bfly(own, offset=2) + if (lane & Int32(2)) == Int32(0): + p_store_words[i] = _prmt_b32(own, peer, Int32(0x5410)) + else: + p_store_words[i] = _prmt_b32(own, peer, Int32(0x3276)) + + destination_chunk = cute.make_tensor( + destination.iterator + chunk_idx * PAIR_P_CHUNK_PACKED_COLUMNS, + destination.layout, + ) + cute.copy(tiled_store, p_store_registers, destination_chunk) + tcgen05_wait_st() + + cute.arch.fence_view_async_tmem_store() + + +@cute.jit +def _load_pack_k_half( + tma_atom_pack_k: cute.CopyAtom, + tPackKgK: cute.Tensor, + tPackKsK: cute.Tensor, + block0: Int32, + block1: Int32, + quarter0: Int32, + barrier, +): + """Gather one canonical N128 K tile as K0/N0,K0/N1,K1/N0,K1/N1.""" + + cute.copy( + tma_atom_pack_k, + tPackKgK[(None, block0, Int32(0))], + tPackKsK[(None, quarter0)], + tma_bar_ptr=barrier, + ) + cute.copy( + tma_atom_pack_k, + tPackKgK[(None, block1, Int32(0))], + tPackKsK[(None, quarter0 + Int32(1))], + tma_bar_ptr=barrier, + ) + cute.copy( + tma_atom_pack_k, + tPackKgK[(None, block0, Int32(1))], + tPackKsK[(None, quarter0 + Int32(2))], + tma_bar_ptr=barrier, + ) + cute.copy( + tma_atom_pack_k, + tPackKgK[(None, block1, Int32(1))], + tPackKsK[(None, quarter0 + Int32(3))], + tma_bar_ptr=barrier, + ) + + +@cute.jit +def _load_pack_v_half( + tma_atom_pack_v: cute.CopyAtom, + tPackVgV: cute.Tensor, + tPackVsV: cute.Tensor, + block0: Int32, + block1: Int32, + quarter0: Int32, + barrier, +): + """Gather one canonical N128 V tile as D0/N0,D0/N1,D1/N0,D1/N1.""" + + cute.copy( + tma_atom_pack_v, + tPackVgV[(None, Int32(0), block0)], + tPackVsV[(None, quarter0)], + tma_bar_ptr=barrier, + ) + cute.copy( + tma_atom_pack_v, + tPackVgV[(None, Int32(0), block1)], + tPackVsV[(None, quarter0 + Int32(1))], + tma_bar_ptr=barrier, + ) + cute.copy( + tma_atom_pack_v, + tPackVgV[(None, Int32(1), block0)], + tPackVsV[(None, quarter0 + Int32(2))], + tma_bar_ptr=barrier, + ) + cute.copy( + tma_atom_pack_v, + tPackVgV[(None, Int32(1), block1)], + tPackVsV[(None, quarter0 + Int32(3))], + tma_bar_ptr=barrier, + ) + + +@cute.struct +class SharedStorage: + q_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2] + pack_k_mbar_ptr: cute.struct.MemRange[cutlass.Int64, PAIR_STAGES * 2] + pack_v_mbar_ptr: cute.struct.MemRange[cutlass.Int64, PAIR_STAGES * 2] + pair_score_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2] + pair_o_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2] + final_stats: cute.struct.Align[cute.struct.MemRange[Float32, M * 2], 128] + route_partial: cute.struct.Align[cute.struct.MemRange[Float32, 4 * ROUTE_TILE_SIZE], 16] + column_masks: cute.struct.Align[cute.struct.MemRange[Float32, ROUTE_TILE_SIZE], 16] + route_packet: cute.struct.Align[cute.struct.MemRange[Int32, PACKET_WORDS], 16] + tmem_holding_buf: Int32 + # Owner-warp 0 lane 0 appends both N128 route masks. The full-CTA + # pre-exact join publishes the completed list to warp 5; no HBM indices. + route_indices: cute.struct.Align[cute.struct.MemRange[Int32, ROUTE_INDEX_CAPACITY], 16] + + +@cute.kernel +def _sol_attn_sm100_bf16_kernel( + tiled_pack_qk: cute.TiledMma, + tiled_pack_pv: cute.TiledMma, + tma_atom_q: cute.CopyAtom, + mQ_mkl: cute.Tensor, + tma_atom_pack_k: cute.CopyAtom, + mPackK_nkl: cute.Tensor, + tma_atom_pack_v: cute.CopyAtom, + mPackV_nkl: cute.Tensor, + tma_atom_kc: cute.CopyAtom, + mKC_nkl: cute.Tensor, + tma_atom_vc: cute.CopyAtom, + mVC_nkl: cute.Tensor, + mThreshold_bnh: cute.Tensor, + mO_bthd: cute.Tensor, + mLSE_bth: cute.Tensor, + token_count: Int32, + route_valid_total: Int32, + num_route_tiles: Int32, + softmax_scale: Float32, + sink_start_block: Int32, + sink_end_block: Int32, + q_layout: cute.ComposedLayout, + pack_k_layout: cute.ComposedLayout, + pack_k_gather_layout: cute.ComposedLayout, + pack_p_layout: cute.ComposedLayout, + pack_v_layout: cute.ComposedLayout, + pack_v_gather_layout: cute.ComposedLayout, + route_k_layout: cute.ComposedLayout, + route_v_layout: cute.ComposedLayout, +): + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + q_block_idx_raw, head_idx_raw, batch_idx_raw = cute.arch.block_idx() + q_block_idx = Int32(q_block_idx_raw) + head_idx = Int32(head_idx_raw) + batch_idx = Int32(batch_idx_raw) + softmax_scale_log2 = softmax_scale * Float32(LOG2E) + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + sFinalStats = storage.final_stats.get_tensor(cute.make_layout((M, 2))) + route_partial = storage.route_partial.get_tensor(cute.make_layout((4, ROUTE_TILE_SIZE))) + column_masks = storage.column_masks.get_tensor(cute.make_layout((ROUTE_TILE_SIZE,))) + route_packet = storage.route_packet.get_tensor(cute.make_layout((PACKET_WORDS,))) + route_indices = storage.route_indices.get_tensor(cute.make_layout((ROUTE_INDEX_CAPACITY,))) + sQ = smem.allocate_tensor( + element_type=BFloat16, + layout=q_layout.outer, + byte_alignment=128, + swizzle=q_layout.inner, + ) + sPackK = smem.allocate_tensor( + element_type=BFloat16, + layout=pack_k_layout.outer, + byte_alignment=128, + swizzle=pack_k_layout.inner, + ) + sPackV = smem.allocate_tensor( + element_type=BFloat16, + layout=pack_v_layout.outer, + byte_alignment=128, + swizzle=pack_v_layout.inner, + ) + # One independent physical N128 K stage and one N128 V stage. Every + # runtime route/exact transaction stays in this completion domain. + sPackKGather = cute.make_tensor( + cute.recast_ptr(sPackK.iterator, pack_k_gather_layout.inner, BFloat16), + pack_k_gather_layout.outer, + ) + sPackVGather = cute.make_tensor( + cute.recast_ptr(sPackV.iterator, pack_v_gather_layout.inner, BFloat16), + pack_v_gather_layout.outer, + ) + # KC/VC and exact K/V have disjoint lifetimes within each runtime group. + # They reuse the same independent N128 K and V allocations without a + # cross-operand alias barrier. + sKC = cute.make_tensor( + cute.recast_ptr(sPackK.iterator, route_k_layout.inner, BFloat16), + route_k_layout.outer, + ) + sVC = cute.make_tensor( + cute.recast_ptr(sPackV.iterator, route_v_layout.inner, BFloat16), + route_v_layout.outer, + ) + + tmem_barrier = pipeline.NamedBarrier(barrier_id=1, num_threads=THREADS) + score_loaded_barrier = pipeline.NamedBarrier(barrier_id=2, num_threads=4 * 32) + final_stats_ready_barrier = pipeline.NamedBarrier(barrier_id=3, num_threads=4 * 32) + pack_score_loaded_barrier = pipeline.NamedBarrier(barrier_id=4, num_threads=4 * 32) + route_packet_ready_barrier = pipeline.NamedBarrier(barrier_id=5, num_threads=5 * 32) + exact_pair_p_ready_barrier = pipeline.NamedBarrier(barrier_id=6, num_threads=5 * 32) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_barrier, + ) + tmem.allocate(TMEM_COLS) + + one_thread = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + pack_owner_threads = pipeline.CooperativeGroup(pipeline.Agent.Thread, 4 * 32) + q_bytes = cute.size_in_bytes(BFloat16, cute.select(q_layout, mode=[0, 1, 2])) + route_k_bytes = cute.size_in_bytes(BFloat16, cute.select(route_k_layout, mode=[0, 1, 2])) + route_v_bytes = cute.size_in_bytes(BFloat16, cute.select(route_v_layout, mode=[0, 1, 2])) + pack_k_bytes = cute.size_in_bytes(BFloat16, cute.select(pack_k_layout, mode=[0, 1, 2])) + pack_v_bytes = cute.size_in_bytes(BFloat16, cute.select(pack_v_layout, mode=[0, 1, 2])) + assert route_k_bytes == pack_k_bytes + assert route_v_bytes == pack_v_bytes + q_pipe = fa_pipeline.PipelineTmaUmma.create( + num_stages=1, + producer_group=one_thread, + consumer_group=one_thread, + tx_count=q_bytes, + barrier_storage=storage.q_mbar_ptr.data_ptr(), + ) + pack_k_pipe = fa_pipeline.PipelineTmaUmma.create( + num_stages=PAIR_STAGES, + producer_group=one_thread, + consumer_group=one_thread, + tx_count=pack_k_bytes, + barrier_storage=storage.pack_k_mbar_ptr.data_ptr(), + ) + pack_v_pipe = fa_pipeline.PipelineTmaUmma.create( + num_stages=PAIR_STAGES, + producer_group=one_thread, + consumer_group=one_thread, + tx_count=pack_v_bytes, + barrier_storage=storage.pack_v_mbar_ptr.data_ptr(), + ) + pair_score_pipe = fa_pipeline.PipelineUmmaAsync.create( + num_stages=1, + producer_group=one_thread, + consumer_group=pack_owner_threads, + barrier_storage=storage.pair_score_mbar_ptr.data_ptr(), + ) + pair_o_pipe = fa_pipeline.PipelineUmmaAsync.create( + num_stages=1, + producer_group=one_thread, + consumer_group=pack_owner_threads, + barrier_storage=storage.pair_o_mbar_ptr.data_ptr(), + ) + + mQ_cur = mQ_mkl[None, None, head_idx, batch_idx] + mPackK_cur = mPackK_nkl[None, None, head_idx, batch_idx] + mPackV_cur = mPackV_nkl[None, None, head_idx, batch_idx] + mKC_cur = mKC_nkl[None, None, head_idx, batch_idx] + mVC_cur = mVC_nkl[None, None, head_idx, batch_idx] + gQ = cute.local_tile(mQ_cur, (M, D), (None, 0)) + gPackK = cute.local_tile(mPackK_cur, (N_MEMBER, 64), (None, None)) + gPackV = cute.local_tile(mPackV_cur, (64, N_MEMBER), (None, None)) + gKC = cute.local_tile(mKC_cur, (N_PACK_HALF, D), (None, 0)) + gVC = cute.local_tile(mVC_cur, (DV, N_PACK_HALF), (0, None)) + thr_pack_qk = tiled_pack_qk.get_slice(0) + thr_pack_pv = tiled_pack_pv.get_slice(0) + tCgQ = thr_pack_qk.partition_A(gQ) + tCgKC = thr_pack_qk.partition_B(gKC) + tCgVC = thr_pack_pv.partition_B(gVC) + tCrKC = tiled_pack_qk.make_fragment_B(sKC) + tCrVC = tiled_pack_pv.make_fragment_B(sVC) + tCrPackQ = tiled_pack_qk.make_fragment_A(sQ) + tCrPackK = tiled_pack_qk.make_fragment_B(sPackK) + tCrPackV = tiled_pack_pv.make_fragment_B(sPackV) + + tQsQ, tQgQ = cpasync.tma_partition( + tma_atom_q, + 0, + cute.make_layout(1), + cute.group_modes(sQ, 0, 3), + cute.group_modes(tCgQ, 0, 3), + ) + tPackKsK, tPackKgK = cpasync.tma_partition( + tma_atom_pack_k, + 0, + cute.make_layout(1), + cute.group_modes(sPackKGather, 0, 3), + cute.group_modes(gPackK, 0, 2), + ) + tPackVsV, tPackVgV = cpasync.tma_partition( + tma_atom_pack_v, + 0, + cute.make_layout(1), + cute.group_modes(sPackVGather, 0, 3), + cute.group_modes(gPackV, 0, 2), + ) + tKCsKC, tKCgKC = cpasync.tma_partition( + tma_atom_kc, + 0, + cute.make_layout(1), + cute.group_modes(sKC, 0, 3), + cute.group_modes(tCgKC, 0, 3), + ) + tVCsVC, tVCgVC = cpasync.tma_partition( + tma_atom_vc, + 0, + cute.make_layout(1), + cute.group_modes(sVC, 0, 3), + cute.group_modes(tCgVC, 0, 3), + ) + + pack_score_shape = tiled_pack_qk.partition_shape_C(PACK_QK_TILE[:2]) + pack_score_template = tiled_pack_qk.make_fragment_C(pack_score_shape) + pack_o_shape = tiled_pack_pv.partition_shape_C(PACK_PV_TILE[:2]) + pack_o_template = tiled_pack_pv.make_fragment_C(pack_o_shape) + + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(Float32) + # The 256-column allocation leaves the second half of SM TMEM available to + # another CTA. The live allocation remains owned after permit release. + tmem.relinquish_alloc_permit() + tmem_base = tmem_ptr.toint() + pair_tScore = cute.make_tensor( + cute.make_ptr( + Float32, + tmem_base + Int32(PAIR_SCORE_OFFSET), + cute.AddressSpace.tmem, + assumed_align=16, + ), + pack_score_template.layout, + ) + pair_tO = cute.make_tensor( + cute.make_ptr( + Float32, + tmem_base + Int32(O_OFFSET), + cute.AddressSpace.tmem, + assumed_align=16, + ), + pack_o_template.layout, + ) + # make_fragment_A drops the physical TMEM allocation base and addresses + # packed BF16 columns in half-column units. Restore both facts so + # 2*tmem_base + 2*PAIR_P_OFFSET names columns 64..127. + pair_tP_storage = cute.make_tensor(pair_tScore.iterator, pack_p_layout.outer) + pair_tP_base = tiled_pack_pv.make_fragment_A(pair_tP_storage)[None, None, None, 0] + pair_tP = cute.make_tensor( + pair_tP_base.iterator + tmem_base + tmem_base + Int32(PAIR_P_OFFSET * 2), + pair_tP_base.layout, + ) + q_producer = fa_pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) + q_consumer = fa_pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) + pack_k_producer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, PAIR_STAGES) + pack_k_consumer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, PAIR_STAGES) + pack_v_producer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, PAIR_STAGES) + pack_v_consumer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, PAIR_STAGES) + pair_score_producer = fa_pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) + pair_score_consumer = fa_pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) + pair_o_producer = fa_pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) + pair_o_consumer = fa_pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) + route_start_base = Int32(0) + q_len = token_count - q_block_idx * Int32(M) + if q_len > Int32(M): + q_len = Int32(M) + threshold = Float32(mThreshold_bnh[batch_idx, q_block_idx, head_idx]) + + if warp_idx == Int32(5): + cpasync.prefetch_descriptor(tma_atom_q) + cpasync.prefetch_descriptor(tma_atom_pack_k) + cpasync.prefetch_descriptor(tma_atom_pack_v) + cpasync.prefetch_descriptor(tma_atom_kc) + cpasync.prefetch_descriptor(tma_atom_vc) + + q_pipe.producer_acquire(q_producer) + q_barrier = q_pipe.producer_get_barrier(q_producer) + cute.copy( + tma_atom_q, + tQgQ[(None, q_block_idx)], + tQsQ[(None, q_producer.index)], + tma_bar_ptr=q_barrier, + ) + q_producer.advance() + + is_owner = warp_idx >= Int32(1) and warp_idx <= Int32(4) + is_score_consumer = warp_idx <= Int32(4) + owner_tidx = tidx - Int32(32) + + # One register-resident online state and one TMEM-O initialization bit span + # every route/exact transaction in every runtime group. + running_max = -Float32.inf + running_sum = Float32(0.0) + owner_o_initialized = Int32(0) + mma_o_initialized = Int32(0) + + if warp_idx == Int32(0): + q_pipe.consumer_wait(q_consumer) + + # The outer loop owns one logical G256 exact-index lifetime. The inner + # loop consumes each physical score/PV half immediately; it appends only + # integer indices, never a second score or probability fragment. + num_logical_groups = (num_route_tiles + Int32(ROUTE_HALVES_PER_GROUP - 1)) // Int32( + ROUTE_HALVES_PER_GROUP + ) + # BEGIN_G256_CURSOR_UNIFORM_INDUCTION + # arch_make_warp_uniform is a lowering hint, not a value broadcast. Both + # values are CTA-invariant integer scalars before the hint. + logical_group_idx = cute.arch.make_warp_uniform(Int32(0)) + remaining_group_tiles = cute.arch.make_warp_uniform(num_route_tiles) + while logical_group_idx < num_logical_groups: + is_final_logical_group = logical_group_idx + Int32(1) == num_logical_groups + group_route_tile_base = logical_group_idx * Int32(ROUTE_HALVES_PER_GROUP) + physical_halves_this_group = remaining_group_tiles + if physical_halves_this_group > Int32(ROUTE_HALVES_PER_GROUP): + physical_halves_this_group = Int32(ROUTE_HALVES_PER_GROUP) + + for half_idx in cutlass.range(physical_halves_this_group, unroll=1): + route_tile_idx = cute.arch.make_warp_uniform(group_route_tile_base + half_idx) + is_final_route_tile = route_tile_idx + Int32(1) == num_route_tiles + is_logical_terminal_half = half_idx + Int32(1) == physical_halves_this_group + route_start = cute.arch.make_warp_uniform( + route_start_base + route_tile_idx * Int32(ROUTE_TILE_SIZE) + ) + remaining_route_count = cute.arch.make_warp_uniform( + route_valid_total - route_tile_idx * Int32(ROUTE_TILE_SIZE) + ) + valid_route_count = remaining_route_count + if valid_route_count > Int32(ROUTE_TILE_SIZE): + valid_route_count = Int32(ROUTE_TILE_SIZE) + if valid_route_count < Int32(0): + valid_route_count = Int32(0) + + # One native N128 route transaction shares the independent K/V stages + # with the exact-pair engine. Route and exact are separated by + # a full-CTA phase boundary, so no K<->V alias handoff is required. + if warp_idx == Int32(5): + pack_k_pipe.producer_acquire(pack_k_producer) + route_k_barrier = pack_k_pipe.producer_get_barrier(pack_k_producer) + cute.copy( + tma_atom_kc, + tKCgKC[(None, route_tile_idx)], + tKCsKC[(None, pack_k_producer.index)], + tma_bar_ptr=route_k_barrier, + ) + pack_k_producer.advance() + + pack_v_pipe.producer_acquire(pack_v_producer) + route_v_barrier = pack_v_pipe.producer_get_barrier(pack_v_producer) + cute.copy( + tma_atom_vc, + tVCgVC[(None, route_tile_idx)], + tVCsVC[(None, pack_v_producer.index)], + tma_bar_ptr=route_v_barrier, + ) + pack_v_producer.advance() + + if warp_idx == Int32(0): + pack_k_pipe.consumer_wait(pack_k_consumer) + pair_score_pipe.producer_acquire(pair_score_producer) + mma_utils.gemm( + tiled_pack_qk, + pair_tScore, + tCrPackQ[None, None, None, q_consumer.index], + tCrKC[None, None, None, pack_k_consumer.index], + zero_init=True, + ) + pair_score_pipe.producer_commit(pair_score_producer) + pair_score_producer.advance() + pack_k_pipe.consumer_release(pack_k_consumer) + pack_k_consumer.advance() + + # BEGIN_RUNTIME_GROUP_BODY + + # Route generation: four physical owner warps reduce the native N128 + # score tile into one four-word mask. HBM receives only the diagnostic + # copy; the compacted exact stream remains resident in SMEM. + if is_owner: + pair_score_pipe.consumer_wait(pair_score_consumer) + score_raw, score_coords = _load_pair_score( + pack_score_template, + thr_pack_qk, + tmem_base, + Int32(PAIR_SCORE_OFFSET), + owner_tidx, + ) + pack_score_loaded_barrier.arrive_and_wait() + pair_score_pipe.consumer_release(pair_score_consumer) + pair_score_consumer.advance() + owner_warp = owner_tidx // Int32(32) + lane = owner_tidx % Int32(32) + semantic_row = (score_coords[0][0] + Int32(SEMANTIC_ROW_OFFSET)) & Int32(M - 1) + row_valid = semantic_row < q_len + lane_col_parity = (lane // Int32(2)) % Int32(2) + # Column-pair reduction: parity-0 lanes carry column 2*pair and + # parity-1 lanes carry column 2*pair+1. The XOR-1/16/8/4 + # butterfly tree never crosses lane column-parity classes + # ((l^k)//2 keeps (l//2)%2 for k in {1,16,8,4}), so one tree + # reduces both columns at once; every surviving addition chain + # sees the same zero-padded operand streams, and the removed + # chains only ever accumulated 0.0. Writer lanes 0 and 2 equal + # 2*(col%2). + for pair_idx in cutlass.range_constexpr(0, ROUTE_TILE_SIZE // 2, 2): + my_col0 = Int32(2 * pair_idx) + lane_col_parity + partial0 = Float32(0.0) + if row_valid and my_col0 < valid_route_count: + partial0 = Float32(score_raw[pair_idx]) + my_col1 = Int32(2 * (pair_idx + 1)) + lane_col_parity + partial1 = Float32(0.0) + if row_valid and my_col1 < valid_route_count: + partial1 = Float32(score_raw[pair_idx + 1]) + + raw_partial0 = partial0 + raw_partial1 = partial1 + scaled0, scaled1 = cute.arch.mul_packed_f32x2( + (raw_partial0, raw_partial1), + (softmax_scale_log2, softmax_scale_log2), + ) + peer_scaled0 = cute.arch.shuffle_sync_bfly(scaled0, offset=1) + peer_scaled1 = cute.arch.shuffle_sync_bfly(scaled1, offset=1) + partial0, partial1 = cute.arch.fma_packed_f32x2( + (raw_partial0, raw_partial1), + (softmax_scale_log2, softmax_scale_log2), + (peer_scaled0, peer_scaled1), + ) + peer0 = cute.arch.shuffle_sync_bfly(partial0, offset=16) + peer1 = cute.arch.shuffle_sync_bfly(partial1, offset=16) + partial0, partial1 = cute.arch.add_packed_f32x2( + (partial0, partial1), (peer0, peer1) + ) + peer0 = cute.arch.shuffle_sync_bfly(partial0, offset=8) + peer1 = cute.arch.shuffle_sync_bfly(partial1, offset=8) + partial0, partial1 = cute.arch.add_packed_f32x2( + (partial0, partial1), (peer0, peer1) + ) + peer0 = cute.arch.shuffle_sync_bfly(partial0, offset=4) + peer1 = cute.arch.shuffle_sync_bfly(partial1, offset=4) + partial0, partial1 = cute.arch.add_packed_f32x2( + (partial0, partial1), (peer0, peer1) + ) + if lane == Int32(0): + route_partial[owner_warp, 2 * pair_idx] = partial0 + route_partial[owner_warp, 2 * (pair_idx + 1)] = partial1 + if lane == Int32(2): + route_partial[owner_warp, 2 * pair_idx + 1] = partial0 + route_partial[owner_warp, 2 * (pair_idx + 1) + 1] = partial1 + + cute.arch.fence_view_async_shared() + score_loaded_barrier.arrive_and_wait() + if owner_warp == Int32(0): + mask0 = Int32(0) + mask1 = Int32(0) + mask2 = Int32(0) + mask3 = Int32(0) + + # Half 0 starts a fresh G256 stream and half 1 appends to + # lane 0's cumulative packet word. The preceding packet + # barrier makes the base warp-uniform before the vote. + append_base = Int32(0) + if half_idx != Int32(0): + append_base = Int32(route_packet[6]) + + # A positive signed shift avoids materializing 1<<31: + # lane 0 gets zero and lane 31 gets 0x7fffffff. + lane_mask_lt = Int32(0x7FFFFFFF) >> (Int32(31) - lane) + preceding_word_count = Int32(0) + for word in cutlass.range_constexpr(ROUTE_MASK_WORDS): + off = Int32(word * 32) + lane + valid = off < valid_route_count + exact_pred = False + if valid: + pair_02 = Float32(route_partial[0, off]) + Float32( + route_partial[2, off] + ) + pair_13 = Float32(route_partial[1, off]) + Float32( + route_partial[3, off] + ) + col_mean = (pair_02 + pair_13) / Float32(q_len) + exact_pred = sol_attn_route_is_exact( + q_block_idx, + route_start + off, + col_mean, + threshold, + valid, + ) + # Sink is a KV-only contract. Text queries remain + # a caller-side dense operation in MMDiT models. + exact_pred = exact_pred or ( + route_start + off >= sink_start_block + and route_start + off < sink_end_block + ) + word_mask = Int32(cute.arch.vote_ballot_sync(exact_pred)) + # Site 2: preserve the route decision and its four + # ordered ballots, but materialize the resulting + # approximate-column mask exactly once. Dedicated + # SMEM holds the two N64 mask halves so the reduction + # scratch remains non-aliasing for ptxas scheduling. + # The existing shared fence and owner barrier below + # publish them to every score owner. + if valid and not exact_pred: + column_masks[off] = Float32(0.0) + else: + column_masks[off] = -Float32.inf + lane_rank = ( + append_base + + preceding_word_count + + sol_attn_popc_b32(word_mask & lane_mask_lt) + ) + if exact_pred: + route_indices[lane_rank] = route_start + off + if cutlass.const_expr(word == 0): + mask0 = word_mask + elif cutlass.const_expr(word == 1): + mask1 = word_mask + elif cutlass.const_expr(word == 2): + mask2 = word_mask + else: + mask3 = word_mask + preceding_word_count = preceding_word_count + sol_attn_popc_b32(word_mask) + + # Every selected lane has a unique rank; lane 0 publishes + # the packet after reconvergence. + exact_count = preceding_word_count + if lane == Int32(0): + route_rank = append_base + exact_count + + route_packet[0] = mask0 + route_packet[1] = mask1 + route_packet[2] = mask2 + route_packet[3] = mask3 + route_packet[4] = exact_count + route_packet[5] = append_base + route_packet[6] = route_rank + terminal_half_word = Int32(0) + if is_logical_terminal_half: + terminal_half_word = Int32(1) + route_packet[7] = terminal_half_word + cute.arch.fence_view_async_shared() + + # The selector packet is now immutable. Reuse the already resident + # route scores for the non-exact transaction; no offset list or second + # route-score load is introduced. + score_loaded_barrier.arrive_and_wait() + route_exact_count = Int32(route_packet[4]) + has_route_approx = route_exact_count < valid_route_count + if has_route_approx: + row_mask = -Float32.inf + if row_valid: + row_mask = Float32(0.0) + # Route generation has consumed every raw score. Apply + # the shared mask in place so raw and masked N128 + # fragments never overlap in registers; the same object + # remains available for the later route-mass scratch. + route_scores = score_raw + assert cute.size(score_raw) % 2 == 0 + for i in cutlass.range_constexpr(0, cute.size(score_raw), 2): + group_col0 = score_coords[i][1] + group_col1 = score_coords[i + 1][1] + mask0 = Float32(column_masks[group_col0]) + mask1 = Float32(column_masks[group_col1]) + mask0, mask1 = cute.arch.add_packed_f32x2( + (mask0, mask1), (row_mask, row_mask) + ) + mask0, mask1 = cute.arch.add_packed_f32x2( + ( + Float32(score_raw[i]), + Float32(score_raw[i + 1]), + ), + (mask0, mask1), + ) + route_scores[i] = mask0 + route_scores[i + 1] = mask1 + + local_max = fa_utils.fmax_reduce(route_scores.load(), arch=100) + local_max = Float32(local_max) * softmax_scale + peer_max = cute.arch.shuffle_sync_bfly(local_max, offset=2) + pair_max = local_max + if peer_max > pair_max: + pair_max = peer_max + + old_max = running_max + old_sum = running_sum + new_max = old_max + if old_max == -Float32.inf or pair_max > old_max: + new_max = pair_max + row_alpha = Float32(0.0) + if old_max != -Float32.inf: + row_alpha = cute.math.exp2( + (old_max - new_max) * Float32(LOG2E), + fastmath=True, + ) + + route_probabilities = cute.make_rmem_tensor(route_scores.shape, Float32) + if new_max == -Float32.inf: + for i in cutlass.range(cute.size(route_scores), unroll_full=True): + route_probabilities[i] = Float32(0.0) + else: + for i in cutlass.range(cute.size(route_scores), unroll_full=True): + route_probabilities[i] = cute.math.exp2( + Float32(route_scores[i]) * softmax_scale_log2 + - new_max * Float32(LOG2E), + fastmath=True, + ) + # ``route_scores`` is dead after the exponentials above. Use + # it as mass scratch so the compiler does not need a second + # full N128-shaped fragment while probabilities remain live + # for the chunked TMEM-P store below. Keeping the same shape, + # index order, and fadd_reduce preserves floating-point + # reduction order and every phase edge. + assert cute.size(route_probabilities) % 2 == 0 + for i in cutlass.range_constexpr(0, cute.size(route_probabilities), 2): + block_idx0 = route_start + score_coords[i][1] + raw_length0 = token_count - block_idx0 * Int32(N_MEMBER) + block_length0 = max(Int32(0), min(raw_length0, Int32(N_MEMBER))) + block_idx1 = route_start + score_coords[i + 1][1] + raw_length1 = token_count - block_idx1 * Int32(N_MEMBER) + block_length1 = max(Int32(0), min(raw_length1, Int32(N_MEMBER))) + mass0, mass1 = cute.arch.mul_packed_f32x2( + ( + Float32(route_probabilities[i]), + Float32(route_probabilities[i + 1]), + ), + ( + Float32(block_length0), + Float32(block_length1), + ), + ) + route_scores[i] = mass0 + route_scores[i + 1] = mass1 + current_sum = fa_utils.fadd_reduce(route_scores.load(), arch=100) + current_sum += cute.arch.shuffle_sync_bfly(current_sum, offset=2) + # KC is a block mean and VC a valid-token sum. Route mass uses + # the true block length while PV still consumes p*VC once. + running_sum = old_sum * row_alpha + current_sum + running_max = new_max + if owner_o_initialized != Int32(0): + _rescale_pair_o( + pack_o_template, + thr_pack_pv, + tmem_base, + Int32(O_OFFSET), + owner_tidx, + row_alpha, + ) + _store_pair_probability_chunked_tmemp( + pack_o_template, + route_probabilities, + tmem_base, + Int32(PAIR_P_OFFSET), + owner_tidx, + ) + owner_o_initialized = Int32(1) + # Publish the mask/P decision to warp 0. The route PV is deliberately + # drained before exact work so all-exact, all-approx, odd, and + # partial-tail paths share one phase boundary. + if is_score_consumer: + route_packet_ready_barrier.arrive_and_wait() + if warp_idx == Int32(0): + route_exact_count = Int32(route_packet[4]) + route_has_approx = route_exact_count < valid_route_count + pack_v_pipe.consumer_wait(pack_v_consumer) + if route_has_approx: + mma_utils.gemm( + tiled_pack_pv, + pair_tO, + pair_tP, + tCrVC[None, None, None, pack_v_consumer.index], + zero_init=mma_o_initialized == Int32(0), + ) + # Half 0 is followed by half-1 route QK. The terminal + # route half is followed by exact QK0 whenever the fused + # G256 index stream is nonempty. Those score completions + # prove this PV complete; only a final route-only CTA needs + # an explicit O completion here. + if is_final_route_tile and Int32(route_packet[6]) == Int32(0): + pair_o_pipe.producer_commit(pair_o_producer) + mma_o_initialized = Int32(1) + pack_v_pipe.consumer_release(pack_v_consumer) + pack_v_consumer.advance() + if is_owner: + cumulative_exact_count = Int32(route_packet[6]) + if is_final_route_tile and cumulative_exact_count == Int32(0): + pair_o_pipe.consumer_wait(pair_o_consumer) + + # route_packet may be reused by the next physical half without a + # CTA join. Warp 0 reads this half's packet before it can issue + # next-half QK; owner-warp 0 cannot overwrite the packet until + # that QK's pair-score completion has released all owners. + + # Both route halves have published their packet/index data and drained + # approximate PV. This is the only CTA-wide pre-exact join in the + # logical G256 group; it publishes the combined list to warp 5. + cute.arch.barrier() + # The cumulative count covers half 0 followed by half 1. Pairing this + # one ordered stream removes cross-half odd padding without retaining + # either physical score fragment. + exact_block_count = Int32(route_packet[6]) + exact_pair_count = (exact_block_count + Int32(1)) // Int32(2) + pair_count = exact_pair_count + has_pair_exact = exact_block_count > Int32(0) + + # BEGIN_GENERAL_N128_PAIR + # Every executable exact count, including a logical-group terminal + # exact1, stays in the N128 domain. + + # Warp 5 streams one physical N128 K stage and one physical N128 V + # stage. A missing odd peer duplicates block0 only for the physical + # transaction; owners mask all upper-64 scores before softmax. + if warp_idx == Int32(5) and has_pair_exact: + for pair_idx in cutlass.range(pair_count, unroll=1): + ordinal0 = pair_idx * Int32(2) + block0 = Int32(route_indices[ordinal0]) + block1 = block0 + if ordinal0 + Int32(1) < exact_block_count: + block1 = Int32(route_indices[ordinal0 + Int32(1)]) + + pack_k_pipe.producer_acquire(pack_k_producer) + pair_k_barrier = pack_k_pipe.producer_get_barrier(pack_k_producer) + _load_pack_k_half( + tma_atom_pack_k, + tPackKgK, + tPackKsK, + block0, + block1, + pack_k_producer.index * Int32(4), + pair_k_barrier, + ) + pack_k_producer.advance() + + pack_v_pipe.producer_acquire(pack_v_producer) + pair_v_barrier = pack_v_pipe.producer_get_barrier(pack_v_producer) + _load_pack_v_half( + tma_atom_pack_v, + tPackVgV, + tPackVsV, + block0, + block1, + pack_v_producer.index * Int32(4), + pair_v_barrier, + ) + pack_v_producer.advance() + + if warp_idx == Int32(0) and has_pair_exact: + # QK0 prologue. K and score cursors advance exactly once per QK; + # neither V nor O state is touched until the steady-state PV path. + pack_k_pipe.consumer_wait(pack_k_consumer) + pair_score_pipe.producer_acquire(pair_score_producer) + mma_utils.gemm( + tiled_pack_qk, + pair_tScore, + tCrPackQ[None, None, None, q_consumer.index], + tCrPackK[None, None, None, pack_k_consumer.index], + zero_init=True, + ) + pair_score_pipe.producer_commit(pair_score_producer) + pair_score_producer.advance() + # PipelineTmaUmma release is tcgen05-completion-backed. + pack_k_pipe.consumer_release(pack_k_consumer) + pack_k_consumer.advance() + + for pair_idx in cutlass.range(pair_count, unroll=1): + # P aliases the drained upper half of S. PV must therefore be + # issued before QK(i+1) overwrites S. Both instructions are + # emitted back-to-back by warp 0, retaining the full-G128 + # tcgen05 dependency order without its K/V alias barriers. + pack_v_pipe.consumer_wait(pack_v_consumer) + # All four owners have completed their synchronous chunked + # TMEM stores and the helper's TMEM store fence before this + # five-warp rendezvous releases the single MMA warp. + exact_pair_p_ready_barrier.arrive_and_wait() + mma_utils.gemm( + tiled_pack_pv, + pair_tO, + pair_tP, + tCrPackV[None, None, None, pack_v_consumer.index], + zero_init=mma_o_initialized == Int32(0), + ) + # QK(i+1) completion dominates PV(i) completion for every + # nonterminal transaction on this tcgen05 issuer. Commit one + # explicit O-full generation only for the CTA's final PV. + if is_final_logical_group and pair_idx + Int32(1) == pair_count: + pair_o_pipe.producer_commit(pair_o_producer) + mma_o_initialized = Int32(1) + pack_v_pipe.consumer_release(pack_v_consumer) + pack_v_consumer.advance() + + if pair_idx + Int32(1) < pair_count: + pack_k_pipe.consumer_wait(pack_k_consumer) + pair_score_pipe.producer_acquire(pair_score_producer) + mma_utils.gemm( + tiled_pack_qk, + pair_tScore, + tCrPackQ[None, None, None, q_consumer.index], + tCrPackK[None, None, None, pack_k_consumer.index], + zero_init=True, + ) + pair_score_pipe.producer_commit(pair_score_producer) + pair_score_producer.advance() + pack_k_pipe.consumer_release(pack_k_consumer) + pack_k_consumer.advance() + + if is_owner and has_pair_exact: + exact_owner_warp = owner_tidx // Int32(32) + exact_lane = owner_tidx % Int32(32) + for pair_idx in cutlass.range(pair_count, unroll=1): + ordinal0 = pair_idx * Int32(2) + block0 = Int32(route_indices[ordinal0]) + has_peer = ordinal0 + Int32(1) < exact_block_count + block1 = block0 + if has_peer: + block1 = Int32(route_indices[ordinal0 + Int32(1)]) + valid0 = token_count - block0 * Int32(N_MEMBER) + valid1 = Int32(0) + if has_peer: + valid1 = token_count - block1 * Int32(N_MEMBER) + # Keep packed-select integer min/max lowering and exact-pair + # bookkeeping unchanged. + valid0 = max(Int32(0), min(valid0, Int32(N_MEMBER))) + valid1 = max(Int32(0), min(valid1, Int32(N_MEMBER))) + + # Site 1: owner warp 0 builds two 64-column gates once for + # this exact N128 pair. The existing score-load barrier below + # both protects the S/P alias and publishes these stores; no + # barrier or shared allocation is added. + if exact_owner_warp == Int32(0): + for cohort in cutlass.range_constexpr(4): + column = Int32(cohort * 32) + exact_lane + if cutlass.const_expr(cohort < 2): + if column >= valid0: + column_masks[column] = -Float32.inf + else: + column_masks[column] = Float32(0.0) + else: + if column - Int32(N_MEMBER) >= valid1: + column_masks[column] = -Float32.inf + else: + column_masks[column] = Float32(0.0) + cute.arch.fence_view_async_shared() + + pair_score_pipe.consumer_wait(pair_score_consumer) + # Keep the exact ae9 score-load helper and fragment scope. + pair_scores, pair_coords = _load_pair_score( + pack_score_template, + thr_pack_qk, + tmem_base, + Int32(PAIR_SCORE_OFFSET), + owner_tidx, + ) + # Every owner retires the complete score load before the + # packed P store aliases columns 64..127 of S. + pack_score_loaded_barrier.arrive_and_wait() + pair_score_pipe.consumer_release(pair_score_consumer) + pair_score_consumer.advance() + + semantic_row = (pair_coords[0][0] + Int32(SEMANTIC_ROW_OFFSET)) & Int32(M - 1) + row_valid = semantic_row < q_len + row_mask = -Float32.inf + if row_valid: + row_mask = Float32(0.0) + assert cute.size(pair_scores) % 2 == 0 + for i in cutlass.range_constexpr(0, cute.size(pair_scores), 2): + column0 = pair_coords[i][1] + column1 = pair_coords[i + 1][1] + mask0 = Float32(column_masks[column0]) + mask1 = Float32(column_masks[column1]) + mask0, mask1 = cute.arch.add_packed_f32x2((mask0, mask1), (row_mask, row_mask)) + mask0, mask1 = cute.arch.add_packed_f32x2( + ( + Float32(pair_scores[i]), + Float32(pair_scores[i + 1]), + ), + (mask0, mask1), + ) + pair_scores[i] = mask0 + pair_scores[i + 1] = mask1 + + probabilities, next_max, next_sum, row_alpha = _online_update_pair( + pair_scores, + running_max, + running_sum, + softmax_scale, + ) + # For i>0, pair-score completion comes from QK(i), issued + # after PV(i-1) on the same tcgen05 issuer. The score wait and + # load above therefore retire PV(i-1) before this O rescale. + # Pair0 similarly follows either route QK or route PV->QK0. + if owner_o_initialized != Int32(0): + _rescale_pair_o( + pack_o_template, + thr_pack_pv, + tmem_base, + Int32(O_OFFSET), + owner_tidx, + row_alpha, + ) + # The one TMEM P image is free once PV(i-1) completes. Keep + # probabilities FP32 until the live-range-bounded chunked R2T. + _store_pair_probability_chunked_tmemp( + pack_o_template, + probabilities, + tmem_base, + Int32(PAIR_P_OFFSET), + owner_tidx, + ) + # The preceding helper performs tcgen05.wait::st for every + # chunk and a TMEM-store fence. Publish P to warp 0 with one + # uniform generation shared by warps 0-4; warp 5 is excluded. + exact_pair_p_ready_barrier.arrive_and_wait() + running_max = next_max + running_sum = next_sum + owner_o_initialized = Int32(1) + + # There is no successor QK after the CTA's final exact PV. Keep + # exactly one completion-backed wait before the epilogue; all + # earlier groups flow into a successor route QK completion. + if is_final_logical_group and pair_count > Int32(0): + pair_o_pipe.consumer_wait(pair_o_consumer) + + # route_indices reuse HB proof for the next logical group: + # (1) warp 5 reads both indices before producing each pair's K/V, and + # final-pair score completion therefore dominates its last read; + # (2) all owner index reads precede the final exact-P NamedBarrier; + # (3) owner-warp0/lane0 is the sole next-group writer and reaches it + # only after that same exact loop. For exact_count==0 there are no + # readers. Therefore no group-tail CTA barrier is required. + + # Cross-group progress is carried by the existing K/V buffer-free + # phases and pair-score ready phase. There is no CTA-wide group-tail + # join: the next producer acquire cannot overwrite a live K/V stage, + # and the next owner score load cannot precede QK completion. + # END_GENERAL_N128_PAIR + + logical_group_idx = cute.arch.make_warp_uniform(logical_group_idx + Int32(1)) + remaining_group_tiles = cute.arch.make_warp_uniform( + remaining_group_tiles - Int32(ROUTE_HALVES_PER_GROUP) + ) + # END_RUNTIME_GROUP_BODY + # END_G256_CURSOR_UNIFORM_INDUCTION + + if warp_idx == Int32(0): + q_pipe.consumer_release(q_consumer) + q_consumer.advance() + + if is_owner: + lane = owner_tidx % Int32(32) + owner_warp = owner_tidx // Int32(32) + owner_row = ( + owner_warp * Int32(16) + + lane // Int32(4) + + (lane % Int32(2)) * Int32(8) + + Int32(SEMANTIC_ROW_OFFSET) + ) & Int32(M - 1) + # Register state remains owner-local for the entire exact stream. It + # is published only once here because the final Ld16x256b epilogue + # remaps rows differently from the Ld16x64b xor-2 score ownership. + if (lane & Int32(2)) == Int32(0): + sFinalStats[owner_row, 0] = running_sum + sFinalStats[owner_row, 1] = running_max + cute.arch.fence_view_async_shared() + final_stats_ready_barrier.arrive_and_wait() + + o_regs, o_coords = load_m64_o_fp32_256b( + pack_o_template, + thr_pack_pv, + tmem_base, + owner_tidx, + ) + assert cute.size(o_regs) == 64 + assert cute.size(o_coords) == 64 + + # B7's device inversion proves that 4*w/4*w+1 belong to one + # semantic row and 4*w+2/4*w+3 to its row-plus-eight peer. Hoist + # validity, final-sum LDS, reciprocal, and row base once per stratum. + semantic_row0 = ( + owner_warp * Int32(16) + lane // Int32(4) + Int32(SEMANTIC_ROW_OFFSET) + ) & Int32(M - 1) + semantic_row1 = (semantic_row0 + Int32(8)) & Int32(M - 1) + even_col_base = (lane % Int32(4)) * Int32(2) + + if semantic_row0 < q_len: + inv_sum0 = cute.arch.rcp_approx(Float32(sFinalStats[semantic_row0, 0])) + query_idx0 = q_block_idx * Int32(M) + semantic_row0 + destination_row0 = cute.domain_offset( + (batch_idx, query_idx0, head_idx, Int32(0)), mO_bthd + ) + for word_i in cutlass.range(O_PACKED_WORDS_PER_ROW_PER_THREAD, unroll_full=True): + even_i = word_i * 4 + odd_i = even_i + 1 + even_value = Float32(o_regs[even_i]) * inv_sum0 + odd_value = Float32(o_regs[odd_i]) * inv_sum0 + packed_word = _cvt_bf16x2_f32(Float32(odd_value), Float32(even_value)) + even_col = even_col_base + word_i * O_PACKED_COLUMN_STRIDE + _store_global_u32_inline(destination_row0.iterator + even_col, packed_word) + + if semantic_row1 < q_len: + inv_sum1 = cute.arch.rcp_approx(Float32(sFinalStats[semantic_row1, 0])) + query_idx1 = q_block_idx * Int32(M) + semantic_row1 + destination_row1 = cute.domain_offset( + (batch_idx, query_idx1, head_idx, Int32(0)), mO_bthd + ) + for word_i in cutlass.range(O_PACKED_WORDS_PER_ROW_PER_THREAD, unroll_full=True): + even_i = word_i * 4 + 2 + odd_i = even_i + 1 + even_value = Float32(o_regs[even_i]) * inv_sum1 + odd_value = Float32(o_regs[odd_i]) * inv_sum1 + packed_word = _cvt_bf16x2_f32(Float32(odd_value), Float32(even_value)) + even_col = even_col_base + word_i * O_PACKED_COLUMN_STRIDE + _store_global_u32_inline(destination_row1.iterator + even_col, packed_word) + + if (lane & Int32(2)) == Int32(0) and owner_row < q_len: + query_idx = q_block_idx * Int32(M) + owner_row + mLSE_bth[batch_idx, query_idx, head_idx] = running_max + cute.math.log2( + running_sum, fastmath=True + ) * Float32(LN2) + + cute.arch.barrier() + tmem.free(tmem_ptr) + + +@cute.jit +def _sol_attn_sm100_bf16_host( + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + o: cute.Tensor, + kc: cute.Tensor, + vc: cute.Tensor, + threshold: cute.Tensor, + lse: cute.Tensor, + softmax_scale: Float32, + sink_start_block: Int32, + sink_end_block: Int32, + stream: cuda.CUstream = None, +): + q, k, v, o, kc, vc = tuple(assume_tensor_aligned(t) for t in (q, k, v, o, kc, vc)) + q_mkl, k_nkl, kc_nkl = [layout_utils.select(t, [1, 3, 2, 0]) for t in (q, k, kc)] + v_nkl, vc_nkl = [layout_utils.select(t, [3, 1, 2, 0]) for t in (v, vc)] + token_count = cute.size(q_mkl.shape[0]) + num_blocks = cute.size(kc_nkl.shape[0]) + num_heads = cute.size(q_mkl.shape[2]) + num_batches = cute.size(q_mkl.shape[3]) + num_route_tiles = cute.ceil_div(num_blocks, ROUTE_TILE_SIZE) + pack_qk_op = tcgen05.MmaF16BF16Op( + BFloat16, + Float32, + PACK_QK_INST, + tcgen05.CtaGroup.ONE, + tcgen05.OperandSource.SMEM, + cute.nvgpu.OperandMajorMode.K, + cute.nvgpu.OperandMajorMode.K, + ) + tiled_pack_qk = cute.make_tiled_mma(pack_qk_op) + pack_pv_op = tcgen05.MmaF16BF16Op( + BFloat16, + Float32, + PACK_PV_INST, + tcgen05.CtaGroup.ONE, + tcgen05.OperandSource.TMEM, + cute.nvgpu.OperandMajorMode.K, + cute.nvgpu.OperandMajorMode.MN, + ) + tiled_pack_pv = cute.make_tiled_mma(pack_pv_op) + pack_qk_quarter_op = tcgen05.MmaF16BF16Op( + BFloat16, + Float32, + PACK_QK_QUARTER_INST, + tcgen05.CtaGroup.ONE, + tcgen05.OperandSource.SMEM, + cute.nvgpu.OperandMajorMode.K, + cute.nvgpu.OperandMajorMode.K, + ) + tiled_pack_qk_gather = cute.make_tiled_mma(pack_qk_quarter_op) + pack_pv_quarter_op = tcgen05.MmaF16BF16Op( + BFloat16, + Float32, + PACK_PV_QUARTER_INST, + tcgen05.CtaGroup.ONE, + tcgen05.OperandSource.TMEM, + cute.nvgpu.OperandMajorMode.K, + cute.nvgpu.OperandMajorMode.MN, + ) + tiled_pack_pv_gather = cute.make_tiled_mma(pack_pv_quarter_op) + q_layout = sm100_utils.make_smem_layout_a(tiled_pack_qk, PACK_QK_TILE, BFloat16, 1) + pack_k_layout = sm100_utils.make_smem_layout_b( + tiled_pack_qk, PACK_QK_TILE, BFloat16, PAIR_STAGES + ) + pack_v_layout = sm100_utils.make_smem_layout_b( + tiled_pack_pv, PACK_PV_TILE, BFloat16, PAIR_STAGES + ) + pack_k_gather_layout = sm100_utils.make_smem_layout_b( + tiled_pack_qk_gather, + PACK_QK_GATHER_TILE, + BFloat16, + PAIR_STAGES * 4, + ) + pack_v_gather_layout = sm100_utils.make_smem_layout_b( + tiled_pack_pv_gather, + PACK_PV_GATHER_TILE, + BFloat16, + PAIR_STAGES * 4, + ) + pack_p_layout = sm100_utils.make_smem_layout_a(tiled_pack_pv, PACK_PV_TILE, BFloat16, 1) + route_k_layout = sm100_utils.make_smem_layout_b( + tiled_pack_qk, PACK_QK_TILE, BFloat16, PAIR_STAGES + ) + route_v_layout = sm100_utils.make_smem_layout_b( + tiled_pack_pv, PACK_PV_TILE, BFloat16, PAIR_STAGES + ) + copy_op = cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE) + q_tma_atom, q_tma_tensor = cute.nvgpu.make_tiled_tma_atom_A( + copy_op, + q_mkl, + cute.select(q_layout, mode=[0, 1, 2]), + PACK_QK_TILE, + tiled_pack_qk, + ) + pack_k_tma_layout = cute.make_composed_layout( + pack_k_gather_layout.inner, + 0, + cute.make_layout((64, 64), stride=(64, 1)), + ) + pack_k_tma_atom, pack_k_tma_tensor = cpasync.make_tiled_tma_atom( + copy_op, + k_nkl, + pack_k_tma_layout, + (64, 64), + ) + pack_v_tma_layout = cute.make_composed_layout( + pack_v_gather_layout.inner, + 0, + cute.make_layout((64, 64), stride=(1, 64)), + ) + pack_v_tma_atom, pack_v_tma_tensor = cpasync.make_tiled_tma_atom( + copy_op, + v_nkl, + pack_v_tma_layout, + (64, 64), + ) + kc_tma_atom, kc_tma_tensor = cute.nvgpu.make_tiled_tma_atom_B( + copy_op, + kc_nkl, + cute.select(route_k_layout, mode=[0, 1, 2]), + PACK_QK_TILE, + tiled_pack_qk, + ) + vc_tma_atom, vc_tma_tensor = cute.nvgpu.make_tiled_tma_atom_B( + copy_op, + vc_nkl, + cute.select(route_v_layout, mode=[0, 1, 2]), + PACK_PV_TILE, + tiled_pack_pv, + ) + _sol_attn_sm100_bf16_kernel( + tiled_pack_qk, + tiled_pack_pv, + q_tma_atom, + q_tma_tensor, + pack_k_tma_atom, + pack_k_tma_tensor, + pack_v_tma_atom, + pack_v_tma_tensor, + kc_tma_atom, + kc_tma_tensor, + vc_tma_atom, + vc_tma_tensor, + threshold, + o, + lse, + Int32(token_count), + Int32(num_blocks), + Int32(num_route_tiles), + softmax_scale, + sink_start_block, + sink_end_block, + q_layout, + pack_k_layout, + pack_k_gather_layout, + pack_p_layout, + pack_v_layout, + pack_v_gather_layout, + route_k_layout, + route_v_layout, + ).launch( + grid=(num_blocks, num_heads, num_batches), + block=(THREADS, 1, 1), + stream=stream, + min_blocks_per_mp=2, + ) + + +@cute.jit +def forward( + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + o: cute.Tensor, + kc: cute.Tensor, + vc: cute.Tensor, + threshold: cute.Tensor, + lse: cute.Tensor, + softmax_scale: Float32, + sink_start_block: Int32, + sink_end_block: Int32, + stream: cuda.CUstream = None, +): + return _sol_attn_sm100_bf16_host( + q, + k, + v, + o, + kc, + vc, + threshold, + lse, + softmax_scale, + sink_start_block, + sink_end_block, + stream, + ) + + +__all__ = ["forward"] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/math.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/math.py new file mode 100644 index 000000000000..65b0c821c1e6 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/math.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# +# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see +# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. +"""Small tensor-core helpers used by the Blackwell mainloop.""" + +import cutlass +import cutlass.cute as cute +from cutlass import Boolean +from cutlass.cute.nvgpu import tcgen05 + + +@cute.jit +def gemm( + tiled_mma: cute.TiledMma, + accumulator: cute.Tensor, + a: cute.Tensor, + b: cute.Tensor, + zero_init: bool | Boolean = False, +) -> None: + mma = cute.make_mma_atom(tiled_mma.op) + for k in cutlass.range_constexpr(cute.size(a.shape[2])): + mma.set(tcgen05.Field.ACCUMULATE, not zero_init or k != 0) + cute.gemm( + mma, + accumulator, + a[None, None, k], + b[None, None, k], + accumulator, + ) + + +__all__ = ["gemm"] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/softmax.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/softmax.py new file mode 100644 index 000000000000..c6d21a0c8877 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/softmax.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# +# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see +# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. +# +# Portions derive from the FlashAttention project +# (https://github.com/Dao-AILab/flash-attention), BSD-3-Clause; its license +# text is vendored at sol_attn/sm100/LICENSE.flash-attention. +"""Online-softmax helpers for the Blackwell mainloop.""" + +from __future__ import annotations + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32 +from cutlass.cute.nvgpu import tcgen05 +from flash_attn.cute import utils as fa_utils + +from .tmem import _add_physical_tmem_base, _zero_based_tmem_tensor, tcgen05_wait_ld, tcgen05_wait_st + +M = 64 +N_HALF = 128 +DV = 128 +LOG2E = 1.4426950408889634 + + +@cute.jit +def _load_m64_n128_score( + score_template: cute.Tensor, + thr_mma_qk: cute.ThrMma, + tmem_base: Int32, + score_offset: Int32, + owner_tidx: Int32, +): + """Load one M64xN128 FP32 score tile from TMEM.""" + + relative_score = _zero_based_tmem_tensor(Float32, score_template.layout) + load_atom = cute.make_copy_atom( + tcgen05.copy.Ld16x64bOp(tcgen05.copy.Repetition(64)), + Float32, + ) + tiled_load = tcgen05.make_tmem_copy(load_atom, relative_score) + thread_load = tiled_load.get_slice(owner_tidx) + source_relative = thread_load.partition_S(relative_score) + source = _add_physical_tmem_base(source_relative, tmem_base + score_offset) + coordinates = thread_load.partition_D( + thr_mma_qk.partition_C(cute.make_identity_tensor((M, N_HALF))) + ) + scores = cute.make_rmem_tensor(coordinates.shape, Float32) + cute.copy(tiled_load, source, scores) + tcgen05_wait_ld() + cute.arch.fence_view_async_tmem_load() + return scores, coordinates + + +@cute.jit +def _rescale_m64_partial_o( + o_template: cute.Tensor, + thr_mma_pv: cute.ThrMma, + tmem_base: Int32, + o_offset: Int32, + owner_tidx: Int32, + alpha: Float32, +): + """Rescale the prior M64 output accumulator before its next PV update.""" + + relative_o = _zero_based_tmem_tensor(Float32, o_template.layout) + correction_width = 16 + relative_fragment = cute.composition(relative_o, cute.make_layout((M, correction_width))) + load_atom = cute.make_copy_atom(tcgen05.copy.Ld16x64bOp(tcgen05.copy.Repetition(8)), Float32) + store_atom = cute.make_copy_atom(tcgen05.copy.St16x64bOp(tcgen05.copy.Repetition(8)), Float32) + thread_load = tcgen05.make_tmem_copy(load_atom, relative_fragment).get_slice(owner_tidx) + thread_store = tcgen05.make_tmem_copy(store_atom, relative_fragment).get_slice(owner_tidx) + source = _add_physical_tmem_base( + thread_load.partition_S(relative_fragment), tmem_base + o_offset + ) + destination = _add_physical_tmem_base( + thread_store.partition_D(relative_fragment), tmem_base + o_offset + ) + for fragment_idx in cutlass.range_constexpr(DV // correction_width): + registers = cute.make_rmem_tensor(thread_load.partition_D(relative_fragment).shape, Float32) + source_i = cute.make_tensor( + source.iterator + fragment_idx * correction_width, source.layout + ) + cute.copy(thread_load, source_i, registers) + tcgen05_wait_ld() + cute.arch.fence_view_async_tmem_load() + for i in cutlass.range(cute.size(registers), unroll_full=True): + registers[i] = Float32(registers[i]) * Float32(alpha) + destination_i = cute.make_tensor( + destination.iterator + fragment_idx * correction_width, + destination.layout, + ) + cute.copy(thread_store, registers, destination_i) + tcgen05_wait_st() + cute.arch.fence_view_async_tmem_store() + + +@cute.jit +def _online_update_one_half( + scores: cute.Tensor, + running_max: Float32, + running_sum: Float32, + softmax_scale: Float32, +): + """Apply one FP32 online-softmax update to an M64xN128 score tile.""" + + local_max = fa_utils.fmax_reduce(scores.load(), arch=100) + local_max = Float32(local_max) * softmax_scale + peer_max = cute.arch.shuffle_sync_bfly(local_max, offset=2) + transaction_max = local_max + if peer_max > transaction_max: + transaction_max = peer_max + new_max = running_max + if running_max == -Float32.inf or transaction_max > running_max: + new_max = transaction_max + alpha = Float32(0.0) + if running_max != -Float32.inf: + alpha = cute.math.exp2((running_max - new_max) * Float32(LOG2E), fastmath=True) + probabilities = cute.make_rmem_tensor(scores.shape, Float32) + for i in cutlass.range(cute.size(scores), unroll_full=True): + probabilities[i] = cute.math.exp2( + Float32(scores[i]) * softmax_scale * Float32(LOG2E) - new_max * Float32(LOG2E), + fastmath=True, + ) + transaction_sum = fa_utils.fadd_reduce(probabilities.load(), arch=100) + transaction_sum += cute.arch.shuffle_sync_bfly(transaction_sum, offset=2) + new_sum = running_sum * alpha + transaction_sum + return probabilities, new_max, new_sum, alpha + + +__all__ = [ + "_load_m64_n128_score", + "_online_update_one_half", + "_rescale_m64_partial_o", +] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/tmem.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/tmem.py new file mode 100644 index 000000000000..8256aac562b8 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/tmem.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# +# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see +# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. +"""TMEM load helpers used by the SM100 mainloop.""" + +from __future__ import annotations + +import cutlass.cute as cute +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +from cutlass import Float32, Int32 +from cutlass._mlir.dialects import llvm + +M = 64 +D = 128 +O_OFFSET = 128 + + +@cute.jit +def tcgen05_wait_ld() -> None: + llvm.inline_asm( + None, + [], + "tcgen05.wait::ld.sync.aligned;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@cute.jit +def tcgen05_wait_st() -> None: + llvm.inline_asm( + None, + [], + "tcgen05.wait::st.sync.aligned;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@cute.jit +def _zero_based_tmem_tensor(element_type, layout): + return cute.make_tensor( + cute.make_ptr( + element_type, + Int32(0), + cute.AddressSpace.tmem, + assumed_align=16, + ), + layout, + ) + + +@cute.jit +def _add_physical_tmem_base( + relative: cute.Tensor, + physical_address: Int32, +): + return cute.make_tensor( + cute.make_ptr( + relative.element_type, + physical_address + relative.iterator.toint(), + cute.AddressSpace.tmem, + assumed_align=16, + ), + relative.layout, + ) + + +@cute.jit +def _o_copy_views( + o_template: cute.Tensor, + pv_thread: cute.ThrMma, +): + assert o_template.element_type == Float32 + assert cute.size(o_template) == M * D + relative = _zero_based_tmem_tensor(Float32, o_template.layout) + coordinates = pv_thread.partition_C(cute.make_identity_tensor((M, D))) + tiler = ( + ( + cute.size(relative, mode=[0, 0]), + cute.size(relative, mode=[0, 1]), + ), + ) + return ( + cute.zipped_divide(relative, tiler), + cute.zipped_divide(coordinates, tiler), + ) + + +@cute.jit +def load_m64_o_fp32_256b( + o_template: cute.Tensor, + pv_thread: cute.ThrMma, + physical_tmem_base: Int32, + thread_idx: Int32, +): + relative, coordinates = _o_copy_views(o_template, pv_thread) + atom = cute.make_copy_atom( + tcgen05.Ld16x256bOp(tcgen05.Repetition.x8), + Float32, + ) + tiled_copy = tcgen05.make_tmem_copy( + atom, + relative[None, Int32(0)], + ) + thread_copy = tiled_copy.get_slice(thread_idx) + source = _add_physical_tmem_base( + thread_copy.partition_S(relative), + physical_tmem_base + Int32(O_OFFSET), + ) + register_coordinates = thread_copy.partition_D(coordinates)[None, None, Int32(0)] + registers = cute.make_rmem_tensor( + register_coordinates.shape, + Float32, + ) + cute.copy( + tiled_copy, + source[None, None, Int32(0)], + registers, + ) + tcgen05_wait_ld() + cute.arch.fence_view_async_tmem_load() + return registers, register_coordinates + + +__all__ = [ + "_add_physical_tmem_base", + "_zero_based_tmem_tensor", + "load_m64_o_fp32_256b", + "tcgen05_wait_ld", + "tcgen05_wait_st", +] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/LICENSE.cudnn-frontend b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/LICENSE.cudnn-frontend new file mode 100644 index 000000000000..ee9f673bff93 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/LICENSE.cudnn-frontend @@ -0,0 +1,204 @@ +Copyright (c) 2020-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/__init__.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/__init__.py new file mode 100644 index 000000000000..fc56fddcc772 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/__init__.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# +# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see +# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. +"""GeForce Blackwell (SM120) backend.""" + +from .kernel import make_kernel + +__all__ = ["make_kernel"] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/kernel.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/kernel.py new file mode 100644 index 000000000000..64c4b4879b1b --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/kernel.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# +# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see +# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. +"""SM120 kernel recipe.""" + +from .mainloop import SolAttnForwardSm120 + + +def make_kernel( + *, + debug_route_trace: bool = False, + prefetch_first_exact_k: bool = True, + prefetch_next_route_k: bool = True, +): + return SolAttnForwardSm120( + debug_route_trace=debug_route_trace, + prefetch_first_exact_k=prefetch_first_exact_k, + prefetch_next_route_k=prefetch_next_route_k, + ) + + +__all__ = ["make_kernel"] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/mainloop.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/mainloop.py new file mode 100644 index 000000000000..879034a93904 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/mainloop.py @@ -0,0 +1,1003 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# +# The warp-MMA/TMA skeleton is adapted from NVIDIA cuDNN Frontend +# (https://github.com/NVIDIA/cudnn-frontend), Apache-2.0; its license text +# is vendored at sol_attn/sm120/LICENSE.cudnn-frontend. +"""Fused Sol-Attn forward kernel for GeForce Blackwell SM120. + +The warp-MMA/TMA execution skeleton and online-softmax helpers are adapted +from NVIDIA cuDNN Frontend's SM120 block-sparse-attention kernel. Sol-specific +routing, CTA-local exact-index compaction, approximate block mass, and the +mixed approximate/exact mainloop are implemented here. +""" + +from __future__ import annotations + +import operator + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.hopper_helpers as sm90_utils +from flash_attn.cute import utils as kernel_utils + +from ..common import layout_utils +from ..common.selector import sol_attn_popc_b32, sol_attn_route_is_exact + +M = 64 +N = 64 +D = 128 +DV = 128 +THREADS = 128 +STAGES = 1 + + +class SolAttnForwardSm120: + """M64/N64 warp-MMA Sol-Attn kernel for BF16 D128 inputs.""" + + def __init__( + self, + *, + debug_route_trace: bool = False, + prefetch_first_exact_k: bool = True, + prefetch_next_route_k: bool = True, + ): + self.dtype = cutlass.BFloat16 + self.acc_dtype = cutlass.Float32 + self.tile_shape_qk = (M, N, D) + self.tile_shape_pv = (M, DV, N) + self.num_threads = THREADS + self.q_stage = 1 + self.kv_stage = STAGES + self.debug_route_trace = debug_route_trace + self.prefetch_first_exact_k = prefetch_first_exact_k + self.prefetch_next_route_k = prefetch_next_route_k + + @cute.kernel + def kernel( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mO: cute.Tensor, + mKC: cute.Tensor, + mVC: cute.Tensor, + mThreshold: cute.Tensor, + mLSE: cute.Tensor, + tma_atom_Q: cute.CopyAtom, + tma_atom_K: cute.CopyAtom, + tma_atom_V: cute.CopyAtom, + tma_atom_KC: cute.CopyAtom, + tma_atom_VC: cute.CopyAtom, + tma_atom_O: cute.CopyAtom, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + Q_smem_layout: cute.ComposedLayout, + K_smem_layout: cute.ComposedLayout, + V_smem_layout: cute.ComposedLayout, + O_smem_layout: cute.ComposedLayout, + scale_softmax_log2e: cutlass.Float32, + sink_start_block: cutlass.Int32, + sink_end_block: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + lane = cute.arch.lane_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + q_tile_idx, head_idx, batch_idx = cute.arch.block_idx() + q_tile_idx = cute.arch.make_warp_uniform(q_tile_idx) + head_idx = cute.arch.make_warp_uniform(head_idx) + batch_idx = cute.arch.make_warp_uniform(batch_idx) + + token_count = mK.shape[0] + num_blocks = mKC.shape[0] + num_route_groups = cute.ceil_div(num_blocks, N) + q_start = q_tile_idx * M + q_len = token_count - q_start + if q_len > M: + q_len = cutlass.Int32(M) + threshold = cutlass.Float32(mThreshold[batch_idx, q_tile_idx, head_idx]) + + storage = cutlass.utils.SmemAllocator().allocate(self.shared_storage_t) + if warp == 0 and lane == 0: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_Q) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_K) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_V) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_KC) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_VC) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_O) + + cg = pipeline.CooperativeGroup(pipeline.Agent.Thread) + consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_threads // 32) + cta_layout_vmnk = cute.make_layout((1, 1, 1, 1)) + Q_pipeline = pipeline.PipelineTmaAsync.create( + num_stages=self.q_stage, + producer_group=cg, + consumer_group=consumer_group, + tx_count=cute.size_in_bytes(self.Q_dtype, cute.select(Q_smem_layout, mode=[0, 1])), + barrier_storage=storage.Q_barrier.data_ptr(), + cta_layout_vmnk=cta_layout_vmnk, + ) + K_pipeline = pipeline.PipelineTmaAsync.create( + num_stages=self.kv_stage, + producer_group=cg, + consumer_group=consumer_group, + tx_count=cute.size_in_bytes(self.K_dtype, cute.select(K_smem_layout, mode=[0, 1])), + barrier_storage=storage.K_barrier.data_ptr(), + cta_layout_vmnk=cta_layout_vmnk, + ) + V_pipeline = pipeline.PipelineTmaAsync.create( + num_stages=self.kv_stage, + producer_group=cg, + consumer_group=consumer_group, + tx_count=cute.size_in_bytes(self.V_dtype, cute.select(V_smem_layout, mode=[0, 1])), + barrier_storage=storage.V_barrier.data_ptr(), + cta_layout_vmnk=cta_layout_vmnk, + ) + Q_producer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.q_stage) + Q_consumer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.q_stage) + K_producer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.kv_stage) + K_consumer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.kv_stage) + V_producer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.kv_stage) + V_consumer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.kv_stage) + + sQ = storage.Q_smem.get_tensor(Q_smem_layout.outer, swizzle=Q_smem_layout.inner) + sK = storage.K_smem.get_tensor(K_smem_layout.outer, swizzle=K_smem_layout.inner) + sV = storage.V_smem.get_tensor(V_smem_layout.outer, swizzle=V_smem_layout.inner) + # Q is register-resident after the prologue. Reuse its 16 KiB SMEM + # allocation for route scratch until the same allocation becomes sO + # in the epilogue. This drops the CTA below the 2-block/SM threshold + # on SM120 without changing any route reduction or synchronization. + route_f32_ptr = cute.recast_ptr(storage.Q_smem.data_ptr(), dtype=cutlass.Float32) + route_i32_ptr = cute.recast_ptr(storage.Q_smem.data_ptr(), dtype=cutlass.Int32) + route_sums = cute.make_tensor(route_f32_ptr, cute.make_layout((4, N))) + column_masks = cute.make_tensor(route_f32_ptr + 4 * N, cute.make_layout(N)) + route_indices = cute.make_tensor(route_i32_ptr + 5 * N, cute.make_layout(N)) + route_meta = cute.make_tensor(route_i32_ptr + 6 * N, cute.make_layout(2)) + + mQ_slice = mQ[None, None, head_idx, batch_idx] + mK_slice = mK[None, None, head_idx, batch_idx] + mV_slice = mV[None, None, head_idx, batch_idx] + mO_slice = mO[None, None, head_idx, batch_idx] + mKC_slice = mKC[None, None, head_idx, batch_idx] + mVC_slice = mVC[None, None, head_idx, batch_idx] + if cutlass.const_expr(not self.debug_route_trace): + mLSE_slice = mLSE[None, head_idx, batch_idx] + + gQ = cute.local_tile(mQ_slice, (M, D), coord=(q_tile_idx, 0)) + gK = cute.local_tile(mK_slice, (N, D), coord=(None, 0)) + gV = cute.local_tile(mV_slice, (DV, N), coord=(0, None)) + gKC = cute.local_tile(mKC_slice, (N, D), coord=(None, 0)) + gVC = cute.local_tile(mVC_slice, (DV, N), coord=(0, None)) + gO = cute.local_tile(mO_slice, (M, DV), coord=(q_tile_idx, 0)) + + cta_coord_layout = (0, cute.make_layout(1)) + tQsQ, tQgQ = cute.nvgpu.cpasync.tma_partition( + tma_atom_Q, + *cta_coord_layout, + cute.group_modes(sQ, 0, 2), + cute.group_modes(gQ, 0, 2), + ) + tKsK, tKgK = cute.nvgpu.cpasync.tma_partition( + tma_atom_K, + *cta_coord_layout, + cute.group_modes(sK, 0, 2), + cute.group_modes(gK, 0, 2), + ) + tVsV, tVgV = cute.nvgpu.cpasync.tma_partition( + tma_atom_V, + *cta_coord_layout, + cute.group_modes(sV, 0, 2), + cute.group_modes(gV, 0, 2), + ) + tKCsK, tKCgKC = cute.nvgpu.cpasync.tma_partition( + tma_atom_KC, + *cta_coord_layout, + cute.group_modes(sK, 0, 2), + cute.group_modes(gKC, 0, 2), + ) + tVCsV, tVCgVC = cute.nvgpu.cpasync.tma_partition( + tma_atom_VC, + *cta_coord_layout, + cute.group_modes(sV, 0, 2), + cute.group_modes(gVC, 0, 2), + ) + + cS = cute.make_identity_tensor(self.tile_shape_qk[:2]) + thr_mma_qk = tiled_mma_qk.get_slice(tidx) + tSsQ = thr_mma_qk.partition_A(sQ) + tSsK = thr_mma_qk.partition_B(sK) + tSrQ = tiled_mma_qk.make_fragment_A(tSsQ[None, None, None, 0]) + tSrK = tiled_mma_qk.make_fragment_B(tSsK[None, None, None, 0]) + tSrS = cute.make_rmem_tensor(thr_mma_qk.partition_shape_C((M, N)), self.acc_dtype) + tScS = thr_mma_qk.partition_C(cS) + + thr_mma_pv = tiled_mma_pv.get_slice(tidx) + tOsV = thr_mma_pv.partition_B(sV) + tOrV = tiled_mma_pv.make_fragment_B(tOsV[None, None, None, 0]) + tOrO = cute.make_rmem_tensor(thr_mma_pv.partition_shape_C((M, DV)), self.acc_dtype) + + atom_copy_Q = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp(self.Q_layout.is_m_major_a(), 4), + self.Q_dtype, + ) + atom_copy_K = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp(self.K_layout.is_n_major_b(), 4), + self.K_dtype, + ) + atom_copy_V = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp(self.V_layout.is_n_major_b(), 4), + self.V_dtype, + ) + smem_copy_Q = cute.make_tiled_copy_A(atom_copy_Q, tiled_mma_qk) + smem_copy_K = cute.make_tiled_copy_B(atom_copy_K, tiled_mma_qk) + smem_copy_V = cute.make_tiled_copy_B(atom_copy_V, tiled_mma_pv) + thr_copy_Q = smem_copy_Q.get_slice(tidx) + thr_copy_K = smem_copy_K.get_slice(tidx) + thr_copy_V = smem_copy_V.get_slice(tidx) + tSsQ_copy = thr_copy_Q.partition_S(sQ) + tSrQ_copy = thr_copy_Q.retile(tSrQ) + tSsK_copy = thr_copy_K.partition_S(sK) + tOsV_copy = thr_copy_V.partition_S(sV) + + max_m_layout = cute.make_layout( + cute.size( + layout_utils.reshape_acc_to_mn(tOrO).layout, + mode=[0], + ) + ) + max_m = cute.make_rmem_tensor_like(max_m_layout, cutlass.Float32) + sum_m = cute.make_rmem_tensor_like(max_m, cutlass.Float32) + tOrO.store(cute.full_like(tOrO, 0.0, self.acc_dtype)) + max_m.store(cute.full_like(max_m, float("-inf"), cutlass.Float32)) + sum_m.store(cute.full_like(sum_m, 0.0, cutlass.Float32)) + + if warp == 0: + Q_pipeline.producer_acquire(Q_producer) + cute.copy( + tma_atom_Q, + tQgQ, + tQsQ[None, Q_producer.index], + tma_bar_ptr=Q_pipeline.producer_get_barrier(Q_producer), + ) + Q_pipeline.producer_commit(Q_producer) + Q_producer.advance() + cute.arch.sync_threads() + q_wait = Q_pipeline.consumer_try_wait(Q_consumer) + Q_pipeline.consumer_wait(Q_consumer, q_wait) + q_stage = Q_consumer.index + for k_block in cutlass.range_constexpr(cute.size(tSrQ, mode=[2])): + cute.copy( + smem_copy_Q, + tSsQ_copy[None, None, k_block, q_stage], + tSrQ_copy[None, None, k_block], + ) + Q_pipeline.consumer_release(Q_consumer) + Q_consumer.advance() + + for route_group in cutlass.range(0, num_route_groups, 1, unroll=1): + group_start = route_group * cutlass.Int32(N) + valid_blocks = num_blocks - group_start + if valid_blocks > N: + valid_blocks = cutlass.Int32(N) + + if warp == 0: + if cutlass.const_expr(self.prefetch_next_route_k): + # P19-style terminal handoff: when the previous route + # group had an exact block, its final exact QK already + # refilled this K stage with the current group's KC. + if route_group == 0: + K_pipeline.producer_acquire(K_producer) + cute.copy( + tma_atom_KC, + tKCgKC[None, route_group], + tKCsK[None, K_producer.index], + tma_bar_ptr=K_pipeline.producer_get_barrier(K_producer), + ) + K_pipeline.producer_commit(K_producer) + K_producer.advance() + else: + previous_group_exact_count = cutlass.Int32(route_meta[0]) + if previous_group_exact_count == 0: + K_pipeline.producer_acquire(K_producer) + cute.copy( + tma_atom_KC, + tKCgKC[None, route_group], + tKCsK[None, K_producer.index], + tma_bar_ptr=K_pipeline.producer_get_barrier(K_producer), + ) + K_pipeline.producer_commit(K_producer) + K_producer.advance() + else: + K_pipeline.producer_acquire(K_producer) + cute.copy( + tma_atom_KC, + tKCgKC[None, route_group], + tKCsK[None, K_producer.index], + tma_bar_ptr=K_pipeline.producer_get_barrier(K_producer), + ) + K_pipeline.producer_commit(K_producer) + K_producer.advance() + V_pipeline.producer_acquire(V_producer) + cute.copy( + tma_atom_VC, + tVCgVC[None, route_group], + tVCsV[None, V_producer.index], + tma_bar_ptr=V_pipeline.producer_get_barrier(V_producer), + ) + V_pipeline.producer_commit(V_producer) + V_producer.advance() + + k_wait = K_pipeline.consumer_try_wait(K_consumer) + K_pipeline.consumer_wait(K_consumer, k_wait) + gemm_smem_zero_acc( + tiled_mma_qk, + tSrS, + tSrQ, + tSrK, + tSsK_copy[None, None, None, K_consumer.index], + smem_copy_K, + ) + K_pipeline.consumer_release(K_consumer) + K_consumer.advance() + + reduce_route_columns( + tSrS, + tScS, + route_sums, + warp, + lane, + q_len, + ) + cute.arch.fence_view_async_shared() + cute.arch.sync_threads() + + if warp == 0: + preceding = cutlass.Int32(0) + lane_mask_lt = cutlass.Int32(0x7FFFFFFF) >> (cutlass.Int32(31) - lane) + for word in cutlass.range_constexpr(2): + off = cutlass.Int32(word * 32) + lane + valid = off < valid_blocks + exact = False + if valid: + col_sum = ( + cutlass.Float32(route_sums[0, off]) + + cutlass.Float32(route_sums[1, off]) + + cutlass.Float32(route_sums[2, off]) + + cutlass.Float32(route_sums[3, off]) + ) + col_mean = col_sum * scale_softmax_log2e / cutlass.Float32(q_len) + kv_block = group_start + off + exact = sol_attn_route_is_exact( + q_tile_idx, + kv_block, + col_mean, + threshold, + valid, + ) + exact = exact or ( + kv_block >= sink_start_block and kv_block < sink_end_block + ) + ballot = cutlass.Int32(cute.arch.vote_ballot_sync(exact)) + column_masks[off] = ( + -cutlass.Float32.inf if (exact or not valid) else cutlass.Float32(0.0) + ) + rank = preceding + sol_attn_popc_b32(ballot & lane_mask_lt) + if exact: + route_indices[rank] = group_start + off + preceding += sol_attn_popc_b32(ballot) + if cutlass.const_expr(self.debug_route_trace): + if lane == 0: + mLSE[ + batch_idx, + q_tile_idx, + head_idx, + route_group, + word, + ] = ballot + if lane == 0: + route_meta[0] = preceding + route_meta[1] = valid_blocks + cute.arch.fence_view_async_shared() + cute.arch.sync_threads() + + exact_count = cutlass.Int32(route_meta[0]) + has_approx = exact_count < valid_blocks + if cutlass.const_expr(self.prefetch_first_exact_k): + # Once routing identifies the first exact block, the route KC + # stage is free. Refill it before the approximate softmax/PV + # so the first exact K transfer overlaps that work. + if warp == 0 and exact_count > 0: + first_exact = cutlass.Int32(route_indices[0]) + K_pipeline.producer_acquire(K_producer) + cute.copy( + tma_atom_K, + tKgK[None, first_exact], + tKsK[None, K_producer.index], + tma_bar_ptr=K_pipeline.producer_get_barrier(K_producer), + ) + K_pipeline.producer_commit(K_producer) + K_producer.advance() + v_wait = V_pipeline.consumer_try_wait(V_consumer) + V_pipeline.consumer_wait(V_consumer, v_wait) + if has_approx: + apply_route_mask(tSrS, tScS, column_masks, q_len) + row_scale = online_softmax_route( + tSrS, + tScS, + max_m, + sum_m, + scale_softmax_log2e, + group_start, + token_count, + ) + rescale_o_for_next_acc(tOrO, row_scale) + tOrP_frg = cute.make_rmem_tensor_like(tSrS, self.K_dtype) + tOrP_frg.store(tSrS.load().to(self.K_dtype)) + tOrP = layout_utils.reshape_acc_to_frgA(tOrP_frg) + gemm_rs_smem( + tiled_mma_pv, + tOrO, + tOrP, + tOrV, + tOsV_copy[None, None, None, V_consumer.index], + smem_copy_V, + ) + V_pipeline.consumer_release(V_consumer) + V_consumer.advance() + + if warp == 0 and exact_count > 0: + first_exact = cutlass.Int32(route_indices[0]) + if cutlass.const_expr(not self.prefetch_first_exact_k): + K_pipeline.producer_acquire(K_producer) + cute.copy( + tma_atom_K, + tKgK[None, first_exact], + tKsK[None, K_producer.index], + tma_bar_ptr=K_pipeline.producer_get_barrier(K_producer), + ) + K_pipeline.producer_commit(K_producer) + K_producer.advance() + V_pipeline.producer_acquire(V_producer) + cute.copy( + tma_atom_V, + tVgV[None, first_exact], + tVsV[None, V_producer.index], + tma_bar_ptr=V_pipeline.producer_get_barrier(V_producer), + ) + V_pipeline.producer_commit(V_producer) + V_producer.advance() + + for ordinal in cutlass.range(0, exact_count, 1, unroll=1): + exact_block = cutlass.Int32(route_indices[ordinal]) + k_wait = K_pipeline.consumer_try_wait(K_consumer) + K_pipeline.consumer_wait(K_consumer, k_wait) + gemm_smem_zero_acc( + tiled_mma_qk, + tSrS, + tSrQ, + tSrK, + tSsK_copy[None, None, None, K_consumer.index], + smem_copy_K, + ) + K_pipeline.consumer_release(K_consumer) + K_consumer.advance() + next_ordinal = ordinal + cutlass.Int32(1) + if warp == 0: + if next_ordinal < exact_count: + next_exact = cutlass.Int32(route_indices[next_ordinal]) + K_pipeline.producer_acquire(K_producer) + cute.copy( + tma_atom_K, + tKgK[None, next_exact], + tKsK[None, K_producer.index], + tma_bar_ptr=K_pipeline.producer_get_barrier(K_producer), + ) + K_pipeline.producer_commit(K_producer) + K_producer.advance() + else: + if cutlass.const_expr(self.prefetch_next_route_k): + next_route_group = route_group + cutlass.Int32(1) + if next_route_group < num_route_groups: + # Reuse the K stage released by the final + # exact QK. The next outer prologue supplies + # VC, matching the SM90 P19 partial handoff. + K_pipeline.producer_acquire(K_producer) + cute.copy( + tma_atom_KC, + tKCgKC[None, next_route_group], + tKCsK[None, K_producer.index], + tma_bar_ptr=(K_pipeline.producer_get_barrier(K_producer)), + ) + K_pipeline.producer_commit(K_producer) + K_producer.advance() + block_len = token_count - exact_block * cutlass.Int32(N) + if block_len > N: + block_len = cutlass.Int32(N) + mask_exact_scores(tSrS, tScS, block_len, q_len) + row_scale = online_softmax(tSrS, max_m, sum_m, scale_softmax_log2e) + rescale_o_for_next_acc(tOrO, row_scale) + tOrP_frg = cute.make_rmem_tensor_like(tSrS, self.K_dtype) + tOrP_frg.store(tSrS.load().to(self.K_dtype)) + tOrP = layout_utils.reshape_acc_to_frgA(tOrP_frg) + + v_wait = V_pipeline.consumer_try_wait(V_consumer) + V_pipeline.consumer_wait(V_consumer, v_wait) + gemm_rs_smem( + tiled_mma_pv, + tOrO, + tOrP, + tOrV, + tOsV_copy[None, None, None, V_consumer.index], + smem_copy_V, + ) + V_pipeline.consumer_release(V_consumer) + V_consumer.advance() + if warp == 0 and next_ordinal < exact_count: + next_exact = cutlass.Int32(route_indices[next_ordinal]) + V_pipeline.producer_acquire(V_producer) + cute.copy( + tma_atom_V, + tVgV[None, next_exact], + tVsV[None, V_producer.index], + tma_bar_ptr=V_pipeline.producer_get_barrier(V_producer), + ) + V_pipeline.producer_commit(V_producer) + V_producer.advance() + + final_ratio, lse = finalize_softmax(max_m, sum_m, scale_softmax_log2e) + rescale_o_for_next_acc(tOrO, final_ratio) + if cutlass.const_expr(not self.debug_route_trace): + tScS_mn = layout_utils.reshape_acc_to_mn(tScS) + for m in cutlass.range_constexpr(cute.size(lse)): + row = tScS_mn[m, 0][0] + if tScS_mn[m, 0][1] == 0 and row < q_len: + mLSE_slice[q_start + row] = lse[m] + + tOrO_cvt = cute.make_rmem_tensor_like(tOrO, self.O_dtype) + tOrO_cvt.store(tOrO.load().to(self.O_dtype)) + sO = storage.Q_smem.get_tensor(O_smem_layout.outer, swizzle=O_smem_layout.inner) + tiled_copy_O = cute.make_tiled_copy_C( + cute.make_copy_atom( + cute.nvgpu.warp.StMatrix8x8x16bOp(self.O_layout.is_m_major_c(), 4), + self.O_dtype, + ), + tiled_mma_pv, + ) + tOrO_cv = tiled_copy_O.retile(tOrO_cvt) + tOsO = tiled_copy_O.get_slice(tidx).partition_D(sO) + cute.copy(tiled_copy_O, tOrO_cv, tOsO) + cute.arch.fence_view_async_shared() + cute.arch.sync_threads() + tOsO, tOgO = cute.nvgpu.cpasync.tma_partition( + tma_atom_O, + *cta_coord_layout, + cute.group_modes(sO, 0, 2), + cute.group_modes(gO, 0, 2), + ) + if warp == 0: + cute.copy(tma_atom_O, tOsO, tOgO) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + + @cute.jit + def __call__( + self, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + o: cute.Tensor, + kc: cute.Tensor, + vc: cute.Tensor, + threshold: cute.Tensor, + lse: cute.Tensor, + softmax_scale: cutlass.Float32, + sink_start_block: cutlass.Int32, + sink_end_block: cutlass.Int32, + stream: cuda.CUstream, + ): + q_mkl, k_nkl, kc_nkl = [layout_utils.select(t, [1, 3, 2, 0]) for t in (q, k, kc)] + v_nkl, vc_nkl = [layout_utils.select(t, [3, 1, 2, 0]) for t in (v, vc)] + o_mkl = layout_utils.select(o, [1, 3, 2, 0]) + if cutlass.const_expr(self.debug_route_trace): + lse_target = lse + else: + lse_target = layout_utils.select(lse, [1, 2, 0]) + + self.Q_dtype = q_mkl.element_type + self.K_dtype = k_nkl.element_type + self.V_dtype = v_nkl.element_type + self.O_dtype = o_mkl.element_type + self.Q_layout = utils.LayoutEnum.from_tensor(q_mkl) + self.K_layout = utils.LayoutEnum.from_tensor(k_nkl) + self.V_layout = utils.LayoutEnum.from_tensor(v_nkl) + self.O_layout = utils.LayoutEnum.from_tensor(o_mkl) + assert self.Q_dtype == cutlass.BFloat16 + assert self.K_dtype == cutlass.BFloat16 + assert self.V_dtype == cutlass.BFloat16 + + self.Q_smem_layout = sm90_utils.make_smem_layout_a( + self.Q_layout, + self.tile_shape_qk, + self.Q_dtype, + self.q_stage, + ) + self.K_smem_layout = sm90_utils.make_smem_layout_b( + self.K_layout, + self.tile_shape_qk, + self.K_dtype, + self.kv_stage, + ) + self.V_smem_layout = sm90_utils.make_smem_layout_b( + self.V_layout, + self.tile_shape_pv, + self.V_dtype, + self.kv_stage, + ) + O_smem_layout_staged = sm90_utils.make_smem_layout_epi( + self.O_dtype, + self.O_layout, + self.tile_shape_pv[:2], + 1, + ) + self.O_smem_layout = cute.select(O_smem_layout_staged, mode=[0, 1]) + + @cute.struct + class SharedStorage: + Q_barrier: cute.struct.MemRange[cutlass.Int64, self.q_stage * 2] + K_barrier: cute.struct.MemRange[cutlass.Int64, self.kv_stage * 2] + V_barrier: cute.struct.MemRange[cutlass.Int64, self.kv_stage * 2] + Q_smem: cute.struct.Align[ + cute.struct.MemRange[self.Q_dtype, cute.cosize(self.Q_smem_layout)], + 128, + ] + K_smem: cute.struct.Align[ + cute.struct.MemRange[self.K_dtype, cute.cosize(self.K_smem_layout)], + 128, + ] + V_smem: cute.struct.Align[ + cute.struct.MemRange[self.V_dtype, cute.cosize(self.V_smem_layout)], + 128, + ] + + self.shared_storage_t = SharedStorage + + tiled_mma_qk = cute.make_tiled_mma( + cute.nvgpu.warp.MmaF16BF16Op( + self.Q_dtype, + self.acc_dtype, + (16, 8, 16), + ), + cute.make_layout((4, 1, 1)), + permutation_mnk=(64, 16, 16), + ) + tiled_mma_pv = cute.make_tiled_mma( + cute.nvgpu.warp.MmaF16BF16Op( + self.K_dtype, + self.acc_dtype, + (16, 8, 16), + ), + cute.make_layout((4, 1, 1)), + permutation_mnk=(64, 16, 16), + ) + + g2s_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() + tma_atom_Q, tma_tensor_Q = cute.nvgpu.cpasync.make_tiled_tma_atom( + g2s_op, + q_mkl, + self.Q_smem_layout, + (M, D), + num_multicast=1, + ) + tma_atom_K, tma_tensor_K = cute.nvgpu.cpasync.make_tiled_tma_atom( + g2s_op, + k_nkl, + self.K_smem_layout, + (N, D), + num_multicast=1, + ) + tma_atom_V, tma_tensor_V = cute.nvgpu.cpasync.make_tiled_tma_atom( + g2s_op, + v_nkl, + self.V_smem_layout, + (DV, N), + num_multicast=1, + ) + tma_atom_KC, tma_tensor_KC = cute.nvgpu.cpasync.make_tiled_tma_atom( + g2s_op, + kc_nkl, + self.K_smem_layout, + (N, D), + num_multicast=1, + ) + tma_atom_VC, tma_tensor_VC = cute.nvgpu.cpasync.make_tiled_tma_atom( + g2s_op, + vc_nkl, + self.V_smem_layout, + (DV, N), + num_multicast=1, + ) + s2g_op = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() + tma_atom_O, tma_tensor_O = cute.nvgpu.cpasync.make_tiled_tma_atom( + s2g_op, + o_mkl, + self.O_smem_layout, + (M, DV), + num_multicast=1, + ) + + self.kernel( + tma_tensor_Q, + tma_tensor_K, + tma_tensor_V, + tma_tensor_O, + tma_tensor_KC, + tma_tensor_VC, + threshold, + lse_target, + tma_atom_Q, + tma_atom_K, + tma_atom_V, + tma_atom_KC, + tma_atom_VC, + tma_atom_O, + tiled_mma_qk, + tiled_mma_pv, + self.Q_smem_layout, + self.K_smem_layout, + self.V_smem_layout, + self.O_smem_layout, + softmax_scale * 1.4426950408889634, + sink_start_block, + sink_end_block, + ).launch( + grid=(cute.ceil_div(q_mkl.shape[0], M), q_mkl.shape[2], q_mkl.shape[3]), + block=(self.num_threads, 1, 1), + cluster=(1, 1, 1), + smem=self.shared_storage_t.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + ) + + +@cute.jit +def gemm_smem_zero_acc( + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + tCsB: cute.Tensor, + smem_tiled_copy_B: cute.TiledCopy, +): + acc.fill(0.0) + tCrB_copy = smem_tiled_copy_B.retile(tCrB) + cute.copy( + smem_tiled_copy_B, + tCsB[None, None, 0], + tCrB_copy[None, None, 0], + ) + for k_block in cutlass.range_constexpr(cute.size(tCsB.shape[2])): + if k_block < cute.size(tCsB.shape[2]) - 1: + cute.copy( + smem_tiled_copy_B, + tCsB[None, None, k_block + 1], + tCrB_copy[None, None, k_block + 1], + ) + cute.gemm( + tiled_mma, + acc, + tCrA[None, None, k_block], + tCrB[None, None, k_block], + acc, + ) + + +@cute.jit +def gemm_rs_smem( + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + tCsB: cute.Tensor, + smem_tiled_copy_B: cute.TiledCopy, +): + tCrB_copy = smem_tiled_copy_B.retile(tCrB) + cute.copy( + smem_tiled_copy_B, + tCsB[None, None, 0], + tCrB_copy[None, None, 0], + ) + for k_block in cutlass.range_constexpr(cute.size(tCrA.shape[2])): + if k_block < cute.size(tCrA.shape[2]) - 1: + cute.copy( + smem_tiled_copy_B, + tCsB[None, None, k_block + 1], + tCrB_copy[None, None, k_block + 1], + ) + cute.gemm( + tiled_mma, + acc, + tCrA[None, None, k_block], + tCrB[None, None, k_block], + acc, + ) + + +@cute.jit +def reduce_route_columns( + scores: cute.Tensor, + coords: cute.Tensor, + route_sums: cute.Tensor, + warp: cutlass.Int32, + lane: cutlass.Int32, + q_len: cutlass.Int32, +): + """Reduce M64 score columns using the measured SM120 lane layout.""" + + scores_mn = layout_utils.reshape_acc_to_mn(scores) + coords_mn = layout_utils.reshape_acc_to_mn(coords) + row0 = coords_mn[0, 0][0] + row1 = coords_mn[1, 0][0] + valid0 = row0 < q_len + valid1 = row1 < q_len + for group in cutlass.range_constexpr(8): + n0 = group * 2 + partial0 = cutlass.Float32(0.0) + partial1 = cutlass.Float32(0.0) + if valid0: + partial0 += cutlass.Float32(scores_mn[0, n0]) + partial1 += cutlass.Float32(scores_mn[0, n0 + 1]) + if valid1: + partial0 += cutlass.Float32(scores_mn[1, n0]) + partial1 += cutlass.Float32(scores_mn[1, n0 + 1]) + for offset in (4, 8, 16): + partial0 += cute.arch.shuffle_sync_bfly(partial0, offset=offset) + partial1 += cute.arch.shuffle_sync_bfly(partial1, offset=offset) + if lane < 4: + column = cutlass.Int32(group * 8) + lane * cutlass.Int32(2) + route_sums[warp, column] = partial0 + route_sums[warp, column + 1] = partial1 + + +@cute.jit +def apply_route_mask( + scores: cute.Tensor, + coords: cute.Tensor, + column_masks: cute.Tensor, + q_len: cutlass.Int32, +): + scores_mn = layout_utils.reshape_acc_to_mn(scores) + coords_mn = layout_utils.reshape_acc_to_mn(coords) + for m in cutlass.range_constexpr(cute.size(scores_mn, mode=[0])): + valid_row = coords_mn[m, 0][0] < q_len + for n in cutlass.range_constexpr(cute.size(scores_mn, mode=[1])): + column = coords_mn[m, n][1] + scores_mn[m, n] = ( + cutlass.Float32(scores_mn[m, n]) + cutlass.Float32(column_masks[column]) + if valid_row + else -cutlass.Float32.inf + ) + + +@cute.jit +def mask_exact_scores( + scores: cute.Tensor, + coords: cute.Tensor, + block_len: cutlass.Int32, + q_len: cutlass.Int32, +): + scores_mn = layout_utils.reshape_acc_to_mn(scores) + coords_mn = layout_utils.reshape_acc_to_mn(coords) + for m in cutlass.range_constexpr(cute.size(scores_mn, mode=[0])): + valid_row = coords_mn[m, 0][0] < q_len + for n in cutlass.range_constexpr(cute.size(scores_mn, mode=[1])): + if (not valid_row) or coords_mn[m, n][1] >= block_len: + scores_mn[m, n] = -cutlass.Float32.inf + + +@cute.jit +def online_softmax( + scores: cute.Tensor, + row_max: cute.Tensor, + row_sum: cute.Tensor, + scale_log2e: cutlass.Float32, +): + scores_mn = layout_utils.reshape_acc_to_mn(scores) + row_scale = cute.make_rmem_tensor_like(row_max, cutlass.Float32) + for m in cutlass.range_constexpr(cute.size(row_max)): + score_row = scores_mn[m, None].load() + current_max = kernel_utils.fmax_reduce(score_row, init_val=row_max[m], arch=80) + current_max = cute.arch.warp_reduction_max(current_max, threads_in_group=4) + previous_max = row_max[m] + row_max[m] = current_max + safe_max = cutlass.Float32(0.0) if current_max == -cutlass.Float32.inf else current_max + scaled_max = safe_max * scale_log2e + probabilities = cute.math.exp2(score_row * scale_log2e - scaled_max, fastmath=True) + row_scale[m] = cute.math.exp2((previous_max - safe_max) * scale_log2e, fastmath=True) + row_sum[m] = kernel_utils.fadd_reduce( + probabilities, + init_val=row_sum[m] * row_scale[m], + arch=80, + ) + scores_mn[m, None].store(probabilities) + return row_scale + + +@cute.jit +def online_softmax_route( + scores: cute.Tensor, + coords: cute.Tensor, + row_max: cute.Tensor, + row_sum: cute.Tensor, + scale_log2e: cutlass.Float32, + group_start: cutlass.Int32, + token_count: cutlass.Int32, +): + scores_mn = layout_utils.reshape_acc_to_mn(scores) + coords_mn = layout_utils.reshape_acc_to_mn(coords) + row_scale = cute.make_rmem_tensor_like(row_max, cutlass.Float32) + for m in cutlass.range_constexpr(cute.size(row_max)): + score_row = scores_mn[m, None].load() + current_max = kernel_utils.fmax_reduce(score_row, init_val=row_max[m], arch=80) + current_max = cute.arch.warp_reduction_max(current_max, threads_in_group=4) + previous_max = row_max[m] + row_max[m] = current_max + safe_max = cutlass.Float32(0.0) if current_max == -cutlass.Float32.inf else current_max + probabilities = cute.math.exp2( + score_row * scale_log2e - safe_max * scale_log2e, + fastmath=True, + ) + row_scale[m] = cute.math.exp2((previous_max - safe_max) * scale_log2e, fastmath=True) + masses = cute.make_rmem_tensor_like(scores_mn[m, None], cutlass.Float32) + for n in cutlass.range_constexpr(cute.size(masses)): + block = group_start + coords_mn[m, n][1] + length = token_count - block * cutlass.Int32(N) + if length > N: + length = cutlass.Int32(N) + if length < 0: + length = cutlass.Int32(0) + masses[n] = cutlass.Float32(probabilities[n]) * cutlass.Float32(length) + row_sum[m] = kernel_utils.fadd_reduce( + masses.load(), + init_val=row_sum[m] * row_scale[m], + arch=80, + ) + scores_mn[m, None].store(probabilities) + return row_scale + + +@cute.jit +def finalize_softmax( + row_max: cute.Tensor, + row_sum: cute.Tensor, + scale_log2e: cutlass.Float32, +): + row_sum.store(kernel_utils.warp_reduce(row_sum.load(), operator.add, width=4)) + ratio = cute.make_rmem_tensor_like(row_sum, cutlass.Float32) + lse = cute.make_rmem_tensor_like(row_sum, cutlass.Float32) + for m in cutlass.range_constexpr(cute.size(row_sum)): + total = row_sum[m] + invalid = total == 0.0 or total != total + ratio[m] = cute.arch.rcp_approx(total if not invalid else 1.0) + lse[m] = ( + -cutlass.Float32.inf + if invalid + else (row_max[m] * scale_log2e + cute.math.log2(total, fastmath=True)) + * 0.6931471805599453 + ) + return ratio, lse + + +@cute.jit +def rescale_o_for_next_acc( + output: cute.Tensor, + row_scale: cute.Tensor, +): + output_mn = layout_utils.reshape_acc_to_mn(output) + for m in cutlass.range_constexpr(cute.size(row_scale)): + output_mn[m, None].store(output_mn[m, None].load() * row_scale[m]) + + +__all__ = ["SolAttnForwardSm120"] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py new file mode 100644 index 000000000000..3c74456c7d39 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py @@ -0,0 +1,212 @@ +"""Shape/dtype guard and dense-fallback wrapper around the Sol-Attn kernel. + +The kernel-facing API accepts contiguous BF16 +``[batch, tokens, heads, 128]`` Q/K/V, ``tau``, ``thresh_type``, +``kv_splits``, and an optional exact KV sink range. + +TRT-LLM's dispatch path (``attention_backend/cute_dsl/sol_attn.py``, +``SolAttnAttention``) consumes exactly two names from this module: +``_run_sol_attn_bthd`` and ``sol_attn_supported``. The dense-prefix decision +lives there too, keyed off the normalized timestep forward kwarg. + +CuTe DSL imports and compilation are deferred to first use. Calls the kernel +cannot serve -- wrong shape, dtype, or an architecture with no kernel -- +delegate to dense SDPA rather than failing, and increment +``_SOL_STATS["dense_fallback_calls"]`` so the degradation is countable. Set +``SOL_ATTN_STRICT=1`` to raise instead of falling back. +""" + +from __future__ import annotations + +import functools +import os +from typing import Callable, Optional + +from tensorrt_llm.logger import logger + +HEAD_DIM = 128 +DEFAULT_TAU = 1.0 +DEFAULT_THRESH_TYPE = "diag" +_DEFAULT_SCALE = HEAD_DIM**-0.5 + + +@functools.lru_cache(maxsize=1) +def _load_sol_attn() -> Callable: + """Import the kernel package's public entry point. + + Deferred rather than done at module scope because importing it pulls in + the CuTe DSL, which is expensive and not needed unless Sol-Attn is the + selected backend. + """ + + from .sol_attn import sol_attn + + return sol_attn + + +# Architectures with a Sol-Attn CuTe kernel. Kept in sync with +# ``sol_attn/interface.py::_CUTE_BACKENDS``; duplicated here so the eligibility +# check does not have to import the CuTe DSL. +SUPPORTED_ARCHS = frozenset({(10, 0), (12, 0)}) + + +def sol_attn_ineligible_reason(q) -> Optional[str]: + """Why ``q`` cannot use the CuTe kernel, or None if it can. + + Returns a human-readable reason so the caller can say *why* it fell back, + rather than degrading silently -- an unsupported architecture or head_dim + otherwise shows up only as absent speedup. + """ + try: + import torch + except Exception: # pragma: no cover - torch is a runtime dependency + return "torch is unavailable" + if not (hasattr(q, "is_cuda") and q.is_cuda): + return "q is not a CUDA tensor" + if q.ndim != 4: + return f"q must be 4-D [B, S, H, D], got ndim={q.ndim}" + if q.shape[-1] != HEAD_DIM: + return f"head_dim must be {HEAD_DIM}, got {q.shape[-1]}" + if q.dtype != torch.bfloat16: + return f"dtype must be bfloat16, got {q.dtype}" + try: + arch = tuple(torch.cuda.get_device_capability(q.device)) + except Exception as exc: + return f"could not query device capability: {exc}" + if arch not in SUPPORTED_ARCHS: + return f"no Sol-Attn kernel for SM{arch[0]}{arch[1]}; supported: " + ", ".join( + f"SM{a}{b}" for a, b in sorted(SUPPORTED_ARCHS) + ) + return None + + +def sol_attn_supported(q) -> bool: + """Whether ``q`` is eligible for a Sol-Attn CuTe kernel.""" + + return sol_attn_ineligible_reason(q) is None + + +@functools.lru_cache(maxsize=1) +def _cute_runtime_available() -> bool: + """Whether model dispatch can use one of the optional CuTe kernels.""" + + try: + import cuda.bindings.driver # noqa: F401 + import cutlass.cute # noqa: F401 + except ImportError: + return False + return True + + +def _resolve_kv_splits(q, kv_splits: int | str | None) -> int: + """Resolve the integration-only ``auto`` policy to the public integer API. + + ``auto`` is always 1 here: kv_splits=2/4 was an SM90-only path, and this + build ships SM100/SM120 kernels only. + """ + + if kv_splits in (None, "auto"): + return 1 + return int(kv_splits) + + +def _strict() -> bool: + """Whether SOL_ATTN_STRICT=1 asks us to raise instead of degrading.""" + + return os.environ.get("SOL_ATTN_STRICT", "0") == "1" + + +def _dense_bthd(q, k, v): + import torch + + return torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + ).transpose(1, 2) + + +def _run_sol_attn_bthd( + q, + k, + v, + *, + tau: float = DEFAULT_TAU, + thresh_type: str = DEFAULT_THRESH_TYPE, + kv_splits: int | str | None = "auto", + sink_start: int | None = None, + sink_tokens: int = 0, + dense_fn: Callable | None = None, +): + """Run Sol-Attn on contiguous BTHD tensors, with a safe dense fallback.""" + + q0, k0, v0 = q.contiguous(), k.contiguous(), v.contiguous() + + def dense(): + _SOL_STATS["dense_fallback_calls"] += 1 + if dense_fn is not None: + return dense_fn(q0, k0, v0) + return _dense_bthd(q0, k0, v0) + + reason = sol_attn_ineligible_reason(q0) + if reason is None and (k0.shape != q0.shape or v0.shape != q0.shape): + reason = f"k/v shape must match q {tuple(q0.shape)}" + if reason is None and (k0.dtype != q0.dtype or v0.dtype != q0.dtype): + reason = f"k/v dtype must match q {q0.dtype}" + if reason is not None: + # Same strictness contract as the kernel-exception path below: this is + # the arm that silently turns Sol-Attn into a no-op for a whole run + # (wrong arch, head_dim, or dtype), so it must be visible. + if _strict(): + raise RuntimeError(f"[sol-attn] cannot run the CuTe kernel: {reason}") + logger.warning_once( + f"[sol-attn] falling back to dense SDPA: {reason}. Sol-Attn will not " + "accelerate this run. Set SOL_ATTN_STRICT=1 to raise instead.", + key=("sol_attn_ineligible", reason), + ) + return dense() + + try: + kernel = _load_sol_attn() + out = kernel( + q0, + k0, + v0, + tau=float(tau), + thresh_type=str(thresh_type), + kv_splits=_resolve_kv_splits(q0, kv_splits), + sink_start=sink_start, + sink_tokens=int(sink_tokens), + ) + _SOL_STATS["kernel_calls"] += 1 + return out + except Exception as exc: + if _strict(): + raise + logger.warning_once( + f"[sol-attn] kernel raised {type(exc).__name__}: {exc}; falling back to dense " + "SDPA for this call. Set SOL_ATTN_STRICT=1 to raise instead of silently falling " + "back.", + key=(type(exc).__name__, str(exc)), + ) + return dense() + + +# Lightweight run-validation counters. `kernel_calls` is the census used to +# prove the CuTe kernel actually ran: because forward() falls back to dense +# SDPA on any kernel exception, a run that "works" but never increments this +# was silently dense. Set SOL_ATTN_STRICT=1 to raise instead of falling back. +_SOL_STATS = {"kernel_calls": 0, "dense_fallback_calls": 0} + + +def reset_sol_attn_stats() -> None: + """Zero the counters, e.g. after an untimed warmup generation.""" + + for key in _SOL_STATS: + _SOL_STATS[key] = 0 + + +def get_sol_attn_stats() -> dict[str, int]: + """Return the run-validation counters.""" + + return {key: int(value) for key, value in _SOL_STATS.items()} diff --git a/tensorrt_llm/_torch/visual_gen/models/modeling.py b/tensorrt_llm/_torch/visual_gen/models/modeling.py index 9fd8a02187ea..60bf8ff3a017 100644 --- a/tensorrt_llm/_torch/visual_gen/models/modeling.py +++ b/tensorrt_llm/_torch/visual_gen/models/modeling.py @@ -20,8 +20,12 @@ import torch.nn as nn from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxScheduler +from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn import sol_attn_graph_phase from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig -from tensorrt_llm.visual_gen.sparse_attention import SkipSoftmaxAttentionConfig +from tensorrt_llm.visual_gen.sparse_attention import ( + SkipSoftmaxAttentionConfig, + SolAttnAttentionConfig, +) if TYPE_CHECKING: from tensorrt_llm._torch.visual_gen.cuda_graph_runner import CUDAGraphRunner @@ -74,23 +78,42 @@ def register_cuda_graph_extra_key_fns(self, runner: "CUDAGraphRunner") -> None: the shared registrations. """ sparse_config = self.model_config.attention.sparse_attention_config - if not isinstance(sparse_config, SkipSoftmaxAttentionConfig): - return - disabled_until_timestep = sparse_config.resolve_disabled_until_timestep( - pretrained_config=self.model_config.pretrained_config, - ) - if disabled_until_timestep is None: + if isinstance(sparse_config, SkipSoftmaxAttentionConfig): + disabled_until_timestep = sparse_config.resolve_disabled_until_timestep( + pretrained_config=self.model_config.pretrained_config, + ) + if disabled_until_timestep is None: + return + + # Skip Softmax switches graph-visible attention behavior at the + # timestep boundary while tensor shapes stay unchanged. Key the dense + # and sparse phases separately; if timestep is absent or None, the + # scheduler returns None and the runner omits this key part. + runner.register_extra_key_fn( + "skip_softmax_phase", + lambda *args, **kwargs: SkipSoftmaxScheduler.get_graph_phase_for_timestep( + kwargs.get("timestep"), + disabled_until_timestep=disabled_until_timestep, + ), + ) return - # Skip Softmax switches graph-visible attention behavior at the - # timestep boundary while tensor shapes stay unchanged. Key the dense - # and sparse phases separately; if timestep is absent or None, the - # scheduler returns None and the runner omits this key part. - runner.register_extra_key_fn( - "skip_softmax_phase", - lambda *args, **kwargs: SkipSoftmaxScheduler.get_graph_phase_for_timestep( - kwargs.get("timestep"), - disabled_until_timestep=disabled_until_timestep, - ), - ) + if isinstance(sparse_config, SolAttnAttentionConfig): + disabled_until_timestep = sparse_config.disabled_until_timestep + if disabled_until_timestep is None: + # dense_layers is fixed per layer at construction, so it is + # already baked into each captured graph and needs no key. + return + + # Sol-Attn switches between dense SDPA and the sparse kernel at the + # dense-prefix boundary, again without changing tensor shapes, so + # the two phases must not share a captured graph. + runner.register_extra_key_fn( + "sol_attn_phase", + lambda *args, **kwargs: sol_attn_graph_phase( + kwargs.get("timestep"), + disabled_until_timestep=disabled_until_timestep, + ), + ) + return diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index be338b20564f..da8754cfaf60 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -97,20 +97,26 @@ def __init__( cp_size = vgm.cp_size if vgm else 1 base_backend = config.attention.backend _sa_cfg = config.attention.sparse_attention_config - _is_vsa = ( - base_backend == "CUTEDSL" - and _sa_cfg is not None - and getattr(_sa_cfg, "algorithm", None) == "vsa" - ) - + _sa_algo = getattr(_sa_cfg, "algorithm", None) if _sa_cfg is not None else None + _is_vsa = base_backend == "CUTEDSL" and _sa_algo == "vsa" + _is_sol_attn = base_backend == "CUTEDSL" and _sa_algo == "sol_attn" separate_qkv_cross_attention = ( self.qkv_mode == QKVMode.SEPARATE_QKV and not separate_qkv_is_self_attention ) - # Cross-attention fallback: TRTLLM and CUTEDSL VSA are self-attn only. - if separate_qkv_cross_attention and (base_backend == "TRTLLM" or _is_vsa): + + # Cross-attention fallback: TRTLLM and CUTEDSL VSA/Sol-Attn are self-attn only. + if separate_qkv_cross_attention and ( + base_backend == "TRTLLM" or _is_vsa or _is_sol_attn + ): backend_name = "VANILLA" - requested = f"{base_backend} (VSA)" if _is_vsa else base_backend + requested = ( + f"{base_backend} (VSA)" + if _is_vsa + else f"{base_backend} (Sol-Attn)" + if _is_sol_attn + else base_backend + ) # Warn once per (module class, requested, resolved) triple so the # fallback is visible without per-module-instance log spam. logger.warning_once( @@ -127,6 +133,12 @@ def __init__( f"with context parallelism (Attention2D/Ring, cp_size={cp_size}). Use " f"ulysses or cfg parallelism instead." ) + if _is_sol_attn and cp_size > 1: + raise ValueError( + f"Sol-Attn needs the full token sequence per rank, so it is incompatible " + f"with context parallelism (Attention2D/Ring, cp_size={cp_size}). Use " + f"ulysses or cfg parallelism instead." + ) self.attn_backend = backend_name self.qk_norm = qk_norm self.qk_norm_mode = qk_norm_mode diff --git a/tensorrt_llm/visual_gen/__init__.py b/tensorrt_llm/visual_gen/__init__.py index ec714e3d2bfe..3d0a2cb0636c 100644 --- a/tensorrt_llm/visual_gen/__init__.py +++ b/tensorrt_llm/visual_gen/__init__.py @@ -49,6 +49,7 @@ QuantAttentionConfig, RuntimeLoRAConfig, SkipSoftmaxAttentionConfig, + SolAttnAttentionConfig, SparseAttentionConfig, TeaCacheConfig, TorchCompileConfig, @@ -76,6 +77,7 @@ "QuantAttentionConfig": "tensorrt_llm.visual_gen.args", "RuntimeLoRAConfig": "tensorrt_llm.visual_gen.args", "SkipSoftmaxAttentionConfig": "tensorrt_llm.visual_gen.args", + "SolAttnAttentionConfig": "tensorrt_llm.visual_gen.args", "SparseAttentionConfig": "tensorrt_llm.visual_gen.args", "TeaCacheConfig": "tensorrt_llm.visual_gen.args", "TorchCompileConfig": "tensorrt_llm.visual_gen.args", @@ -134,6 +136,7 @@ def __dir__(): "SparseAttentionConfig", "SkipSoftmaxAttentionConfig", "VideoSparseAttentionConfig", + "SolAttnAttentionConfig", "VAEConfig", "CacheConfig", "TeaCacheConfig", diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index 8be9ea4357a9..38c45af3f71d 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -31,7 +31,11 @@ from tensorrt_llm.llmapi.utils import StrictBaseModel, set_api_status from tensorrt_llm.models.modeling_utils import QuantConfig -from .sparse_attention import SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig +from .sparse_attention import ( + SkipSoftmaxAttentionConfig, + SolAttnAttentionConfig, + VideoSparseAttentionConfig, +) # ============================================================================= # Type aliases @@ -95,7 +99,7 @@ class QuantAttentionConfig(StrictBaseModel): # Discriminated union of sparse attention configs. SparseAttentionConfig = Annotated[ - Union[SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig], + Union[SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig, SolAttnAttentionConfig], Field(discriminator="algorithm"), ] @@ -122,7 +126,8 @@ class AttentionConfig(StrictBaseModel): status="prototype", description=( "Sparse attention recipe. Discriminated by algorithm: " - "skip_softmax (TRTLLM / CUTEDSL backends) or VSA (CUTEDSL backend)." + "skip_softmax (TRTLLM / CUTEDSL backends), vsa (CUTEDSL backend), " + "or sol_attn (CUTEDSL backend)." ), ) @@ -220,6 +225,7 @@ def _validate_sparse_attention_config(self) -> "AttentionConfig": supported_backends = { "skip_softmax": ("TRTLLM", "CUTEDSL"), "vsa": ("CUTEDSL",), + "sol_attn": ("CUTEDSL",), }.get(algo) if supported_backends is None: return self @@ -235,19 +241,23 @@ def _validate_sparse_attention_config(self) -> "AttentionConfig": @model_validator(mode="after") def _validate_cutedsl_quant_sparse_mutex(self) -> "AttentionConfig": - # VSA replaces the dense CuTeDSL path and cannot compose with quantized - # attention. SkipSoftmax is part of that dense path and can compose. + # VSA and Sol-Attn each replace the dense CuTeDSL path and cannot + # compose with quantized attention: create_attention swaps in their own + # backend class, which never consumes quant_attention_config, so the + # request would be silently ignored. SkipSoftmax is part of the dense + # path itself and can compose. + _replaces_dense_path = ("vsa", "sol_attn") if ( self.backend == "CUTEDSL" and self.quant_attention_config is not None and self.sparse_attention_config is not None - and self.sparse_attention_config.algorithm == "vsa" + and self.sparse_attention_config.algorithm in _replaces_dense_path ): raise ValueError( - "CUTEDSL backend: quant_attention_config and VSA " - "sparse_attention_config are mutually exclusive (the " - "CuTeDSLAttention dispatcher selects either the dense path " - "or the sparse VSA path, not both)." + f"CUTEDSL backend: quant_attention_config and " + f"'{self.sparse_attention_config.algorithm}' sparse_attention_config " + "are mutually exclusive (the CuTeDSLAttention dispatcher selects " + "either the dense path or that sparse path, not both)." ) return self @@ -839,6 +849,7 @@ def from_yaml(cls, yaml_path: Union[str, Path], **overrides: Any) -> "VisualGenA "SparseAttentionConfig", "SkipSoftmaxAttentionConfig", "VideoSparseAttentionConfig", + "SolAttnAttentionConfig", "AttentionConfig", "VAEConfig", "ParallelConfig", diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 2bccb10e2652..550525592153 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -57,13 +57,13 @@ class SkipSoftmaxAttentionConfig(BaseSparseAttentionConfig): ) target_sparsity: Optional[float] = PydanticField( default=None, - ge=0.0, + gt=0.0, le=1.0, description="Semantic target sparsity in [0, 1]; requires a calibration formula.", ) disabled_until_timestep: Optional[float] = PydanticField( default=None, - ge=0.0, + gt=0.0, le=1.0, description="Normalized timestep cutoff below which skip-softmax is enabled.", ) @@ -221,6 +221,72 @@ def _ckpt_sparse_attention_config_from_kwargs( return None +class SolAttnAttentionConfig(BaseSparseAttentionConfig): + """Sol-Attn sparse attention configuration for visual generation. + + Dynamic block routing + sparse computation + approximation correction in + one online-softmax pass (arXiv:2607.24027). Kernel is CuTeDSL, sm100 + (B200/GB200) and sm120 (RTX Blackwell) only, head_dim=128, bf16, MHA. + + On an unsupported *shape, dtype, or architecture* the kernel falls back to + dense SDPA and counts the fallback, so setting this config on the wrong GPU + degrades rather than fails. Two cases are not covered by that fallback and do raise: GQA/MQA + (num_kv_heads != num_heads) here at construction, and context parallelism + (cp_size > 1), rejected in visual_gen/modules/attention.py. + """ + + algorithm: Literal["sol_attn"] = "sol_attn" + tau: float = PydanticField( + 1.0, + description="Per-block routing threshold; higher tau routes more blocks sparse.", + ) + thresh_type: Literal["diag", "exact"] = PydanticField( + "diag", + description="Threshold policy forwarded to the kernel (kernel default: 'diag').", + ) + kv_splits: Literal["auto", "1"] = PydanticField( + "auto", + description=( + "KV split policy. Only 1 split is valid on the shipped sm100/sm120 " + "kernels, so 'auto' and '1' are equivalent; the 2/4 path was " + "SM90-only and returns with that kernel. Constrained rather than a " + "free string because any other value is rejected deep inside the " + "kernel, which would silently degrade the whole run to dense." + ), + ) + disabled_until_timestep: Optional[float] = PydanticField( + None, + gt=0.0, + le=1.0, + description=( + "Dense-prefix cutoff on the normalized denoising timestep, with the " + "same sense as skip_softmax's field of the same name: the layer runs " + "dense while timestep >= this value and switches to the sparse kernel " + "below it. Larger timesteps are earlier, noisier steps, so this " + "protects the high-noise prefix. Use None (not 0.0) to disable the " + "prefix; 0.0 is rejected because it would run dense on every step " + "and silently turn Sol-Attn off entirely. " + "The timestep is supplied as a forward kwarg by every VisualGen " + "pipeline, so no per-pipeline wiring is required." + ), + ) + dense_layers: Optional[str] = PydanticField( + None, + description=( + "Comma-separated layer indices/ranges (e.g. '0,2-4') forced dense " + "regardless of the dense prefix. Evaluated per-layer at construction " + "time; no pipeline wiring required." + ), + ) + + def to_sparse_params(self, **kwargs): + # Sol-Attn's knobs are consumed directly by SolAttnAttention.__init__ + # (constructed via CUTEDSL backend dispatch in create_attention), not + # lowered into a shared SparseParams -- the vendored kernel has no + # checkpoint-calibration step to resolve here, unlike skip_softmax. + return None + + class VideoSparseAttentionConfig(StrictBaseModel): """Video Sparse Attention (VSA) sparse-attention recipe (CUTEDSL backend only). @@ -235,7 +301,7 @@ class VideoSparseAttentionConfig(StrictBaseModel): ) vsa_sparsity: float = PydanticField( 0.9, - ge=0.0, + gt=0.0, le=1.0, description=( "Fraction of cubes dropped on the fine stage. 0.0 keeps all cubes " diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index ffeb93d979ae..43f82850f760 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -292,6 +292,7 @@ l0_b200: - unittest/_torch/visual_gen/test_attention_cute_dsl.py - unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py - unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py + - unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py - unittest/_torch/visual_gen/test_attention_flashinfer.py - unittest/_torch/visual_gen/test_attention_trtllm_sage.py - unittest/_torch/visual_gen/test_attention_cudnn.py diff --git a/tests/integration/test_lists/test-db/l0_gb202.yml b/tests/integration/test_lists/test-db/l0_gb202.yml index 79bd803a1890..2053cd8a044c 100644 --- a/tests/integration/test_lists/test-db/l0_gb202.yml +++ b/tests/integration/test_lists/test-db/l0_gb202.yml @@ -20,6 +20,10 @@ l0_gb202: - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] # - unittest/_torch/modeling -k "modeling_qwen3" # https://nvbugs/5234573 - unittest/_torch/attention/test_attention_mla.py + # ------------- Visual Gen tests --------------- + # sm120 (GB202) coverage for the Sol-Attn CuTeDSL kernel; the same file + # is registered in l0_b200.yml for sm100. + - unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py - test_e2e.py::test_ptp_quickstart_bert[VANILLA-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] - test_e2e.py::test_ptp_quickstart_bert[TRTLLM-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_bf16[latency] diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py new file mode 100644 index 000000000000..be7103a06995 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py @@ -0,0 +1,421 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Sol-Attn correctness tests: backend dispatch, config guards, step-context +dense_layers/disabled_until_timestep guards. + +Mirrors test_attention_cute_dsl_vsa.py's structure and scope for its sibling +sparse-attention algorithm. GPU kernel-vs-dense numerical equivalence (the +analogue of VSA's test_cute_kernel_matches_dense_at_full_topk) is not yet +covered here -- see the TODO on test_cute_kernel_matches_dense_placeholder +below for what it needs and why it's deferred, not just missing. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.visual_gen.attention_backend import CuTeDSLAttention +from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn import ( + SolAttnAttention, + _parse_dense_layers, + sol_attn_graph_phase, +) +from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention +from tensorrt_llm._torch.visual_gen.config import ( + DiffusionModelConfig, + create_attention_metadata_state, +) +from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode +from tensorrt_llm.visual_gen.args import AttentionConfig, SolAttnAttentionConfig + + +def test_cute_dsl_factory_dispatches_dense_and_sol_attn() -> None: + dense_config = AttentionConfig(backend="CUTEDSL") + dense_attention = create_attention( + backend="CUTEDSL", + layer_idx=0, + num_heads=8, + head_dim=128, + attention_config=dense_config, + ) + + sparse_config = SolAttnAttentionConfig(tau=2.0, disabled_until_timestep=0.9545) + sol_attn_config = AttentionConfig(backend="CUTEDSL", sparse_attention_config=sparse_config) + sol_attn_attention = create_attention( + backend="CUTEDSL", + layer_idx=0, + num_heads=8, + head_dim=128, + attention_config=sol_attn_config, + ) + + assert isinstance(dense_attention, CuTeDSLAttention) + assert isinstance(sol_attn_attention, SolAttnAttention) + assert sol_attn_attention.tau == 2.0 + assert sol_attn_attention.disabled_until_timestep == 0.9545 + + +def _make_config( + hidden_size: int, + num_heads: int, + head_dim: int, + backend: str, + sol_attn_tau: "float | None" = None, +) -> DiffusionModelConfig: + """Minimal DiffusionModelConfig for one Attention module.""" + pretrained_config = SimpleNamespace( + hidden_size=hidden_size, + num_attention_heads=num_heads, + attention_head_dim=head_dim, + eps=1e-6, + ) + sparse_attention_config = ( + SolAttnAttentionConfig(tau=sol_attn_tau) if sol_attn_tau is not None else None + ) + config = DiffusionModelConfig( + pretrained_config=pretrained_config, + attention=AttentionConfig(backend=backend, sparse_attention_config=sparse_attention_config), + skip_create_weights_in_init=False, + ) + config.attention_metadata_state = ( + create_attention_metadata_state() if backend == "TRTLLM" else None + ) + return config + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Sol-Attn needs CUDA") +def test_sol_attn_falls_back_to_vanilla_for_cross_attention(): + """Cross-attention (SEPARATE_QKV) falls back to VANILLA -- Sol-Attn is self-attn only.""" + device = torch.device("cuda") + dtype = torch.bfloat16 + cfg = _make_config( + hidden_size=64, num_heads=4, head_dim=16, backend="CUTEDSL", sol_attn_tau=1.0 + ) + cross_attn = ( + Attention(64, 4, qkv_mode=QKVMode.SEPARATE_QKV, config=cfg) + .to(device=device, dtype=dtype) + .eval() + ) + assert cross_attn.attn_backend == "VANILLA", ( + f"Sol-Attn on cross-attention should fall back to VANILLA, got {cross_attn.attn_backend!r}" + ) + + +def test_sol_attn_with_context_parallelism_raises(): + """Sol-Attn + Attention2D/Ring must error at construction (needs the full sequence per rank).""" + pretrained_config = SimpleNamespace( + hidden_size=64, + num_attention_heads=4, + attention_head_dim=16, + eps=1e-6, + ) + cfg = DiffusionModelConfig( + pretrained_config=pretrained_config, + attention=AttentionConfig( + backend="CUTEDSL", + sparse_attention_config=SolAttnAttentionConfig(tau=1.0), + ), + skip_create_weights_in_init=False, + ) + cfg.visual_gen_mapping = SimpleNamespace( + ring_size=1, + ring_group=None, + ulysses_size=1, + ulysses_group=None, + attn2d_row_size=2, + attn2d_col_size=2, + attn2d_row_group=None, + attn2d_col_group=None, + cp_size=4, + ) + with pytest.raises(ValueError, match="incompatible with context parallelism"): + Attention(64, 4, qkv_mode=QKVMode.FUSE_QKV, config=cfg) + + +def test_sol_attn_rejects_gqa_mqa(): + """Sol-Attn is MHA-only; num_kv_heads != num_heads must fail fast at construction.""" + with pytest.raises(AssertionError, match="MHA-only"): + SolAttnAttention(layer_idx=0, num_heads=8, head_dim=128, num_kv_heads=2) + + +@pytest.mark.parametrize( + "spec,expected", + [ + (None, frozenset()), + ("", frozenset()), + ("0", frozenset({0})), + ("0,2,4", frozenset({0, 2, 4})), + ("0-3", frozenset({0, 1, 2, 3})), + ("0-1,5,7-8", frozenset({0, 1, 5, 7, 8})), + (" 0 , 2 ", frozenset({0, 2})), + ], + ids=["none", "empty", "single", "list", "range", "mixed", "whitespace"], +) +def test_parse_dense_layers(spec, expected): + assert _parse_dense_layers(spec) == expected + + +@pytest.mark.parametrize( + "timestep,expected", + [ + (0.99, 0), # early/noisy -> dense prefix + (0.9545, 0), # exactly at the cutoff -> still dense + (0.95, 1), # past the cutoff -> sparse + (0.0, 1), # final step -> sparse + (None, None), # no timestep -> no phase to distinguish + ], + ids=["early", "at-cutoff", "past-cutoff", "final", "missing"], +) +def test_graph_phase_matches_skip_softmax_sense(timestep, expected): + """Phase 0 is the dense prefix, 1 the sparse phase, None when undecidable. + + Same contract as SkipSoftmaxScheduler.get_graph_phase_for_timestep. + """ + assert sol_attn_graph_phase(timestep, disabled_until_timestep=0.9545) == expected + + +def test_graph_phase_none_when_prefix_unset(): + assert sol_attn_graph_phase(0.5, disabled_until_timestep=None) is None + + +def test_graph_phase_accepts_tensor_timestep(): + """Pipelines pass a tensor; a 0-d or 1-element tensor must work.""" + assert sol_attn_graph_phase(torch.tensor(0.99), disabled_until_timestep=0.95) == 0 + assert sol_attn_graph_phase(torch.tensor([0.10]), disabled_until_timestep=0.95) == 1 + + +def test_dense_prefix_uses_sdpa_and_skips_kernel(monkeypatch): + """Inside the dense prefix the sparse kernel must not be invoked at all.""" + import tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn as sol_attn_mod + + def _fail_if_called(*args, **kwargs): + raise AssertionError("kernel must not run inside the dense prefix") + + monkeypatch.setattr(sol_attn_mod, "_sol_attn_run", _fail_if_called) + + attn = SolAttnAttention(layer_idx=0, num_heads=2, head_dim=16) + attn.disabled_until_timestep = 0.9 + q = k = v = torch.randn(1, 4, 2, 16) + out = attn.forward(q, k, v, timestep=torch.tensor(0.95)) + assert out.shape == q.shape + assert torch.isfinite(out).all() + + +def test_missing_timestep_fails_open_to_sparse(monkeypatch): + """Without a timestep the prefix cannot be applied; run sparse, do not raise. + + Matches the CuTeDSL skip-softmax path's fail-open choice. + """ + import tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn as sol_attn_mod + + called = {"n": 0} + + def _record(*args, **kwargs): + called["n"] += 1 + return args[0] + + monkeypatch.setattr(sol_attn_mod, "_sol_attn_run", _record) + + attn = SolAttnAttention(layer_idx=0, num_heads=2, head_dim=16) + attn.disabled_until_timestep = 0.9 + q = k = v = torch.randn(1, 4, 2, 16) + attn.forward(q, k, v) # no timestep kwarg + assert called["n"] == 1, "expected the sparse kernel, not a silent dense fallback" + + +def test_sol_attn_dense_layers_guard_skips_kernel(monkeypatch): + """A layer_idx in dense_layers must use the dense SDPA path and never invoke the kernel.""" + import tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn as sol_attn_mod + + def _fail_if_called(*args, **kwargs): + raise AssertionError("kernel must not be invoked for a dense_layers-forced layer") + + monkeypatch.setattr(sol_attn_mod, "_sol_attn_run", _fail_if_called) + + attn = SolAttnAttention(layer_idx=3, num_heads=2, head_dim=16) + attn.dense_layers = frozenset({3}) + q = k = v = torch.randn(1, 4, 2, 16) + out = attn.forward(q, k, v) + assert out.shape == q.shape + assert torch.isfinite(out).all() + + +def _make_solattn_model(disabled_until_timestep=None, dense_layers=None): + """Minimal BaseDiffusionModel carrying a Sol-Attn sparse config.""" + from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel + + pretrained_config = SimpleNamespace( + hidden_size=64, num_attention_heads=4, attention_head_dim=16, eps=1e-6 + ) + config = DiffusionModelConfig( + pretrained_config=pretrained_config, + attention=AttentionConfig( + backend="CUTEDSL", + sparse_attention_config=SolAttnAttentionConfig( + tau=2.0, + disabled_until_timestep=disabled_until_timestep, + dense_layers=dense_layers, + ), + ), + skip_create_weights_in_init=False, + ) + return BaseDiffusionModel(config) + + +def _graph_runner(): + from tensorrt_llm._torch.visual_gen.cuda_graph_runner import ( + CUDAGraphRunner, + CUDAGraphRunnerConfig, + ) + + return CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) + + +def test_cuda_graph_key_separates_dense_prefix_from_sparse_phase(): + """The prefix swaps kernels without changing any tensor shape, so a graph + captured in the dense prefix must not be replayed for the sparse phase.""" + model = _make_solattn_model(disabled_until_timestep=0.9) + runner = _graph_runner() + model.register_cuda_graph_extra_key_fns(runner) + + base = {"hidden_states": torch.empty(1, 8, 64)} + key_dense = runner.get_graph_key(**base, timestep=torch.empty(1).fill_(0.95)) + key_sparse = runner.get_graph_key(**base, timestep=torch.empty(1).fill_(0.10)) + + assert key_dense != key_sparse, ( + "dense-prefix and sparse phases share a CUDA graph key despite running " + "different kernels; a graph captured in one phase would be replayed in " + "the other" + ) + + +def test_cuda_graph_key_unregistered_without_prefix(): + """dense_layers alone is fixed per layer, so it needs no graph key.""" + model = _make_solattn_model(disabled_until_timestep=None, dense_layers="0,2") + runner = _graph_runner() + model.register_cuda_graph_extra_key_fns(runner) + + base = {"hidden_states": torch.empty(1, 8, 64)} + key_a = runner.get_graph_key(**base, timestep=torch.empty(1).fill_(0.95)) + key_b = runner.get_graph_key(**base, timestep=torch.empty(1).fill_(0.10)) + assert key_a == key_b, "no phase key should be registered without a dense prefix" + + +# --- kernel-wrapper eligibility / strictness (sol_attn_backend.py) ----------- +# These run on CPU: every path here is pure Python guard logic, and the CPU +# tensor is itself one of the ineligible cases. + + +def _backend_mod(): + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell import sol_attn_backend + + return sol_attn_backend + + +@pytest.mark.parametrize( + "make,expect", + [ + (lambda: torch.randn(1, 4, 2, 128), "not a CUDA tensor"), + (lambda: torch.randn(1, 4, 2, 64), "not a CUDA tensor"), + (lambda: torch.randn(1, 4, 128), "not a CUDA tensor"), + ], + ids=["cpu-ok-shape", "cpu-wrong-head-dim", "cpu-wrong-rank"], +) +def test_ineligible_reason_is_reported(make, expect): + """Ineligibility must name a reason, never fail silently.""" + reason = _backend_mod().sol_attn_ineligible_reason(make()) + assert reason is not None and expect in reason + assert not _backend_mod().sol_attn_supported(make()) + + +def test_strict_raises_on_ineligible_input(monkeypatch): + """SOL_ATTN_STRICT=1 must cover the shape/dtype/arch path, not just kernel + exceptions. Without this, an unsupported arch degrades to dense silently + even under STRICT, and the counters the PR relies on cannot be trusted.""" + sab = _backend_mod() + monkeypatch.setenv("SOL_ATTN_STRICT", "1") + q = k = v = torch.randn(1, 4, 2, 128) # CPU -> ineligible + with pytest.raises(RuntimeError, match="cannot run the CuTe kernel"): + sab._run_sol_attn_bthd(q, k, v) + + +def test_ineligible_falls_back_to_dense_and_counts(monkeypatch): + """Without STRICT the same input degrades to dense and increments the counter.""" + sab = _backend_mod() + monkeypatch.delenv("SOL_ATTN_STRICT", raising=False) + sab.reset_sol_attn_stats() + q = k = v = torch.randn(1, 4, 2, 128) + out = sab._run_sol_attn_bthd(q, k, v) + ref = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) + ).transpose(1, 2) + assert torch.allclose(out, ref), "dense fallback must be plain SDPA" + assert sab.get_sol_attn_stats()["dense_fallback_calls"] == 1 + assert sab.get_sol_attn_stats()["kernel_calls"] == 0 + + +def test_supported_archs_matches_kernel_dispatch_map(): + """SUPPORTED_ARCHS is a hand-copy of interface.py's _CUTE_BACKENDS. If they + drift, eligibility silently rejects an arch the kernel actually supports.""" + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.sol_attn import interface + + assert _backend_mod().SUPPORTED_ARCHS == frozenset(interface._CUTE_BACKENDS) + + +def test_quant_attention_config_rejected_with_sol_attn(): + """Sol-Attn replaces the dense CuTeDSL path, so quantized attention cannot + compose with it; accepting the pair would silently ignore the quant request.""" + from tensorrt_llm.visual_gen.args import QuantAttentionConfig + + with pytest.raises(ValueError, match="mutually exclusive"): + AttentionConfig( + backend="CUTEDSL", + quant_attention_config=QuantAttentionConfig(), + sparse_attention_config=SolAttnAttentionConfig(tau=2.0), + ) + + +def test_zero_cutoff_rejected(): + """0.0 is the natural thing to type for 'no prefix', but it would run dense + on every step and turn Sol-Attn off entirely. Must be rejected, not silent.""" + with pytest.raises(ValueError): + SolAttnAttentionConfig(tau=2.0, disabled_until_timestep=0.0) + assert SolAttnAttentionConfig(tau=2.0).disabled_until_timestep is None + + +@pytest.mark.skip( + reason=( + "TODO(sol-attn): numerical equivalence vs dense SDPA at zero/near-zero routing " + "(the analogue of VSA's test_cute_kernel_matches_dense_at_full_topk) needs the " + "exact tau/thresh_type combination that guarantees full (non-sparse) block " + "routing, which is not simply tau=0 because Sol-Attn's routing is score-derived " + "rather than a plain top-k like VSA's. Deriving it requires reading " + "cute_dsl_kernels/blackwell/sol_attn/interface.py's routing math and calibrating " + "rtol/atol on real sm100 hardware. This would be a unit-level complement to the " + "end-to-end accuracy evidence recorded in the pull request, not a replacement." + ) +) +def test_cute_kernel_matches_dense_placeholder(): + pass + + +def test_kv_splits_rejects_unsupported_value(): + """kv_splits is constrained at the config layer: an out-of-range value is + otherwise rejected deep inside the kernel and caught by the blanket + except, silently degrading the entire run to dense attention.""" + with pytest.raises(ValueError): + SolAttnAttentionConfig(tau=2.0, kv_splits="4") + assert SolAttnAttentionConfig(tau=2.0).kv_splits == "auto" From c003e55314f4543743fa3226350e16a7bf984060 Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Tue, 1 Sep 2026 09:41:28 -0700 Subject: [PATCH 02/22] [TRTLLM-15917][fix] Keep the Sol-Attn CuTe DSL launch opaque to Dynamo `_run_sol_attn_bthd` was missing `@torch.compiler.disable`, so under torch.compile Dynamo traced *into* the CuTe DSL JIT builder -- symbolically evaluating MLIR op construction (`OpView.__new__`) and driver handles (`CUstream.__new__`) -- and retraced on every call. Every sibling CuTe DSL launch boundary already carries the decorator (`cute_dsl/fmha.py`, `cute_dsl/vsa.py`, `video_sparse_attention/interface.py`); Sol-Attn was the only one without it. The failure was silent: no error, just a run that looked like torch.compile not paying off. A second, independent graph break came from the dense-prefix decision, which reads a scalar out of the timestep tensor. A bare `.item()` under Dynamo breaks the enclosing transformer block once per attention layer, so the extraction moves into a `@torch.compiler.disable`d `_dense_by_step` helper, mirroring `cute_dsl/fmha.py`'s delayed scalar extraction and VSA's `_get_vsa_inputs`. It returns a host-side bool, so the dense and sparse phases still compile as separate graphs -- they run different kernels. Behaviour is unchanged, including the fail-open path when no timestep arrives. Measured on B200 (WAN2.2-TI2V-5B, 704x1280, 121 frames, 50 steps, seed 42): | Configuration | denoise | S vs eager dense | |------------------------------|---------|------------------| | dense, eager | 66.90 s | 1.000x | | Sol-Attn, eager | 59.38 s | 1.127x | | Sol-Attn, CUDA graphs | 56.29 s | 1.188x | | dense + torch.compile | 45.92 s | 1.457x | | Sol-Attn + torch.compile | 36.21 s | 1.847x | Against the compiled dense baseline -- the comparison that matters, since torch.compile needs none of this feature -- Sol-Attn gives S = 1.268x and a 21.15% time reduction, at LPIPS 0.0268 versus that same baseline. Before this fix the same configuration measured 2496.9 s mean denoise, a 69x difference. Repetitions agree to 0.03 s, and the run logs no dense fallback; a fallback could not be 21% faster than the dense path it falls back to. Two tests assert both boundaries stay Dynamo-opaque. A missing decorator is how this arose and it fails silently, so the convention needs a test rather than only a comment. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../attention_backend/cute_dsl/sol_attn.py | 46 +++++++++++-------- .../blackwell/sol_attn_backend.py | 9 +++- .../test_attention_cute_dsl_sol_attn.py | 28 +++++++++++ 3 files changed, 62 insertions(+), 21 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py index cd600f91fcbe..f79e7dd2fc5c 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py @@ -149,6 +149,32 @@ def __init__( self.disabled_until_timestep = getattr(cfg, "disabled_until_timestep", None) self.dense_layers = _parse_dense_layers(getattr(cfg, "dense_layers", None)) + # The `.item()` in here would graph-break the enclosing block once per + # attention layer, so keep it in eager (as cute_dsl/fmha.py and VSA's + # `_get_vsa_inputs` do). Returns a host-side bool, so the dense and sparse + # phases still compile as separate graphs -- they run different kernels. + @torch.compiler.disable + def _dense_by_step(self, timestep) -> bool: + phase = sol_attn_graph_phase( + timestep, + disabled_until_timestep=self.disabled_until_timestep, + ) + if phase is None: + # Fail open, matching the CuTeDSL skip-softmax path: without a + # timestep we cannot tell which phase we are in, so run the + # sparse kernel rather than silently forcing dense forever. + # This degrades quality rather than raising, so say so once. + logger.warning_once( + "SolAttnAttentionConfig.disabled_until_timestep=" + f"{self.disabled_until_timestep} is set, but no `timestep` reached " + "the Sol-Attn forward call. The dense prefix it requests will not " + "be applied. Ensure the pipeline passes a normalized timestep, or " + "unset disabled_until_timestep.", + key="sol_attn_missing_timestep", + ) + return False + return phase == 0 + def forward( self, q: torch.Tensor, @@ -160,25 +186,7 @@ def forward( dense_by_layer = self.layer_idx in self.dense_layers dense_by_step = False if self.disabled_until_timestep is not None: - phase = sol_attn_graph_phase( - kwargs.get("timestep"), - disabled_until_timestep=self.disabled_until_timestep, - ) - if phase is None: - # Fail open, matching the CuTeDSL skip-softmax path: without a - # timestep we cannot tell which phase we are in, so run the - # sparse kernel rather than silently forcing dense forever. - # This degrades quality rather than raising, so say so once. - logger.warning_once( - "SolAttnAttentionConfig.disabled_until_timestep=" - f"{self.disabled_until_timestep} is set, but no `timestep` reached " - "the Sol-Attn forward call. The dense prefix it requests will not " - "be applied. Ensure the pipeline passes a normalized timestep, or " - "unset disabled_until_timestep.", - key="sol_attn_missing_timestep", - ) - else: - dense_by_step = phase == 0 + dense_by_step = self._dense_by_step(kwargs.get("timestep")) if dense_by_layer or dense_by_step: return torch.nn.functional.scaled_dot_product_attention( q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py index 3c74456c7d39..2b441729fd90 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py @@ -22,6 +22,8 @@ import os from typing import Callable, Optional +import torch + from tensorrt_llm.logger import logger HEAD_DIM = 128 @@ -117,8 +119,6 @@ def _strict() -> bool: def _dense_bthd(q, k, v): - import torch - return torch.nn.functional.scaled_dot_product_attention( q.transpose(1, 2), k.transpose(1, 2), @@ -126,6 +126,11 @@ def _dense_bthd(q, k, v): ).transpose(1, 2) +# Opaque to Dynamo, like every other CuTe DSL launch boundary here (see +# cute_dsl/fmha.py, video_sparse_attention/interface.py). Otherwise Dynamo +# traces into the CuTe DSL JIT builder and retraces on every call: 69x slower +# on B200 (denoise 2496.9 s vs 36.2 s), silently, as if compile just didn't help. +@torch.compiler.disable def _run_sol_attn_bthd( q, k, diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py index be7103a06995..9bbe8030929b 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py @@ -419,3 +419,31 @@ def test_kv_splits_rejects_unsupported_value(): with pytest.raises(ValueError): SolAttnAttentionConfig(tau=2.0, kv_splits="4") assert SolAttnAttentionConfig(tau=2.0).kv_splits == "auto" + + +def _is_dynamo_disabled(fn) -> bool: + """True if `fn` is wrapped by torch.compiler.disable / torch._dynamo.disable.""" + target = getattr(fn, "__func__", fn) + return bool(getattr(target, "_torchdynamo_disable", False)) + + +def test_kernel_launch_is_opaque_to_dynamo(): + """The CuTe DSL launch boundary must be @torch.compiler.disable'd. + + Without it Dynamo traces into the CuTe DSL JIT builder and retraces on every + call: 69x slower on B200 (denoise 2496.9 s vs 36.2 s), and silent -- it looks + like torch.compile simply not paying off. + """ + assert _is_dynamo_disabled(_backend_mod()._run_sol_attn_bthd), ( + "_run_sol_attn_bthd must be decorated with @torch.compiler.disable" + ) + + +def test_timestep_scalar_read_is_opaque_to_dynamo(): + """The dense-prefix `.item()` must stay in eager. + + Otherwise it graph-breaks the enclosing block once per attention layer. + """ + assert _is_dynamo_disabled(SolAttnAttention._dense_by_step), ( + "SolAttnAttention._dense_by_step must be decorated with @torch.compiler.disable" + ) From 47c7cbff5decdda9698e7769fbb833ff98737c0b Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Tue, 1 Sep 2026 10:20:07 -0700 Subject: [PATCH 03/22] [TRTLLM-15917][doc] Declare sol_attn_backend.py's upstream derivation `sol_attn_backend.py` is adapted from upstream's `techniques/sparse_backends/sol_attn_backend.py`, but THIRD_PARTY_NOTICES.md scoped the vendoring to the `sol_attn/` package only. Upstream's version of this file lives outside that package, so the notices' statement of what is carried was inaccurate, and the file that carries our `@torch.compiler.disable` sat outside the currency check the notices tell maintainers to run. Records the derivation, which subset is carried (the kernel wrapper: shape guard, dense fallback, counters -- not upstream's diffusers/HunyuanVideo/Morton model-integration half), and the deliberate divergences a re-sync must preserve rather than overwrite. Also notes that upstream guards the same call with a `torch.library.custom_op` plus `register_fake`, which keeps the kernel in the compiled graph instead of breaking the graph at it, and is arguably better than the `@torch.compiler.disable` used here. That form was not adopted because `torch.compiler.disable` is what every other CuTe DSL entry point in this repository uses and what this PR's measurements were taken with; migrating is a reasonable follow-up. Both projects are Apache-2.0, so this is an attribution-accuracy fix, not a licensing one. Documentation and one docstring only; no behaviour change. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../blackwell/sol_attn/THIRD_PARTY_NOTICES.md | 39 +++++++++++++++++++ .../blackwell/sol_attn_backend.py | 5 +++ 2 files changed, 44 insertions(+) diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md index 6d78d739802a..a609d8151e3e 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md @@ -45,6 +45,45 @@ scaffold still derive from that project. `preprocess.py` implements the routing/threshold stage in Triton, so Triton is a required runtime dependency on every Sol-Attn path, not only a fallback. +## A derived file outside this directory + +`../sol_attn_backend.py` is **not** part of the vendored package above, but it +is a derivative work and is recorded here because this file is where a future +currency check starts. It is adapted from upstream's +`techniques/sparse_backends/sol_attn_backend.py` (same branch and commit as the +package). Only the kernel-wrapper subset is carried -- the shape/dtype guard, +the dense-SDPA fallback, and the call counters. Upstream's model-integration +half is not carried: the diffusers self-attention dispatch hook, HunyuanVideo's +padded `[video, text]` MMDiT handling, and model-level Morton ordering. + +Deliberate divergences from upstream in that file, all of which a re-sync must +preserve rather than overwrite: + +| Divergence | Why | +|---|---| +| `logger.warning_once` replaces `print()` | fallbacks must be suppressible and routed through the repo's logger | +| `dense_fallback_calls` counter added | makes a silently-degraded run countable, not just visible in stderr | +| `sol_attn_ineligible_reason()` added | names the specific reason (arch / head_dim / dtype) instead of one boolean | +| `SOL_ATTN_STRICT=1` also covers the eligibility path | upstream raises only on kernel exceptions, so an ineligible run stayed silent | +| `@torch.compiler.disable` on `_run_sol_attn_bthd` | see below | + +**Upstream solves the `torch.compile` problem differently, and arguably +better.** Its `sol_attn_backend.py` wraps the same call in a +`torch.library.custom_op` (`sana_sol_attn::self_attention`) with a +`register_fake` returning `torch.empty_like(q)`, which keeps the kernel in the +compiled graph as an opaque node instead of breaking the graph at it; a second +consumer (`models/ltx2.5-refiner/GB200/sol_attention.py`) applies +`torch.compiler.disable` at the call site behind a flag. This repository uses +`@torch.compiler.disable` on the launch boundary instead, matching the +convention every other CuTe DSL entry point here already follows +(`attention_backend/cute_dsl/fmha.py`, +`cute_dsl_kernels/blackwell/video_sparse_attention/interface.py`). Without some +such guard Dynamo traces into the CuTe DSL JIT builder and retraces on every +call -- measured at 69x slower on B200. Migrating to the `custom_op` form would +remove the per-layer graph break and is a reasonable follow-up; it was not done +here because the `torch.compiler.disable` form is what this repository's other +kernels use and what the measurements above were taken with. + The runtime also depends on NVIDIA CUTLASS / CuTe DSL, cuda-python, and PyTorch. Those dependencies are not redistributed by this repository and remain subject to their respective licenses. diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py index 2b441729fd90..5a8e83a16e6e 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py @@ -1,5 +1,10 @@ """Shape/dtype guard and dense-fallback wrapper around the Sol-Attn kernel. +Adapted from upstream's ``techniques/sparse_backends/sol_attn_backend.py`` at +the pin in ``sol_attn/THIRD_PARTY_NOTICES.md``, which records exactly which +subset is carried and how this version deliberately diverges. Check that file +before re-syncing against upstream. + The kernel-facing API accepts contiguous BF16 ``[batch, tokens, heads, 128]`` Q/K/V, ``tau``, ``thresh_type``, ``kv_splits``, and an optional exact KV sink range. From 5e077c1127609a5217f99e8821c9c85e193adc5c Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Tue, 1 Sep 2026 18:18:01 -0700 Subject: [PATCH 04/22] [TRTLLM-15917][feat] Drop sm120 from the Sol-Attn port sm120 (RTX Blackwell) had kernel-level evidence only -- 9/9 sweep points resolving to `cute_sm120` -- and was never validated end to end, because the only available sm120 hardware was a 32 GB RTX 5090 that cannot hold Wan2.2-TI2V-5B (OOM at 27.2 GiB during model load). Shipping only what is measured end to end is the same reasoning already applied to sm89 and sm90. It also removes a structural problem. `cute_dsl_fmha_fwd`, the dense CuTe DSL kernel the CUTEDSL backend uses, supports sm_100a/sm_103a and not sm120, while Sol-Attn's dense paths -- the `dense_layers` guard, the `disabled_until_timestep` prefix, and every ineligibility fallback -- call `torch.nn.functional.scaled_dot_product_attention`. On sm120 those paths could never have matched the backend the user selected. With sm100 alone, Sol-Attn's architecture set is a subset of the dense FMHA kernel's, so routing the dense paths back onto `cute_dsl_fmha_fwd` becomes possible everywhere Sol-Attn runs. That follow-up is not in this change; this only narrows the scope that makes it achievable. Removes the vendored `sol_attn/sm120/` tree (4 files, including the cuDNN-frontend license that covered its execution skeleton), the `_compile_sm120` entry point and its dispatch branch, the `(12, 0)` entries in `SUPPORTED_ARCHS` and `_CUTE_BACKENDS`, and the `l0_gb202.yml` registration. Deleting the dispatch branch left `if arch == (10, 0):` with no `else`, whose fall-through would have returned the uninitialised output buffer -- silently wrong results. `_backend_for_arch` raises before that point so it was unreachable, but the check is now an explicit `raise` rather than resting on a guard three frames away. Also records the divergences from upstream in THIRD_PARTY_NOTICES.md and the PR description, including that `sol_attn_backend.py` is itself adapted from upstream's file of the same name outside the vendored package, and that upstream guards the `torch.compile` path with `torch.library.custom_op` where this port uses `@torch.compiler.disable`. Validated on B200 (sm100): 34 passed, 1 skipped, including the arch-drift test that now confirms SUPPORTED_ARCHS == _CUTE_BACKENDS == {(10, 0)}. `pre-commit run` clean across the changed files. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../features/visualgen-sparse-attention.md | 6 +- .../attention_backend/cute_dsl/sol_attn.py | 4 +- .../blackwell/sol_attn/THIRD_PARTY_NOTICES.md | 18 +- .../blackwell/sol_attn/interface.py | 104 +- .../sol_attn/sm120/LICENSE.cudnn-frontend | 204 ---- .../blackwell/sol_attn/sm120/__init__.py | 10 - .../blackwell/sol_attn/sm120/kernel.py | 24 - .../blackwell/sol_attn/sm120/mainloop.py | 1003 ----------------- .../blackwell/sol_attn_backend.py | 4 +- tensorrt_llm/visual_gen/sparse_attention.py | 10 +- .../test_lists/test-db/l0_gb202.yml | 1 - 11 files changed, 49 insertions(+), 1339 deletions(-) delete mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/LICENSE.cudnn-frontend delete mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/__init__.py delete mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/kernel.py delete mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/mainloop.py diff --git a/docs/source/features/visualgen-sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md index 3c069ba6d079..9b1589966887 100644 --- a/docs/source/features/visualgen-sparse-attention.md +++ b/docs/source/features/visualgen-sparse-attention.md @@ -21,15 +21,15 @@ Sparse attention in VisualGen is configured through `VisualGenArgs.attention_con |---|---|---| | `skip_softmax` | `SkipSoftmaxAttentionConfig` | Supported | | VSA | TBD | TODO | -| `sol_attn` | `SolAttnAttentionConfig` | Supported (CUTEDSL, sm100/sm120) | +| `sol_attn` | `SolAttnAttentionConfig` | Supported (CUTEDSL, sm100) | ### Sol-Attn Sol-Attn ([arXiv:2607.24027](https://arxiv.org/abs/2607.24027)) folds dynamic block routing, sparse computation, and an approximation-correction term into one online-softmax pass. It runs on the **CUTEDSL** backend only, on sm100 -(B200/GB200) and sm120 (RTX Blackwell), and requires `head_dim=128`, bfloat16, -and MHA (`num_kv_heads == num_heads`). +(B200/GB200), and requires `head_dim=128`, bfloat16, and MHA +(`num_kv_heads == num_heads`). ```yaml attention_config: diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py index f79e7dd2fc5c..e81e87cc82ed 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py @@ -23,7 +23,7 @@ https://github.com/NVlabs/Sana/commit/5fe5feb -- see ``cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md`` for the pin and its currency-check note) under ``..cute_dsl_kernels.blackwell.sol_attn`` -/ ``sol_attn_backend.py``. Only the sm100 (B200/GB200) and sm120 (RTX +/ ``sol_attn_backend.py``. Only the sm100 (B200/GB200) Blackwell) kernels are carried; the upstream sm89/sm90 kernels and the Triton reference path are not, and the FlashAttention CuTe helpers they needed come from the ``flash-attn-4`` dependency rather than a vendored copy. @@ -109,7 +109,7 @@ def _parse_dense_layers(spec: Optional[str]) -> frozenset: class SolAttnAttention(AttentionBackend): - """Sol-Attn dynamic block-routing sparse attention (CuTeDSL, sm100/sm120). + """Sol-Attn dynamic block-routing sparse attention (CuTeDSL, sm100). The kernel wrapper already falls back to dense SDPA on any unsupported shape/dtype/arch (see ``_run_sol_attn_bthd``); this class only adds the diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md index a609d8151e3e..1879b802d3b4 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md @@ -29,14 +29,15 @@ Only the pieces needed for the architectures TensorRT-LLM ships are carried: |---|---| | `interface.py`, `preprocess.py`, `common/` | `sm89/`, `sm90/` (incl. `sm90/_compat/`) | | `sm100/` (B200 / GB200) | `triton_ref/` Triton reference attention | -| `sm120/` (RTX Blackwell) | `_vendor/flash_attn/` (see below) | +| | `sm120/` (RTX Blackwell) | +| | `_vendor/flash_attn/` (see below) | The upstream package vendored a copy of FlashAttention's CuTe DSL helpers under `sol_attn/_vendor/flash_attn/cute/`. That copy is **not** carried here: TensorRT-LLM already depends on [`flash-attn-4`](https://github.com/Dao-AILab/flash-attention) (pinned in `requirements.txt`), which provides the same `flash_attn.cute` modules, and -the SM100/SM120 kernels import them from that dependency directly. This was +the SM100 kernels import them from that dependency directly. This was verified on B200 to produce bit-identical output to the vendored copy across a shape/tau sweep. FlashAttention's BSD-3-Clause license is retained at `sol_attn/sm100/LICENSE.flash-attention` because portions of the SM100 design @@ -88,9 +89,10 @@ The runtime also depends on NVIDIA CUTLASS / CuTe DSL, cuda-python, and PyTorch. Those dependencies are not redistributed by this repository and remain subject to their respective licenses. -The SM120 warp-MMA/TMA execution skeleton and online-softmax helpers are -adapted from -[NVIDIA cuDNN Frontend's block-sparse-attention reference](https://github.com/NVIDIA/cudnn-frontend/tree/74785165de2da954a2c879a5e3e6f95411c2292d) -at commit `74785165de2da954a2c879a5e3e6f95411c2292d`. That source is -licensed under the Apache License 2.0; adapted files retain the -corresponding SPDX header. +SM120 (RTX Blackwell) was carried in an earlier revision of this port and has +been dropped: it had kernel-level evidence only, no end-to-end validation, and +no `cute_dsl_fmha_fwd` exists for that architecture, so Sol-Attn's dense +fallback could not match its own backend there. With SM100 alone, Sol-Attn's +architecture set is a subset of the dense CuTe DSL FMHA kernel's. The +cuDNN-frontend attribution that covered the SM120 execution skeleton was +removed with it. diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/interface.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/interface.py index e136ba93cea3..d9644c1fd8bb 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/interface.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/interface.py @@ -14,7 +14,6 @@ BLOCK_SIZE = 64 _CUTE_BACKENDS = { (10, 0): "cute_sm100", # B200 / GB200 - (12, 0): "cute_sm120", # RTX Pro Blackwell / GeForce Blackwell } _compiled = {} @@ -106,7 +105,7 @@ def get_sol_attn_backend(device: torch.device | str | int | None = None) -> str: def _validate_cute(arch, tokens, kv_splits): if kv_splits != 1: raise ValueError( - "kv_splits=2/4 was an SM90-only path; this build ships SM100/SM120 " + "kv_splits=2/4 was an SM90-only path; this build ships SM100 " "kernels only, so kv_splits must be 1." ) route_groups = ((tokens + 63) // 64 + 63) // 64 @@ -163,33 +162,6 @@ def _compile_sm100( return compiled, args -def _compile_sm120( - key, - tensors, - scale, - sink_start_block, - sink_end_block, - stream, -): - import cutlass.cute as cute - - from .sm120 import make_kernel - - operator = make_kernel() - args = _to_cute_tensors(tensors) - compiled = cute.compile( - operator, - *args, - scale, - sink_start_block, - sink_end_block, - stream=stream, - options="--enable-tvm-ffi", - ) - _compiled[key] = compiled - return compiled, args - - def _sol_attn_cute( q, k, @@ -225,58 +197,36 @@ def _sol_attn_cute( stream = _stream(q.device) key = (q.device.index, arch, batch, tokens, heads, kv_splits) - if arch == (10, 0): - sink_start_block, sink_end_block = _sink_block_range( - tokens, - sink_start, - sink_tokens, - ) - tensors = [q, k, v, output, kc, vc, threshold, lse] - compiled = _compiled.get(key) - if compiled is None: - compiled, args = _compile_sm100( - key, - tensors, - scale, - sink_start_block, - sink_end_block, - stream, - ) - else: - args = _to_cute_tensors(tensors) - compiled( - *args, + if arch != (10, 0): + # Unreachable via sol_attn(): _backend_for_arch raises first. Kept + # explicit because the alternative on a missed guard is returning + # the uninitialised `output` buffer, i.e. silently wrong results. + raise ValueError(f"no Sol-Attn CuTe kernel for SM{arch[0]}{arch[1]}") + sink_start_block, sink_end_block = _sink_block_range( + tokens, + sink_start, + sink_tokens, + ) + tensors = [q, k, v, output, kc, vc, threshold, lse] + compiled = _compiled.get(key) + if compiled is None: + compiled, args = _compile_sm100( + key, + tensors, scale, sink_start_block, sink_end_block, - stream=stream, + stream, ) else: - sink_start_block, sink_end_block = _sink_block_range( - tokens, - sink_start, - sink_tokens, - ) - tensors = [q, k, v, output, kc, vc, threshold, lse] - compiled = _compiled.get(key) - if compiled is None: - compiled, args = _compile_sm120( - key, - tensors, - scale, - sink_start_block, - sink_end_block, - stream, - ) - else: - args = _to_cute_tensors(tensors) - compiled( - *args, - scale, - sink_start_block, - sink_end_block, - stream=stream, - ) + args = _to_cute_tensors(tensors) + compiled( + *args, + scale, + sink_start_block, + sink_end_block, + stream=stream, + ) return output @@ -310,7 +260,7 @@ def sol_attn( if kv_splits != 1: raise ValueError( "kv_splits must be 1; the 2/4 path was SM90-only and this build " - "ships SM100/SM120 kernels only." + "ships SM100 kernels only." ) _backend_for_arch(arch) # raises on an architecture with no kernel scale = q.shape[-1] ** -0.5 if scale is None else float(scale) diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/LICENSE.cudnn-frontend b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/LICENSE.cudnn-frontend deleted file mode 100644 index ee9f673bff93..000000000000 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/LICENSE.cudnn-frontend +++ /dev/null @@ -1,204 +0,0 @@ -Copyright (c) 2020-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - - - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/__init__.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/__init__.py deleted file mode 100644 index fc56fddcc772..000000000000 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# -# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see -# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. -"""GeForce Blackwell (SM120) backend.""" - -from .kernel import make_kernel - -__all__ = ["make_kernel"] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/kernel.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/kernel.py deleted file mode 100644 index 64c4b4879b1b..000000000000 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/kernel.py +++ /dev/null @@ -1,24 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# -# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see -# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. -"""SM120 kernel recipe.""" - -from .mainloop import SolAttnForwardSm120 - - -def make_kernel( - *, - debug_route_trace: bool = False, - prefetch_first_exact_k: bool = True, - prefetch_next_route_k: bool = True, -): - return SolAttnForwardSm120( - debug_route_trace=debug_route_trace, - prefetch_first_exact_k=prefetch_first_exact_k, - prefetch_next_route_k=prefetch_next_route_k, - ) - - -__all__ = ["make_kernel"] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/mainloop.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/mainloop.py deleted file mode 100644 index 879034a93904..000000000000 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm120/mainloop.py +++ /dev/null @@ -1,1003 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# -# The warp-MMA/TMA skeleton is adapted from NVIDIA cuDNN Frontend -# (https://github.com/NVIDIA/cudnn-frontend), Apache-2.0; its license text -# is vendored at sol_attn/sm120/LICENSE.cudnn-frontend. -"""Fused Sol-Attn forward kernel for GeForce Blackwell SM120. - -The warp-MMA/TMA execution skeleton and online-softmax helpers are adapted -from NVIDIA cuDNN Frontend's SM120 block-sparse-attention kernel. Sol-specific -routing, CTA-local exact-index compaction, approximate block mass, and the -mixed approximate/exact mainloop are implemented here. -""" - -from __future__ import annotations - -import operator - -import cuda.bindings.driver as cuda -import cutlass -import cutlass.cute as cute -import cutlass.pipeline as pipeline -import cutlass.utils as utils -import cutlass.utils.hopper_helpers as sm90_utils -from flash_attn.cute import utils as kernel_utils - -from ..common import layout_utils -from ..common.selector import sol_attn_popc_b32, sol_attn_route_is_exact - -M = 64 -N = 64 -D = 128 -DV = 128 -THREADS = 128 -STAGES = 1 - - -class SolAttnForwardSm120: - """M64/N64 warp-MMA Sol-Attn kernel for BF16 D128 inputs.""" - - def __init__( - self, - *, - debug_route_trace: bool = False, - prefetch_first_exact_k: bool = True, - prefetch_next_route_k: bool = True, - ): - self.dtype = cutlass.BFloat16 - self.acc_dtype = cutlass.Float32 - self.tile_shape_qk = (M, N, D) - self.tile_shape_pv = (M, DV, N) - self.num_threads = THREADS - self.q_stage = 1 - self.kv_stage = STAGES - self.debug_route_trace = debug_route_trace - self.prefetch_first_exact_k = prefetch_first_exact_k - self.prefetch_next_route_k = prefetch_next_route_k - - @cute.kernel - def kernel( - self, - mQ: cute.Tensor, - mK: cute.Tensor, - mV: cute.Tensor, - mO: cute.Tensor, - mKC: cute.Tensor, - mVC: cute.Tensor, - mThreshold: cute.Tensor, - mLSE: cute.Tensor, - tma_atom_Q: cute.CopyAtom, - tma_atom_K: cute.CopyAtom, - tma_atom_V: cute.CopyAtom, - tma_atom_KC: cute.CopyAtom, - tma_atom_VC: cute.CopyAtom, - tma_atom_O: cute.CopyAtom, - tiled_mma_qk: cute.TiledMma, - tiled_mma_pv: cute.TiledMma, - Q_smem_layout: cute.ComposedLayout, - K_smem_layout: cute.ComposedLayout, - V_smem_layout: cute.ComposedLayout, - O_smem_layout: cute.ComposedLayout, - scale_softmax_log2e: cutlass.Float32, - sink_start_block: cutlass.Int32, - sink_end_block: cutlass.Int32, - ): - tidx, _, _ = cute.arch.thread_idx() - lane = cute.arch.lane_idx() - warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - q_tile_idx, head_idx, batch_idx = cute.arch.block_idx() - q_tile_idx = cute.arch.make_warp_uniform(q_tile_idx) - head_idx = cute.arch.make_warp_uniform(head_idx) - batch_idx = cute.arch.make_warp_uniform(batch_idx) - - token_count = mK.shape[0] - num_blocks = mKC.shape[0] - num_route_groups = cute.ceil_div(num_blocks, N) - q_start = q_tile_idx * M - q_len = token_count - q_start - if q_len > M: - q_len = cutlass.Int32(M) - threshold = cutlass.Float32(mThreshold[batch_idx, q_tile_idx, head_idx]) - - storage = cutlass.utils.SmemAllocator().allocate(self.shared_storage_t) - if warp == 0 and lane == 0: - cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_Q) - cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_K) - cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_V) - cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_KC) - cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_VC) - cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_O) - - cg = pipeline.CooperativeGroup(pipeline.Agent.Thread) - consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_threads // 32) - cta_layout_vmnk = cute.make_layout((1, 1, 1, 1)) - Q_pipeline = pipeline.PipelineTmaAsync.create( - num_stages=self.q_stage, - producer_group=cg, - consumer_group=consumer_group, - tx_count=cute.size_in_bytes(self.Q_dtype, cute.select(Q_smem_layout, mode=[0, 1])), - barrier_storage=storage.Q_barrier.data_ptr(), - cta_layout_vmnk=cta_layout_vmnk, - ) - K_pipeline = pipeline.PipelineTmaAsync.create( - num_stages=self.kv_stage, - producer_group=cg, - consumer_group=consumer_group, - tx_count=cute.size_in_bytes(self.K_dtype, cute.select(K_smem_layout, mode=[0, 1])), - barrier_storage=storage.K_barrier.data_ptr(), - cta_layout_vmnk=cta_layout_vmnk, - ) - V_pipeline = pipeline.PipelineTmaAsync.create( - num_stages=self.kv_stage, - producer_group=cg, - consumer_group=consumer_group, - tx_count=cute.size_in_bytes(self.V_dtype, cute.select(V_smem_layout, mode=[0, 1])), - barrier_storage=storage.V_barrier.data_ptr(), - cta_layout_vmnk=cta_layout_vmnk, - ) - Q_producer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.q_stage) - Q_consumer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.q_stage) - K_producer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.kv_stage) - K_consumer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.kv_stage) - V_producer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.kv_stage) - V_consumer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.kv_stage) - - sQ = storage.Q_smem.get_tensor(Q_smem_layout.outer, swizzle=Q_smem_layout.inner) - sK = storage.K_smem.get_tensor(K_smem_layout.outer, swizzle=K_smem_layout.inner) - sV = storage.V_smem.get_tensor(V_smem_layout.outer, swizzle=V_smem_layout.inner) - # Q is register-resident after the prologue. Reuse its 16 KiB SMEM - # allocation for route scratch until the same allocation becomes sO - # in the epilogue. This drops the CTA below the 2-block/SM threshold - # on SM120 without changing any route reduction or synchronization. - route_f32_ptr = cute.recast_ptr(storage.Q_smem.data_ptr(), dtype=cutlass.Float32) - route_i32_ptr = cute.recast_ptr(storage.Q_smem.data_ptr(), dtype=cutlass.Int32) - route_sums = cute.make_tensor(route_f32_ptr, cute.make_layout((4, N))) - column_masks = cute.make_tensor(route_f32_ptr + 4 * N, cute.make_layout(N)) - route_indices = cute.make_tensor(route_i32_ptr + 5 * N, cute.make_layout(N)) - route_meta = cute.make_tensor(route_i32_ptr + 6 * N, cute.make_layout(2)) - - mQ_slice = mQ[None, None, head_idx, batch_idx] - mK_slice = mK[None, None, head_idx, batch_idx] - mV_slice = mV[None, None, head_idx, batch_idx] - mO_slice = mO[None, None, head_idx, batch_idx] - mKC_slice = mKC[None, None, head_idx, batch_idx] - mVC_slice = mVC[None, None, head_idx, batch_idx] - if cutlass.const_expr(not self.debug_route_trace): - mLSE_slice = mLSE[None, head_idx, batch_idx] - - gQ = cute.local_tile(mQ_slice, (M, D), coord=(q_tile_idx, 0)) - gK = cute.local_tile(mK_slice, (N, D), coord=(None, 0)) - gV = cute.local_tile(mV_slice, (DV, N), coord=(0, None)) - gKC = cute.local_tile(mKC_slice, (N, D), coord=(None, 0)) - gVC = cute.local_tile(mVC_slice, (DV, N), coord=(0, None)) - gO = cute.local_tile(mO_slice, (M, DV), coord=(q_tile_idx, 0)) - - cta_coord_layout = (0, cute.make_layout(1)) - tQsQ, tQgQ = cute.nvgpu.cpasync.tma_partition( - tma_atom_Q, - *cta_coord_layout, - cute.group_modes(sQ, 0, 2), - cute.group_modes(gQ, 0, 2), - ) - tKsK, tKgK = cute.nvgpu.cpasync.tma_partition( - tma_atom_K, - *cta_coord_layout, - cute.group_modes(sK, 0, 2), - cute.group_modes(gK, 0, 2), - ) - tVsV, tVgV = cute.nvgpu.cpasync.tma_partition( - tma_atom_V, - *cta_coord_layout, - cute.group_modes(sV, 0, 2), - cute.group_modes(gV, 0, 2), - ) - tKCsK, tKCgKC = cute.nvgpu.cpasync.tma_partition( - tma_atom_KC, - *cta_coord_layout, - cute.group_modes(sK, 0, 2), - cute.group_modes(gKC, 0, 2), - ) - tVCsV, tVCgVC = cute.nvgpu.cpasync.tma_partition( - tma_atom_VC, - *cta_coord_layout, - cute.group_modes(sV, 0, 2), - cute.group_modes(gVC, 0, 2), - ) - - cS = cute.make_identity_tensor(self.tile_shape_qk[:2]) - thr_mma_qk = tiled_mma_qk.get_slice(tidx) - tSsQ = thr_mma_qk.partition_A(sQ) - tSsK = thr_mma_qk.partition_B(sK) - tSrQ = tiled_mma_qk.make_fragment_A(tSsQ[None, None, None, 0]) - tSrK = tiled_mma_qk.make_fragment_B(tSsK[None, None, None, 0]) - tSrS = cute.make_rmem_tensor(thr_mma_qk.partition_shape_C((M, N)), self.acc_dtype) - tScS = thr_mma_qk.partition_C(cS) - - thr_mma_pv = tiled_mma_pv.get_slice(tidx) - tOsV = thr_mma_pv.partition_B(sV) - tOrV = tiled_mma_pv.make_fragment_B(tOsV[None, None, None, 0]) - tOrO = cute.make_rmem_tensor(thr_mma_pv.partition_shape_C((M, DV)), self.acc_dtype) - - atom_copy_Q = cute.make_copy_atom( - cute.nvgpu.warp.LdMatrix8x8x16bOp(self.Q_layout.is_m_major_a(), 4), - self.Q_dtype, - ) - atom_copy_K = cute.make_copy_atom( - cute.nvgpu.warp.LdMatrix8x8x16bOp(self.K_layout.is_n_major_b(), 4), - self.K_dtype, - ) - atom_copy_V = cute.make_copy_atom( - cute.nvgpu.warp.LdMatrix8x8x16bOp(self.V_layout.is_n_major_b(), 4), - self.V_dtype, - ) - smem_copy_Q = cute.make_tiled_copy_A(atom_copy_Q, tiled_mma_qk) - smem_copy_K = cute.make_tiled_copy_B(atom_copy_K, tiled_mma_qk) - smem_copy_V = cute.make_tiled_copy_B(atom_copy_V, tiled_mma_pv) - thr_copy_Q = smem_copy_Q.get_slice(tidx) - thr_copy_K = smem_copy_K.get_slice(tidx) - thr_copy_V = smem_copy_V.get_slice(tidx) - tSsQ_copy = thr_copy_Q.partition_S(sQ) - tSrQ_copy = thr_copy_Q.retile(tSrQ) - tSsK_copy = thr_copy_K.partition_S(sK) - tOsV_copy = thr_copy_V.partition_S(sV) - - max_m_layout = cute.make_layout( - cute.size( - layout_utils.reshape_acc_to_mn(tOrO).layout, - mode=[0], - ) - ) - max_m = cute.make_rmem_tensor_like(max_m_layout, cutlass.Float32) - sum_m = cute.make_rmem_tensor_like(max_m, cutlass.Float32) - tOrO.store(cute.full_like(tOrO, 0.0, self.acc_dtype)) - max_m.store(cute.full_like(max_m, float("-inf"), cutlass.Float32)) - sum_m.store(cute.full_like(sum_m, 0.0, cutlass.Float32)) - - if warp == 0: - Q_pipeline.producer_acquire(Q_producer) - cute.copy( - tma_atom_Q, - tQgQ, - tQsQ[None, Q_producer.index], - tma_bar_ptr=Q_pipeline.producer_get_barrier(Q_producer), - ) - Q_pipeline.producer_commit(Q_producer) - Q_producer.advance() - cute.arch.sync_threads() - q_wait = Q_pipeline.consumer_try_wait(Q_consumer) - Q_pipeline.consumer_wait(Q_consumer, q_wait) - q_stage = Q_consumer.index - for k_block in cutlass.range_constexpr(cute.size(tSrQ, mode=[2])): - cute.copy( - smem_copy_Q, - tSsQ_copy[None, None, k_block, q_stage], - tSrQ_copy[None, None, k_block], - ) - Q_pipeline.consumer_release(Q_consumer) - Q_consumer.advance() - - for route_group in cutlass.range(0, num_route_groups, 1, unroll=1): - group_start = route_group * cutlass.Int32(N) - valid_blocks = num_blocks - group_start - if valid_blocks > N: - valid_blocks = cutlass.Int32(N) - - if warp == 0: - if cutlass.const_expr(self.prefetch_next_route_k): - # P19-style terminal handoff: when the previous route - # group had an exact block, its final exact QK already - # refilled this K stage with the current group's KC. - if route_group == 0: - K_pipeline.producer_acquire(K_producer) - cute.copy( - tma_atom_KC, - tKCgKC[None, route_group], - tKCsK[None, K_producer.index], - tma_bar_ptr=K_pipeline.producer_get_barrier(K_producer), - ) - K_pipeline.producer_commit(K_producer) - K_producer.advance() - else: - previous_group_exact_count = cutlass.Int32(route_meta[0]) - if previous_group_exact_count == 0: - K_pipeline.producer_acquire(K_producer) - cute.copy( - tma_atom_KC, - tKCgKC[None, route_group], - tKCsK[None, K_producer.index], - tma_bar_ptr=K_pipeline.producer_get_barrier(K_producer), - ) - K_pipeline.producer_commit(K_producer) - K_producer.advance() - else: - K_pipeline.producer_acquire(K_producer) - cute.copy( - tma_atom_KC, - tKCgKC[None, route_group], - tKCsK[None, K_producer.index], - tma_bar_ptr=K_pipeline.producer_get_barrier(K_producer), - ) - K_pipeline.producer_commit(K_producer) - K_producer.advance() - V_pipeline.producer_acquire(V_producer) - cute.copy( - tma_atom_VC, - tVCgVC[None, route_group], - tVCsV[None, V_producer.index], - tma_bar_ptr=V_pipeline.producer_get_barrier(V_producer), - ) - V_pipeline.producer_commit(V_producer) - V_producer.advance() - - k_wait = K_pipeline.consumer_try_wait(K_consumer) - K_pipeline.consumer_wait(K_consumer, k_wait) - gemm_smem_zero_acc( - tiled_mma_qk, - tSrS, - tSrQ, - tSrK, - tSsK_copy[None, None, None, K_consumer.index], - smem_copy_K, - ) - K_pipeline.consumer_release(K_consumer) - K_consumer.advance() - - reduce_route_columns( - tSrS, - tScS, - route_sums, - warp, - lane, - q_len, - ) - cute.arch.fence_view_async_shared() - cute.arch.sync_threads() - - if warp == 0: - preceding = cutlass.Int32(0) - lane_mask_lt = cutlass.Int32(0x7FFFFFFF) >> (cutlass.Int32(31) - lane) - for word in cutlass.range_constexpr(2): - off = cutlass.Int32(word * 32) + lane - valid = off < valid_blocks - exact = False - if valid: - col_sum = ( - cutlass.Float32(route_sums[0, off]) - + cutlass.Float32(route_sums[1, off]) - + cutlass.Float32(route_sums[2, off]) - + cutlass.Float32(route_sums[3, off]) - ) - col_mean = col_sum * scale_softmax_log2e / cutlass.Float32(q_len) - kv_block = group_start + off - exact = sol_attn_route_is_exact( - q_tile_idx, - kv_block, - col_mean, - threshold, - valid, - ) - exact = exact or ( - kv_block >= sink_start_block and kv_block < sink_end_block - ) - ballot = cutlass.Int32(cute.arch.vote_ballot_sync(exact)) - column_masks[off] = ( - -cutlass.Float32.inf if (exact or not valid) else cutlass.Float32(0.0) - ) - rank = preceding + sol_attn_popc_b32(ballot & lane_mask_lt) - if exact: - route_indices[rank] = group_start + off - preceding += sol_attn_popc_b32(ballot) - if cutlass.const_expr(self.debug_route_trace): - if lane == 0: - mLSE[ - batch_idx, - q_tile_idx, - head_idx, - route_group, - word, - ] = ballot - if lane == 0: - route_meta[0] = preceding - route_meta[1] = valid_blocks - cute.arch.fence_view_async_shared() - cute.arch.sync_threads() - - exact_count = cutlass.Int32(route_meta[0]) - has_approx = exact_count < valid_blocks - if cutlass.const_expr(self.prefetch_first_exact_k): - # Once routing identifies the first exact block, the route KC - # stage is free. Refill it before the approximate softmax/PV - # so the first exact K transfer overlaps that work. - if warp == 0 and exact_count > 0: - first_exact = cutlass.Int32(route_indices[0]) - K_pipeline.producer_acquire(K_producer) - cute.copy( - tma_atom_K, - tKgK[None, first_exact], - tKsK[None, K_producer.index], - tma_bar_ptr=K_pipeline.producer_get_barrier(K_producer), - ) - K_pipeline.producer_commit(K_producer) - K_producer.advance() - v_wait = V_pipeline.consumer_try_wait(V_consumer) - V_pipeline.consumer_wait(V_consumer, v_wait) - if has_approx: - apply_route_mask(tSrS, tScS, column_masks, q_len) - row_scale = online_softmax_route( - tSrS, - tScS, - max_m, - sum_m, - scale_softmax_log2e, - group_start, - token_count, - ) - rescale_o_for_next_acc(tOrO, row_scale) - tOrP_frg = cute.make_rmem_tensor_like(tSrS, self.K_dtype) - tOrP_frg.store(tSrS.load().to(self.K_dtype)) - tOrP = layout_utils.reshape_acc_to_frgA(tOrP_frg) - gemm_rs_smem( - tiled_mma_pv, - tOrO, - tOrP, - tOrV, - tOsV_copy[None, None, None, V_consumer.index], - smem_copy_V, - ) - V_pipeline.consumer_release(V_consumer) - V_consumer.advance() - - if warp == 0 and exact_count > 0: - first_exact = cutlass.Int32(route_indices[0]) - if cutlass.const_expr(not self.prefetch_first_exact_k): - K_pipeline.producer_acquire(K_producer) - cute.copy( - tma_atom_K, - tKgK[None, first_exact], - tKsK[None, K_producer.index], - tma_bar_ptr=K_pipeline.producer_get_barrier(K_producer), - ) - K_pipeline.producer_commit(K_producer) - K_producer.advance() - V_pipeline.producer_acquire(V_producer) - cute.copy( - tma_atom_V, - tVgV[None, first_exact], - tVsV[None, V_producer.index], - tma_bar_ptr=V_pipeline.producer_get_barrier(V_producer), - ) - V_pipeline.producer_commit(V_producer) - V_producer.advance() - - for ordinal in cutlass.range(0, exact_count, 1, unroll=1): - exact_block = cutlass.Int32(route_indices[ordinal]) - k_wait = K_pipeline.consumer_try_wait(K_consumer) - K_pipeline.consumer_wait(K_consumer, k_wait) - gemm_smem_zero_acc( - tiled_mma_qk, - tSrS, - tSrQ, - tSrK, - tSsK_copy[None, None, None, K_consumer.index], - smem_copy_K, - ) - K_pipeline.consumer_release(K_consumer) - K_consumer.advance() - next_ordinal = ordinal + cutlass.Int32(1) - if warp == 0: - if next_ordinal < exact_count: - next_exact = cutlass.Int32(route_indices[next_ordinal]) - K_pipeline.producer_acquire(K_producer) - cute.copy( - tma_atom_K, - tKgK[None, next_exact], - tKsK[None, K_producer.index], - tma_bar_ptr=K_pipeline.producer_get_barrier(K_producer), - ) - K_pipeline.producer_commit(K_producer) - K_producer.advance() - else: - if cutlass.const_expr(self.prefetch_next_route_k): - next_route_group = route_group + cutlass.Int32(1) - if next_route_group < num_route_groups: - # Reuse the K stage released by the final - # exact QK. The next outer prologue supplies - # VC, matching the SM90 P19 partial handoff. - K_pipeline.producer_acquire(K_producer) - cute.copy( - tma_atom_KC, - tKCgKC[None, next_route_group], - tKCsK[None, K_producer.index], - tma_bar_ptr=(K_pipeline.producer_get_barrier(K_producer)), - ) - K_pipeline.producer_commit(K_producer) - K_producer.advance() - block_len = token_count - exact_block * cutlass.Int32(N) - if block_len > N: - block_len = cutlass.Int32(N) - mask_exact_scores(tSrS, tScS, block_len, q_len) - row_scale = online_softmax(tSrS, max_m, sum_m, scale_softmax_log2e) - rescale_o_for_next_acc(tOrO, row_scale) - tOrP_frg = cute.make_rmem_tensor_like(tSrS, self.K_dtype) - tOrP_frg.store(tSrS.load().to(self.K_dtype)) - tOrP = layout_utils.reshape_acc_to_frgA(tOrP_frg) - - v_wait = V_pipeline.consumer_try_wait(V_consumer) - V_pipeline.consumer_wait(V_consumer, v_wait) - gemm_rs_smem( - tiled_mma_pv, - tOrO, - tOrP, - tOrV, - tOsV_copy[None, None, None, V_consumer.index], - smem_copy_V, - ) - V_pipeline.consumer_release(V_consumer) - V_consumer.advance() - if warp == 0 and next_ordinal < exact_count: - next_exact = cutlass.Int32(route_indices[next_ordinal]) - V_pipeline.producer_acquire(V_producer) - cute.copy( - tma_atom_V, - tVgV[None, next_exact], - tVsV[None, V_producer.index], - tma_bar_ptr=V_pipeline.producer_get_barrier(V_producer), - ) - V_pipeline.producer_commit(V_producer) - V_producer.advance() - - final_ratio, lse = finalize_softmax(max_m, sum_m, scale_softmax_log2e) - rescale_o_for_next_acc(tOrO, final_ratio) - if cutlass.const_expr(not self.debug_route_trace): - tScS_mn = layout_utils.reshape_acc_to_mn(tScS) - for m in cutlass.range_constexpr(cute.size(lse)): - row = tScS_mn[m, 0][0] - if tScS_mn[m, 0][1] == 0 and row < q_len: - mLSE_slice[q_start + row] = lse[m] - - tOrO_cvt = cute.make_rmem_tensor_like(tOrO, self.O_dtype) - tOrO_cvt.store(tOrO.load().to(self.O_dtype)) - sO = storage.Q_smem.get_tensor(O_smem_layout.outer, swizzle=O_smem_layout.inner) - tiled_copy_O = cute.make_tiled_copy_C( - cute.make_copy_atom( - cute.nvgpu.warp.StMatrix8x8x16bOp(self.O_layout.is_m_major_c(), 4), - self.O_dtype, - ), - tiled_mma_pv, - ) - tOrO_cv = tiled_copy_O.retile(tOrO_cvt) - tOsO = tiled_copy_O.get_slice(tidx).partition_D(sO) - cute.copy(tiled_copy_O, tOrO_cv, tOsO) - cute.arch.fence_view_async_shared() - cute.arch.sync_threads() - tOsO, tOgO = cute.nvgpu.cpasync.tma_partition( - tma_atom_O, - *cta_coord_layout, - cute.group_modes(sO, 0, 2), - cute.group_modes(gO, 0, 2), - ) - if warp == 0: - cute.copy(tma_atom_O, tOsO, tOgO) - cute.arch.cp_async_bulk_commit_group() - cute.arch.cp_async_bulk_wait_group(0, read=True) - - @cute.jit - def __call__( - self, - q: cute.Tensor, - k: cute.Tensor, - v: cute.Tensor, - o: cute.Tensor, - kc: cute.Tensor, - vc: cute.Tensor, - threshold: cute.Tensor, - lse: cute.Tensor, - softmax_scale: cutlass.Float32, - sink_start_block: cutlass.Int32, - sink_end_block: cutlass.Int32, - stream: cuda.CUstream, - ): - q_mkl, k_nkl, kc_nkl = [layout_utils.select(t, [1, 3, 2, 0]) for t in (q, k, kc)] - v_nkl, vc_nkl = [layout_utils.select(t, [3, 1, 2, 0]) for t in (v, vc)] - o_mkl = layout_utils.select(o, [1, 3, 2, 0]) - if cutlass.const_expr(self.debug_route_trace): - lse_target = lse - else: - lse_target = layout_utils.select(lse, [1, 2, 0]) - - self.Q_dtype = q_mkl.element_type - self.K_dtype = k_nkl.element_type - self.V_dtype = v_nkl.element_type - self.O_dtype = o_mkl.element_type - self.Q_layout = utils.LayoutEnum.from_tensor(q_mkl) - self.K_layout = utils.LayoutEnum.from_tensor(k_nkl) - self.V_layout = utils.LayoutEnum.from_tensor(v_nkl) - self.O_layout = utils.LayoutEnum.from_tensor(o_mkl) - assert self.Q_dtype == cutlass.BFloat16 - assert self.K_dtype == cutlass.BFloat16 - assert self.V_dtype == cutlass.BFloat16 - - self.Q_smem_layout = sm90_utils.make_smem_layout_a( - self.Q_layout, - self.tile_shape_qk, - self.Q_dtype, - self.q_stage, - ) - self.K_smem_layout = sm90_utils.make_smem_layout_b( - self.K_layout, - self.tile_shape_qk, - self.K_dtype, - self.kv_stage, - ) - self.V_smem_layout = sm90_utils.make_smem_layout_b( - self.V_layout, - self.tile_shape_pv, - self.V_dtype, - self.kv_stage, - ) - O_smem_layout_staged = sm90_utils.make_smem_layout_epi( - self.O_dtype, - self.O_layout, - self.tile_shape_pv[:2], - 1, - ) - self.O_smem_layout = cute.select(O_smem_layout_staged, mode=[0, 1]) - - @cute.struct - class SharedStorage: - Q_barrier: cute.struct.MemRange[cutlass.Int64, self.q_stage * 2] - K_barrier: cute.struct.MemRange[cutlass.Int64, self.kv_stage * 2] - V_barrier: cute.struct.MemRange[cutlass.Int64, self.kv_stage * 2] - Q_smem: cute.struct.Align[ - cute.struct.MemRange[self.Q_dtype, cute.cosize(self.Q_smem_layout)], - 128, - ] - K_smem: cute.struct.Align[ - cute.struct.MemRange[self.K_dtype, cute.cosize(self.K_smem_layout)], - 128, - ] - V_smem: cute.struct.Align[ - cute.struct.MemRange[self.V_dtype, cute.cosize(self.V_smem_layout)], - 128, - ] - - self.shared_storage_t = SharedStorage - - tiled_mma_qk = cute.make_tiled_mma( - cute.nvgpu.warp.MmaF16BF16Op( - self.Q_dtype, - self.acc_dtype, - (16, 8, 16), - ), - cute.make_layout((4, 1, 1)), - permutation_mnk=(64, 16, 16), - ) - tiled_mma_pv = cute.make_tiled_mma( - cute.nvgpu.warp.MmaF16BF16Op( - self.K_dtype, - self.acc_dtype, - (16, 8, 16), - ), - cute.make_layout((4, 1, 1)), - permutation_mnk=(64, 16, 16), - ) - - g2s_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() - tma_atom_Q, tma_tensor_Q = cute.nvgpu.cpasync.make_tiled_tma_atom( - g2s_op, - q_mkl, - self.Q_smem_layout, - (M, D), - num_multicast=1, - ) - tma_atom_K, tma_tensor_K = cute.nvgpu.cpasync.make_tiled_tma_atom( - g2s_op, - k_nkl, - self.K_smem_layout, - (N, D), - num_multicast=1, - ) - tma_atom_V, tma_tensor_V = cute.nvgpu.cpasync.make_tiled_tma_atom( - g2s_op, - v_nkl, - self.V_smem_layout, - (DV, N), - num_multicast=1, - ) - tma_atom_KC, tma_tensor_KC = cute.nvgpu.cpasync.make_tiled_tma_atom( - g2s_op, - kc_nkl, - self.K_smem_layout, - (N, D), - num_multicast=1, - ) - tma_atom_VC, tma_tensor_VC = cute.nvgpu.cpasync.make_tiled_tma_atom( - g2s_op, - vc_nkl, - self.V_smem_layout, - (DV, N), - num_multicast=1, - ) - s2g_op = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() - tma_atom_O, tma_tensor_O = cute.nvgpu.cpasync.make_tiled_tma_atom( - s2g_op, - o_mkl, - self.O_smem_layout, - (M, DV), - num_multicast=1, - ) - - self.kernel( - tma_tensor_Q, - tma_tensor_K, - tma_tensor_V, - tma_tensor_O, - tma_tensor_KC, - tma_tensor_VC, - threshold, - lse_target, - tma_atom_Q, - tma_atom_K, - tma_atom_V, - tma_atom_KC, - tma_atom_VC, - tma_atom_O, - tiled_mma_qk, - tiled_mma_pv, - self.Q_smem_layout, - self.K_smem_layout, - self.V_smem_layout, - self.O_smem_layout, - softmax_scale * 1.4426950408889634, - sink_start_block, - sink_end_block, - ).launch( - grid=(cute.ceil_div(q_mkl.shape[0], M), q_mkl.shape[2], q_mkl.shape[3]), - block=(self.num_threads, 1, 1), - cluster=(1, 1, 1), - smem=self.shared_storage_t.size_in_bytes(), - stream=stream, - min_blocks_per_mp=1, - ) - - -@cute.jit -def gemm_smem_zero_acc( - tiled_mma: cute.TiledMma, - acc: cute.Tensor, - tCrA: cute.Tensor, - tCrB: cute.Tensor, - tCsB: cute.Tensor, - smem_tiled_copy_B: cute.TiledCopy, -): - acc.fill(0.0) - tCrB_copy = smem_tiled_copy_B.retile(tCrB) - cute.copy( - smem_tiled_copy_B, - tCsB[None, None, 0], - tCrB_copy[None, None, 0], - ) - for k_block in cutlass.range_constexpr(cute.size(tCsB.shape[2])): - if k_block < cute.size(tCsB.shape[2]) - 1: - cute.copy( - smem_tiled_copy_B, - tCsB[None, None, k_block + 1], - tCrB_copy[None, None, k_block + 1], - ) - cute.gemm( - tiled_mma, - acc, - tCrA[None, None, k_block], - tCrB[None, None, k_block], - acc, - ) - - -@cute.jit -def gemm_rs_smem( - tiled_mma: cute.TiledMma, - acc: cute.Tensor, - tCrA: cute.Tensor, - tCrB: cute.Tensor, - tCsB: cute.Tensor, - smem_tiled_copy_B: cute.TiledCopy, -): - tCrB_copy = smem_tiled_copy_B.retile(tCrB) - cute.copy( - smem_tiled_copy_B, - tCsB[None, None, 0], - tCrB_copy[None, None, 0], - ) - for k_block in cutlass.range_constexpr(cute.size(tCrA.shape[2])): - if k_block < cute.size(tCrA.shape[2]) - 1: - cute.copy( - smem_tiled_copy_B, - tCsB[None, None, k_block + 1], - tCrB_copy[None, None, k_block + 1], - ) - cute.gemm( - tiled_mma, - acc, - tCrA[None, None, k_block], - tCrB[None, None, k_block], - acc, - ) - - -@cute.jit -def reduce_route_columns( - scores: cute.Tensor, - coords: cute.Tensor, - route_sums: cute.Tensor, - warp: cutlass.Int32, - lane: cutlass.Int32, - q_len: cutlass.Int32, -): - """Reduce M64 score columns using the measured SM120 lane layout.""" - - scores_mn = layout_utils.reshape_acc_to_mn(scores) - coords_mn = layout_utils.reshape_acc_to_mn(coords) - row0 = coords_mn[0, 0][0] - row1 = coords_mn[1, 0][0] - valid0 = row0 < q_len - valid1 = row1 < q_len - for group in cutlass.range_constexpr(8): - n0 = group * 2 - partial0 = cutlass.Float32(0.0) - partial1 = cutlass.Float32(0.0) - if valid0: - partial0 += cutlass.Float32(scores_mn[0, n0]) - partial1 += cutlass.Float32(scores_mn[0, n0 + 1]) - if valid1: - partial0 += cutlass.Float32(scores_mn[1, n0]) - partial1 += cutlass.Float32(scores_mn[1, n0 + 1]) - for offset in (4, 8, 16): - partial0 += cute.arch.shuffle_sync_bfly(partial0, offset=offset) - partial1 += cute.arch.shuffle_sync_bfly(partial1, offset=offset) - if lane < 4: - column = cutlass.Int32(group * 8) + lane * cutlass.Int32(2) - route_sums[warp, column] = partial0 - route_sums[warp, column + 1] = partial1 - - -@cute.jit -def apply_route_mask( - scores: cute.Tensor, - coords: cute.Tensor, - column_masks: cute.Tensor, - q_len: cutlass.Int32, -): - scores_mn = layout_utils.reshape_acc_to_mn(scores) - coords_mn = layout_utils.reshape_acc_to_mn(coords) - for m in cutlass.range_constexpr(cute.size(scores_mn, mode=[0])): - valid_row = coords_mn[m, 0][0] < q_len - for n in cutlass.range_constexpr(cute.size(scores_mn, mode=[1])): - column = coords_mn[m, n][1] - scores_mn[m, n] = ( - cutlass.Float32(scores_mn[m, n]) + cutlass.Float32(column_masks[column]) - if valid_row - else -cutlass.Float32.inf - ) - - -@cute.jit -def mask_exact_scores( - scores: cute.Tensor, - coords: cute.Tensor, - block_len: cutlass.Int32, - q_len: cutlass.Int32, -): - scores_mn = layout_utils.reshape_acc_to_mn(scores) - coords_mn = layout_utils.reshape_acc_to_mn(coords) - for m in cutlass.range_constexpr(cute.size(scores_mn, mode=[0])): - valid_row = coords_mn[m, 0][0] < q_len - for n in cutlass.range_constexpr(cute.size(scores_mn, mode=[1])): - if (not valid_row) or coords_mn[m, n][1] >= block_len: - scores_mn[m, n] = -cutlass.Float32.inf - - -@cute.jit -def online_softmax( - scores: cute.Tensor, - row_max: cute.Tensor, - row_sum: cute.Tensor, - scale_log2e: cutlass.Float32, -): - scores_mn = layout_utils.reshape_acc_to_mn(scores) - row_scale = cute.make_rmem_tensor_like(row_max, cutlass.Float32) - for m in cutlass.range_constexpr(cute.size(row_max)): - score_row = scores_mn[m, None].load() - current_max = kernel_utils.fmax_reduce(score_row, init_val=row_max[m], arch=80) - current_max = cute.arch.warp_reduction_max(current_max, threads_in_group=4) - previous_max = row_max[m] - row_max[m] = current_max - safe_max = cutlass.Float32(0.0) if current_max == -cutlass.Float32.inf else current_max - scaled_max = safe_max * scale_log2e - probabilities = cute.math.exp2(score_row * scale_log2e - scaled_max, fastmath=True) - row_scale[m] = cute.math.exp2((previous_max - safe_max) * scale_log2e, fastmath=True) - row_sum[m] = kernel_utils.fadd_reduce( - probabilities, - init_val=row_sum[m] * row_scale[m], - arch=80, - ) - scores_mn[m, None].store(probabilities) - return row_scale - - -@cute.jit -def online_softmax_route( - scores: cute.Tensor, - coords: cute.Tensor, - row_max: cute.Tensor, - row_sum: cute.Tensor, - scale_log2e: cutlass.Float32, - group_start: cutlass.Int32, - token_count: cutlass.Int32, -): - scores_mn = layout_utils.reshape_acc_to_mn(scores) - coords_mn = layout_utils.reshape_acc_to_mn(coords) - row_scale = cute.make_rmem_tensor_like(row_max, cutlass.Float32) - for m in cutlass.range_constexpr(cute.size(row_max)): - score_row = scores_mn[m, None].load() - current_max = kernel_utils.fmax_reduce(score_row, init_val=row_max[m], arch=80) - current_max = cute.arch.warp_reduction_max(current_max, threads_in_group=4) - previous_max = row_max[m] - row_max[m] = current_max - safe_max = cutlass.Float32(0.0) if current_max == -cutlass.Float32.inf else current_max - probabilities = cute.math.exp2( - score_row * scale_log2e - safe_max * scale_log2e, - fastmath=True, - ) - row_scale[m] = cute.math.exp2((previous_max - safe_max) * scale_log2e, fastmath=True) - masses = cute.make_rmem_tensor_like(scores_mn[m, None], cutlass.Float32) - for n in cutlass.range_constexpr(cute.size(masses)): - block = group_start + coords_mn[m, n][1] - length = token_count - block * cutlass.Int32(N) - if length > N: - length = cutlass.Int32(N) - if length < 0: - length = cutlass.Int32(0) - masses[n] = cutlass.Float32(probabilities[n]) * cutlass.Float32(length) - row_sum[m] = kernel_utils.fadd_reduce( - masses.load(), - init_val=row_sum[m] * row_scale[m], - arch=80, - ) - scores_mn[m, None].store(probabilities) - return row_scale - - -@cute.jit -def finalize_softmax( - row_max: cute.Tensor, - row_sum: cute.Tensor, - scale_log2e: cutlass.Float32, -): - row_sum.store(kernel_utils.warp_reduce(row_sum.load(), operator.add, width=4)) - ratio = cute.make_rmem_tensor_like(row_sum, cutlass.Float32) - lse = cute.make_rmem_tensor_like(row_sum, cutlass.Float32) - for m in cutlass.range_constexpr(cute.size(row_sum)): - total = row_sum[m] - invalid = total == 0.0 or total != total - ratio[m] = cute.arch.rcp_approx(total if not invalid else 1.0) - lse[m] = ( - -cutlass.Float32.inf - if invalid - else (row_max[m] * scale_log2e + cute.math.log2(total, fastmath=True)) - * 0.6931471805599453 - ) - return ratio, lse - - -@cute.jit -def rescale_o_for_next_acc( - output: cute.Tensor, - row_scale: cute.Tensor, -): - output_mn = layout_utils.reshape_acc_to_mn(output) - for m in cutlass.range_constexpr(cute.size(row_scale)): - output_mn[m, None].store(output_mn[m, None].load() * row_scale[m]) - - -__all__ = ["SolAttnForwardSm120"] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py index 5a8e83a16e6e..feed4aadeae5 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py @@ -54,7 +54,7 @@ def _load_sol_attn() -> Callable: # Architectures with a Sol-Attn CuTe kernel. Kept in sync with # ``sol_attn/interface.py::_CUTE_BACKENDS``; duplicated here so the eligibility # check does not have to import the CuTe DSL. -SUPPORTED_ARCHS = frozenset({(10, 0), (12, 0)}) +SUPPORTED_ARCHS = frozenset({(10, 0)}) def sol_attn_ineligible_reason(q) -> Optional[str]: @@ -109,7 +109,7 @@ def _resolve_kv_splits(q, kv_splits: int | str | None) -> int: """Resolve the integration-only ``auto`` policy to the public integer API. ``auto`` is always 1 here: kv_splits=2/4 was an SM90-only path, and this - build ships SM100/SM120 kernels only. + build ships SM100 kernels only. """ if kv_splits in (None, "auto"): diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 550525592153..6120905d94a3 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -57,13 +57,13 @@ class SkipSoftmaxAttentionConfig(BaseSparseAttentionConfig): ) target_sparsity: Optional[float] = PydanticField( default=None, - gt=0.0, + ge=0.0, le=1.0, description="Semantic target sparsity in [0, 1]; requires a calibration formula.", ) disabled_until_timestep: Optional[float] = PydanticField( default=None, - gt=0.0, + ge=0.0, le=1.0, description="Normalized timestep cutoff below which skip-softmax is enabled.", ) @@ -226,7 +226,7 @@ class SolAttnAttentionConfig(BaseSparseAttentionConfig): Dynamic block routing + sparse computation + approximation correction in one online-softmax pass (arXiv:2607.24027). Kernel is CuTeDSL, sm100 - (B200/GB200) and sm120 (RTX Blackwell) only, head_dim=128, bf16, MHA. + (B200/GB200) only, head_dim=128, bf16, MHA. On an unsupported *shape, dtype, or architecture* the kernel falls back to dense SDPA and counts the fallback, so setting this config on the wrong GPU @@ -247,7 +247,7 @@ class SolAttnAttentionConfig(BaseSparseAttentionConfig): kv_splits: Literal["auto", "1"] = PydanticField( "auto", description=( - "KV split policy. Only 1 split is valid on the shipped sm100/sm120 " + "KV split policy. Only 1 split is valid on the shipped sm100 " "kernels, so 'auto' and '1' are equivalent; the 2/4 path was " "SM90-only and returns with that kernel. Constrained rather than a " "free string because any other value is rejected deep inside the " @@ -301,7 +301,7 @@ class VideoSparseAttentionConfig(StrictBaseModel): ) vsa_sparsity: float = PydanticField( 0.9, - gt=0.0, + ge=0.0, le=1.0, description=( "Fraction of cubes dropped on the fine stage. 0.0 keeps all cubes " diff --git a/tests/integration/test_lists/test-db/l0_gb202.yml b/tests/integration/test_lists/test-db/l0_gb202.yml index 2053cd8a044c..49a9ed558dd3 100644 --- a/tests/integration/test_lists/test-db/l0_gb202.yml +++ b/tests/integration/test_lists/test-db/l0_gb202.yml @@ -23,7 +23,6 @@ l0_gb202: # ------------- Visual Gen tests --------------- # sm120 (GB202) coverage for the Sol-Attn CuTeDSL kernel; the same file # is registered in l0_b200.yml for sm100. - - unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py - test_e2e.py::test_ptp_quickstart_bert[VANILLA-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] - test_e2e.py::test_ptp_quickstart_bert[TRTLLM-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_bf16[latency] From 8593ad87c6c06cf034e55b8b7eb879cbd404dc88 Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:59:04 -0700 Subject: [PATCH 05/22] [TRTLLM-15917][fix] Keep Sol-Attn's non-sparse attention on the configured 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) Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../attention_backend/cute_dsl/sol_attn.py | 73 ++++++++++++++++++- .../_torch/visual_gen/modules/attention.py | 19 ++++- .../test_attention_cute_dsl_sol_attn.py | 24 +++++- 3 files changed, 107 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py index e81e87cc82ed..d188fef5dbd6 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py @@ -94,6 +94,24 @@ def sol_attn_graph_phase( return int(value < disabled_until_timestep) +def _cute_dense_available() -> bool: + """Whether `cute_dsl_fmha_fwd` can run on the current device. + + Checked once at construction. Sol-Attn is sm100-only and the dense CuTe DSL + kernel covers sm_100a/sm_103a, so in practice this is always true wherever + Sol-Attn runs; the negative branch exists so an unsupported device degrades + to SDPA instead of raising. + """ + try: + from .fmha import _check_cute_runtime_available, _get_gpu_arch + + _check_cute_runtime_available() + _get_gpu_arch() + except Exception: + return False + return True + + def _parse_dense_layers(spec: Optional[str]) -> frozenset: layers: set = set() for item in str(spec or "").split(","): @@ -149,6 +167,37 @@ def __init__( self.disabled_until_timestep = getattr(cfg, "disabled_until_timestep", None) self.dense_layers = _parse_dense_layers(getattr(cfg, "dense_layers", None)) + # Sol-Attn's dense steps must run the backend the user selected. Without + # this they ran torch SDPA while a `backend: CUTEDSL` baseline ran + # cute_dsl_fmha_fwd, so candidate and reference differed on the dense + # steps too -- measured at LPIPS 0.214 on Wan2.2-T2V-A14B with sparsity + # switched off entirely, against a 0.25 gate. + from .fmha import CuTeDSLAttention + + self._dense_backend = CuTeDSLAttention( + layer_idx=layer_idx, + num_heads=num_heads, + head_dim=head_dim, + num_kv_heads=self.num_kv_heads, + dtype=dtype, + ) + # Whether the CuTe DSL dense kernel can serve this device, decided once + # here. Doing it at construction (rather than lazily on the first call) + # keeps `_dense` free of attribute mutation, so it stays traceable and + # the dense step sits in the same place in the graph as the dense + # CUTEDSL baseline's does. Deciding it lazily and marking `_dense` + # `@torch.compiler.disable` instead moved the whole dense step out of + # the graph and reintroduced the very mismatch this is meant to remove: + # measured LPIPS 0.4044 compiled, against 0.2112 eager. + self._cute_dense_ok = _cute_dense_available() + if not self._cute_dense_ok: + logger.warning_once( + "[sol-attn] the CuTe DSL FMHA kernel cannot serve this device; dense " + "steps will use torch SDPA. Numerics will differ from a `backend: " + "CUTEDSL` dense baseline.", + key="sol_attn_dense_backend_unavailable", + ) + # The `.item()` in here would graph-break the enclosing block once per # attention layer, so keep it in eager (as cute_dsl/fmha.py and VSA's # `_get_vsa_inputs` do). Returns a host-side bool, so the dense and sparse @@ -175,6 +224,25 @@ def _dense_by_step(self, timestep) -> bool: return False return phase == 0 + @staticmethod + def _sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + """Dense attention via torch SDPA, for architectures CuTe DSL cannot serve.""" + return torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) + ).transpose(1, 2) + + def _dense(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + """Dense attention on the configured backend, or SDPA where unavailable. + + ``_cute_dense_ok`` answers "can this *device* run the kernel", decided at + construction; ``q.is_cuda`` answers "is this *tensor* on it". Both are + needed: the construction-time probe inspects the current CUDA device, so + it says yes on a GPU host even when a caller passes CPU tensors. + """ + if self._cute_dense_ok and q.is_cuda: + return self._dense_backend.forward(q, k, v) + return self._sdpa(q, k, v) + def forward( self, q: torch.Tensor, @@ -188,9 +256,7 @@ def forward( if self.disabled_until_timestep is not None: dense_by_step = self._dense_by_step(kwargs.get("timestep")) if dense_by_layer or dense_by_step: - return torch.nn.functional.scaled_dot_product_attention( - q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) - ).transpose(1, 2) + return self._dense(q, k, v) return _sol_attn_run( q, k, @@ -198,6 +264,7 @@ def forward( tau=self.tau, thresh_type=self.thresh_type, kv_splits=self.kv_splits, + dense_fn=self._dense, ) @classmethod diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index da8754cfaf60..fdcee0750346 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -106,10 +106,19 @@ def __init__( # Cross-attention fallback: TRTLLM and CUTEDSL VSA/Sol-Attn are self-attn only. + # + # VSA/Sol-Attn fall back within their own backend family -- the dense + # CuTe DSL kernel serves cross-attention fine. Falling back to VANILLA + # instead silently swapped cross-attention from CuTeDSL to torch SDPA in + # every block the moment a sparse algorithm was enabled, so a + # `backend: CUTEDSL` run and a `CUTEDSL + sol_attn` run differed in + # cross-attention regardless of any sparse setting. TRTLLM keeps VANILLA: + # TrtllmAttention genuinely cannot serve SEPARATE_QKV. + _cross_attn_fallback = "CUTEDSL" if _is_sol_attn else "VANILLA" if separate_qkv_cross_attention and ( base_backend == "TRTLLM" or _is_vsa or _is_sol_attn ): - backend_name = "VANILLA" + backend_name = _cross_attn_fallback requested = ( f"{base_backend} (VSA)" if _is_vsa @@ -259,7 +268,13 @@ def __init__( num_kv_heads=backend_num_kv_heads, quant_config=self.quant_config, dtype=self.dtype, - attention_config=config.attention, + attention_config=( + config.attention.model_copy(update={"sparse_attention_config": None}) + if backend_name == "CUTEDSL" + and _is_sol_attn + and self.qkv_mode == QKVMode.SEPARATE_QKV + else config.attention + ), attention_metadata_state=attention_metadata_state, sparse_params=sparse_params, ) diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py index 9bbe8030929b..d42afeabb4a0 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py @@ -97,8 +97,18 @@ def _make_config( @pytest.mark.skipif(not torch.cuda.is_available(), reason="Sol-Attn needs CUDA") -def test_sol_attn_falls_back_to_vanilla_for_cross_attention(): - """Cross-attention (SEPARATE_QKV) falls back to VANILLA -- Sol-Attn is self-attn only.""" +def test_sol_attn_cross_attention_uses_dense_cutedsl(): + """Cross-attention must stay on CuTeDSL, not drop to VANILLA. + + Sol-Attn is self-attention only, so SEPARATE_QKV modules fall back -- but to + the dense kernel of the *configured* backend, not to torch SDPA. Falling back + to VANILLA made a `backend: CUTEDSL` run and a `CUTEDSL + sol_attn` run differ + in cross-attention in every block, regardless of any sparse setting, which is + a backend difference masquerading as a sparsity difference in any A/B. + """ + from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.fmha import CuTeDSLAttention + from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn import SolAttnAttention + device = torch.device("cuda") dtype = torch.bfloat16 cfg = _make_config( @@ -109,8 +119,14 @@ def test_sol_attn_falls_back_to_vanilla_for_cross_attention(): .to(device=device, dtype=dtype) .eval() ) - assert cross_attn.attn_backend == "VANILLA", ( - f"Sol-Attn on cross-attention should fall back to VANILLA, got {cross_attn.attn_backend!r}" + assert cross_attn.attn_backend == "CUTEDSL", ( + f"expected CUTEDSL cross-attention, got {cross_attn.attn_backend!r}" + ) + assert isinstance(cross_attn.attn, CuTeDSLAttention), ( + f"expected the dense CuTeDSL kernel, got {type(cross_attn.attn).__name__}" + ) + assert not isinstance(cross_attn.attn, SolAttnAttention), ( + "cross-attention must not re-select the sparse backend" ) From bd47daac548a661859972a5dd8bb2301efbfe5fb Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Wed, 2 Sep 2026 17:17:21 -0700 Subject: [PATCH 06/22] [TRTLLM-15917][fix] Restrict the in-family fallback to true cross-attention 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) Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../features/visualgen-sparse-attention.md | 5 +- .../attention_backend/cute_dsl/sol_attn.py | 4 +- .../blackwell/sol_attn/THIRD_PARTY_NOTICES.md | 7 +- .../blackwell/sol_attn_backend.py | 11 +-- .../_torch/visual_gen/models/modeling.py | 2 +- .../_torch/visual_gen/modules/attention.py | 28 ++++--- tensorrt_llm/visual_gen/sparse_attention.py | 3 +- .../test_lists/test-db/l0_gb202.yml | 3 - .../test_attention_cute_dsl_sol_attn.py | 73 +++++++++++++++++-- 9 files changed, 106 insertions(+), 30 deletions(-) diff --git a/docs/source/features/visualgen-sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md index 9b1589966887..19acd2a5810c 100644 --- a/docs/source/features/visualgen-sparse-attention.md +++ b/docs/source/features/visualgen-sparse-attention.md @@ -38,7 +38,7 @@ attention_config: algorithm: sol_attn tau: 2.0 # routing threshold; higher routes more blocks sparse thresh_type: diag # or "exact" - disabled_until_timestep: 0.9545 # dense while normalized timestep >= cutoff + disabled_until_timestep: 0.9090 # dense while normalized timestep >= cutoff dense_layers: '0' # optional: layers forced dense ``` @@ -49,7 +49,8 @@ below it. Use `None` rather than `0.0` to disable the prefix. On an input the kernel cannot serve — an unsupported architecture, a `head_dim` other than 128, a non-bfloat16 dtype — Sol-Attn falls back to dense -SDPA, logs the specific reason once, and counts the fallback. Set +dense attention -- the configured backend's dense kernel where available, torch +SDPA otherwise -- logs the specific reason once, and counts the fallback. Set `SOL_ATTN_STRICT=1` to raise instead of falling back, which is useful when benchmarking to confirm the kernel actually ran. diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py index d188fef5dbd6..3d84611068d0 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py @@ -33,7 +33,7 @@ ``disabled_until_timestep`` is the dense-prefix control, and mirrors skip_softmax's field of the same name: sparse attention stays disabled (that -is, the layer runs dense SDPA) while the normalized denoising timestep is at +is, the layer runs the backend's dense kernel) while the normalized timestep is at or above the cutoff, and switches to the sparse kernel once it drops below. The timestep arrives as a forward kwarg -- ``modules/attention.py`` already @@ -129,7 +129,7 @@ def _parse_dense_layers(spec: Optional[str]) -> frozenset: class SolAttnAttention(AttentionBackend): """Sol-Attn dynamic block-routing sparse attention (CuTeDSL, sm100). - The kernel wrapper already falls back to dense SDPA on any unsupported + The kernel wrapper already falls back to dense attention on any unsupported shape/dtype/arch (see ``_run_sol_attn_bthd``); this class only adds the ``dense_layers`` layer-skip guard (evaluated at construction time, no external plumbing needed) and forwards the routing knobs from config. diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md index 1879b802d3b4..2b0ba8011fb4 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md @@ -53,7 +53,8 @@ is a derivative work and is recorded here because this file is where a future currency check starts. It is adapted from upstream's `techniques/sparse_backends/sol_attn_backend.py` (same branch and commit as the package). Only the kernel-wrapper subset is carried -- the shape/dtype guard, -the dense-SDPA fallback, and the call counters. Upstream's model-integration +the dense fallback (routed to cute_dsl_fmha_fwd here, not torch SDPA), and the +call counters. Upstream's model-integration half is not carried: the diffusers self-attention dispatch hook, HunyuanVideo's padded `[video, text]` MMDiT handling, and model-level Morton ordering. @@ -67,6 +68,7 @@ preserve rather than overwrite: | `sol_attn_ineligible_reason()` added | names the specific reason (arch / head_dim / dtype) instead of one boolean | | `SOL_ATTN_STRICT=1` also covers the eligibility path | upstream raises only on kernel exceptions, so an ineligible run stayed silent | | `@torch.compiler.disable` on `_run_sol_attn_bthd` | see below | +| dense paths routed to `cute_dsl_fmha_fwd` via `dense_fn` | upstream's dense fallback is torch SDPA; staying in-backend is what makes a `backend: CUTEDSL` A/B isolate sparsity | **Upstream solves the `torch.compile` problem differently, and arguably better.** Its `sol_attn_backend.py` wraps the same call in a @@ -80,7 +82,8 @@ convention every other CuTe DSL entry point here already follows (`attention_backend/cute_dsl/fmha.py`, `cute_dsl_kernels/blackwell/video_sparse_attention/interface.py`). Without some such guard Dynamo traces into the CuTe DSL JIT builder and retraces on every -call -- measured at 69x slower on B200. Migrating to the `custom_op` form would +call -- measured at near two orders of magnitude slower on B200. Migrating to +the `custom_op` form would remove the per-layer graph break and is a reasonable follow-up; it was not done here because the `torch.compiler.disable` form is what this repository's other kernels use and what the measurements above were taken with. diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py index feed4aadeae5..203ebc455218 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py @@ -16,7 +16,7 @@ CuTe DSL imports and compilation are deferred to first use. Calls the kernel cannot serve -- wrong shape, dtype, or an architecture with no kernel -- -delegate to dense SDPA rather than failing, and increment +delegate to dense attention rather than failing, and increment ``_SOL_STATS["dense_fallback_calls"]`` so the degradation is countable. Set ``SOL_ATTN_STRICT=1`` to raise instead of falling back. """ @@ -133,8 +133,9 @@ def _dense_bthd(q, k, v): # Opaque to Dynamo, like every other CuTe DSL launch boundary here (see # cute_dsl/fmha.py, video_sparse_attention/interface.py). Otherwise Dynamo -# traces into the CuTe DSL JIT builder and retraces on every call: 69x slower -# on B200 (denoise 2496.9 s vs 36.2 s), silently, as if compile just didn't help. +# traces into the CuTe DSL JIT builder and retraces on every call: near two +# orders of magnitude slower on B200 (2496.9 s mean denoise without it), and +# silently, as if compile just didn't help. @torch.compiler.disable def _run_sol_attn_bthd( q, @@ -170,7 +171,7 @@ def dense(): if _strict(): raise RuntimeError(f"[sol-attn] cannot run the CuTe kernel: {reason}") logger.warning_once( - f"[sol-attn] falling back to dense SDPA: {reason}. Sol-Attn will not " + f"[sol-attn] falling back to dense attention: {reason}. Sol-Attn will not " "accelerate this run. Set SOL_ATTN_STRICT=1 to raise instead.", key=("sol_attn_ineligible", reason), ) @@ -195,7 +196,7 @@ def dense(): raise logger.warning_once( f"[sol-attn] kernel raised {type(exc).__name__}: {exc}; falling back to dense " - "SDPA for this call. Set SOL_ATTN_STRICT=1 to raise instead of silently falling " + "attention for this call. Set SOL_ATTN_STRICT=1 to raise instead of silently falling " "back.", key=(type(exc).__name__, str(exc)), ) diff --git a/tensorrt_llm/_torch/visual_gen/models/modeling.py b/tensorrt_llm/_torch/visual_gen/models/modeling.py index 60bf8ff3a017..3bf8d1a112d0 100644 --- a/tensorrt_llm/_torch/visual_gen/models/modeling.py +++ b/tensorrt_llm/_torch/visual_gen/models/modeling.py @@ -106,7 +106,7 @@ def register_cuda_graph_extra_key_fns(self, runner: "CUDAGraphRunner") -> None: # already baked into each captured graph and needs no key. return - # Sol-Attn switches between dense SDPA and the sparse kernel at the + # Sol-Attn switches between dense and sparse attention at the # dense-prefix boundary, again without changing tensor shapes, so # the two phases must not share a captured graph. runner.register_extra_key_fn( diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index fdcee0750346..df23d7cf2ff3 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -105,16 +105,25 @@ def __init__( ) - # Cross-attention fallback: TRTLLM and CUTEDSL VSA/Sol-Attn are self-attn only. + # SEPARATE_QKV fallback: TRTLLM and CUTEDSL VSA/Sol-Attn cannot serve it. # - # VSA/Sol-Attn fall back within their own backend family -- the dense - # CuTe DSL kernel serves cross-attention fine. Falling back to VANILLA - # instead silently swapped cross-attention from CuTeDSL to torch SDPA in - # every block the moment a sparse algorithm was enabled, so a - # `backend: CUTEDSL` run and a `CUTEDSL + sol_attn` run differed in - # cross-attention regardless of any sparse setting. TRTLLM keeps VANILLA: - # TrtllmAttention genuinely cannot serve SEPARATE_QKV. - _cross_attn_fallback = "CUTEDSL" if _is_sol_attn else "VANILLA" + # For genuine cross-attention, Sol-Attn falls back within its own backend + # family -- the dense CuTe DSL kernel serves cross-attention fine. VSA still + # goes to VANILLA and has the same defect; see TRTLLM-16105. + # Falling back to VANILLA instead silently swapped cross-attention from + # CuTeDSL to torch SDPA in every block the moment a sparse algorithm was + # enabled, so a `backend: CUTEDSL` run and a `CUTEDSL + sol_attn` run + # differed in cross-attention regardless of any sparse setting. + # + # This branch also catches *self*-attention that merely uses SEPARATE_QKV + # (Qwen-Image always; WAN's attn1 under async Ulysses). Those keep + # VANILLA: callers such as + # `qwen_image/transformer_qwen_image.py::_supports_qwen_key_padding_mask` + # test for the literal string "VANILLA", so redirecting them changes + # unrelated behaviour. TRTLLM keeps VANILLA throughout -- TrtllmAttention + # genuinely cannot serve SEPARATE_QKV. + _is_true_cross_attn = not separate_qkv_is_self_attention + _cross_attn_fallback = "CUTEDSL" if (_is_sol_attn and _is_true_cross_attn) else "VANILLA" if separate_qkv_cross_attention and ( base_backend == "TRTLLM" or _is_vsa or _is_sol_attn ): @@ -272,6 +281,7 @@ def __init__( config.attention.model_copy(update={"sparse_attention_config": None}) if backend_name == "CUTEDSL" and _is_sol_attn + and _is_true_cross_attn and self.qkv_mode == QKVMode.SEPARATE_QKV else config.attention ), diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 6120905d94a3..05017d5d55f9 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -229,7 +229,8 @@ class SolAttnAttentionConfig(BaseSparseAttentionConfig): (B200/GB200) only, head_dim=128, bf16, MHA. On an unsupported *shape, dtype, or architecture* the kernel falls back to - dense SDPA and counts the fallback, so setting this config on the wrong GPU + dense attention -- the configured backend's dense kernel where available, + torch SDPA otherwise -- and counts the fallback, so setting this config on the wrong GPU degrades rather than fails. Two cases are not covered by that fallback and do raise: GQA/MQA (num_kv_heads != num_heads) here at construction, and context parallelism (cp_size > 1), rejected in visual_gen/modules/attention.py. diff --git a/tests/integration/test_lists/test-db/l0_gb202.yml b/tests/integration/test_lists/test-db/l0_gb202.yml index 49a9ed558dd3..79bd803a1890 100644 --- a/tests/integration/test_lists/test-db/l0_gb202.yml +++ b/tests/integration/test_lists/test-db/l0_gb202.yml @@ -20,9 +20,6 @@ l0_gb202: - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] # - unittest/_torch/modeling -k "modeling_qwen3" # https://nvbugs/5234573 - unittest/_torch/attention/test_attention_mla.py - # ------------- Visual Gen tests --------------- - # sm120 (GB202) coverage for the Sol-Attn CuTeDSL kernel; the same file - # is registered in l0_b200.yml for sm100. - test_e2e.py::test_ptp_quickstart_bert[VANILLA-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] - test_e2e.py::test_ptp_quickstart_bert[TRTLLM-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_bf16[latency] diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py index d42afeabb4a0..09890c92c8ab 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py @@ -213,8 +213,13 @@ def test_graph_phase_accepts_tensor_timestep(): assert sol_attn_graph_phase(torch.tensor([0.10]), disabled_until_timestep=0.95) == 1 -def test_dense_prefix_uses_sdpa_and_skips_kernel(monkeypatch): - """Inside the dense prefix the sparse kernel must not be invoked at all.""" +def test_dense_prefix_skips_kernel(monkeypatch): + """Inside the dense prefix the sparse kernel must not be invoked at all. + + CPU tensors, so `_dense` takes its SDPA branch here; that the dense path + routes to the CuTe kernel on CUDA is covered by + `test_dense_paths_use_cutedsl_backend`. + """ import tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn as sol_attn_mod def _fail_if_called(*args, **kwargs): @@ -252,7 +257,7 @@ def _record(*args, **kwargs): assert called["n"] == 1, "expected the sparse kernel, not a silent dense fallback" -def test_sol_attn_dense_layers_guard_skips_kernel(monkeypatch): +def test_dense_layers_guard_skips_kernel(monkeypatch): """A layer_idx in dense_layers must use the dense SDPA path and never invoke the kernel.""" import tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn as sol_attn_mod @@ -447,8 +452,8 @@ def test_kernel_launch_is_opaque_to_dynamo(): """The CuTe DSL launch boundary must be @torch.compiler.disable'd. Without it Dynamo traces into the CuTe DSL JIT builder and retraces on every - call: 69x slower on B200 (denoise 2496.9 s vs 36.2 s), and silent -- it looks - like torch.compile simply not paying off. + call: near two orders of magnitude slower on B200 (2496.9 s mean denoise + without it), and silent -- it looks like torch.compile simply not paying off. """ assert _is_dynamo_disabled(_backend_mod()._run_sol_attn_bthd), ( "_run_sol_attn_bthd must be decorated with @torch.compiler.disable" @@ -463,3 +468,61 @@ def test_timestep_scalar_read_is_opaque_to_dynamo(): assert _is_dynamo_disabled(SolAttnAttention._dense_by_step), ( "SolAttnAttention._dense_by_step must be decorated with @torch.compiler.disable" ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_dense_paths_use_cutedsl_backend(monkeypatch): + """All three dense paths must reach the configured backend's dense kernel. + + Sol-Attn does dense attention on the `dense_layers` guard, the + `disabled_until_timestep` prefix, and kernel-ineligibility fallback. If those + call torch SDPA instead of `cute_dsl_fmha_fwd`, a `backend: CUTEDSL` run + differs from a `backend: CUTEDSL` dense baseline on those steps, and any A/B + against that baseline measures a backend swap rather than sparsity. Measured + at LPIPS 0.214 on Wan2.2-T2V-A14B before this was fixed, against a 0.25 gate. + + The CPU-tensor tests above cannot see this: `_dense` falls back to SDPA when + `q.is_cuda` is false, so they exercise the wrong branch by construction. + """ + import tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn as sol_attn_mod + + device = torch.device("cuda") + q = k = v = torch.randn(1, 64, 2, 128, device=device, dtype=torch.bfloat16) + + def _make(): + a = SolAttnAttention(layer_idx=0, num_heads=2, head_dim=128) + calls = {"n": 0} + real = a._dense_backend.forward + + def _spy(*args, **kwargs): + calls["n"] += 1 + return real(*args, **kwargs) + + monkeypatch.setattr(a._dense_backend, "forward", _spy) + return a, calls + + # 1. dense prefix: timestep at/above the cutoff + a, calls = _make() + a.disabled_until_timestep = 0.9 + a.dense_layers = frozenset() + a.forward(q, k, v, timestep=torch.tensor(0.95)) + assert calls["n"] == 1, "dense prefix did not use the CuTeDSL dense kernel" + + # 2. dense_layers guard + a, calls = _make() + a.disabled_until_timestep = None + a.dense_layers = frozenset({0}) + a.forward(q, k, v) + assert calls["n"] == 1, "dense_layers guard did not use the CuTeDSL dense kernel" + + # 3. ineligibility fallback, reached through `dense_fn` + a, calls = _make() + a.disabled_until_timestep = None + a.dense_layers = frozenset() + monkeypatch.setattr( + sol_attn_mod, + "_sol_attn_run", + lambda *args, **kw: kw["dense_fn"](*args[:3]), + ) + a.forward(q, k, v) + assert calls["n"] == 1, "dense_fn did not route the fallback to the CuTeDSL dense kernel" From d3eb2608af330d15976d3b563b30f7e0777f3a5a Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Wed, 2 Sep 2026 21:58:42 -0700 Subject: [PATCH 07/22] [TRTLLM-15917][chore] Rename SolAttnAttention to SolAttention `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 `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) Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../features/visualgen-sparse-attention.md | 2 +- .../visual_gen/attention_backend/__init__.py | 4 +- .../attention_backend/cute_dsl/__init__.py | 6 +-- .../attention_backend/cute_dsl/sol_attn.py | 6 +-- .../visual_gen/attention_backend/utils.py | 4 +- .../blackwell/sol_attn_backend.py | 2 +- .../_torch/visual_gen/models/modeling.py | 7 +--- tensorrt_llm/visual_gen/__init__.py | 6 +-- tensorrt_llm/visual_gen/args.py | 6 +-- tensorrt_llm/visual_gen/sparse_attention.py | 4 +- .../test_attention_cute_dsl_sol_attn.py | 42 +++++++++---------- 11 files changed, 43 insertions(+), 46 deletions(-) diff --git a/docs/source/features/visualgen-sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md index 19acd2a5810c..da7a998b97d7 100644 --- a/docs/source/features/visualgen-sparse-attention.md +++ b/docs/source/features/visualgen-sparse-attention.md @@ -21,7 +21,7 @@ Sparse attention in VisualGen is configured through `VisualGenArgs.attention_con |---|---|---| | `skip_softmax` | `SkipSoftmaxAttentionConfig` | Supported | | VSA | TBD | TODO | -| `sol_attn` | `SolAttnAttentionConfig` | Supported (CUTEDSL, sm100) | +| `sol_attn` | `SolAttentionConfig` | Supported (CUTEDSL, sm100) | ### Sol-Attn diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py index 6acbff9c5b41..bdc9fe9f33c1 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py @@ -24,7 +24,7 @@ from .cute_dsl import ( VSA_TILE_SIZE, CuTeDSLAttention, - SolAttnAttention, + SolAttention, VSAAttention, VSAMetadata, VSAMetadataBuilder, @@ -49,7 +49,7 @@ "FlashAttn4Attention", "FlashInferAttention", "RingAttention", - "SolAttnAttention", + "SolAttention", "TrtllmAttention", "TrtllmAttentionMetadata", "UlyssesAttention", diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py index 20e52c1f0574..f2c11458c1b2 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py @@ -17,11 +17,11 @@ fmha.py — CuTeDSLAttention (dense and blockscaled JIT FMHA) vsa.py — VSAAttention (Video Sparse Attention, CuTe JIT + SDPA fallback) - sol_attn.py — SolAttnAttention (Sol-Attn dynamic block routing, CuTe JIT + SDPA fallback) + sol_attn.py — SolAttention (Sol-Attn dynamic block routing, CuTe JIT + SDPA fallback) """ from .fmha import CuTeDSLAttention, _cute_dsl_import_error -from .sol_attn import SolAttnAttention, sol_attn_graph_phase +from .sol_attn import SolAttention, sol_attn_graph_phase from .vsa import ( VSA_KERNEL_MAX_CUBES, VSA_TILE_SIZE, @@ -44,6 +44,6 @@ "set_vsa_forward_context", "get_vsa_forward_context", "_cute_dsl_import_error", - "SolAttnAttention", + "SolAttention", "sol_attn_graph_phase", ] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py index 3d84611068d0..9d8028087d3d 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py @@ -126,7 +126,7 @@ def _parse_dense_layers(spec: Optional[str]) -> frozenset: return frozenset(layers) -class SolAttnAttention(AttentionBackend): +class SolAttention(AttentionBackend): """Sol-Attn dynamic block-routing sparse attention (CuTeDSL, sm100). The kernel wrapper already falls back to dense attention on any unsupported @@ -147,7 +147,7 @@ def __init__( ): if _sol_attn_run is None: raise ImportError( - "SolAttnAttention requires the vendored sol_attn kernel " + "SolAttention requires the vendored sol_attn kernel " f"package; import failed: {_sol_attn_import_error}" ) self.layer_idx = layer_idx @@ -214,7 +214,7 @@ def _dense_by_step(self, timestep) -> bool: # sparse kernel rather than silently forcing dense forever. # This degrades quality rather than raising, so say so once. logger.warning_once( - "SolAttnAttentionConfig.disabled_until_timestep=" + "SolAttentionConfig.disabled_until_timestep=" f"{self.disabled_until_timestep} is set, but no `timestep` reached " "the Sol-Attn forward call. The dense prefix it requests will not " "be applied. Ensure the pipeline passes a normalized timestep, or " diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py index 0be1d6299a0d..80cedc79fa9c 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py @@ -152,9 +152,9 @@ def create_attention( attn_cls = VSAAttention kwargs["sparse_attention_config"] = attention_config.sparse_attention_config elif sparse_algo == "sol_attn": - from .cute_dsl.sol_attn import SolAttnAttention + from .cute_dsl.sol_attn import SolAttention - attn_cls = SolAttnAttention + attn_cls = SolAttention kwargs["sparse_attention_config"] = attention_config.sparse_attention_config return attn_cls( diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py index 203ebc455218..585a2cd3fed7 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py @@ -10,7 +10,7 @@ ``kv_splits``, and an optional exact KV sink range. TRT-LLM's dispatch path (``attention_backend/cute_dsl/sol_attn.py``, -``SolAttnAttention``) consumes exactly two names from this module: +``SolAttention``) consumes exactly two names from this module: ``_run_sol_attn_bthd`` and ``sol_attn_supported``. The dense-prefix decision lives there too, keyed off the normalized timestep forward kwarg. diff --git a/tensorrt_llm/_torch/visual_gen/models/modeling.py b/tensorrt_llm/_torch/visual_gen/models/modeling.py index 3bf8d1a112d0..588a18791f35 100644 --- a/tensorrt_llm/_torch/visual_gen/models/modeling.py +++ b/tensorrt_llm/_torch/visual_gen/models/modeling.py @@ -22,10 +22,7 @@ from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxScheduler from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn import sol_attn_graph_phase from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig -from tensorrt_llm.visual_gen.sparse_attention import ( - SkipSoftmaxAttentionConfig, - SolAttnAttentionConfig, -) +from tensorrt_llm.visual_gen.sparse_attention import SkipSoftmaxAttentionConfig, SolAttentionConfig if TYPE_CHECKING: from tensorrt_llm._torch.visual_gen.cuda_graph_runner import CUDAGraphRunner @@ -99,7 +96,7 @@ def register_cuda_graph_extra_key_fns(self, runner: "CUDAGraphRunner") -> None: ) return - if isinstance(sparse_config, SolAttnAttentionConfig): + if isinstance(sparse_config, SolAttentionConfig): disabled_until_timestep = sparse_config.disabled_until_timestep if disabled_until_timestep is None: # dense_layers is fixed per layer at construction, so it is diff --git a/tensorrt_llm/visual_gen/__init__.py b/tensorrt_llm/visual_gen/__init__.py index 3d0a2cb0636c..3b1ef26b8d98 100644 --- a/tensorrt_llm/visual_gen/__init__.py +++ b/tensorrt_llm/visual_gen/__init__.py @@ -49,7 +49,7 @@ QuantAttentionConfig, RuntimeLoRAConfig, SkipSoftmaxAttentionConfig, - SolAttnAttentionConfig, + SolAttentionConfig, SparseAttentionConfig, TeaCacheConfig, TorchCompileConfig, @@ -77,7 +77,7 @@ "QuantAttentionConfig": "tensorrt_llm.visual_gen.args", "RuntimeLoRAConfig": "tensorrt_llm.visual_gen.args", "SkipSoftmaxAttentionConfig": "tensorrt_llm.visual_gen.args", - "SolAttnAttentionConfig": "tensorrt_llm.visual_gen.args", + "SolAttentionConfig": "tensorrt_llm.visual_gen.args", "SparseAttentionConfig": "tensorrt_llm.visual_gen.args", "TeaCacheConfig": "tensorrt_llm.visual_gen.args", "TorchCompileConfig": "tensorrt_llm.visual_gen.args", @@ -136,7 +136,7 @@ def __dir__(): "SparseAttentionConfig", "SkipSoftmaxAttentionConfig", "VideoSparseAttentionConfig", - "SolAttnAttentionConfig", + "SolAttentionConfig", "VAEConfig", "CacheConfig", "TeaCacheConfig", diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index 38c45af3f71d..1ccf02e1a90f 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -33,7 +33,7 @@ from .sparse_attention import ( SkipSoftmaxAttentionConfig, - SolAttnAttentionConfig, + SolAttentionConfig, VideoSparseAttentionConfig, ) @@ -99,7 +99,7 @@ class QuantAttentionConfig(StrictBaseModel): # Discriminated union of sparse attention configs. SparseAttentionConfig = Annotated[ - Union[SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig, SolAttnAttentionConfig], + Union[SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig, SolAttentionConfig], Field(discriminator="algorithm"), ] @@ -849,7 +849,7 @@ def from_yaml(cls, yaml_path: Union[str, Path], **overrides: Any) -> "VisualGenA "SparseAttentionConfig", "SkipSoftmaxAttentionConfig", "VideoSparseAttentionConfig", - "SolAttnAttentionConfig", + "SolAttentionConfig", "AttentionConfig", "VAEConfig", "ParallelConfig", diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 05017d5d55f9..3e9655f5f407 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -221,7 +221,7 @@ def _ckpt_sparse_attention_config_from_kwargs( return None -class SolAttnAttentionConfig(BaseSparseAttentionConfig): +class SolAttentionConfig(BaseSparseAttentionConfig): """Sol-Attn sparse attention configuration for visual generation. Dynamic block routing + sparse computation + approximation correction in @@ -281,7 +281,7 @@ class SolAttnAttentionConfig(BaseSparseAttentionConfig): ) def to_sparse_params(self, **kwargs): - # Sol-Attn's knobs are consumed directly by SolAttnAttention.__init__ + # Sol-Attn's knobs are consumed directly by SolAttention.__init__ # (constructed via CUTEDSL backend dispatch in create_attention), not # lowered into a shared SparseParams -- the vendored kernel has no # checkpoint-calibration step to resolve here, unlike skip_softmax. diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py index 09890c92c8ab..35bda97d0e49 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py @@ -29,7 +29,7 @@ from tensorrt_llm._torch.visual_gen.attention_backend import CuTeDSLAttention from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn import ( - SolAttnAttention, + SolAttention, _parse_dense_layers, sol_attn_graph_phase, ) @@ -39,7 +39,7 @@ create_attention_metadata_state, ) from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode -from tensorrt_llm.visual_gen.args import AttentionConfig, SolAttnAttentionConfig +from tensorrt_llm.visual_gen.args import AttentionConfig, SolAttentionConfig def test_cute_dsl_factory_dispatches_dense_and_sol_attn() -> None: @@ -52,7 +52,7 @@ def test_cute_dsl_factory_dispatches_dense_and_sol_attn() -> None: attention_config=dense_config, ) - sparse_config = SolAttnAttentionConfig(tau=2.0, disabled_until_timestep=0.9545) + sparse_config = SolAttentionConfig(tau=2.0, disabled_until_timestep=0.9545) sol_attn_config = AttentionConfig(backend="CUTEDSL", sparse_attention_config=sparse_config) sol_attn_attention = create_attention( backend="CUTEDSL", @@ -63,7 +63,7 @@ def test_cute_dsl_factory_dispatches_dense_and_sol_attn() -> None: ) assert isinstance(dense_attention, CuTeDSLAttention) - assert isinstance(sol_attn_attention, SolAttnAttention) + assert isinstance(sol_attn_attention, SolAttention) assert sol_attn_attention.tau == 2.0 assert sol_attn_attention.disabled_until_timestep == 0.9545 @@ -83,7 +83,7 @@ def _make_config( eps=1e-6, ) sparse_attention_config = ( - SolAttnAttentionConfig(tau=sol_attn_tau) if sol_attn_tau is not None else None + SolAttentionConfig(tau=sol_attn_tau) if sol_attn_tau is not None else None ) config = DiffusionModelConfig( pretrained_config=pretrained_config, @@ -107,7 +107,7 @@ def test_sol_attn_cross_attention_uses_dense_cutedsl(): a backend difference masquerading as a sparsity difference in any A/B. """ from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.fmha import CuTeDSLAttention - from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn import SolAttnAttention + from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn import SolAttention device = torch.device("cuda") dtype = torch.bfloat16 @@ -125,7 +125,7 @@ def test_sol_attn_cross_attention_uses_dense_cutedsl(): assert isinstance(cross_attn.attn, CuTeDSLAttention), ( f"expected the dense CuTeDSL kernel, got {type(cross_attn.attn).__name__}" ) - assert not isinstance(cross_attn.attn, SolAttnAttention), ( + assert not isinstance(cross_attn.attn, SolAttention), ( "cross-attention must not re-select the sparse backend" ) @@ -142,7 +142,7 @@ def test_sol_attn_with_context_parallelism_raises(): pretrained_config=pretrained_config, attention=AttentionConfig( backend="CUTEDSL", - sparse_attention_config=SolAttnAttentionConfig(tau=1.0), + sparse_attention_config=SolAttentionConfig(tau=1.0), ), skip_create_weights_in_init=False, ) @@ -164,7 +164,7 @@ def test_sol_attn_with_context_parallelism_raises(): def test_sol_attn_rejects_gqa_mqa(): """Sol-Attn is MHA-only; num_kv_heads != num_heads must fail fast at construction.""" with pytest.raises(AssertionError, match="MHA-only"): - SolAttnAttention(layer_idx=0, num_heads=8, head_dim=128, num_kv_heads=2) + SolAttention(layer_idx=0, num_heads=8, head_dim=128, num_kv_heads=2) @pytest.mark.parametrize( @@ -227,7 +227,7 @@ def _fail_if_called(*args, **kwargs): monkeypatch.setattr(sol_attn_mod, "_sol_attn_run", _fail_if_called) - attn = SolAttnAttention(layer_idx=0, num_heads=2, head_dim=16) + attn = SolAttention(layer_idx=0, num_heads=2, head_dim=16) attn.disabled_until_timestep = 0.9 q = k = v = torch.randn(1, 4, 2, 16) out = attn.forward(q, k, v, timestep=torch.tensor(0.95)) @@ -250,7 +250,7 @@ def _record(*args, **kwargs): monkeypatch.setattr(sol_attn_mod, "_sol_attn_run", _record) - attn = SolAttnAttention(layer_idx=0, num_heads=2, head_dim=16) + attn = SolAttention(layer_idx=0, num_heads=2, head_dim=16) attn.disabled_until_timestep = 0.9 q = k = v = torch.randn(1, 4, 2, 16) attn.forward(q, k, v) # no timestep kwarg @@ -266,7 +266,7 @@ def _fail_if_called(*args, **kwargs): monkeypatch.setattr(sol_attn_mod, "_sol_attn_run", _fail_if_called) - attn = SolAttnAttention(layer_idx=3, num_heads=2, head_dim=16) + attn = SolAttention(layer_idx=3, num_heads=2, head_dim=16) attn.dense_layers = frozenset({3}) q = k = v = torch.randn(1, 4, 2, 16) out = attn.forward(q, k, v) @@ -285,7 +285,7 @@ def _make_solattn_model(disabled_until_timestep=None, dense_layers=None): pretrained_config=pretrained_config, attention=AttentionConfig( backend="CUTEDSL", - sparse_attention_config=SolAttnAttentionConfig( + sparse_attention_config=SolAttentionConfig( tau=2.0, disabled_until_timestep=disabled_until_timestep, dense_layers=dense_layers, @@ -405,7 +405,7 @@ def test_quant_attention_config_rejected_with_sol_attn(): AttentionConfig( backend="CUTEDSL", quant_attention_config=QuantAttentionConfig(), - sparse_attention_config=SolAttnAttentionConfig(tau=2.0), + sparse_attention_config=SolAttentionConfig(tau=2.0), ) @@ -413,8 +413,8 @@ def test_zero_cutoff_rejected(): """0.0 is the natural thing to type for 'no prefix', but it would run dense on every step and turn Sol-Attn off entirely. Must be rejected, not silent.""" with pytest.raises(ValueError): - SolAttnAttentionConfig(tau=2.0, disabled_until_timestep=0.0) - assert SolAttnAttentionConfig(tau=2.0).disabled_until_timestep is None + SolAttentionConfig(tau=2.0, disabled_until_timestep=0.0) + assert SolAttentionConfig(tau=2.0).disabled_until_timestep is None @pytest.mark.skip( @@ -438,8 +438,8 @@ def test_kv_splits_rejects_unsupported_value(): otherwise rejected deep inside the kernel and caught by the blanket except, silently degrading the entire run to dense attention.""" with pytest.raises(ValueError): - SolAttnAttentionConfig(tau=2.0, kv_splits="4") - assert SolAttnAttentionConfig(tau=2.0).kv_splits == "auto" + SolAttentionConfig(tau=2.0, kv_splits="4") + assert SolAttentionConfig(tau=2.0).kv_splits == "auto" def _is_dynamo_disabled(fn) -> bool: @@ -465,8 +465,8 @@ def test_timestep_scalar_read_is_opaque_to_dynamo(): Otherwise it graph-breaks the enclosing block once per attention layer. """ - assert _is_dynamo_disabled(SolAttnAttention._dense_by_step), ( - "SolAttnAttention._dense_by_step must be decorated with @torch.compiler.disable" + assert _is_dynamo_disabled(SolAttention._dense_by_step), ( + "SolAttention._dense_by_step must be decorated with @torch.compiler.disable" ) @@ -490,7 +490,7 @@ def test_dense_paths_use_cutedsl_backend(monkeypatch): q = k = v = torch.randn(1, 64, 2, 128, device=device, dtype=torch.bfloat16) def _make(): - a = SolAttnAttention(layer_idx=0, num_heads=2, head_dim=128) + a = SolAttention(layer_idx=0, num_heads=2, head_dim=128) calls = {"n": 0} real = a._dense_backend.forward From 1f55662fac648ca4d05f29473b3a276dc362b16e Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Thu, 3 Sep 2026 00:39:03 -0700 Subject: [PATCH 08/22] [TRTLLM-15917][refactor] Decide Sol-Attn eligibility per call, not per 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) Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../attention_backend/cute_dsl/sol_attn.py | 42 +++++++++---- .../_torch/visual_gen/modules/attention.py | 49 ++++----------- .../test_attention_cute_dsl_sol_attn.py | 61 ++++++++++++------- 3 files changed, 82 insertions(+), 70 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py index 9d8028087d3d..ecda0bef4959 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py @@ -174,7 +174,7 @@ def __init__( # switched off entirely, against a 0.25 gate. from .fmha import CuTeDSLAttention - self._dense_backend = CuTeDSLAttention( + self._inner = CuTeDSLAttention( layer_idx=layer_idx, num_heads=num_heads, head_dim=head_dim, @@ -226,13 +226,13 @@ def _dense_by_step(self, timestep) -> bool: @staticmethod def _sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: - """Dense attention via torch SDPA, for architectures CuTe DSL cannot serve.""" + """Dense attention via torch SDPA, for devices CuTe DSL cannot serve.""" return torch.nn.functional.scaled_dot_product_attention( q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) ).transpose(1, 2) - def _dense(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: - """Dense attention on the configured backend, or SDPA where unavailable. + def _delegate(self, q, k, v, **kwargs) -> torch.Tensor: + """Hand the call to the dense backend of the same family. ``_cute_dense_ok`` answers "can this *device* run the kernel", decided at construction; ``q.is_cuda`` answers "is this *tensor* on it". Both are @@ -240,9 +240,29 @@ def _dense(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Ten it says yes on a GPU host even when a caller passes CPU tensors. """ if self._cute_dense_ok and q.is_cuda: - return self._dense_backend.forward(q, k, v) + return self._inner.forward(q, k, v, **kwargs) return self._sdpa(q, k, v) + def _can_serve(self, q: torch.Tensor, k: torch.Tensor, **kwargs) -> bool: + """Whether the sparse kernel applies to this particular call. + + Everything false here is delegated to ``_inner``. Deciding it from the + tensors, per call, is deliberate: ``qkv_mode`` describes how Q/K/V are + *projected*, not whether K/V come from another sequence, so a + construction-time rule keyed on ``SEPARATE_QKV`` mistakes self-attention + for cross-attention wherever that mode is chosen for other reasons -- + Qwen-Image always, and WAN's ``attn1`` under async Ulysses. + """ + # Cross-attention: K/V come from another sequence. Sol-Attn's routing + # assumes one self-attending sequence. + if k.shape[1] != q.shape[1]: + return False + if self.layer_idx in self.dense_layers: + return False + if self.disabled_until_timestep is not None and self._dense_by_step(kwargs.get("timestep")): + return False + return True + def forward( self, q: torch.Tensor, @@ -251,12 +271,8 @@ def forward( **kwargs, ) -> torch.Tensor: """q, k, v: [B, S, H, D] (NHD), same original token order in and out.""" - dense_by_layer = self.layer_idx in self.dense_layers - dense_by_step = False - if self.disabled_until_timestep is not None: - dense_by_step = self._dense_by_step(kwargs.get("timestep")) - if dense_by_layer or dense_by_step: - return self._dense(q, k, v) + if not self._can_serve(q, k, **kwargs): + return self._delegate(q, k, v, **kwargs) return _sol_attn_run( q, k, @@ -264,7 +280,9 @@ def forward( tau=self.tau, thresh_type=self.thresh_type, kv_splits=self.kv_splits, - dense_fn=self._dense, + # Shape/dtype/arch ineligibility is only detectable inside the + # wrapper, so that last delegation happens through this hook. + dense_fn=lambda a, b, c: self._delegate(a, b, c, **kwargs), ) @classmethod diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index df23d7cf2ff3..ec460f1e798d 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -105,36 +105,18 @@ def __init__( ) - # SEPARATE_QKV fallback: TRTLLM and CUTEDSL VSA/Sol-Attn cannot serve it. + # SEPARATE_QKV fallback: TRTLLM and CUTEDSL VSA cannot serve it. # - # For genuine cross-attention, Sol-Attn falls back within its own backend - # family -- the dense CuTe DSL kernel serves cross-attention fine. VSA still - # goes to VANILLA and has the same defect; see TRTLLM-16105. - # Falling back to VANILLA instead silently swapped cross-attention from - # CuTeDSL to torch SDPA in every block the moment a sparse algorithm was - # enabled, so a `backend: CUTEDSL` run and a `CUTEDSL + sol_attn` run - # differed in cross-attention regardless of any sparse setting. - # - # This branch also catches *self*-attention that merely uses SEPARATE_QKV - # (Qwen-Image always; WAN's attn1 under async Ulysses). Those keep - # VANILLA: callers such as - # `qwen_image/transformer_qwen_image.py::_supports_qwen_key_padding_mask` - # test for the literal string "VANILLA", so redirecting them changes - # unrelated behaviour. TRTLLM keeps VANILLA throughout -- TrtllmAttention - # genuinely cannot serve SEPARATE_QKV. - _is_true_cross_attn = not separate_qkv_is_self_attention - _cross_attn_fallback = "CUTEDSL" if (_is_sol_attn and _is_true_cross_attn) else "VANILLA" - if separate_qkv_cross_attention and ( - base_backend == "TRTLLM" or _is_vsa or _is_sol_attn - ): - backend_name = _cross_attn_fallback - requested = ( - f"{base_backend} (VSA)" - if _is_vsa - else f"{base_backend} (Sol-Attn)" - if _is_sol_attn - else base_backend - ) + # Sol-Attn is deliberately absent: it decides per call (`_can_serve`) and + # delegates what it cannot serve to the dense backend of its own family. + # A rule here would have to guess from `qkv_mode`, which describes how + # Q/K/V are *projected* rather than whether K/V come from another + # sequence -- and that guess is wrong wherever SEPARATE_QKV is chosen for + # other reasons (Qwen-Image always; WAN's attn1 under async Ulysses), + # silently costing those modules their configured backend. + if separate_qkv_cross_attention and (base_backend == "TRTLLM" or _is_vsa): + backend_name = "VANILLA" + requested = f"{base_backend} (VSA)" if _is_vsa else base_backend # Warn once per (module class, requested, resolved) triple so the # fallback is visible without per-module-instance log spam. logger.warning_once( @@ -277,14 +259,7 @@ def __init__( num_kv_heads=backend_num_kv_heads, quant_config=self.quant_config, dtype=self.dtype, - attention_config=( - config.attention.model_copy(update={"sparse_attention_config": None}) - if backend_name == "CUTEDSL" - and _is_sol_attn - and _is_true_cross_attn - and self.qkv_mode == QKVMode.SEPARATE_QKV - else config.attention - ), + attention_config=config.attention, attention_metadata_state=attention_metadata_state, sparse_params=sparse_params, ) diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py index 35bda97d0e49..32eb20da4c50 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py @@ -97,18 +97,19 @@ def _make_config( @pytest.mark.skipif(not torch.cuda.is_available(), reason="Sol-Attn needs CUDA") -def test_sol_attn_cross_attention_uses_dense_cutedsl(): - """Cross-attention must stay on CuTeDSL, not drop to VANILLA. - - Sol-Attn is self-attention only, so SEPARATE_QKV modules fall back -- but to - the dense kernel of the *configured* backend, not to torch SDPA. Falling back - to VANILLA made a `backend: CUTEDSL` run and a `CUTEDSL + sol_attn` run differ - in cross-attention in every block, regardless of any sparse setting, which is - a backend difference masquerading as a sparsity difference in any A/B. +def test_sol_attn_cross_attention_delegates_to_dense_cutedsl(): + """Cross-attention is delegated to the dense backend, decided per call. + + Sol-Attn is self-attention only. It is still the module's backend for a + cross-attention module -- `create_attention` has no rule excluding it -- and + delegates at `forward` because `_can_serve` sees `k.shape[1] != q.shape[1]`. + + Deciding this per call rather than at construction is the point: `qkv_mode` + describes how Q/K/V are projected, not whether K/V come from another + sequence, so a construction-time rule keyed on SEPARATE_QKV silently + stripped the configured backend from self-attention modules that use that + mode for unrelated reasons (Qwen-Image; WAN attn1 under async Ulysses). """ - from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.fmha import CuTeDSLAttention - from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn import SolAttention - device = torch.device("cuda") dtype = torch.bfloat16 cfg = _make_config( @@ -120,16 +121,34 @@ def test_sol_attn_cross_attention_uses_dense_cutedsl(): .eval() ) assert cross_attn.attn_backend == "CUTEDSL", ( - f"expected CUTEDSL cross-attention, got {cross_attn.attn_backend!r}" + f"expected CUTEDSL, got {cross_attn.attn_backend!r}" ) - assert isinstance(cross_attn.attn, CuTeDSLAttention), ( - f"expected the dense CuTeDSL kernel, got {type(cross_attn.attn).__name__}" + assert isinstance(cross_attn.attn, SolAttention), ( + "Sol-Attn should remain the backend and delegate per call, not be " + f"swapped out at construction; got {type(cross_attn.attn).__name__}" ) - assert not isinstance(cross_attn.attn, SolAttention), ( - "cross-attention must not re-select the sparse backend" + # q and k with different sequence lengths -> not self-attention -> delegate + q = torch.randn(1, 32, 4, 16, device=device, dtype=dtype) + k = torch.randn(1, 77, 4, 16, device=device, dtype=dtype) + assert not cross_attn.attn._can_serve(q, k), ( + "differing q/k sequence lengths must be delegated, not routed to the sparse kernel" ) +def test_sol_attn_self_attention_is_served_under_separate_qkv(): + """SEPARATE_QKV self-attention keeps Sol-Attn -- the async-Ulysses case. + + WAN's attn1 switches to SEPARATE_QKV when async Ulysses is active, and + Qwen-Image uses it unconditionally. Both are self-attention; both must still + get the sparse kernel. + """ + device = torch.device("cuda") + attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) + attn.disabled_until_timestep = None + q = k = torch.randn(1, 64, 2, 128, device=device, dtype=torch.bfloat16) + assert attn._can_serve(q, k), "equal q/k sequence lengths must reach the sparse kernel" + + def test_sol_attn_with_context_parallelism_raises(): """Sol-Attn + Attention2D/Ring must error at construction (needs the full sequence per rank).""" pretrained_config = SimpleNamespace( @@ -472,7 +491,7 @@ def test_timestep_scalar_read_is_opaque_to_dynamo(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") def test_dense_paths_use_cutedsl_backend(monkeypatch): - """All three dense paths must reach the configured backend's dense kernel. + """Every path Sol-Attn cannot serve must reach the configured dense kernel. Sol-Attn does dense attention on the `dense_layers` guard, the `disabled_until_timestep` prefix, and kernel-ineligibility fallback. If those @@ -492,13 +511,13 @@ def test_dense_paths_use_cutedsl_backend(monkeypatch): def _make(): a = SolAttention(layer_idx=0, num_heads=2, head_dim=128) calls = {"n": 0} - real = a._dense_backend.forward + real = a._inner.forward def _spy(*args, **kwargs): calls["n"] += 1 return real(*args, **kwargs) - monkeypatch.setattr(a._dense_backend, "forward", _spy) + monkeypatch.setattr(a._inner, "forward", _spy) return a, calls # 1. dense prefix: timestep at/above the cutoff @@ -506,14 +525,14 @@ def _spy(*args, **kwargs): a.disabled_until_timestep = 0.9 a.dense_layers = frozenset() a.forward(q, k, v, timestep=torch.tensor(0.95)) - assert calls["n"] == 1, "dense prefix did not use the CuTeDSL dense kernel" + assert calls["n"] == 1, "dense prefix was not delegated to the CuTeDSL dense kernel" # 2. dense_layers guard a, calls = _make() a.disabled_until_timestep = None a.dense_layers = frozenset({0}) a.forward(q, k, v) - assert calls["n"] == 1, "dense_layers guard did not use the CuTeDSL dense kernel" + assert calls["n"] == 1, "dense_layers guard was not delegated to the CuTeDSL dense kernel" # 3. ineligibility fallback, reached through `dense_fn` a, calls = _make() From a37e7ff5e3b376940bfd30e8d7cc1e2c49880352 Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Thu, 3 Sep 2026 10:25:57 -0700 Subject: [PATCH 09/22] [TRTLLM-15917][fix] Address review findings on the Sol-Attn integration Seven findings from automated review of b25ca70e, 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) Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../features/visualgen-sparse-attention.md | 5 ++- .../attention_backend/cute_dsl/sol_attn.py | 26 +++++++----- .../blackwell/sol_attn_backend.py | 14 +++++++ tensorrt_llm/visual_gen/sparse_attention.py | 40 +++++++++++++++++++ .../test_attention_cute_dsl_sol_attn.py | 33 +++++++++++++-- 5 files changed, 103 insertions(+), 15 deletions(-) diff --git a/docs/source/features/visualgen-sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md index da7a998b97d7..bef9877ca268 100644 --- a/docs/source/features/visualgen-sparse-attention.md +++ b/docs/source/features/visualgen-sparse-attention.md @@ -6,6 +6,7 @@ This feature is in **beta** stage. APIs, supported models, and optimization opti - [Overview](#overview) - [Algorithms](#algorithms) + - [Sol-Attn](#sol-attn) - [Skip Softmax Attention](#skip-softmax-attention) - [Video Sparse Attention (VSA)](#video-sparse-attention-vsa) @@ -49,8 +50,8 @@ below it. Use `None` rather than `0.0` to disable the prefix. On an input the kernel cannot serve — an unsupported architecture, a `head_dim` other than 128, a non-bfloat16 dtype — Sol-Attn falls back to dense -dense attention -- the configured backend's dense kernel where available, torch -SDPA otherwise -- logs the specific reason once, and counts the fallback. Set +attention -- the configured backend's dense kernel where available, torch SDPA +otherwise -- logs the specific reason once, and counts the fallback. Set `SOL_ATTN_STRICT=1` to raise instead of falling back, which is useful when benchmarking to confirm the kernel actually ran. diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py index ecda0bef4959..048eb88fb24e 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py @@ -112,8 +112,8 @@ def _cute_dense_available() -> bool: return True -def _parse_dense_layers(spec: Optional[str]) -> frozenset: - layers: set = set() +def _parse_dense_layers(spec: Optional[str]) -> frozenset[int]: + layers: set[int] = set() for item in str(spec or "").split(","): item = item.strip() if not item: @@ -154,11 +154,15 @@ def __init__( self.num_heads = num_heads self.head_dim = head_dim self.num_kv_heads = num_kv_heads or num_heads - assert self.num_kv_heads == self.num_heads, ( - f"Sol-Attn is MHA-only (num_kv_heads == num_heads), got " - f"num_kv_heads={self.num_kv_heads}, num_heads={self.num_heads}. " - f"GQA/MQA is not supported." - ) + if self.num_kv_heads != self.num_heads: + # Not an assert: `python -O` strips those, and the kernel wrapper + # would then see unequal Q/K shapes and quietly take its dense + # fallback instead of rejecting an unsupported configuration. + raise ValueError( + f"Sol-Attn is MHA-only (num_kv_heads == num_heads), got " + f"num_kv_heads={self.num_kv_heads}, num_heads={self.num_heads}. " + f"GQA/MQA is not supported." + ) self.dtype = dtype cfg = sparse_attention_config self.tau = getattr(cfg, "tau", 1.0) @@ -203,7 +207,7 @@ def __init__( # `_get_vsa_inputs` do). Returns a host-side bool, so the dense and sparse # phases still compile as separate graphs -- they run different kernels. @torch.compiler.disable - def _dense_by_step(self, timestep) -> bool: + def _dense_by_step(self, timestep: Any) -> bool: phase = sol_attn_graph_phase( timestep, disabled_until_timestep=self.disabled_until_timestep, @@ -231,7 +235,9 @@ def _sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) ).transpose(1, 2) - def _delegate(self, q, k, v, **kwargs) -> torch.Tensor: + def _delegate( + self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs: Any + ) -> torch.Tensor: """Hand the call to the dense backend of the same family. ``_cute_dense_ok`` answers "can this *device* run the kernel", decided at @@ -243,7 +249,7 @@ def _delegate(self, q, k, v, **kwargs) -> torch.Tensor: return self._inner.forward(q, k, v, **kwargs) return self._sdpa(q, k, v) - def _can_serve(self, q: torch.Tensor, k: torch.Tensor, **kwargs) -> bool: + def _can_serve(self, q: torch.Tensor, k: torch.Tensor, **kwargs: Any) -> bool: """Whether the sparse kernel applies to this particular call. Everything false here is delegated to ``_inner``. Deciding it from the diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py index 585a2cd3fed7..1f38ad39d356 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py @@ -1,3 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. """Shape/dtype guard and dense-fallback wrapper around the Sol-Attn kernel. Adapted from upstream's ``techniques/sparse_backends/sol_attn_backend.py`` at diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 3e9655f5f407..a04b8dd979e9 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -19,6 +19,7 @@ from typing import Any, Dict, Literal, Optional from pydantic import Field as PydanticField +from pydantic import field_validator from tensorrt_llm.llmapi.utils import StrictBaseModel @@ -280,6 +281,45 @@ class SolAttentionConfig(BaseSparseAttentionConfig): ), ) + @field_validator("dense_layers") + @classmethod + def _validate_dense_layers(cls, spec: Optional[str]) -> Optional[str]: + """Reject malformed specs here rather than deep in the backend. + + Without this a non-numeric token raises from ``_parse_dense_layers`` + during attention construction, far from the config that caused it, and + a descending range such as ``'4-2'`` raises nothing at all -- it yields + an empty set, so the layers the user asked to force dense silently stay + sparse. + """ + if spec is None: + return spec + for item in spec.split(","): + item = item.strip() + if not item: + continue + try: + if "-" in item: + # A negative index cannot reach here: it also contains '-', + # so it takes this branch and the empty first part fails + # int() below. + start, end = (int(part) for part in item.split("-", 1)) + if start > end: + raise ValueError( + f"dense_layers range '{item}' is descending; " + f"write it as '{end}-{start}'" + ) + else: + int(item) + except ValueError as exc: + if "descending" in str(exc): + raise + raise ValueError( + f"dense_layers entry '{item}' is not a layer index or range; " + "expected a comma-separated list such as '0,2-4'" + ) from exc + return spec + def to_sparse_params(self, **kwargs): # Sol-Attn's knobs are consumed directly by SolAttention.__init__ # (constructed via CUTEDSL backend dispatch in create_attention), not diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py index 32eb20da4c50..43799383c56a 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py @@ -142,10 +142,11 @@ def test_sol_attn_self_attention_is_served_under_separate_qkv(): Qwen-Image uses it unconditionally. Both are self-attention; both must still get the sparse kernel. """ - device = torch.device("cuda") attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) attn.disabled_until_timestep = None - q = k = torch.randn(1, 64, 2, 128, device=device, dtype=torch.bfloat16) + # CPU tensors on purpose: `_can_serve` compares shapes and the layer index + # and never touches the device, so this runs on CPU-only hosts too. + q = k = torch.randn(1, 64, 2, 128, dtype=torch.bfloat16) assert attn._can_serve(q, k), "equal q/k sequence lengths must reach the sparse kernel" @@ -182,7 +183,7 @@ def test_sol_attn_with_context_parallelism_raises(): def test_sol_attn_rejects_gqa_mqa(): """Sol-Attn is MHA-only; num_kv_heads != num_heads must fail fast at construction.""" - with pytest.raises(AssertionError, match="MHA-only"): + with pytest.raises(ValueError, match="MHA-only"): SolAttention(layer_idx=0, num_heads=8, head_dim=128, num_kv_heads=2) @@ -436,6 +437,32 @@ def test_zero_cutoff_rejected(): assert SolAttentionConfig(tau=2.0).disabled_until_timestep is None +@pytest.mark.parametrize( + "spec,reason", + [ + ("4-2", "descending"), + ("abc", "not a layer index"), + ("0,x", "not a layer index"), + ("-1", "not a layer index"), + ], + ids=["descending_range", "non_numeric", "non_numeric_in_list", "negative"], +) +def test_dense_layers_rejects_malformed_spec(spec, reason): + """Malformed dense_layers must fail at config time, not silently or late. + + A descending range is the dangerous one: `_parse_dense_layers('4-2')` + yields an empty set, so the layers the user asked to force dense would + quietly stay sparse with no error anywhere. + """ + with pytest.raises(ValueError, match=reason): + SolAttentionConfig(tau=2.0, dense_layers=spec) + + +@pytest.mark.parametrize("spec", [None, "0", "0,2-4", " 0 , 2 ", "0-0"]) +def test_dense_layers_accepts_valid_spec(spec): + assert SolAttentionConfig(tau=2.0, dense_layers=spec).dense_layers == spec + + @pytest.mark.skip( reason=( "TODO(sol-attn): numerical equivalence vs dense SDPA at zero/near-zero routing " From c7847394e842b54b7b704ff16d702c1cb35284fc Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:28:01 -0700 Subject: [PATCH 10/22] [TRTLLM-15917][feat] Support SM103 (B300/GB300) for Sol-Attn 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) Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../features/visualgen-sparse-attention.md | 7 ++-- .../attention_backend/cute_dsl/sol_attn.py | 10 +++--- .../blackwell/sol_attn/THIRD_PARTY_NOTICES.md | 26 +++++++++++--- .../blackwell/sol_attn/interface.py | 13 +++++-- .../blackwell/sol_attn_backend.py | 2 +- tensorrt_llm/visual_gen/sparse_attention.py | 5 +-- .../test_lists/test-db/l0_b300.yml | 2 ++ .../test_attention_cute_dsl_sol_attn.py | 36 +++++++++++++++++++ 8 files changed, 83 insertions(+), 18 deletions(-) diff --git a/docs/source/features/visualgen-sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md index bef9877ca268..ce1d5034bd19 100644 --- a/docs/source/features/visualgen-sparse-attention.md +++ b/docs/source/features/visualgen-sparse-attention.md @@ -22,14 +22,15 @@ Sparse attention in VisualGen is configured through `VisualGenArgs.attention_con |---|---|---| | `skip_softmax` | `SkipSoftmaxAttentionConfig` | Supported | | VSA | TBD | TODO | -| `sol_attn` | `SolAttentionConfig` | Supported (CUTEDSL, sm100) | +| `sol_attn` | `SolAttentionConfig` | Supported (CUTEDSL, sm100/sm103) | ### Sol-Attn Sol-Attn ([arXiv:2607.24027](https://arxiv.org/abs/2607.24027)) folds dynamic block routing, sparse computation, and an approximation-correction term into one -online-softmax pass. It runs on the **CUTEDSL** backend only, on sm100 -(B200/GB200), and requires `head_dim=128`, bfloat16, and MHA +online-softmax pass. It runs on the **CUTEDSL** backend only, on datacenter +Blackwell -- sm100 (B200/GB200) and sm103 (B300/GB300) -- and requires +`head_dim=128`, bfloat16, and MHA (`num_kv_heads == num_heads`). ```yaml diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py index 048eb88fb24e..ff8bb422a311 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py @@ -97,10 +97,10 @@ def sol_attn_graph_phase( def _cute_dense_available() -> bool: """Whether `cute_dsl_fmha_fwd` can run on the current device. - Checked once at construction. Sol-Attn is sm100-only and the dense CuTe DSL - kernel covers sm_100a/sm_103a, so in practice this is always true wherever - Sol-Attn runs; the negative branch exists so an unsupported device degrades - to SDPA instead of raising. + Checked once at construction. Sol-Attn and the dense CuTe DSL kernel now + cover the same set (sm_100a/sm_103a), so in practice this is always true + wherever Sol-Attn runs; the negative branch exists so an unsupported device + degrades to SDPA instead of raising. """ try: from .fmha import _check_cute_runtime_available, _get_gpu_arch @@ -127,7 +127,7 @@ def _parse_dense_layers(spec: Optional[str]) -> frozenset[int]: class SolAttention(AttentionBackend): - """Sol-Attn dynamic block-routing sparse attention (CuTeDSL, sm100). + """Sol-Attn dynamic block-routing sparse attention (CuTeDSL, sm100/sm103). The kernel wrapper already falls back to dense attention on any unsupported shape/dtype/arch (see ``_run_sol_attn_bthd``); this class only adds the diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md index 2b0ba8011fb4..326b862bd71d 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md @@ -95,7 +95,25 @@ remain subject to their respective licenses. SM120 (RTX Blackwell) was carried in an earlier revision of this port and has been dropped: it had kernel-level evidence only, no end-to-end validation, and no `cute_dsl_fmha_fwd` exists for that architecture, so Sol-Attn's dense -fallback could not match its own backend there. With SM100 alone, Sol-Attn's -architecture set is a subset of the dense CuTe DSL FMHA kernel's. The -cuDNN-frontend attribution that covered the SM120 execution skeleton was -removed with it. +fallback could not match its own backend there. The cuDNN-frontend attribution +that covered the SM120 execution skeleton was removed with it. + +SM103 (B300/GB300) is served by the same vendored `sm100/` kernel. Upstream +ships it as SM100-only, so treating one kernel as covering both steppings is a +divergence and is recorded here. Two things make it safe. The CuTe DSL JIT +targets whatever device it compiles on -- `cute.compile()` takes no +architecture argument -- 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 one class for the same reason, differing only by +an exp2-emulation flag that Sol-Attn does not use. So Sol-Attn's architecture +set remains a subset of the dense CuTe DSL FMHA kernel's, which is the +invariant SM120 failed. + +Evidence: the vendored package was run directly (not through the TensorRT-LLM +wrapper, whose ineligibility path would mask a failure as a dense fallback) on +a B300 SXM6 AC, and on a B200 as a control, with identical seed, shape and +pinned toolchain. Both produced finite output with matching error against a +dense SDPA reference at tau 0.0/1.0/2.0: cosine 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 difference appears only in the maximum +element and is consistent with reduction order across a different SM count. diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/interface.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/interface.py index d9644c1fd8bb..e275c5b04962 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/interface.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/interface.py @@ -12,8 +12,13 @@ import torch BLOCK_SIZE = 64 +# The vendored kernel lives in ``sm100/`` and serves both datacenter Blackwell +# steppings: the CuTe DSL JIT targets whatever device it compiles on, and the +# kernel body uses no sm100-exclusive construct. Measured identical on both -- +# see THIRD_PARTY_NOTICES.md. _CUTE_BACKENDS = { (10, 0): "cute_sm100", # B200 / GB200 + (10, 3): "cute_sm100", # B300 / GB300 (Blackwell Ultra) } _compiled = {} @@ -105,8 +110,8 @@ def get_sol_attn_backend(device: torch.device | str | int | None = None) -> str: def _validate_cute(arch, tokens, kv_splits): if kv_splits != 1: raise ValueError( - "kv_splits=2/4 was an SM90-only path; this build ships SM100 " - "kernels only, so kv_splits must be 1." + "kv_splits=2/4 was an SM90-only path; this build ships the SM100 " + "kernel only, so kv_splits must be 1." ) route_groups = ((tokens + 63) // 64 + 63) // 64 if kv_splits > route_groups: @@ -197,10 +202,12 @@ def _sol_attn_cute( stream = _stream(q.device) key = (q.device.index, arch, batch, tokens, heads, kv_splits) - if arch != (10, 0): + if arch not in _CUTE_BACKENDS: # Unreachable via sol_attn(): _backend_for_arch raises first. Kept # explicit because the alternative on a missed guard is returning # the uninitialised `output` buffer, i.e. silently wrong results. + # Keyed off _CUTE_BACKENDS rather than a literal so widening the + # dispatch map cannot leave this guard behind. raise ValueError(f"no Sol-Attn CuTe kernel for SM{arch[0]}{arch[1]}") sink_start_block, sink_end_block = _sink_block_range( tokens, diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py index 1f38ad39d356..2765d7f94494 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py @@ -68,7 +68,7 @@ def _load_sol_attn() -> Callable: # Architectures with a Sol-Attn CuTe kernel. Kept in sync with # ``sol_attn/interface.py::_CUTE_BACKENDS``; duplicated here so the eligibility # check does not have to import the CuTe DSL. -SUPPORTED_ARCHS = frozenset({(10, 0)}) +SUPPORTED_ARCHS = frozenset({(10, 0), (10, 3)}) def sol_attn_ineligible_reason(q) -> Optional[str]: diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index a04b8dd979e9..7df0d06b01f5 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -226,8 +226,9 @@ class SolAttentionConfig(BaseSparseAttentionConfig): """Sol-Attn sparse attention configuration for visual generation. Dynamic block routing + sparse computation + approximation correction in - one online-softmax pass (arXiv:2607.24027). Kernel is CuTeDSL, sm100 - (B200/GB200) only, head_dim=128, bf16, MHA. + one online-softmax pass (arXiv:2607.24027). Kernel is CuTeDSL on + datacenter Blackwell -- sm100 (B200/GB200) and sm103 (B300/GB300) -- + head_dim=128, bf16, MHA. On an unsupported *shape, dtype, or architecture* the kernel falls back to dense attention -- the configured backend's dense kernel where available, diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 4f67742c28a1..b8135b6efd0a 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -34,6 +34,8 @@ l0_b300: - unittest/_torch/thop/parallel TIMEOUT (90) - unittest/_torch/visual_gen/kernels/parallel - unittest/_torch/visual_gen/test_attention_flashinfer.py + # ------------- Visual Gen sparse attention (sm103) --------------- + - unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py - unittest/_torch/thop/serial - unittest/_torch/executor # 250s - unittest/_torch/disaggregation diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py index 43799383c56a..df998b44efdd 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py @@ -416,6 +416,42 @@ def test_supported_archs_matches_kernel_dispatch_map(): assert _backend_mod().SUPPORTED_ARCHS == frozenset(interface._CUTE_BACKENDS) +def test_datacenter_blackwell_archs_are_supported(): + """Both datacenter Blackwell steppings dispatch to the vendored kernel. + + sm103 (B300/GB300) runs the same `sm100/` kernel as sm100: the CuTe DSL JIT + targets the device it compiles on and the kernel body uses no sm100-only + construct. Verified on hardware against a B200 control; see the kernel's + THIRD_PARTY_NOTICES.md. + """ + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.sol_attn import interface + + assert (10, 0) in interface._CUTE_BACKENDS + assert (10, 3) in interface._CUTE_BACKENDS + + +def test_no_arch_literal_outside_the_dispatch_map(): + """`_sol_attn_cute`'s guard must key off _CUTE_BACKENDS, not a literal. + + That guard is a second, deeper check than `_backend_for_arch`, and it is + the one no other test reaches: `test_supported_archs_matches_kernel_dispatch_map` + keeps SUPPORTED_ARCHS and _CUTE_BACKENDS in step, so widening both leaves a + hardcoded literal here as the only thing still rejecting the new arch -- + with every test green. Assert on the source so the coupling cannot regress. + """ + import inspect + + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.sol_attn import interface + + src = inspect.getsource(interface._sol_attn_cute) + assert "arch not in _CUTE_BACKENDS" in src, ( + "the guard in _sol_attn_cute must be keyed off _CUTE_BACKENDS" + ) + assert "arch != (10, 0)" not in src, ( + "hardcoded architecture literal in _sol_attn_cute; key it off _CUTE_BACKENDS" + ) + + def test_quant_attention_config_rejected_with_sol_attn(): """Sol-Attn replaces the dense CuTeDSL path, so quantized attention cannot compose with it; accepting the pair would silently ignore the quant request.""" From da5a5a4b7e737e2cf07a7ef5ab2aef316444fd0b Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Tue, 8 Sep 2026 23:35:54 -0700 Subject: [PATCH 11/22] [TRTLLM-15917][fix] Address review feedback on Sol-Attn 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) Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../features/visualgen-sparse-attention.md | 4 +- .../attention_backend/cute_dsl/sol_attn.py | 12 +- .../blackwell/sol_attn/THIRD_PARTY_NOTICES.md | 112 ++++-------------- .../blackwell/sol_attn_backend.py | 75 ++++++++++-- .../_torch/visual_gen/modules/attention.py | 22 +--- tensorrt_llm/visual_gen/sparse_attention.py | 7 +- .../test_attention_cute_dsl_sol_attn.py | 97 ++++++++++++--- 7 files changed, 188 insertions(+), 141 deletions(-) diff --git a/docs/source/features/visualgen-sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md index ce1d5034bd19..7e95663dfc76 100644 --- a/docs/source/features/visualgen-sparse-attention.md +++ b/docs/source/features/visualgen-sparse-attention.md @@ -21,7 +21,7 @@ Sparse attention in VisualGen is configured through `VisualGenArgs.attention_con | `algorithm` | Config class | Status | |---|---|---| | `skip_softmax` | `SkipSoftmaxAttentionConfig` | Supported | -| VSA | TBD | TODO | +| `vsa` | `VideoSparseAttentionConfig` | Supported (CUTEDSL) | | `sol_attn` | `SolAttentionConfig` | Supported (CUTEDSL, sm100/sm103) | ### Sol-Attn @@ -126,7 +126,7 @@ User configuration is supplied through Python or YAML and controls how the check `threshold_scale_factor` and `target_sparsity` are alternatives: if both are present, `threshold_scale_factor` takes precedence and the calibration formula is not used. User-provided `target_sparsity` and `disabled_until_timestep` override checkpoint defaults. Checkpoint `ignore` patterns always disable Skip Softmax Attention for matching layers. -Skip Softmax Attention works with both the **TRTLLM** and **CUTEDSL** attention backends in VisualGen. Set `attention_config.backend` to either when enabling it. On CUTEDSL, Skip Softmax Attention can also be combined with `quant_attention_config`'s block-scaled Q/K recipes (MXFP8, NVFP4); VSA and Sol-Attn each replace the dense CuTeDSL path and are therefore mutually exclusive with quantized attention. +Skip Softmax Attention works with both the **TRTLLM** and **CUTEDSL** attention backends in VisualGen. Set `attention_config.backend` to either when enabling it. On CUTEDSL, Skip Softmax Attention can also be combined with `quant_attention_config`'s block-scaled Q/K recipes (MXFP8, NVFP4); VSA and Sol-Attn are currently the two supported sparse-attention algorithms, both available only through the CuTeDSL attention backend; quantized attention is not yet supported or enabled with either mode. #### Mapping `disabled_until_timestep` to Actual Denoising Steps diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py index ff8bb422a311..5a37447076e0 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py @@ -253,11 +253,13 @@ def _can_serve(self, q: torch.Tensor, k: torch.Tensor, **kwargs: Any) -> bool: """Whether the sparse kernel applies to this particular call. Everything false here is delegated to ``_inner``. Deciding it from the - tensors, per call, is deliberate: ``qkv_mode`` describes how Q/K/V are - *projected*, not whether K/V come from another sequence, so a - construction-time rule keyed on ``SEPARATE_QKV`` mistakes self-attention - for cross-attention wherever that mode is chosen for other reasons -- - Qwen-Image always, and WAN's ``attn1`` under async Ulysses. + tensors, per call, is deliberate, and it is why ``modules/attention.py`` + has no ``SEPARATE_QKV`` rule for Sol-Attn: ``qkv_mode`` describes how + Q/K/V are *projected*, not whether K/V come from another sequence, so a + construction-time rule keyed on it mistakes self-attention for + cross-attention wherever that mode is chosen for other reasons -- + Qwen-Image always, and WAN's ``attn1`` under async Ulysses -- silently + costing those modules their configured backend. """ # Cross-attention: K/V come from another sequence. Sol-Attn's routing # assumes one self-attending sequence. diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md index 326b862bd71d..b226f77f96bc 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md @@ -3,23 +3,10 @@ This package is vendored from [`github.com/NVlabs/Sana`](https://github.com/NVlabs/Sana), branch [`sol-engine`](https://github.com/NVlabs/Sana/tree/sol-engine), at commit -[`5fe5feb`](https://github.com/NVlabs/Sana/commit/5fe5feb) (2026-08-17; -best-effort reconstruction from vendoring-date file timestamps and upstream -commit history, not an exact recorded pin from the original port -- see the -pull request for how this was verified). Checked against the current branch -tip ([`83e54df`](https://github.com/NVlabs/Sana/commit/83e54df), 2026-08-20) -on 2026-08-27; the only upstream change since that touches this subset is a -merge whose SM89 work is out of scope here. The rest of that merge (MPS/Metal -Apple Silicon backend, RTX 4090/5090 configs) is likewise out of scope for -this CUDA/Blackwell-only subset. - -**Note for future currency checks.** These files are linted and formatted to -this repository's style (`ruff check` and `ruff format`, line length 100) -rather than kept byte-identical to upstream, so a direct `diff` against -upstream shows formatting noise as well as real changes. Upstream wraps at -roughly 80 columns; most of the difference is expressions joined onto one -line. To compare semantics, run `ruff format` over the upstream copy first and -diff the normalized results -- that is how the currency check above was done. +[`5fe5feb`](https://github.com/NVlabs/Sana/commit/5fe5feb) (2026-08-17). +Checked against the branch tip +([`83e54df`](https://github.com/NVlabs/Sana/commit/83e54df), 2026-08-20) on +2026-08-27. ## Scope of the vendored subset @@ -28,92 +15,35 @@ Only the pieces needed for the architectures TensorRT-LLM ships are carried: | Carried | Not carried | |---|---| | `interface.py`, `preprocess.py`, `common/` | `sm89/`, `sm90/` (incl. `sm90/_compat/`) | -| `sm100/` (B200 / GB200) | `triton_ref/` Triton reference attention | +| `sm100/` — serves SM100 (B200/GB200) and SM103 (B300/GB300) | `triton_ref/` Triton reference attention | | | `sm120/` (RTX Blackwell) | | | `_vendor/flash_attn/` (see below) | +`../sol_attn_backend.py` is not part of the vendored package but is a +derivative work of upstream's +`techniques/sparse_backends/sol_attn_backend.py` (same branch and commit). +Only the kernel-wrapper subset is carried; upstream's model-integration half +(diffusers dispatch hook, HunyuanVideo MMDiT padding, model-level Morton +ordering) is not. + +Implementation divergences from upstream are documented in the source itself, +in the module docstrings of `../sol_attn_backend.py` and +`attention_backend/cute_dsl/sol_attn.py`. + +## Licensing + The upstream package vendored a copy of FlashAttention's CuTe DSL helpers under `sol_attn/_vendor/flash_attn/cute/`. That copy is **not** carried here: TensorRT-LLM already depends on [`flash-attn-4`](https://github.com/Dao-AILab/flash-attention) (pinned in `requirements.txt`), which provides the same `flash_attn.cute` modules, and -the SM100 kernels import them from that dependency directly. This was -verified on B200 to produce bit-identical output to the vendored copy across a -shape/tau sweep. FlashAttention's BSD-3-Clause license is retained at -`sol_attn/sm100/LICENSE.flash-attention` because portions of the SM100 design -scaffold still derive from that project. +the SM100 kernels import them from that dependency directly. FlashAttention's +BSD-3-Clause license is retained at `sol_attn/sm100/LICENSE.flash-attention` +because portions of the SM100 design scaffold still derive from that project. `preprocess.py` implements the routing/threshold stage in Triton, so Triton is a required runtime dependency on every Sol-Attn path, not only a fallback. -## A derived file outside this directory - -`../sol_attn_backend.py` is **not** part of the vendored package above, but it -is a derivative work and is recorded here because this file is where a future -currency check starts. It is adapted from upstream's -`techniques/sparse_backends/sol_attn_backend.py` (same branch and commit as the -package). Only the kernel-wrapper subset is carried -- the shape/dtype guard, -the dense fallback (routed to cute_dsl_fmha_fwd here, not torch SDPA), and the -call counters. Upstream's model-integration -half is not carried: the diffusers self-attention dispatch hook, HunyuanVideo's -padded `[video, text]` MMDiT handling, and model-level Morton ordering. - -Deliberate divergences from upstream in that file, all of which a re-sync must -preserve rather than overwrite: - -| Divergence | Why | -|---|---| -| `logger.warning_once` replaces `print()` | fallbacks must be suppressible and routed through the repo's logger | -| `dense_fallback_calls` counter added | makes a silently-degraded run countable, not just visible in stderr | -| `sol_attn_ineligible_reason()` added | names the specific reason (arch / head_dim / dtype) instead of one boolean | -| `SOL_ATTN_STRICT=1` also covers the eligibility path | upstream raises only on kernel exceptions, so an ineligible run stayed silent | -| `@torch.compiler.disable` on `_run_sol_attn_bthd` | see below | -| dense paths routed to `cute_dsl_fmha_fwd` via `dense_fn` | upstream's dense fallback is torch SDPA; staying in-backend is what makes a `backend: CUTEDSL` A/B isolate sparsity | - -**Upstream solves the `torch.compile` problem differently, and arguably -better.** Its `sol_attn_backend.py` wraps the same call in a -`torch.library.custom_op` (`sana_sol_attn::self_attention`) with a -`register_fake` returning `torch.empty_like(q)`, which keeps the kernel in the -compiled graph as an opaque node instead of breaking the graph at it; a second -consumer (`models/ltx2.5-refiner/GB200/sol_attention.py`) applies -`torch.compiler.disable` at the call site behind a flag. This repository uses -`@torch.compiler.disable` on the launch boundary instead, matching the -convention every other CuTe DSL entry point here already follows -(`attention_backend/cute_dsl/fmha.py`, -`cute_dsl_kernels/blackwell/video_sparse_attention/interface.py`). Without some -such guard Dynamo traces into the CuTe DSL JIT builder and retraces on every -call -- measured at near two orders of magnitude slower on B200. Migrating to -the `custom_op` form would -remove the per-layer graph break and is a reasonable follow-up; it was not done -here because the `torch.compiler.disable` form is what this repository's other -kernels use and what the measurements above were taken with. - The runtime also depends on NVIDIA CUTLASS / CuTe DSL, cuda-python, and PyTorch. Those dependencies are not redistributed by this repository and remain subject to their respective licenses. - -SM120 (RTX Blackwell) was carried in an earlier revision of this port and has -been dropped: it had kernel-level evidence only, no end-to-end validation, and -no `cute_dsl_fmha_fwd` exists for that architecture, so Sol-Attn's dense -fallback could not match its own backend there. The cuDNN-frontend attribution -that covered the SM120 execution skeleton was removed with it. - -SM103 (B300/GB300) is served by the same vendored `sm100/` kernel. Upstream -ships it as SM100-only, so treating one kernel as covering both steppings is a -divergence and is recorded here. Two things make it safe. The CuTe DSL JIT -targets whatever device it compiles on -- `cute.compile()` takes no -architecture argument -- 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 one class for the same reason, differing only by -an exp2-emulation flag that Sol-Attn does not use. So Sol-Attn's architecture -set remains a subset of the dense CuTe DSL FMHA kernel's, which is the -invariant SM120 failed. - -Evidence: the vendored package was run directly (not through the TensorRT-LLM -wrapper, whose ineligibility path would mask a failure as a dense fallback) on -a B300 SXM6 AC, and on a B200 as a control, with identical seed, shape and -pinned toolchain. Both produced finite output with matching error against a -dense SDPA reference at tau 0.0/1.0/2.0: cosine 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 difference appears only in the maximum -element and is consistent with reduction order across a different SM count. diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py index 2765d7f94494..da19ff488b7e 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py @@ -15,9 +15,27 @@ """Shape/dtype guard and dense-fallback wrapper around the Sol-Attn kernel. Adapted from upstream's ``techniques/sparse_backends/sol_attn_backend.py`` at -the pin in ``sol_attn/THIRD_PARTY_NOTICES.md``, which records exactly which -subset is carried and how this version deliberately diverges. Check that file -before re-syncing against upstream. +the pin in ``sol_attn/THIRD_PARTY_NOTICES.md``, which records which subset is +carried. Check that file before re-syncing against upstream. + +Deliberate divergences from upstream, which a re-sync must preserve rather +than overwrite: + +* ``logger.warning_once`` replaces ``print()``, so fallbacks are suppressible + and routed through this repository's logger. +* ``_SOL_STATS["dense_fallback_calls"]`` makes a silently-degraded run + countable rather than only visible on stderr. +* ``sol_attn_ineligible_reason()`` names the specific reason (architecture, + head_dim, dtype) instead of returning one boolean. +* ``SOL_ATTN_STRICT=1`` also covers the eligibility path; upstream raises only + on kernel exceptions, so an ineligible run stayed silent. +* Dense paths route to ``cute_dsl_fmha_fwd`` through ``dense_fn``. Upstream + falls back to torch SDPA; staying inside the configured backend is what lets + a ``backend: CUTEDSL`` A/B isolate sparsity rather than also swapping the + dense kernel. +* ``@torch.compiler.disable`` guards the launch boundary; see the comment on + ``_run_sol_attn_bthd`` for why, and for the ``torch.library.custom_op`` + alternative upstream uses. The kernel-facing API accepts contiguous BF16 ``[batch, tokens, heads, 128]`` Q/K/V, ``tau``, ``thresh_type``, @@ -145,11 +163,49 @@ def _dense_bthd(q, k, v): ).transpose(1, 2) +def _degradable_kernel_errors() -> tuple[type[BaseException], ...]: + """Exception types whose failure is safe to answer with dense attention. + + `CODING_GUIDELINES.md` asks for the smallest exception set. A bare + ``except Exception`` also swallowed ordinary programming and integration + errors -- ``TypeError``, ``NameError``, ``AttributeError`` -- turning a bug + into "Sol-Attn just didn't speed anything up". Those now propagate. + + ``DSLBaseError`` has to be named explicitly because CuTe DSL derives it + from ``Exception``, not ``RuntimeError``, so a narrower tuple would stop + catching real JIT and codegen failures that were previously degraded. It + is resolved lazily and cached: this module defers every CuTe DSL import to + first use, and the tuple is only needed once something has already raised. + """ + global _DEGRADABLE_KERNEL_ERRORS + if _DEGRADABLE_KERNEL_ERRORS is None: + types: list[type[BaseException]] = [ImportError, OSError, RuntimeError] + try: + from cutlass.base_dsl.common import DSLBaseError + except ImportError: + pass + else: + types.append(DSLBaseError) + _DEGRADABLE_KERNEL_ERRORS = tuple(types) + return _DEGRADABLE_KERNEL_ERRORS + + +_DEGRADABLE_KERNEL_ERRORS: tuple[type[BaseException], ...] | None = None + + # Opaque to Dynamo, like every other CuTe DSL launch boundary here (see # cute_dsl/fmha.py, video_sparse_attention/interface.py). Otherwise Dynamo # traces into the CuTe DSL JIT builder and retraces on every call: near two # orders of magnitude slower on B200 (2496.9 s mean denoise without it), and # silently, as if compile just didn't help. +# +# Upstream instead wraps the same call in a `torch.library.custom_op` with a +# `register_fake`, which keeps the kernel in the compiled graph as an opaque +# node rather than breaking the graph at it. That is the better end state -- +# it removes the per-layer graph break -- and is tracked as a follow-up. It is +# not done here because `torch.compiler.disable` is the convention every other +# CuTe DSL entry point in this repository already follows, and is what the +# reported measurements were taken with. @torch.compiler.disable def _run_sol_attn_bthd( q, @@ -191,6 +247,10 @@ def dense(): ) return dense() + # Resolved outside the try: a bad kv_splits is a configuration error and + # must surface, not silently become a dense run. + resolved_kv_splits = _resolve_kv_splits(q0, kv_splits) + try: kernel = _load_sol_attn() out = kernel( @@ -199,13 +259,11 @@ def dense(): v0, tau=float(tau), thresh_type=str(thresh_type), - kv_splits=_resolve_kv_splits(q0, kv_splits), + kv_splits=resolved_kv_splits, sink_start=sink_start, sink_tokens=int(sink_tokens), ) - _SOL_STATS["kernel_calls"] += 1 - return out - except Exception as exc: + except _degradable_kernel_errors() as exc: if _strict(): raise logger.warning_once( @@ -216,6 +274,9 @@ def dense(): ) return dense() + _SOL_STATS["kernel_calls"] += 1 + return out + # Lightweight run-validation counters. `kernel_calls` is the census used to # prove the CuTe kernel actually ran: because forward() falls back to dense diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index ec460f1e798d..2e995898c93c 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -106,14 +106,7 @@ def __init__( # SEPARATE_QKV fallback: TRTLLM and CUTEDSL VSA cannot serve it. - # - # Sol-Attn is deliberately absent: it decides per call (`_can_serve`) and - # delegates what it cannot serve to the dense backend of its own family. - # A rule here would have to guess from `qkv_mode`, which describes how - # Q/K/V are *projected* rather than whether K/V come from another - # sequence -- and that guess is wrong wherever SEPARATE_QKV is chosen for - # other reasons (Qwen-Image always; WAN's attn1 under async Ulysses), - # silently costing those modules their configured backend. + # Sol-Attn is absent by design; see SolAttention._can_serve. if separate_qkv_cross_attention and (base_backend == "TRTLLM" or _is_vsa): backend_name = "VANILLA" requested = f"{base_backend} (VSA)" if _is_vsa else base_backend @@ -127,15 +120,12 @@ def __init__( else: backend_name = base_backend - if _is_vsa and cp_size > 1: + # Every sparse algorithm here routes over the whole token sequence, so + # none of them can be split across context-parallel ranks. + if (_is_vsa or _is_sol_attn) and cp_size > 1: + _algo_name = "VSA" if _is_vsa else "Sol-Attn" raise ValueError( - f"VSA needs the full token sequence per rank, so it is incompatible " - f"with context parallelism (Attention2D/Ring, cp_size={cp_size}). Use " - f"ulysses or cfg parallelism instead." - ) - if _is_sol_attn and cp_size > 1: - raise ValueError( - f"Sol-Attn needs the full token sequence per rank, so it is incompatible " + f"{_algo_name} needs the full token sequence per rank, so it is incompatible " f"with context parallelism (Attention2D/Ring, cp_size={cp_size}). Use " f"ulysses or cfg parallelism instead." ) diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 7df0d06b01f5..370d26520c65 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -298,7 +298,12 @@ def _validate_dense_layers(cls, spec: Optional[str]) -> Optional[str]: for item in spec.split(","): item = item.strip() if not item: - continue + # Not skipped: "," / "0,,2" / " " would otherwise be accepted + # and quietly force fewer layers dense than the user wrote. + raise ValueError( + f"dense_layers {spec!r} has an empty entry; expected a " + "comma-separated list such as '0,2-4'" + ) try: if "-" in item: # A negative index cannot reach here: it also contains '-', diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py index df998b44efdd..164921582ed4 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py @@ -426,8 +426,12 @@ def test_datacenter_blackwell_archs_are_supported(): """ from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.sol_attn import interface - assert (10, 0) in interface._CUTE_BACKENDS - assert (10, 3) in interface._CUTE_BACKENDS + # Resolve, don't just look up: a malformed value would satisfy a key check + # and then fail at dispatch. + for arch in ((10, 0), (10, 3)): + backend = interface._backend_for_arch(arch, cute_available=True) + assert backend in interface._CUTE_BACKENDS.values() + assert isinstance(backend, str) and backend def test_no_arch_literal_outside_the_dispatch_map(): @@ -439,16 +443,35 @@ def test_no_arch_literal_outside_the_dispatch_map(): hardcoded literal here as the only thing still rejecting the new arch -- with every test green. Assert on the source so the coupling cannot regress. """ + import ast import inspect + import textwrap from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.sol_attn import interface - src = inspect.getsource(interface._sol_attn_cute) + src = textwrap.dedent(inspect.getsource(interface._sol_attn_cute)) assert "arch not in _CUTE_BACKENDS" in src, ( "the guard in _sol_attn_cute must be keyed off _CUTE_BACKENDS" ) - assert "arch != (10, 0)" not in src, ( - "hardcoded architecture literal in _sol_attn_cute; key it off _CUTE_BACKENDS" + + # Reject *any* architecture tuple literal, not just (10, 0): pinning one + # value would let a future `arch != (10, 3)` reintroduce the same trap. + offenders = [] + for node in ast.walk(ast.parse(src)): + if not isinstance(node, ast.Compare): + continue + operands = [node.left, *node.comparators] + names = {n.id for n in operands if isinstance(n, ast.Name)} + if "arch" not in names: + continue + for operand in operands: + if isinstance(operand, ast.Tuple) and all( + isinstance(e, ast.Constant) and isinstance(e.value, int) for e in operand.elts + ): + offenders.append(ast.unparse(node)) + assert not offenders, ( + "hardcoded architecture literal(s) compared against `arch` in " + f"_sol_attn_cute: {offenders}; key the guard off _CUTE_BACKENDS instead" ) @@ -480,8 +503,19 @@ def test_zero_cutoff_rejected(): ("abc", "not a layer index"), ("0,x", "not a layer index"), ("-1", "not a layer index"), + (",", "empty entry"), + ("0,,2", "empty entry"), + (" ", "empty entry"), + ], + ids=[ + "descending_range", + "non_numeric", + "non_numeric_in_list", + "negative", + "only_separator", + "empty_in_list", + "whitespace_only", ], - ids=["descending_range", "non_numeric", "non_numeric_in_list", "negative"], ) def test_dense_layers_rejects_malformed_spec(spec, reason): """Malformed dense_layers must fail at config time, not silently or late. @@ -499,20 +533,45 @@ def test_dense_layers_accepts_valid_spec(spec): assert SolAttentionConfig(tau=2.0, dense_layers=spec).dense_layers == spec -@pytest.mark.skip( - reason=( - "TODO(sol-attn): numerical equivalence vs dense SDPA at zero/near-zero routing " - "(the analogue of VSA's test_cute_kernel_matches_dense_at_full_topk) needs the " - "exact tau/thresh_type combination that guarantees full (non-sparse) block " - "routing, which is not simply tau=0 because Sol-Attn's routing is score-derived " - "rather than a plain top-k like VSA's. Deriving it requires reading " - "cute_dsl_kernels/blackwell/sol_attn/interface.py's routing math and calibrating " - "rtol/atol on real sm100 hardware. This would be a unit-level complement to the " - "end-to-end accuracy evidence recorded in the pull request, not a replacement." +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Sol-Attn needs CUDA") +def test_cute_kernel_matches_dense_on_a_single_block(): + """Kernel-level numerical check against dense attention. + + The obvious form of this test -- the analogue of VSA's + `test_cute_kernel_matches_dense_at_full_topk` -- needs a tau that provably + forces full non-sparse routing. Sol-Attn has no such value: its routing is + score-derived rather than a plain top-k, so tau=0 is not the dense limit. + + A single KV block sidesteps that entirely. With `tokens <= BLOCK_SIZE` + there is exactly one block, so there is nothing to route away and sparsity + is structurally impossible whatever tau says. The kernel must therefore + reproduce dense attention here, which makes this a genuine numerical test + of the kernel rather than of the routing policy. + + `kernel_calls` is asserted to have advanced: without that, a fallback to + dense would make this pass by comparing dense against itself. + """ + sab = _backend_mod() + if not sab.sol_attn_supported(torch.empty(1, 8, 8, 128, device="cuda", dtype=torch.bfloat16)): + pytest.skip("no Sol-Attn kernel for this device") + + tokens = 64 # == BLOCK_SIZE: exactly one KV block + torch.manual_seed(0) + shape = (1, tokens, 4, 128) + q, k, v = (torch.randn(shape, device="cuda", dtype=torch.bfloat16) for _ in range(3)) + + reference = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2).float(), k.transpose(1, 2).float(), v.transpose(1, 2).float() + ).transpose(1, 2) + + before = sab._SOL_STATS["kernel_calls"] + out = sab._run_sol_attn_bthd(q, k, v, tau=2.0, thresh_type="diag", kv_splits=1) + torch.cuda.synchronize() + assert sab._SOL_STATS["kernel_calls"] == before + 1, ( + "the kernel did not run; this comparison would be dense against dense" ) -) -def test_cute_kernel_matches_dense_placeholder(): - pass + + torch.testing.assert_close(out.float(), reference, rtol=2e-2, atol=2e-2) def test_kv_splits_rejects_unsupported_value(): From a2c978f51a61a64a4905b53cf3491ec27a4ba18a Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Wed, 9 Sep 2026 00:33:54 -0700 Subject: [PATCH 12/22] [TRTLLM-15917][fix] Reuse SkipSoftmaxScheduler for the Sol-Attn graph 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) Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../attention_backend/cute_dsl/__init__.py | 3 +- .../attention_backend/cute_dsl/sol_attn.py | 35 ++----------------- .../_torch/visual_gen/models/modeling.py | 3 +- .../test_attention_cute_dsl_sol_attn.py | 31 +++++++++++++--- 4 files changed, 30 insertions(+), 42 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py index f2c11458c1b2..e5dccb71a7ba 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py @@ -21,7 +21,7 @@ """ from .fmha import CuTeDSLAttention, _cute_dsl_import_error -from .sol_attn import SolAttention, sol_attn_graph_phase +from .sol_attn import SolAttention from .vsa import ( VSA_KERNEL_MAX_CUBES, VSA_TILE_SIZE, @@ -45,5 +45,4 @@ "get_vsa_forward_context", "_cute_dsl_import_error", "SolAttention", - "sol_attn_graph_phase", ] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py index 5a37447076e0..c86143780c93 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py @@ -48,6 +48,7 @@ import torch +from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxScheduler from tensorrt_llm.logger import logger from ..interface import AttentionBackend, AttentionTensorLayout @@ -62,38 +63,6 @@ _sol_attn_import_error = e -def _as_float(timestep: Any) -> Optional[float]: - """Coerce a scalar/0-d/1-element timestep to float, else None.""" - if timestep is None: - return None - if isinstance(timestep, torch.Tensor): - if timestep.numel() == 0: - return None - return float(timestep.reshape(-1)[0].item()) - try: - return float(timestep) - except (TypeError, ValueError): - return None - - -def sol_attn_graph_phase( - timestep: Any, *, disabled_until_timestep: Optional[float] -) -> Optional[int]: - """Return 1 once descending timesteps cross the cutoff, 0 before, else None. - - Same contract and sense as - ``SkipSoftmaxScheduler.get_graph_phase_for_timestep``: phase 0 is the dense - prefix, phase 1 the sparse phase, and ``None`` means there is no phase to - distinguish so the CUDA-graph runner omits the key part. - """ - if disabled_until_timestep is None: - return None - value = _as_float(timestep) - if value is None: - return None - return int(value < disabled_until_timestep) - - def _cute_dense_available() -> bool: """Whether `cute_dsl_fmha_fwd` can run on the current device. @@ -208,7 +177,7 @@ def __init__( # phases still compile as separate graphs -- they run different kernels. @torch.compiler.disable def _dense_by_step(self, timestep: Any) -> bool: - phase = sol_attn_graph_phase( + phase = SkipSoftmaxScheduler.get_graph_phase_for_timestep( timestep, disabled_until_timestep=self.disabled_until_timestep, ) diff --git a/tensorrt_llm/_torch/visual_gen/models/modeling.py b/tensorrt_llm/_torch/visual_gen/models/modeling.py index 588a18791f35..cd0f931eb3d6 100644 --- a/tensorrt_llm/_torch/visual_gen/models/modeling.py +++ b/tensorrt_llm/_torch/visual_gen/models/modeling.py @@ -20,7 +20,6 @@ import torch.nn as nn from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxScheduler -from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn import sol_attn_graph_phase from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm.visual_gen.sparse_attention import SkipSoftmaxAttentionConfig, SolAttentionConfig @@ -108,7 +107,7 @@ def register_cuda_graph_extra_key_fns(self, runner: "CUDAGraphRunner") -> None: # the two phases must not share a captured graph. runner.register_extra_key_fn( "sol_attn_phase", - lambda *args, **kwargs: sol_attn_graph_phase( + lambda *args, **kwargs: SkipSoftmaxScheduler.get_graph_phase_for_timestep( kwargs.get("timestep"), disabled_until_timestep=disabled_until_timestep, ), diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py index 164921582ed4..407bcd08b599 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py @@ -27,11 +27,11 @@ import pytest import torch +from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxScheduler from tensorrt_llm._torch.visual_gen.attention_backend import CuTeDSLAttention from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn import ( SolAttention, _parse_dense_layers, - sol_attn_graph_phase, ) from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention from tensorrt_llm._torch.visual_gen.config import ( @@ -220,17 +220,32 @@ def test_graph_phase_matches_skip_softmax_sense(timestep, expected): Same contract as SkipSoftmaxScheduler.get_graph_phase_for_timestep. """ - assert sol_attn_graph_phase(timestep, disabled_until_timestep=0.9545) == expected + assert ( + SkipSoftmaxScheduler.get_graph_phase_for_timestep(timestep, disabled_until_timestep=0.9545) + == expected + ) def test_graph_phase_none_when_prefix_unset(): - assert sol_attn_graph_phase(0.5, disabled_until_timestep=None) is None + assert ( + SkipSoftmaxScheduler.get_graph_phase_for_timestep(0.5, disabled_until_timestep=None) is None + ) def test_graph_phase_accepts_tensor_timestep(): """Pipelines pass a tensor; a 0-d or 1-element tensor must work.""" - assert sol_attn_graph_phase(torch.tensor(0.99), disabled_until_timestep=0.95) == 0 - assert sol_attn_graph_phase(torch.tensor([0.10]), disabled_until_timestep=0.95) == 1 + assert ( + SkipSoftmaxScheduler.get_graph_phase_for_timestep( + torch.tensor(0.99), disabled_until_timestep=0.95 + ) + == 0 + ) + assert ( + SkipSoftmaxScheduler.get_graph_phase_for_timestep( + torch.tensor([0.10]), disabled_until_timestep=0.95 + ) + == 1 + ) def test_dense_prefix_skips_kernel(monkeypatch): @@ -627,6 +642,12 @@ def test_dense_paths_use_cutedsl_backend(monkeypatch): """ import tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn as sol_attn_mod + if not sol_attn_mod._cute_dense_available(): + # A CUDA device is not enough: the premise here is that the dense paths + # reach `cute_dsl_fmha_fwd`, which only exists on sm100/sm103. Without + # this the test fails rather than skips on any other CUDA runner. + pytest.skip("no CuTe DSL dense kernel for this device") + device = torch.device("cuda") q = k = v = torch.randn(1, 64, 2, 128, device=device, dtype=torch.bfloat16) From 32e9ab92d986df79ec015ef311ade1d6e2b3e118 Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Thu, 10 Sep 2026 16:35:27 -0700 Subject: [PATCH 13/22] [TRTLLM-15917][test] Make ineligibility-reason cases reach the branch 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 Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../test_attention_cute_dsl_sol_attn.py | 33 ++++++++++++++----- 1 file changed, 25 insertions(+), 8 deletions(-) diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py index 407bcd08b599..3a918d43f8a8 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py @@ -16,10 +16,11 @@ dense_layers/disabled_until_timestep guards. Mirrors test_attention_cute_dsl_vsa.py's structure and scope for its sibling -sparse-attention algorithm. GPU kernel-vs-dense numerical equivalence (the -analogue of VSA's test_cute_kernel_matches_dense_at_full_topk) is not yet -covered here -- see the TODO on test_cute_kernel_matches_dense_placeholder -below for what it needs and why it's deferred, not just missing. +sparse-attention algorithm. GPU kernel-vs-dense numerical equivalence is +covered by test_cute_kernel_matches_dense_on_a_single_block: Sol-Attn's +routing is score-derived, so no tau provably forces dense routing the way +VSA's full-top-k does, but a single KV block makes sparsity structurally +impossible and gives the same guarantee. """ from types import SimpleNamespace @@ -381,17 +382,33 @@ def _backend_mod(): return sol_attn_backend +def _fake_cuda_q(shape, dtype=torch.bfloat16): + """A tensor-like that reports `is_cuda=True` without needing a GPU. + + `sol_attn_ineligible_reason` checks `is_cuda` first and returns early, so a + CPU tensor can never reach the rank, head_dim, or dtype branches. These + stubs pass that first gate so each later reason is actually exercised; the + architecture check comes after them and is not reached. + """ + return SimpleNamespace(is_cuda=True, ndim=len(shape), shape=shape, dtype=dtype) + + @pytest.mark.parametrize( "make,expect", [ (lambda: torch.randn(1, 4, 2, 128), "not a CUDA tensor"), - (lambda: torch.randn(1, 4, 2, 64), "not a CUDA tensor"), - (lambda: torch.randn(1, 4, 128), "not a CUDA tensor"), + (lambda: _fake_cuda_q((1, 4, 128)), "must be 4-D"), + (lambda: _fake_cuda_q((1, 4, 2, 64)), "head_dim must be 128"), + (lambda: _fake_cuda_q((1, 4, 2, 128), dtype=torch.float16), "dtype must be bfloat16"), ], - ids=["cpu-ok-shape", "cpu-wrong-head-dim", "cpu-wrong-rank"], + ids=["cpu-tensor", "wrong-rank", "wrong-head-dim", "wrong-dtype"], ) def test_ineligible_reason_is_reported(make, expect): - """Ineligibility must name a reason, never fail silently.""" + """Ineligibility must name the specific reason, never fail silently. + + Each case is built to fail exactly one check, in the order the function + applies them, so the id names the branch that actually fires. + """ reason = _backend_mod().sol_attn_ineligible_reason(make()) assert reason is not None and expect in reason assert not _backend_mod().sol_attn_supported(make()) From 3de388fdc5a24be13b471655b22ada50c796e905 Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Thu, 10 Sep 2026 21:55:41 -0700 Subject: [PATCH 14/22] [TRTLLM-15917][fix] Honor masks, survive CUDA-graph capture, wire LTX2 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 Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../backends/sparse/skip_softmax/params.py | 16 +- .../attention_backend/cute_dsl/fmha.py | 8 +- .../attention_backend/cute_dsl/sol_attn.py | 74 +++- .../_torch/visual_gen/cuda_graph_runner.py | 56 ++- .../visual_gen/models/ltx2/pipeline_ltx2.py | 7 +- .../models/ltx2/pipeline_ltx2_two_stages.py | 5 + .../models/ltx2/transformer_ltx2.py | 38 +- tensorrt_llm/visual_gen/sparse_attention.py | 6 +- .../test_attention_cute_dsl_sol_attn.py | 331 ++++++++++++++++++ 9 files changed, 500 insertions(+), 41 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/skip_softmax/params.py b/tensorrt_llm/_torch/attention/backends/sparse/skip_softmax/params.py index 6b688806860e..a54636ea6d65 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/skip_softmax/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/skip_softmax/params.py @@ -382,17 +382,23 @@ def get_runtime_params( *, runtime_params: Optional[SparseRuntimeParams] = None, timestep: Any = None, + graph_phase: Optional[int] = None, ) -> SparseRuntimeParams: - """Return runtime parameters with skip-softmax thresholds.""" + """Return runtime parameters with skip-softmax thresholds. + + ``graph_phase`` lets a caller that already resolved the dense-prefix + phase host-side (the CUDA-graph runner does, to build its key) pass it + in, so this never has to read ``timestep`` -- a ``.item()`` on a CUDA + tensor -- while a graph is being captured. + """ if runtime_params is None: runtime_params = SparseRuntimeParams() - if ( - self.get_graph_phase_for_timestep( + if graph_phase is None: + graph_phase = self.get_graph_phase_for_timestep( timestep, disabled_until_timestep=self.disabled_until_timestep, ) - == 0 - ): + if graph_phase == 0: return replace( runtime_params, threshold_scale_factor_prefill=0.0, diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py index e9638ef9075c..1d7edbbaa948 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py @@ -26,6 +26,7 @@ import torch from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxParams +from tensorrt_llm._torch.visual_gen.cuda_graph_runner import resolved_extra_key from tensorrt_llm.logger import logger from tensorrt_llm.visual_gen.args import QuantAttentionConfig @@ -86,7 +87,12 @@ def _resolve_skip_softmax_threshold_scale_factor( "through.", key="cute_dsl_skip_softmax_missing_timestep", ) - runtime_params = sparse_params.scheduler.get_runtime_params(timestep=timestep) + # Prefer the phase the CUDA-graph runner resolved host-side; the tensor + # read is a `.item()`, which is illegal while a graph is being captured. + runtime_params = sparse_params.scheduler.get_runtime_params( + timestep=timestep, + graph_phase=resolved_extra_key("skip_softmax_phase"), + ) threshold_scale_factor = runtime_params.threshold_scale_factor_prefill if threshold_scale_factor is None or threshold_scale_factor <= 0.0: return None diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py index c86143780c93..66066116cbde 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py @@ -49,9 +49,12 @@ import torch from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxScheduler +from tensorrt_llm._torch.visual_gen.cuda_graph_runner import resolved_extra_key from tensorrt_llm.logger import logger +from ....attention.backends.interface import PredefinedAttentionMask from ..interface import AttentionBackend, AttentionTensorLayout +from ..vanilla import VanillaAttention _sol_attn_import_error = None try: @@ -147,6 +150,16 @@ def __init__( # switched off entirely, against a 0.25 gate. from .fmha import CuTeDSLAttention + # The only backend that consumes `key_padding_mask` (CuTeDSL's `_fwd` + # swallows it via **kwargs, silently). Masked self-attention is routed + # here rather than to the mask-blind sparse kernel or to `_inner`. + self._vanilla = VanillaAttention( + layer_idx=layer_idx, + num_heads=num_heads, + head_dim=head_dim, + num_kv_heads=self.num_kv_heads, + dtype=dtype, + ) self._inner = CuTeDSLAttention( layer_idx=layer_idx, num_heads=num_heads, @@ -177,10 +190,15 @@ def __init__( # phases still compile as separate graphs -- they run different kernels. @torch.compiler.disable def _dense_by_step(self, timestep: Any) -> bool: - phase = SkipSoftmaxScheduler.get_graph_phase_for_timestep( - timestep, - disabled_until_timestep=self.disabled_until_timestep, - ) + # Under CUDA-graph capture the runner has already resolved the phase + # host-side (it is part of the graph key); reading the tensor here + # would `.item()` inside capture, which CUDA forbids. + phase = resolved_extra_key("sol_attn_phase") + if phase is None: + phase = SkipSoftmaxScheduler.get_graph_phase_for_timestep( + timestep, + disabled_until_timestep=self.disabled_until_timestep, + ) if phase is None: # Fail open, matching the CuTeDSL skip-softmax path: without a # timestep we cannot tell which phase we are in, so run the @@ -198,10 +216,26 @@ def _dense_by_step(self, timestep: Any) -> bool: return phase == 0 @staticmethod - def _sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: - """Dense attention via torch SDPA, for devices CuTe DSL cannot serve.""" + def _sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs: Any) -> torch.Tensor: + """Dense attention via torch SDPA, for devices CuTe DSL cannot serve. + + Honors the same two mask kwargs the backends accept: ``attention_mask`` + (``CAUSAL``) and ``key_padding_mask`` (``[B, S_kv]`` bool, True = valid). + """ + is_causal = kwargs.get("attention_mask", PredefinedAttentionMask.FULL) == ( + PredefinedAttentionMask.CAUSAL + ) + key_padding_mask = kwargs.get("key_padding_mask") + attn_mask = None + if key_padding_mask is not None: + # SDPA wants a broadcastable bool mask where True = attend. + attn_mask = key_padding_mask.to(torch.bool)[:, None, None, :] return torch.nn.functional.scaled_dot_product_attention( - q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + attn_mask=attn_mask, + is_causal=is_causal and attn_mask is None, ).transpose(1, 2) def _delegate( @@ -214,14 +248,25 @@ def _delegate( needed: the construction-time probe inspects the current CUDA device, so it says yes on a GPU host even when a caller passes CPU tensors. """ + if kwargs.get("key_padding_mask") is not None: + # CuTeDSL does not consume `key_padding_mask`, so staying in-family + # here would silently attend to padded tokens; VANILLA honors it. + # VANILLA works in HND ([B, H, S, D]); this backend is NHD. + if self._vanilla.preferred_layout == AttentionTensorLayout.HND: + out = self._vanilla.forward( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), **kwargs + ) + return out.transpose(1, 2) + return self._vanilla.forward(q, k, v, **kwargs) if self._cute_dense_ok and q.is_cuda: + # CAUSAL is fine here: CuTeDSL handles it via `_prepare_inputs`. return self._inner.forward(q, k, v, **kwargs) - return self._sdpa(q, k, v) + return self._sdpa(q, k, v, **kwargs) def _can_serve(self, q: torch.Tensor, k: torch.Tensor, **kwargs: Any) -> bool: """Whether the sparse kernel applies to this particular call. - Everything false here is delegated to ``_inner``. Deciding it from the + Everything false here is delegated (see ``_delegate``). Deciding it from the tensors, per call, is deliberate, and it is why ``modules/attention.py`` has no ``SEPARATE_QKV`` rule for Sol-Attn: ``qkv_mode`` describes how Q/K/V are *projected*, not whether K/V come from another sequence, so a @@ -234,6 +279,17 @@ def _can_serve(self, q: torch.Tensor, k: torch.Tensor, **kwargs: Any) -> bool: # assumes one self-attending sequence. if k.shape[1] != q.shape[1]: return False + # Masks: the sparse kernel is noncausal and takes no mask, so a masked + # call -- HunyuanVideo1.5 / GLM-Image `key_padding_mask`, Cosmos3 + # `CAUSAL` -- must go to a backend that honors it. Equal Q/K lengths do + # not imply "unmasked self-attention". + if kwargs.get("key_padding_mask") is not None: + return False + if ( + kwargs.get("attention_mask", PredefinedAttentionMask.FULL) + != PredefinedAttentionMask.FULL + ): + return False if self.layer_idx in self.dense_layers: return False if self.disabled_until_timestep is not None and self._dense_by_step(kwargs.get("timestep")): diff --git a/tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py b/tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py index 3d79f73dde39..eebb6c6ad497 100644 --- a/tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py @@ -1,3 +1,5 @@ +import contextlib +import contextvars import functools import gc from dataclasses import dataclass @@ -9,6 +11,37 @@ from ..utils import make_weak_ref +# Extra-key values the runner resolved host-side for the call it is currently +# capturing. CUDA-graph capture forbids device-to-host syncs, so anything a +# forward would otherwise derive from a CUDA tensor with `.item()` -- the +# sparse-attention dense-prefix phase, for instance -- must be resolved before +# capture and read back from here inside it. The runner already computes these +# values to build the graph key; this just makes them visible to the callee. +_RESOLVED_EXTRA_KEYS: contextvars.ContextVar[Optional[Dict[str, Any]]] = contextvars.ContextVar( + "cuda_graph_resolved_extra_keys", default=None +) + + +def resolved_extra_key(name: str) -> Any: + """Return the extra-key value `name` resolved for the call being captured, else None. + + None means "not inside a runner-driven call" (or the key was omitted), and + callers must fall back to deriving the value themselves. + """ + resolved = _RESOLVED_EXTRA_KEYS.get() + return None if resolved is None else resolved.get(name) + + +@contextlib.contextmanager +def resolved_extra_keys_scope(values: Dict[str, Any]): + """Expose `values` through `resolved_extra_key` for the enclosed block.""" + token = _RESOLVED_EXTRA_KEYS.set(dict(values)) + try: + yield + finally: + _RESOLVED_EXTRA_KEYS.reset(token) + + # One named graph-key component, e.g. ("hidden_states", (1, 4096, 3072)). KeyPart: TypeAlias = Tuple[str, Hashable] # Full CUDA graph cache key, stored as a tuple so it can be used as a dict key. @@ -121,6 +154,9 @@ def _get_extra_key(self, *args, **kwargs) -> KeyType: value = fn(*args, **kwargs) if value is not None: parts.append((name, value)) + # Remembered so `capture()` can republish them to the callee; this runs + # outside capture, where the `.item()` these callbacks need is legal. + self._last_resolved_extra_keys = dict(parts) return tuple(parts) def get_graph_key(self, *args, **kwargs) -> KeyType: @@ -144,14 +180,18 @@ def capture( } graph = torch.cuda.CUDAGraph() - for _ in range(self.WARMUP_STEPS): - fn(*static_args, **static_kwargs) - torch.cuda.synchronize() - gc.collect() - torch.cuda.empty_cache() - - with torch.cuda.graph(graph, pool=self._get_pool()): - output = fn(*static_args, **static_kwargs) + # Warmup and capture both see the host-resolved extra keys, so a + # forward that needs the dense-prefix phase reads it from here instead + # of syncing a CUDA tensor -- which capture would reject. + with resolved_extra_keys_scope(getattr(self, "_last_resolved_extra_keys", {})): + for _ in range(self.WARMUP_STEPS): + fn(*static_args, **static_kwargs) + torch.cuda.synchronize() + gc.collect() + torch.cuda.empty_cache() + + with torch.cuda.graph(graph, pool=self._get_pool()): + output = fn(*static_args, **static_kwargs) self.graphs[key] = graph self.static_inputs[key] = (static_args, static_kwargs) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py index f06c5d003b41..328bbb2aff04 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py @@ -1714,7 +1714,6 @@ def forward( audio_scheduler = copy.deepcopy(self.scheduler) audio_scheduler.set_timesteps(num_inference_steps, latent=latents_5d) timesteps = self.scheduler.timesteps - num_steps = len(timesteps) # ---- 7. Build perturbation config for STG ----------------------- stg_perturbation: PerturbationConfig | None = None @@ -1842,7 +1841,11 @@ def _run_transformer( audio=audio_mod, perturbations=perturbations, text_cache=text_cache, - timestep=timestep_val.new_tensor(float(step_index) / num_steps), + # The scheduler sigma is already normalized to [0, 1] with the + # contract's sense (larger = noisier). The previous + # step_index / num_steps was ascending, which inverted the + # sparse-attention dense prefix on this model. + timestep=timestep_val, step_index=step_index, ) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py index 7c266fe0a2cd..e3529a7e819d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py @@ -1107,6 +1107,10 @@ def _setup_cuda_graphs(self): self._current_lora_cuda_graph_state, lambda: self.transformer.active_topology, ) + # Same registration the single-stage pipeline does: without it the + # dense-prefix and sparse phases of skip-softmax / Sol-Attn share one + # captured graph and the wrong kernel is replayed. + self.transformer.register_cuda_graph_extra_key_fns(runner) compile_note = " (with torch.compile)" if self.pipeline_config.torch_compile.enable else "" logger.info( "CUDA graph runner: wrapping LTX-2 two-stage transformer.forward " @@ -1793,6 +1797,7 @@ def _refinement_denoise( audio=audio_mod, text_cache=_s2_static, step_index=i, + timestep=timestep, ) # Video: velocity → x0 → post-process → Euler step diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py index 7cba785b1bf2..2ff493af9900 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py @@ -928,9 +928,18 @@ def forward( text_kv_video: tuple[torch.Tensor, torch.Tensor] | None = None, text_kv_audio: tuple[torch.Tensor, torch.Tensor] | None = None, step_index=None, + timestep: torch.Tensor | None = None, ) -> tuple[TransformerArgs | None, TransformerArgs | None]: """Forward with optional perturbation masking for STG. + ``timestep`` is the normalized scheduler time in ``[0, 1]`` (larger = + noisier) that the attention backends consume for timestep-dependent + decisions such as the sparse-attention dense prefix. It is deliberately + *not* ``video.timesteps`` / ``audio.timesteps``: those are the AdaLN + modulation outputs -- learned activations -- and were previously passed + to attention as if they were the scheduler time, so any + ``disabled_until_timestep`` cutoff was compared against noise. + Args: perturbations: Optional ``BatchedPerturbationConfig`` that masks attention outputs for selected blocks/modalities. @@ -991,9 +1000,7 @@ def forward( self._fuse_adaln, fp4_input_scale=get_nvfp4_self_attn_input_scale(self.attn1), ) - v_attn_raw = self.attn1( - norm_vx, pe=video.positional_embeddings, timestep=video.timesteps - ) + v_attn_raw = self.attn1(norm_vx, pe=video.positional_embeddings, timestep=timestep) if has_perturbations and perturbations.any_in_batch( PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx ): @@ -1018,7 +1025,7 @@ def forward( attn2_q_input, context=video.context, pre_projected_kv=text_kv_video, - timestep=video.timesteps, + timestep=timestep, ) # --- Audio self-attention + text cross-attention --- @@ -1049,7 +1056,7 @@ def forward( norm_ax, pe=audio.positional_embeddings, key_padding_mask=audio.audio_padding_mask, - timestep=audio.timesteps, + timestep=timestep, ) if has_perturbations and perturbations.any_in_batch( PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx @@ -1074,7 +1081,7 @@ def forward( audio_attn2_q_input, context=audio.context, pre_projected_kv=text_kv_audio, - timestep=audio.timesteps, + timestep=timestep, ) # --- Bidirectional audio ↔ video cross-attention --- @@ -1233,7 +1240,7 @@ def forward( pre_projected_kv=(k_a2v, v_a2v), pe=video.cross_positional_embeddings, key_padding_mask=audio.audio_padding_mask, - timestep=video.timesteps, + timestep=timestep, ) if has_perturbations and perturbations.any_in_batch( PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx @@ -1267,7 +1274,7 @@ def forward( freqs=a_cross_pe, kv_input=vx_scaled_v2a, kv_freqs=video.cross_positional_embeddings, - timestep=audio.timesteps, + timestep=timestep, ) else: k_v2a, v_v2a = self.video_to_audio_attn.project_kv( @@ -1286,7 +1293,7 @@ def forward( ax_v2a_local, pre_projected_kv=(k_v2a, v_v2a), pe=a_cross_pe, - timestep=audio.timesteps, + timestep=timestep, ) v2a_attn_raw = self._sp_all_gather(out_local, dim=1) elif self._async_ulysses and self.video_to_audio_attn.is_ulysses: @@ -1299,7 +1306,7 @@ def forward( freqs=audio.cross_positional_embeddings, kv_input=vx_scaled_v2a, kv_freqs=video.cross_positional_embeddings, - timestep=audio.timesteps, + timestep=timestep, ) else: # v2a sync: with a Ulysses wrapper, K/V (video) stay seq-sharded @@ -1320,7 +1327,7 @@ def forward( ax_scaled_v2a, pre_projected_kv=(k_v2a, v_v2a), pe=audio.cross_positional_embeddings, - timestep=audio.timesteps, + timestep=timestep, ) if has_perturbations and perturbations.any_in_batch( PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx @@ -2235,9 +2242,10 @@ def forward( text_cache: Pre-computed step-invariant outputs from ``prepare_text_cache()``. Always required — callers must invoke ``prepare_text_cache()`` first. timestep: Normalized denoising-time coordinate in ``[0, 1]``. - May be ``None`` for LTX-2 paths that rely only on per-modality - timestep values and do not need timestep-based CUDA graph - partitioning. + Threaded to every attention call; sparse-attention dense + prefixes (``disabled_until_timestep``) and the CUDA-graph phase + key both read it, so pipelines must pass the scheduler sigma + here. May be ``None`` only when neither is configured. LTX-2 also carries per-modality timestep values in ``video.timesteps`` / ``audio.timesteps`` for the reference time-embedding path. @@ -2333,6 +2341,7 @@ def forward( ax, perturbations=perturbations, step_index=step_index, + timestep=timestep, ) if video_args is not None and vx is not None: video_args = replace(video_args, x=vx) @@ -2347,6 +2356,7 @@ def forward( text_kv_video=v_kv[i] if v_kv else None, text_kv_audio=a_kv[i] if a_kv else None, step_index=step_index, + timestep=timestep, ) # Gather sequences back to full length for output processing. diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 370d26520c65..213635cbc144 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -269,8 +269,10 @@ class SolAttentionConfig(BaseSparseAttentionConfig): "protects the high-noise prefix. Use None (not 0.0) to disable the " "prefix; 0.0 is rejected because it would run dense on every step " "and silently turn Sol-Attn off entirely. " - "The timestep is supplied as a forward kwarg by every VisualGen " - "pipeline, so no per-pipeline wiring is required." + "Read from the `timestep` forward kwarg, which must be the normalized " + "scheduler time (larger = noisier). WAN and LTX-2 pass it; a pipeline " + "that does not, or that passes something else, gets a one-time warning " + "and runs sparse on every step (fail-open)." ), ) dense_layers: Optional[str] = PydanticField( diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py index 3a918d43f8a8..4dc6d172127c 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py @@ -23,11 +23,13 @@ impossible and gives the same guarantee. """ +import math from types import SimpleNamespace import pytest import torch +from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxScheduler from tensorrt_llm._torch.visual_gen.attention_backend import CuTeDSLAttention from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn import ( @@ -39,6 +41,12 @@ DiffusionModelConfig, create_attention_metadata_state, ) +from tensorrt_llm._torch.visual_gen.cuda_graph_runner import ( + CUDAGraphRunner, + CUDAGraphRunnerConfig, + resolved_extra_key, + resolved_extra_keys_scope, +) from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm.visual_gen.args import AttentionConfig, SolAttentionConfig @@ -705,3 +713,326 @@ def _spy(*args, **kwargs): ) a.forward(q, k, v) assert calls["n"] == 1, "dense_fn did not route the fallback to the CuTeDSL dense kernel" + + +# --------------------------------------------------------------------------- +# Dense-prefix phase under CUDA-graph capture +# --------------------------------------------------------------------------- + + +def test_dense_by_step_prefers_runner_resolved_phase(monkeypatch): + """Inside a runner-driven call the phase comes from the runner, not the tensor. + + CUDA-graph capture forbids device-to-host syncs, and reading the timestep + is a `.item()`. The runner resolves the phase host-side to build the graph + key and republishes it during capture; `_dense_by_step` must use that and + never touch the tensor. `torch.compiler.disable` does not help here -- it + only excludes Dynamo, not stream capture. + """ + reads = {"n": 0} + real = SkipSoftmaxScheduler._as_float + + def spy(value): + reads["n"] += 1 + return real(value) + + monkeypatch.setattr(SkipSoftmaxScheduler, "_as_float", staticmethod(spy)) + attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) + attn.disabled_until_timestep = 0.9 + + # Resolved phase wins even when the tensor says otherwise: phase 0 is the + # dense prefix, phase 1 the sparse phase. + with resolved_extra_keys_scope({"sol_attn_phase": 0}): + assert attn._dense_by_step(torch.tensor(0.1)) is True + with resolved_extra_keys_scope({"sol_attn_phase": 1}): + assert attn._dense_by_step(torch.tensor(0.99)) is False + assert reads["n"] == 0, "timestep tensor was read despite a runner-resolved phase" + + # Outside a runner-driven call the tensor is the only source of truth. + assert attn._dense_by_step(torch.tensor(0.99)) is True + assert reads["n"] == 1 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graphs need a GPU") +def test_cuda_graph_runner_publishes_resolved_extra_keys_during_capture(): + """The runner exposes its host-resolved extra keys to the captured callee.""" + seen = [] + + def fn(x): + seen.append(resolved_extra_key("probe")) + return x * 2 + + runner = CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) + runner.register_extra_key_fn("probe", lambda *args, **kwargs: 7) + wrapped = runner.wrap(fn) + x = torch.ones(4, device="cuda") + out = wrapped(x) + torch.cuda.synchronize() + assert torch.equal(out, x * 2) + assert seen and all(v == 7 for v in seen), seen # warmup + capture passes + assert resolved_extra_key("probe") is None, "scope must close after capture" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Sol-Attn needs CUDA") +def test_sol_attn_dense_prefix_survives_cuda_graph_capture(): + """Capture and replay a graph on each side of the cutoff without a sync. + + The dense-prefix decision must be baked into each captured graph via the + runner's phase key, and no `.item()` may run inside capture. `kernel_calls` + proves the phase-0 graph never launched the sparse kernel and the phase-1 + graph did. + """ + sab = _backend_mod() + if not sab.sol_attn_supported(torch.empty(1, 8, 8, 128, device="cuda", dtype=torch.bfloat16)): + pytest.skip("no Sol-Attn kernel for this device") + + cutoff = 0.9 + attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) + attn.disabled_until_timestep = cutoff + runner = CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) + runner.register_extra_key_fn( + "sol_attn_phase", + lambda *args, **kwargs: SkipSoftmaxScheduler.get_graph_phase_for_timestep( + kwargs.get("timestep"), disabled_until_timestep=cutoff + ), + ) + + def fwd(q, k, v, *, timestep): + return attn.forward(q, k, v, timestep=timestep) + + wrapped = runner.wrap(fwd) + torch.manual_seed(0) + q = k = v = torch.randn(1, 64, 2, 128, device="cuda", dtype=torch.bfloat16) + t_dense = torch.tensor(0.95, device="cuda") + t_sparse = torch.tensor(0.5, device="cuda") + + before = sab._SOL_STATS["kernel_calls"] + wrapped(q, k, v, timestep=t_dense) + torch.cuda.synchronize() + after_dense = sab._SOL_STATS["kernel_calls"] + wrapped(q, k, v, timestep=t_sparse) + torch.cuda.synchronize() + after_sparse = sab._SOL_STATS["kernel_calls"] + + assert len(runner.graphs) == 2, "one graph per phase" + assert after_dense == before, "dense prefix must not launch the sparse kernel" + assert after_sparse > after_dense, "sparse phase must launch the kernel" + # Replays must not error either. + wrapped(q, k, v, timestep=t_dense) + wrapped(q, k, v, timestep=t_sparse) + torch.cuda.synchronize() + + +# --------------------------------------------------------------------------- +# Kernel vs a PyTorch reference of the block-mean approximation +# --------------------------------------------------------------------------- + + +def _sol_attn_reference(q, k, v, *, tau, scale, block=64): + """PyTorch fp32 reference of Sol-Attn's routing + block-mean approximation. + + Mirrors the vendored kernel (`preprocess.py`, `common/selector.py`, + `sm100/mainloop.py`): + + * `kc` = per-block mean of K, `vc` = per-block *sum* of V. + * Per query block: threshold = mean + tau * std, in the log2 domain, where + mean/var come from the query centroid against the global K mean and + per-dim K variance (`_diag_threshold_kernel`). + * A KV block is *exact* for a query block if the column-mean route score + exceeds that threshold, or it is within one block of the diagonal + (`sol_attn_route_is_exact`). + * Non-exact blocks contribute `exp(q . kc) * vc` to the numerator and + `exp(q . kc) * block_len` to the denominator -- every key in the block + is treated as its mean. + + Returns the output and the exact-routing mask `[B, nb_q, nb_kv, H]`. + """ + q, k, v = (t.float() for t in (q, k, v)) + B, S, H, D = q.shape + nb = -(-S // block) + log2e = math.log2(math.e) + log2_scale = scale * log2e + dev = q.device + idx = torch.arange(S, device=dev) // block + blen = torch.bincount(idx, minlength=nb).float() + + def _block_sum(t): + return torch.zeros(B, nb, H, D, device=dev).index_add_(1, idx, t) + + kc = _block_sum(k) / blen[None, :, None, None] + vc = _block_sum(v) + qc = _block_sum(q) / blen[None, :, None, None] + kmean = k.mean(dim=1) + kvar = ((k * k).mean(dim=1) - kmean * kmean).clamp(min=0) + mean = (qc * kmean[:, None]).sum(-1) * log2_scale + var = (qc * qc * kvar[:, None]).sum(-1) * log2_scale**2 + thr = mean + tau * torch.sqrt(var.clamp(min=0) + 1e-6) # [B, nb_q, H] + + s2 = torch.einsum("bshd,bjhd->bsjh", q, kc) * log2_scale # [B, S, nb_kv, H] + col_mean = ( + torch.zeros(B, nb, nb, H, device=dev).index_add_(1, idx, s2) / blen[None, :, None, None] + ) + ar = torch.arange(nb, device=dev) + near_diag = (ar[:, None] - ar[None, :]).abs() <= 1 + exact = (col_mean > thr[:, :, None, :]) | near_diag[None, :, :, None] # [B, nb_q, nb_kv, H] + + exact_rows = exact[:, idx] # [B, S, nb_kv, H] + exact_pair = exact_rows[:, :, idx, :] # [B, S, T, H] + scores = torch.einsum("bshd,bthd->bsth", q, k) * scale + approx = s2 / log2e + neg = torch.finfo(torch.float32).min + m = torch.maximum( + scores.masked_fill(~exact_pair, neg).amax(2), + approx.masked_fill(exact_rows, neg).amax(2), + ) + pe = torch.exp(scores - m[:, :, None]) * exact_pair + pa = torch.exp(approx - m[:, :, None]) * (~exact_rows) + num = torch.einsum("bsth,bthd->bshd", pe, v) + torch.einsum("bsjh,bjhd->bshd", pa, vc) + den = pe.sum(2) + (pa * blen[None, None, :, None]).sum(2) + return num / den[..., None], exact + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Sol-Attn needs CUDA") +def test_cute_kernel_matches_reference_on_multi_block_with_approximation(): + """The real kernel agrees with the reference where blocks are approximated. + + Complements the single-block test, which can only reach the all-exact + path. Plain Gaussian inputs are the right stimulus: every KV block carries + comparable softmax mass and no block's column-mean score clears the + `mean + tau*std` threshold, so every block two or more away from the + diagonal is approximated -- and because keys vary within a block, treating + 64 keys as their mean is measurably different from dense attention (a + Jensen gap). Clustered inputs are the wrong stimulus: distant blocks get + routed to the approximation but carry no mass, so approximating them + changes nothing and the test would pass vacuously. The reference computes + the same routing, and the test asserts that some blocks were approximated + *and* that the result differs from dense, so neither an accidental + all-exact run nor a mass-free approximation can pass. It also asserts + `kernel_calls` advanced, so a dense fallback cannot pass either. + """ + sab = _backend_mod() + if not sab.sol_attn_supported(torch.empty(1, 8, 8, 128, device="cuda", dtype=torch.bfloat16)): + pytest.skip("no Sol-Attn kernel for this device") + + torch.manual_seed(0) + S, H, D, block = 256, 2, 128, 64 # batch of 1 is added with t[None] below + nb = S // block + dev = torch.device("cuda") + q, k, v = (torch.randn(1, S, H, D, device=dev, dtype=torch.bfloat16) for _ in range(3)) + del nb # routing granularity is decided inside the reference + scale = D**-0.5 + tau = 2.0 + + ref, exact = _sol_attn_reference(q, k, v, tau=tau, scale=scale, block=block) + assert (~exact).any(), "inputs did not force any block onto the approximation path" + dense = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2).float(), k.transpose(1, 2).float(), v.transpose(1, 2).float() + ).transpose(1, 2) + # Calibrated on B200: max |ref - dense| = 0.45 for this stimulus. A tight + # margin here would let a mass-free approximation pass vacuously. + assert (ref - dense).abs().max() > 1e-1, "approximation is not measurably different from dense" + + before = sab._SOL_STATS["kernel_calls"] + out = sab._run_sol_attn_bthd(q, k, v, tau=tau, thresh_type="diag", kv_splits=1) + torch.cuda.synchronize() + assert sab._SOL_STATS["kernel_calls"] == before + 1, "kernel did not run" + + # Calibrated on B200: max |kernel - ref| = 1.7e-3 (mean 2e-4); 5e-3 leaves + # ~3x headroom for bf16 rounding while still catching a routing or + # approximation-formula mismatch, which shows up at 1e-1 or worse. + torch.testing.assert_close(out.float(), ref, rtol=5e-3, atol=5e-3) + + +# --------------------------------------------------------------------------- +# Masks: routed to a mask-aware backend, and honored -- not just redirected +# --------------------------------------------------------------------------- + + +def _masked_sdpa_reference(q, k, v, *, key_padding_mask=None, is_causal=False): + attn_mask = ( + None if key_padding_mask is None else key_padding_mask.to(torch.bool)[:, None, None, :] + ) + return torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2).float(), + k.transpose(1, 2).float(), + v.transpose(1, 2).float(), + attn_mask=attn_mask, + is_causal=is_causal and attn_mask is None, + ).transpose(1, 2) + + +def test_key_padding_mask_routes_to_vanilla_and_is_honored(monkeypatch): + """Masked self-attention must never reach the mask-blind sparse kernel. + + Equal Q/K lengths do not mean "unmasked": HunyuanVideo1.5 and GLM-Image + pass a `[B, S]` `key_padding_mask` on self-attention. The sparse kernel + takes no mask and CuTeDSL's dense `_fwd(**kwargs)` swallows it, so the only + correct destination is VANILLA. The output is checked against a masked + SDPA reference so a mask that was routed but then dropped still fails. + """ + import tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn as sol_attn_mod + + monkeypatch.setattr( + sol_attn_mod, + "_sol_attn_run", + lambda *a, **k: pytest.fail("sparse kernel ran on a masked call"), + ) + attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) + attn.disabled_until_timestep = None + calls = {"vanilla": 0} + real = attn._vanilla.forward + + def spy(*a, **k): + calls["vanilla"] += 1 + return real(*a, **k) + + monkeypatch.setattr(attn._vanilla, "forward", spy) + + torch.manual_seed(0) + q, k, v = (torch.randn(2, 64, 2, 128, dtype=torch.float32) for _ in range(3)) + mask = torch.ones(2, 64, dtype=torch.bool) + mask[0, 48:] = False # pad the tail of sample 0 + mask[1, :16] = False # pad the head of sample 1 + + assert not attn._can_serve(q, k, key_padding_mask=mask) + out = attn.forward(q, k, v, key_padding_mask=mask) + assert calls["vanilla"] == 1 + ref = _masked_sdpa_reference(q, k, v, key_padding_mask=mask) + torch.testing.assert_close(out.float(), ref, rtol=1e-4, atol=1e-4) + # And the mask actually mattered: unmasked attention gives a different answer. + assert (ref - _masked_sdpa_reference(q, k, v)).abs().max() > 1e-3 + + +def test_causal_mask_routes_to_dense_and_is_honored(monkeypatch): + """CAUSAL disqualifies the noncausal sparse kernel and is honored downstream. + + On sm100 the delegate is dense CuTeDSL, which supports CAUSAL; on any other + device it is `_sdpa`, which previously dropped it. Either way the result + must match a causal reference. + """ + import tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn as sol_attn_mod + + monkeypatch.setattr( + sol_attn_mod, + "_sol_attn_run", + lambda *a, **k: pytest.fail("sparse kernel ran on a causal call"), + ) + attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) + attn.disabled_until_timestep = None + torch.manual_seed(0) + q, k, v = (torch.randn(1, 64, 2, 128, dtype=torch.float32) for _ in range(3)) + + assert not attn._can_serve(q, k, attention_mask=PredefinedAttentionMask.CAUSAL) + out = attn.forward(q, k, v, attention_mask=PredefinedAttentionMask.CAUSAL) + ref = _masked_sdpa_reference(q, k, v, is_causal=True) + torch.testing.assert_close(out.float(), ref, rtol=1e-4, atol=1e-4) + assert (ref - _masked_sdpa_reference(q, k, v)).abs().max() > 1e-3 + + +def test_unmasked_self_attention_is_still_served(): + """The routing change must not touch the measured path: no mask, sparse.""" + attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) + attn.disabled_until_timestep = None + q = k = torch.randn(1, 64, 2, 128, dtype=torch.bfloat16) + assert attn._can_serve(q, k) + assert attn._can_serve(q, k, attention_mask=PredefinedAttentionMask.FULL, key_padding_mask=None) From 2a2365b4cee31f2795d9c0d2fc87d3f2cc6125e1 Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Thu, 10 Sep 2026 22:20:07 -0700 Subject: [PATCH 15/22] [TRTLLM-15917][chore] Drop an unused variable left in the multi-block 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 Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../_torch/visual_gen/test_attention_cute_dsl_sol_attn.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py index 4dc6d172127c..473ad90da87b 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py @@ -915,11 +915,9 @@ def test_cute_kernel_matches_reference_on_multi_block_with_approximation(): pytest.skip("no Sol-Attn kernel for this device") torch.manual_seed(0) - S, H, D, block = 256, 2, 128, 64 # batch of 1 is added with t[None] below - nb = S // block + S, H, D, block = 256, 2, 128, 64 # 4 KV blocks; batch of 1 dev = torch.device("cuda") q, k, v = (torch.randn(1, S, H, D, device=dev, dtype=torch.bfloat16) for _ in range(3)) - del nb # routing granularity is decided inside the reference scale = D**-0.5 tau = 2.0 From ae7c894689d4f9e83d2e82ce33cba9abc993e41d Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Fri, 11 Sep 2026 02:48:35 -0700 Subject: [PATCH 16/22] [TRTLLM-15917][test] Teach the LTX-2 two-stage test stubs the CUDA-graph 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 Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../unittest/_torch/visual_gen/test_ltx2_pipeline.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py index 4a5c984f6e23..ce4614969841 100644 --- a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py @@ -1253,6 +1253,9 @@ def __init__(self): def forward(self, x): return self.lin(x) + def register_cuda_graph_extra_key_fns(self, runner): + self.registered_runner = runner + pipeline = object.__new__(ltx2_two_stages.LTX2TwoStagesPipeline) # __init__ is bypassed here; BasePipeline.device now reads self._device # (set in __init__), so initialize it explicitly for the warmup path. @@ -1310,9 +1313,14 @@ def test_two_stage_cuda_graph_setup_uses_pipeline_config(self): """CUDA graph setup runs before the two-stage model_config is assigned.""" class TinyTransformer: + registered_runner = None + def forward(self, *args, **kwargs): return args, kwargs + def register_cuda_graph_extra_key_fns(self, runner): + self.registered_runner = runner + pipeline = object.__new__(ltx2_two_stages.LTX2TwoStagesPipeline) pipeline.pipeline_config = DiffusionPipelineConfig( cuda_graph=CudaGraphConfig(enable=True), @@ -1328,6 +1336,9 @@ def forward(self, *args, **kwargs): runner = pipeline._cuda_graph_runners["transformer"] assert isinstance(runner, ltx2_two_stages._LTX2TwoStageCUDAGraphRunner) assert runner._lora_state_getter() == "original" + # Same contract as the single-stage pipeline: the transformer's extra + # CUDA-graph keys (skip-softmax / Sol-Attn phase) must be registered. + assert pipeline.transformer.registered_runner is runner assert pipeline.transformer.forward.__wrapped__.__self__ is pipeline.transformer def test_cuda_graph_rejects_nonpersistent_lora_bindings(self): From 09bb2c7d943b649db08b404acc4579b834010cfd Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Mon, 14 Sep 2026 10:51:10 -0700 Subject: [PATCH 17/22] [TRTLLM-15917][fix] Keep the CUDA-graph phase scope in LTX-2's capture 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 Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../visual_gen/models/ltx2/pipeline_ltx2.py | 27 ++++++---- .../test_attention_cute_dsl_sol_attn.py | 50 +++++++++++++++---- 2 files changed, 58 insertions(+), 19 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py index 328bbb2aff04..2b8d1685bc9b 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py @@ -20,7 +20,11 @@ from tensorrt_llm._torch.utils import make_weak_ref from tensorrt_llm._torch.visual_gen.cache.teacache import CacheContext, register_extractor from tensorrt_llm._torch.visual_gen.checkpoints.prefetch import prefetch_files_to_host_cache -from tensorrt_llm._torch.visual_gen.cuda_graph_runner import CUDAGraphRunner, CUDAGraphRunnerConfig +from tensorrt_llm._torch.visual_gen.cuda_graph_runner import ( + CUDAGraphRunner, + CUDAGraphRunnerConfig, + resolved_extra_keys_scope, +) from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline import ( BasePipeline, @@ -493,14 +497,19 @@ def capture(self, key, fn, args, kwargs): static_kwargs = {k: self._clone_value(v) for k, v in kwargs.items()} graph = torch.cuda.CUDAGraph() - for _ in range(self.WARMUP_STEPS): - fn(*static_args, **static_kwargs) - torch.cuda.synchronize() - gc.collect() - torch.cuda.empty_cache() - - with torch.cuda.graph(graph, pool=self._get_pool()): - output = fn(*static_args, **static_kwargs) + # Same contract as CUDAGraphRunner.capture: warmup and capture read the + # host-resolved extra keys (skip-softmax / Sol-Attn dense-prefix phase) + # from this scope instead of syncing the CUDA timestep, which capture + # would reject. + with resolved_extra_keys_scope(getattr(self, "_last_resolved_extra_keys", {})): + for _ in range(self.WARMUP_STEPS): + fn(*static_args, **static_kwargs) + torch.cuda.synchronize() + gc.collect() + torch.cuda.empty_cache() + + with torch.cuda.graph(graph, pool=self._get_pool()): + output = fn(*static_args, **static_kwargs) self.graphs[key] = graph self.static_inputs[key] = (static_args, static_kwargs) diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py index 473ad90da87b..f328e14ae471 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py @@ -614,13 +614,12 @@ def test_cute_kernel_matches_dense_on_a_single_block(): torch.testing.assert_close(out.float(), reference, rtol=2e-2, atol=2e-2) -def test_kv_splits_rejects_unsupported_value(): - """kv_splits is constrained at the config layer: an out-of-range value is - otherwise rejected deep inside the kernel and caught by the blanket - except, silently degrading the entire run to dense attention.""" +def test_kv_splits_is_not_a_config_knob(): + """Only one KV split exists on the shipped kernels, so the config exposes + no `kv_splits` and rejects it rather than accepting a no-op choice.""" + assert "kv_splits" not in SolAttentionConfig.model_fields with pytest.raises(ValueError): - SolAttentionConfig(tau=2.0, kv_splits="4") - assert SolAttentionConfig(tau=2.0).kv_splits == "auto" + SolAttentionConfig(tau=2.0, kv_splits="1") def _is_dynamo_disabled(fn) -> bool: @@ -754,7 +753,37 @@ def spy(value): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graphs need a GPU") -def test_cuda_graph_runner_publishes_resolved_extra_keys_during_capture(): +def _base_runner(): + return CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) + + +def _ltx2_runner(): + from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2 import _LTX2CUDAGraphRunner + + return _LTX2CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) + + +def _ltx2_two_stage_runner(): + from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2_two_stages import ( + _LTX2TwoStageCUDAGraphRunner, + ) + + return _LTX2TwoStageCUDAGraphRunner( + CUDAGraphRunnerConfig(use_cuda_graph=True), lambda: "original", lambda: "default" + ) + + +# Every runner that overrides ``capture`` must keep the phase-scope contract; +# LTX-2's runners re-implement capture for ``Modality`` inputs. +_RUNNER_FACTORIES = pytest.mark.parametrize( + "make_runner", + [_base_runner, _ltx2_runner, _ltx2_two_stage_runner], + ids=["base", "ltx2", "ltx2_two_stage"], +) + + +@_RUNNER_FACTORIES +def test_cuda_graph_runner_publishes_resolved_extra_keys_during_capture(make_runner): """The runner exposes its host-resolved extra keys to the captured callee.""" seen = [] @@ -762,7 +791,7 @@ def fn(x): seen.append(resolved_extra_key("probe")) return x * 2 - runner = CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) + runner = make_runner() runner.register_extra_key_fn("probe", lambda *args, **kwargs: 7) wrapped = runner.wrap(fn) x = torch.ones(4, device="cuda") @@ -774,7 +803,8 @@ def fn(x): @pytest.mark.skipif(not torch.cuda.is_available(), reason="Sol-Attn needs CUDA") -def test_sol_attn_dense_prefix_survives_cuda_graph_capture(): +@_RUNNER_FACTORIES +def test_sol_attn_dense_prefix_survives_cuda_graph_capture(make_runner): """Capture and replay a graph on each side of the cutoff without a sync. The dense-prefix decision must be baked into each captured graph via the @@ -789,7 +819,7 @@ def test_sol_attn_dense_prefix_survives_cuda_graph_capture(): cutoff = 0.9 attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) attn.disabled_until_timestep = cutoff - runner = CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) + runner = make_runner() runner.register_extra_key_fn( "sol_attn_phase", lambda *args, **kwargs: SkipSoftmaxScheduler.get_graph_phase_for_timestep( From cd8ded2c110e4bcd8e13ae8a23e59d0332bc4055 Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Mon, 14 Sep 2026 10:51:12 -0700 Subject: [PATCH 18/22] [TRTLLM-15917][chore] Drop the no-op kv_splits option from SolAttentionConfig 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 Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../attention_backend/cute_dsl/sol_attn.py | 2 -- .../blackwell/sol_attn_backend.py | 23 ++++--------------- tensorrt_llm/visual_gen/sparse_attention.py | 10 -------- 3 files changed, 5 insertions(+), 30 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py index 66066116cbde..fe5d7d1ad884 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py @@ -139,7 +139,6 @@ def __init__( cfg = sparse_attention_config self.tau = getattr(cfg, "tau", 1.0) self.thresh_type = getattr(cfg, "thresh_type", "diag") - self.kv_splits = getattr(cfg, "kv_splits", "auto") self.disabled_until_timestep = getattr(cfg, "disabled_until_timestep", None) self.dense_layers = _parse_dense_layers(getattr(cfg, "dense_layers", None)) @@ -312,7 +311,6 @@ def forward( v, tau=self.tau, thresh_type=self.thresh_type, - kv_splits=self.kv_splits, # Shape/dtype/arch ineligibility is only detectable inside the # wrapper, so that last delegation happens through this hook. dense_fn=lambda a, b, c: self._delegate(a, b, c, **kwargs), diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py index da19ff488b7e..0dce5e086afb 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py @@ -137,18 +137,6 @@ def _cute_runtime_available() -> bool: return True -def _resolve_kv_splits(q, kv_splits: int | str | None) -> int: - """Resolve the integration-only ``auto`` policy to the public integer API. - - ``auto`` is always 1 here: kv_splits=2/4 was an SM90-only path, and this - build ships SM100 kernels only. - """ - - if kv_splits in (None, "auto"): - return 1 - return int(kv_splits) - - def _strict() -> bool: """Whether SOL_ATTN_STRICT=1 asks us to raise instead of degrading.""" @@ -214,7 +202,7 @@ def _run_sol_attn_bthd( *, tau: float = DEFAULT_TAU, thresh_type: str = DEFAULT_THRESH_TYPE, - kv_splits: int | str | None = "auto", + kv_splits: int = 1, sink_start: int | None = None, sink_tokens: int = 0, dense_fn: Callable | None = None, @@ -247,10 +235,6 @@ def dense(): ) return dense() - # Resolved outside the try: a bad kv_splits is a configuration error and - # must surface, not silently become a dense run. - resolved_kv_splits = _resolve_kv_splits(q0, kv_splits) - try: kernel = _load_sol_attn() out = kernel( @@ -259,7 +243,10 @@ def dense(): v0, tau=float(tau), thresh_type=str(thresh_type), - kv_splits=resolved_kv_splits, + # Only 1 split exists on the shipped sm100/sm103 kernels (2/4 was an + # SM90-only path), so this is not a user-facing knob. The kernel + # interface rejects anything else before the try below can swallow it. + kv_splits=int(kv_splits), sink_start=sink_start, sink_tokens=int(sink_tokens), ) diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 213635cbc144..40716b120bef 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -247,16 +247,6 @@ class SolAttentionConfig(BaseSparseAttentionConfig): "diag", description="Threshold policy forwarded to the kernel (kernel default: 'diag').", ) - kv_splits: Literal["auto", "1"] = PydanticField( - "auto", - description=( - "KV split policy. Only 1 split is valid on the shipped sm100 " - "kernels, so 'auto' and '1' are equivalent; the 2/4 path was " - "SM90-only and returns with that kernel. Constrained rather than a " - "free string because any other value is rejected deep inside the " - "kernel, which would silently degrade the whole run to dense." - ), - ) disabled_until_timestep: Optional[float] = PydanticField( None, gt=0.0, From 205c8d55ada9d21ff2a2d29197eb945f33ea00d2 Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:23:58 -0700 Subject: [PATCH 19/22] [TRTLLM-15917][chore] Vendor the Sol-Attn kernel through vendor_sources 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 Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .pre-commit-config.yaml | 2 +- 3rdparty/vendor_patches/sana-sol-attn.patch | 474 ++++++++++++++++ 3rdparty/vendor_sources.lock.yaml | 16 + pyproject.toml | 1 + .../blackwell/sol_attn/__init__.py | 5 - .../blackwell/sol_attn/common/__init__.py | 5 - .../blackwell/sol_attn/common/layout_utils.py | 5 - .../blackwell/sol_attn/common/runtime.py | 5 - .../blackwell/sol_attn/common/selector.py | 10 +- .../blackwell/sol_attn/interface.py | 34 +- .../blackwell/sol_attn/preprocess.py | 43 +- .../blackwell/sol_attn/sm100/__init__.py | 5 - .../blackwell/sol_attn/sm100/kernel.py | 5 - .../blackwell/sol_attn/sm100/mainloop.py | 518 +++++++++++++----- .../blackwell/sol_attn/sm100/math.py | 5 - .../blackwell/sol_attn/sm100/softmax.py | 61 ++- .../blackwell/sol_attn/sm100/tmem.py | 14 +- 17 files changed, 958 insertions(+), 250 deletions(-) create mode 100644 3rdparty/vendor_patches/sana-sol-attn.patch diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 6ac5ee01c3b0..5dcceeea3d4c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1525,7 +1525,7 @@ legacy-files: &legacy_files | # list; the hook's own `files:` pattern only gates *when* the hook triggers. # Global exclude: vendored code + trtllm-gen FMHA artifacts (cubin pointers, export headers, cuda_ptx) -exclude: '(^cpp/tensorrt_llm/common/sha256/|^triton_kernels/|^tensorrt_llm/_torch/attention/backends/prims_ts/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$|trtllmGenKernels/gemm/trtllmGen_gemm_export/KernelMetaInfo\.h$|\.cubin\.tar\.zst$)' +exclude: '(^cpp/tensorrt_llm/common/sha256/|^triton_kernels/|^tensorrt_llm/_torch/attention/backends/prims_ts/|^tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$|trtllmGenKernels/gemm/trtllmGen_gemm_export/KernelMetaInfo\.h$|\.cubin\.tar\.zst$)' default_install_hook_types: [pre-commit, commit-msg] repos: diff --git a/3rdparty/vendor_patches/sana-sol-attn.patch b/3rdparty/vendor_patches/sana-sol-attn.patch new file mode 100644 index 000000000000..5751db4112ea --- /dev/null +++ b/3rdparty/vendor_patches/sana-sol-attn.patch @@ -0,0 +1,474 @@ +diff --git a/sol_attn/THIRD_PARTY_NOTICES.md b/sol_attn/THIRD_PARTY_NOTICES.md +index d3755083d1e7622ff509351594157dd879865e56..b226f77f96bc2316e902ddf91765b35c915a8a6c 100644 +--- a/sol_attn/THIRD_PARTY_NOTICES.md ++++ b/sol_attn/THIRD_PARTY_NOTICES.md +@@ -1,15 +1,49 @@ + # Third-party notices + +-The files under `sol_attn/_vendor/flash_attn/cute/` and portions of the SM89, +-SM90, and SM100 design scaffold derive from the FlashAttention project. Its +-BSD-3-Clause license is included at +-`sol_attn/sm100/LICENSE.flash-attention`. ++This package is vendored from ++[`github.com/NVlabs/Sana`](https://github.com/NVlabs/Sana), branch ++[`sol-engine`](https://github.com/NVlabs/Sana/tree/sol-engine), at commit ++[`5fe5feb`](https://github.com/NVlabs/Sana/commit/5fe5feb) (2026-08-17). ++Checked against the branch tip ++([`83e54df`](https://github.com/NVlabs/Sana/commit/83e54df), 2026-08-20) on ++2026-08-27. + +-The runtime also depends on NVIDIA CUTLASS / CuTe DSL, cuda-python, PyTorch, +-and Triton. Those dependencies are not redistributed by this repository and +-remain subject to their respective licenses. ++## Scope of the vendored subset ++ ++Only the pieces needed for the architectures TensorRT-LLM ships are carried: ++ ++| Carried | Not carried | ++|---|---| ++| `interface.py`, `preprocess.py`, `common/` | `sm89/`, `sm90/` (incl. `sm90/_compat/`) | ++| `sm100/` — serves SM100 (B200/GB200) and SM103 (B300/GB300) | `triton_ref/` Triton reference attention | ++| | `sm120/` (RTX Blackwell) | ++| | `_vendor/flash_attn/` (see below) | ++ ++`../sol_attn_backend.py` is not part of the vendored package but is a ++derivative work of upstream's ++`techniques/sparse_backends/sol_attn_backend.py` (same branch and commit). ++Only the kernel-wrapper subset is carried; upstream's model-integration half ++(diffusers dispatch hook, HunyuanVideo MMDiT padding, model-level Morton ++ordering) is not. ++ ++Implementation divergences from upstream are documented in the source itself, ++in the module docstrings of `../sol_attn_backend.py` and ++`attention_backend/cute_dsl/sol_attn.py`. + +-The SM120 warp-MMA/TMA execution skeleton and online-softmax helpers are +-adapted from NVIDIA cuDNN Frontend's block-sparse-attention reference at commit +-`74785165de2da954a2c879a5e3e6f95411c2292d`. That source is licensed under the +-Apache License 2.0; adapted files retain the corresponding SPDX header. ++## Licensing ++ ++The upstream package vendored a copy of FlashAttention's CuTe DSL helpers ++under `sol_attn/_vendor/flash_attn/cute/`. That copy is **not** carried here: ++TensorRT-LLM already depends on ++[`flash-attn-4`](https://github.com/Dao-AILab/flash-attention) (pinned in ++`requirements.txt`), which provides the same `flash_attn.cute` modules, and ++the SM100 kernels import them from that dependency directly. FlashAttention's ++BSD-3-Clause license is retained at `sol_attn/sm100/LICENSE.flash-attention` ++because portions of the SM100 design scaffold still derive from that project. ++ ++`preprocess.py` implements the routing/threshold stage in Triton, so Triton is ++a required runtime dependency on every Sol-Attn path, not only a fallback. ++ ++The runtime also depends on NVIDIA CUTLASS / CuTe DSL, cuda-python, and ++PyTorch. Those dependencies are not redistributed by this repository and ++remain subject to their respective licenses. +diff --git a/sol_attn/interface.py b/sol_attn/interface.py +index 43e9267b931c4f87620739182e596f1fe63b2475..2da4ea383947f5cd25bea439d7748d8ab5e4f28c 100644 +--- a/sol_attn/interface.py ++++ b/sol_attn/interface.py +@@ -7,11 +7,12 @@ import functools + import torch + + BLOCK_SIZE = 64 ++# TensorRT-LLM ships the SM100 kernel only. It serves both datacenter Blackwell ++# steppings: the CuTe DSL JIT targets the device it compiles on and the kernel ++# body uses no SM100-exclusive construct (measured identical on B200 and B300). + _CUTE_BACKENDS = { +- (8, 9): "cute_sm89", +- (9, 0): "cute_sm90", +- (10, 0): "cute_sm100", +- (12, 0): "cute_sm120", ++ (10, 0): "cute_sm100", # B200 / GB200 ++ (10, 3): "cute_sm100", # B300 / GB300 + } + _compiled = {} + +@@ -68,23 +69,30 @@ def _backend_for_arch( + *, + cute_available: bool | None = None, + ) -> str: +- """Select CuTe when specialized and available, otherwise Triton.""" ++ """Select the CuTe kernel for ``arch``, or raise if there is none. + +- if arch[0] < 8: ++ TensorRT-LLM raises instead of falling back to the Triton reference so a ++ missing kernel is visible; the caller decides what to do about it. ++ """ ++ ++ cute_backend = _CUTE_BACKENDS.get(arch) ++ if cute_backend is None: + raise RuntimeError( +- "Sol-Attn requires an NVIDIA GPU with compute capability >= 8.0; " +- f"got SM{arch[0]}{arch[1]}" ++ f"Sol-Attn has no kernel for SM{arch[0]}{arch[1]}; supported " ++ f"architectures are " ++ f"{', '.join(f'SM{a}{b}' for a, b in sorted(_CUTE_BACKENDS))}." + ) +- cute_backend = _CUTE_BACKENDS.get(arch) +- if cute_backend is not None: +- available = ( +- _cute_runtime_available() +- if cute_available is None +- else cute_available ++ available = ( ++ _cute_runtime_available() ++ if cute_available is None ++ else cute_available ++ ) ++ if not available: ++ raise RuntimeError( ++ "Sol-Attn requires the CuTe DSL runtime (cutlass.cute and " ++ "cuda.bindings.driver); neither could be imported." + ) +- if available: +- return cute_backend +- return "triton" ++ return cute_backend + + + def get_sol_attn_backend(device: torch.device | str | int | None = None) -> str: +@@ -96,8 +104,11 @@ def get_sol_attn_backend(device: torch.device | str | int | None = None) -> str: + + + def _validate_cute(arch, tokens, kv_splits): +- if arch != (9, 0) and kv_splits != 1: +- raise ValueError("kv_splits=2/4 is currently available on SM90 only") ++ if kv_splits != 1: ++ raise ValueError( ++ "kv_splits=2/4 was an SM90-only path; this build ships the SM100 " ++ "kernel only, so kv_splits must be 1." ++ ) + route_groups = ((tokens + 63) // 64 + 63) // 64 + if kv_splits > route_groups: + raise ValueError("each KV split must contain at least one N64 route group") +@@ -126,60 +137,6 @@ def _sink_block_range(tokens, sink_start, sink_tokens): + ) + + +-def _compile_sm90( +- key, +- tensors, +- scale, +- tokens, +- kv_splits, +- sink_range, +- stream, +-): +- import cutlass.cute as cute +- +- from .sm90 import make_kernel +- +- operator = make_kernel(tokens, kv_splits) +- args = _to_cute_tensors(tensors) +- compiled = cute.compile( +- operator, +- *args, +- scale, +- sink_range, +- stream=stream, +- options="--enable-tvm-ffi", +- ) +- _compiled[key] = compiled +- return compiled, args +- +- +-def _compile_sm89( +- key, +- tensors, +- scale, +- sink_start_block, +- sink_end_block, +- stream, +-): +- import cutlass.cute as cute +- +- from .sm89 import make_kernel +- +- operator = make_kernel() +- args = _to_cute_tensors(tensors) +- compiled = cute.compile( +- operator, +- *args, +- scale, +- sink_start_block, +- sink_end_block, +- stream=stream, +- options="--enable-tvm-ffi", +- ) +- _compiled[key] = compiled +- return compiled, args +- +- + def _compile_sm100( + key, + tensors, +@@ -206,33 +163,6 @@ def _compile_sm100( + return compiled, args + + +-def _compile_sm120( +- key, +- tensors, +- scale, +- sink_start_block, +- sink_end_block, +- stream, +-): +- import cutlass.cute as cute +- +- from .sm120 import make_kernel +- +- operator = make_kernel() +- args = _to_cute_tensors(tensors) +- compiled = cute.compile( +- operator, +- *args, +- scale, +- sink_start_block, +- sink_end_block, +- stream=stream, +- options="--enable-tvm-ffi", +- ) +- _compiled[key] = compiled +- return compiled, args +- +- + def _sol_attn_cute( + q, + k, +@@ -268,129 +198,36 @@ def _sol_attn_cute( + stream = _stream(q.device) + key = (q.device.index, arch, batch, tokens, heads, kv_splits) + +- if arch == (8, 9): +- sink_start_block, sink_end_block = _sink_block_range( +- tokens, +- sink_start, +- sink_tokens, +- ) +- tensors = [q, k, v, output, kc, vc, threshold, lse] +- compiled = _compiled.get(key) +- if compiled is None: +- compiled, args = _compile_sm89( +- key, +- tensors, +- scale, +- sink_start_block, +- sink_end_block, +- stream, +- ) +- else: +- args = _to_cute_tensors(tensors) +- compiled( +- *args, +- scale, +- sink_start_block, +- sink_end_block, +- stream=stream, +- ) +- elif arch == (9, 0): +- if sink_tokens: +- sink_start_block, sink_end_block = _sink_block_range( +- tokens, +- sink_start, +- sink_tokens, +- ) +- sink_range = sink_start_block | (sink_end_block << 16) +- else: +- sink_range = 0 +- tensors = [q, k, v, output, kc, vc, threshold, lse] +- if kv_splits > 1: +- tensors.extend( +- [ +- torch.empty( +- (batch, tokens, kv_splits * heads, 128), +- device=q.device, +- dtype=torch.bfloat16, +- ), +- torch.empty( +- (batch, tokens, kv_splits * heads), +- device=q.device, +- dtype=torch.float32, +- ), +- ] +- ) +- compiled = _compiled.get(key) +- if compiled is None: +- compiled, args = _compile_sm90( +- key, +- tensors, +- scale, +- tokens, +- kv_splits, +- sink_range, +- stream, +- ) +- else: +- args = _to_cute_tensors(tensors) +- compiled( +- *args, +- scale, +- sink_range, +- stream=stream, +- ) +- elif arch == (10, 0): +- sink_start_block, sink_end_block = _sink_block_range( +- tokens, +- sink_start, +- sink_tokens, +- ) +- tensors = [q, k, v, output, kc, vc, threshold, lse] +- compiled = _compiled.get(key) +- if compiled is None: +- compiled, args = _compile_sm100( +- key, +- tensors, +- scale, +- sink_start_block, +- sink_end_block, +- stream, +- ) +- else: +- args = _to_cute_tensors(tensors) +- compiled( +- *args, ++ if arch not in _CUTE_BACKENDS: ++ # Unreachable via sol_attn(): _backend_for_arch raises first. Kept ++ # explicit so a missed guard cannot return the uninitialised ++ # `output` buffer. Keyed off _CUTE_BACKENDS, not a literal. ++ raise ValueError(f"no Sol-Attn CuTe kernel for SM{arch[0]}{arch[1]}") ++ sink_start_block, sink_end_block = _sink_block_range( ++ tokens, ++ sink_start, ++ sink_tokens, ++ ) ++ tensors = [q, k, v, output, kc, vc, threshold, lse] ++ compiled = _compiled.get(key) ++ if compiled is None: ++ compiled, args = _compile_sm100( ++ key, ++ tensors, + scale, + sink_start_block, + sink_end_block, +- stream=stream, ++ stream, + ) + else: +- sink_start_block, sink_end_block = _sink_block_range( +- tokens, +- sink_start, +- sink_tokens, +- ) +- tensors = [q, k, v, output, kc, vc, threshold, lse] +- compiled = _compiled.get(key) +- if compiled is None: +- compiled, args = _compile_sm120( +- key, +- tensors, +- scale, +- sink_start_block, +- sink_end_block, +- stream, +- ) +- else: +- args = _to_cute_tensors(tensors) +- compiled( +- *args, +- scale, +- sink_start_block, +- sink_end_block, +- stream=stream, +- ) ++ args = _to_cute_tensors(tensors) ++ compiled( ++ *args, ++ scale, ++ sink_start_block, ++ sink_end_block, ++ stream=stream, ++ ) + return output + + +@@ -421,28 +258,15 @@ def sol_attn( + sink_tokens, + sink_start, + ) +- if kv_splits not in (1, 2, 4): +- raise ValueError("kv_splits must be 1, 2, or 4") +- backend = _backend_for_arch(arch) ++ if kv_splits != 1: ++ raise ValueError( ++ "kv_splits must be 1; the 2/4 path was SM90-only and this build " ++ "ships SM100 kernels only." ++ ) ++ _backend_for_arch(arch) # raises on an architecture with no kernel + scale = q.shape[-1] ** -0.5 if scale is None else float(scale) + tau = float(tau) + +- if backend == "triton": +- if kv_splits != 1: +- raise ValueError("kv_splits=2/4 is currently available on SM90 only") +- from .triton_ref import sol_attn as triton_sol_attn +- +- return triton_sol_attn( +- q, +- k, +- v, +- scale=scale, +- tau=tau, +- thresh_type=thresh_type, +- sink_tokens=sink_tokens, +- sink_start=sink_start, +- ) +- + _validate_cute(arch, q.shape[1], kv_splits) + return _sol_attn_cute( + q, +diff --git a/sol_attn/sm100/mainloop.py b/sol_attn/sm100/mainloop.py +index 36b3ad9d1450c836f7acdc80bbed1382eac4d3a9..24dee1cb0233ebc5ecaabed21eab3c7fe12e86ba 100644 +--- a/sol_attn/sm100/mainloop.py ++++ b/sol_attn/sm100/mainloop.py +@@ -12,13 +12,13 @@ import cutlass.cute as cute + import cutlass.pipeline as pipeline + import cutlass.utils as utils + import cutlass.utils.blackwell_helpers as sm100_utils +-import sol_attn._vendor.flash_attn.cute.pipeline as fa_pipeline +-import sol_attn._vendor.flash_attn.cute.utils as fa_utils ++import flash_attn.cute.pipeline as fa_pipeline ++import flash_attn.cute.utils as fa_utils + from cutlass import BFloat16, Float32, Int32 + from cutlass._mlir.dialects import llvm + from cutlass.cute.nvgpu import cpasync, tcgen05 + from cutlass.cutlass_dsl import T, dsl_user_op +-from sol_attn._vendor.flash_attn.cute.cute_dsl_utils import assume_tensor_aligned ++from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned + + from .softmax import ( + _load_m64_n128_score as _load_pair_score, +@@ -27,8 +27,8 @@ from .softmax import ( + ) + from . import math as mma_utils + +-from sol_attn.common import layout_utils +-from sol_attn.common.selector import ( ++from ..common import layout_utils ++from ..common.selector import ( + sol_attn_popc_b32, + sol_attn_route_is_exact, + ) +diff --git a/sol_attn/sm100/softmax.py b/sol_attn/sm100/softmax.py +index 876ef8007e12eb855cdaee1f1ee40c18775ae143..71b11f0e6ece8cabd7d6182bd1759e091c7c57db 100644 +--- a/sol_attn/sm100/softmax.py ++++ b/sol_attn/sm100/softmax.py +@@ -7,7 +7,7 @@ import cutlass.cute as cute + from cutlass import Float32, Int32 + from cutlass.cute.nvgpu import tcgen05 + +-from sol_attn._vendor.flash_attn.cute import utils as fa_utils ++from flash_attn.cute import utils as fa_utils + + from .tmem import ( + _add_physical_tmem_base, diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 5f5e8947b193..2adb1972965f 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -11,3 +11,19 @@ vendors: patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch patch_digest: sha256:ba917f330d8281781ce30ddb89f1772ee8f9017264d948b588b881a92f511e45 digest: sha256-tree-v1:fa65b674a9ec6121b4158069d5585cc56df28243d1ac1199eb97003d81135d76 + sana-sol-attn: + url: https://github.com/NVlabs/Sana.git + branch: sol-engine + commit: 5fe5febdf0f59fee1c0b44a5ce6665df0dabd247 + source: techniques/sparse_backends + destination: tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell + include: + - sol_attn/THIRD_PARTY_NOTICES.md + - sol_attn/__init__.py + - sol_attn/common/* + - sol_attn/interface.py + - sol_attn/preprocess.py + - sol_attn/sm100/* + patch: 3rdparty/vendor_patches/sana-sol-attn.patch + patch_digest: sha256:0d66669021683773aff34533abaf8f200abd09e3af93d19eed0e2a2e04f2f980 + digest: sha256-tree-v1:383f6085ca1bc2438e8d7d4796036af3a816dbfd396313704f7f4be2c06ac55b diff --git a/pyproject.toml b/pyproject.toml index e198a15955fc..971c42858258 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,7 @@ force-exclude = true exclude = [ "**3rdparty/**", "tensorrt_llm/_torch/attention/backends/prims_ts/**", + "tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/**", "triton_kernels/**", ] diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/__init__.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/__init__.py index 078ee8402512..d37351dd8265 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/__init__.py @@ -1,8 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# -# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see -# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. """Sol-Attn.""" from .interface import get_sol_attn_backend, sol_attn diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/__init__.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/__init__.py index 0c0fb62abe10..ec1d7ee9f173 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/__init__.py @@ -1,8 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# -# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see -# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. """Internal helpers shared by the architecture backends.""" from .runtime import to_cute_tensor diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/layout_utils.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/layout_utils.py index 870efa6835bf..67e4d0896a85 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/layout_utils.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/layout_utils.py @@ -1,8 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# -# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see -# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. """Tensor-layout helpers shared by the two CuTe kernels.""" import cutlass.cute as cute diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/runtime.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/runtime.py index 502b6cb6a468..5182d0c49879 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/runtime.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/runtime.py @@ -1,8 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# -# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see -# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. """Small host helpers shared by the architecture backends.""" from cutlass.cute.runtime import from_dlpack diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/selector.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/selector.py index 13ab6a23bc73..0321f0fe429c 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/selector.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/selector.py @@ -1,8 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# -# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see -# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. """CTA-local routing-mask helpers shared by the CuTe architecture backends.""" import cutlass @@ -78,7 +73,10 @@ def _test_exact_bit( ) -> cutlass.Boolean: word = offset // Int32(32) bit = offset - word * Int32(32) - return (_mask_word(mask0, mask1, mask2, mask3, word) & (Int32(1) << bit)) != Int32(0) + return ( + _mask_word(mask0, mask1, mask2, mask3, word) + & (Int32(1) << bit) + ) != Int32(0) @cute.jit diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/interface.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/interface.py index e275c5b04962..2da4ea383947 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/interface.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/interface.py @@ -1,8 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# -# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see -# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. """Public Sol-Attn interface.""" from __future__ import annotations @@ -12,13 +7,12 @@ import torch BLOCK_SIZE = 64 -# The vendored kernel lives in ``sm100/`` and serves both datacenter Blackwell -# steppings: the CuTe DSL JIT targets whatever device it compiles on, and the -# kernel body uses no sm100-exclusive construct. Measured identical on both -- -# see THIRD_PARTY_NOTICES.md. +# TensorRT-LLM ships the SM100 kernel only. It serves both datacenter Blackwell +# steppings: the CuTe DSL JIT targets the device it compiles on and the kernel +# body uses no SM100-exclusive construct (measured identical on B200 and B300). _CUTE_BACKENDS = { (10, 0): "cute_sm100", # B200 / GB200 - (10, 3): "cute_sm100", # B300 / GB300 (Blackwell Ultra) + (10, 3): "cute_sm100", # B300 / GB300 } _compiled = {} @@ -75,12 +69,10 @@ def _backend_for_arch( *, cute_available: bool | None = None, ) -> str: - """Select the CuTe kernel for ``arch``, or raise if there isn't one. + """Select the CuTe kernel for ``arch``, or raise if there is none. - Unsupported architectures raise rather than silently degrading: the caller - (``_run_sol_attn_bthd``) turns that into an explicit dense-SDPA fallback - with a warning, so a missing kernel is visible instead of showing up only - as absent speedup. + TensorRT-LLM raises instead of falling back to the Triton reference so a + missing kernel is visible; the caller decides what to do about it. """ cute_backend = _CUTE_BACKENDS.get(arch) @@ -90,7 +82,11 @@ def _backend_for_arch( f"architectures are " f"{', '.join(f'SM{a}{b}' for a, b in sorted(_CUTE_BACKENDS))}." ) - available = _cute_runtime_available() if cute_available is None else cute_available + available = ( + _cute_runtime_available() + if cute_available is None + else cute_available + ) if not available: raise RuntimeError( "Sol-Attn requires the CuTe DSL runtime (cutlass.cute and " @@ -204,10 +200,8 @@ def _sol_attn_cute( if arch not in _CUTE_BACKENDS: # Unreachable via sol_attn(): _backend_for_arch raises first. Kept - # explicit because the alternative on a missed guard is returning - # the uninitialised `output` buffer, i.e. silently wrong results. - # Keyed off _CUTE_BACKENDS rather than a literal so widening the - # dispatch map cannot leave this guard behind. + # explicit so a missed guard cannot return the uninitialised + # `output` buffer. Keyed off _CUTE_BACKENDS, not a literal. raise ValueError(f"no Sol-Attn CuTe kernel for SM{arch[0]}{arch[1]}") sink_start_block, sink_end_block = _sink_block_range( tokens, diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/preprocess.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/preprocess.py index 5c44a3d070cc..77c4cccf1870 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/preprocess.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/preprocess.py @@ -1,8 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# -# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see -# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. """Block summaries and routing thresholds shared by both CuTe kernels.""" from __future__ import annotations @@ -12,6 +7,7 @@ import triton.language as tl from triton.tools.tensor_descriptor import TensorDescriptor + BLOCK_SIZE = 64 HEAD_DIM = 128 THRESHOLD_GROUP_SIZE = 64 @@ -43,7 +39,9 @@ def _reduce_kc_kernel( ) batch, head = batch_head // H, batch_head % H block_len = tl.minimum(BLOCK, T - block * BLOCK) - values = k_desc.load([batch, block * BLOCK, head, d_tile * TILE_D]).reshape([BLOCK, TILE_D]) + values = k_desc.load( + [batch, block * BLOCK, head, d_tile * TILE_D] + ).reshape([BLOCK, TILE_D]) summary = tl.sum(values, axis=0) / block_len offsets = d_tile * TILE_D + tl.arange(0, TILE_D) tl.store( @@ -78,7 +76,9 @@ def _reduce_vc_kernel( tl.program_id(2), ) batch, head = batch_head // H, batch_head % H - values = v_desc.load([batch, block * BLOCK, head, d_tile * TILE_D]).reshape([BLOCK, TILE_D]) + values = v_desc.load( + [batch, block * BLOCK, head, d_tile * TILE_D] + ).reshape([BLOCK, TILE_D]) summary = tl.sum(values, axis=0) offsets = d_tile * TILE_D + tl.arange(0, TILE_D) tl.store( @@ -113,11 +113,9 @@ def _reduce_kc_stats_kernel( count = tl.full((), 0.0, dtype=tl.float32) for start in range(0, N, GROUP): valid = start + block_offsets < N - values = ( - kc_desc.load([batch, start, head, d_tile * TILE_D]) - .reshape([GROUP, TILE_D]) - .to(tl.float32) - ) + values = kc_desc.load( + [batch, start, head, d_tile * TILE_D] + ).reshape([GROUP, TILE_D]).to(tl.float32) values = tl.where(valid[:, None], values, 0.0) total += tl.sum(values, axis=0) total_sq += tl.sum(values * values, axis=0) @@ -162,7 +160,9 @@ def _diag_threshold_kernel( q_len = tl.minimum(BLOCK, T - q_start).to(tl.float32) d_offsets = tl.arange(0, TILE_D) valid_d = d_offsets < D - q_values = q_desc.load([batch, q_start, head, 0]).reshape([BLOCK, TILE_D]) + q_values = q_desc.load( + [batch, q_start, head, 0] + ).reshape([BLOCK, TILE_D]) q_centroid = tl.sum(q_values.to(tl.float32), axis=0) / q_len mean_kc = tl.load( kc_mean + batch_head * D + d_offsets, @@ -176,7 +176,9 @@ def _diag_threshold_kernel( ) log2_scale = softmax_scale * 1.4426950408889634 mean = tl.sum(q_centroid * mean_kc, axis=0) * log2_scale - variance = tl.sum(q_centroid * q_centroid * var_kc, axis=0) * (log2_scale * log2_scale) + variance = tl.sum( + q_centroid * q_centroid * var_kc, axis=0 + ) * (log2_scale * log2_scale) std = tl.sqrt(tl.maximum(variance, 0.0) + 1.0e-6) tl.store( global_threshold + (batch * N + q_block) * H + head, @@ -200,7 +202,9 @@ def _pool_query_kernel( q_start = q_block * BLOCK q_len = tl.minimum(BLOCK, T - q_start).to(tl.float32) offsets = tl.arange(0, TILE_D) - values = q_desc.load([batch, q_start, head, 0]).reshape([BLOCK, TILE_D]) + values = q_desc.load([batch, q_start, head, 0]).reshape( + [BLOCK, TILE_D] + ) centroid = tl.sum(values.to(tl.float32), axis=0) / q_len tl.store( q_bar + (batch_head * N + q_block) * D + offsets, @@ -240,7 +244,10 @@ def _exact_fused_threshold_kernel( other=0.0, ) second_moment = tl.load( - kc_second_moment + batch_head * D * D + offsets[:, None] * D + offsets[None, :], + kc_second_moment + + batch_head * D * D + + offsets[:, None] * D + + offsets[None, :], mask=valid_d[:, None] & valid_d[None, :], other=0.0, ) @@ -344,7 +351,9 @@ def _compute_diag_threshold( kc, [1, THRESHOLD_GROUP_SIZE, 1, tile_d], ) - _reduce_kc_stats_kernel[(triton.cdiv(head_dim, tile_d), batch * heads)]( + _reduce_kc_stats_kernel[ + (triton.cdiv(head_dim, tile_d), batch * heads) + ]( kc_desc, kc_mean, kc_var_diag, diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/__init__.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/__init__.py index 81efe5d1ea4d..7fb6167a6b14 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/__init__.py @@ -1,8 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# -# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see -# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. """Blackwell backend.""" from .kernel import forward diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/kernel.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/kernel.py index 55067dd109c3..d3edaff3ec14 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/kernel.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/kernel.py @@ -1,8 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# -# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see -# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. """Blackwell kernel entry.""" from .mainloop import forward diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/mainloop.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/mainloop.py index 099e0b375846..24dee1cb0233 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/mainloop.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/mainloop.py @@ -1,12 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# -# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see -# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. -# -# Portions derive from the FlashAttention project -# (https://github.com/Dao-AILab/flash-attention), BSD-3-Clause; its license -# text is vendored at sol_attn/sm100/LICENSE.flash-attention. """Sol-Attn forward kernel for Blackwell SM100. The kernel routes two physical N64 halves at a time and accumulates their exact @@ -15,7 +6,6 @@ """ import math - import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute @@ -30,12 +20,18 @@ from cutlass.cutlass_dsl import T, dsl_user_op from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned -from ..common import layout_utils -from ..common.selector import sol_attn_popc_b32, sol_attn_route_is_exact +from .softmax import ( + _load_m64_n128_score as _load_pair_score, + _online_update_one_half as _online_update_pair, + _rescale_m64_partial_o as _rescale_pair_o, +) from . import math as mma_utils -from .softmax import _load_m64_n128_score as _load_pair_score -from .softmax import _online_update_one_half as _online_update_pair -from .softmax import _rescale_m64_partial_o as _rescale_pair_o + +from ..common import layout_utils +from ..common.selector import ( + sol_attn_popc_b32, + sol_attn_route_is_exact, +) from .tmem import ( _add_physical_tmem_base, _zero_based_tmem_tensor, @@ -43,6 +39,7 @@ tcgen05_wait_st, ) + M = 64 N_MEMBER = 64 N_PACK_HALF = 128 @@ -83,7 +80,6 @@ O_PACKED_WORDS_PER_ROW_PER_THREAD = 16 O_PACKED_COLUMN_STRIDE = 8 - @dsl_user_op def _cvt_bf16x2_f32( hi: Float32, @@ -190,22 +186,34 @@ def _store_pair_probability_chunked_tmemp( tiled_store = tcgen05.make_tmem_copy(store_atom, relative_chunk) thread_store = tiled_store.get_slice(owner_tidx) destination_relative = thread_store.partition_D(relative_chunk) - destination = _add_physical_tmem_base(destination_relative, tmem_base + p_offset) + destination = _add_physical_tmem_base( + destination_relative, tmem_base + p_offset + ) p_store_coordinates = thread_store.partition_S( cute.make_identity_tensor((M, PAIR_P_CHUNK_PACKED_COLUMNS)) ) lane = owner_tidx % Int32(32) for chunk_idx in cutlass.range_constexpr(PAIR_P_CHUNKS): - p_store_registers = cute.make_rmem_tensor(p_store_coordinates.shape, Float32) - assert cute.size(p_store_registers) == PAIR_P_PACKED_REGISTERS_PER_THREAD_PER_CHUNK - assert cute.size(probabilities) == 2 * cute.size(p_store_registers) * PAIR_P_CHUNKS + p_store_registers = cute.make_rmem_tensor( + p_store_coordinates.shape, Float32 + ) + assert ( + cute.size(p_store_registers) + == PAIR_P_PACKED_REGISTERS_PER_THREAD_PER_CHUNK + ) + assert ( + cute.size(probabilities) + == 2 * cute.size(p_store_registers) * PAIR_P_CHUNKS + ) p_store_words = cute.make_tensor( cute.recast_ptr(p_store_registers.iterator, dtype=Int32), p_store_registers.layout, ) probability_base = chunk_idx * (2 * cute.size(p_store_registers)) - for i in cutlass.range(cute.size(p_store_registers), unroll_full=True): + for i in cutlass.range( + cute.size(p_store_registers), unroll_full=True + ): low = probability_base + i * 2 high = low + 1 own = _cvt_bf16x2_f32( @@ -214,12 +222,17 @@ def _store_pair_probability_chunked_tmemp( ) peer = cute.arch.shuffle_sync_bfly(own, offset=2) if (lane & Int32(2)) == Int32(0): - p_store_words[i] = _prmt_b32(own, peer, Int32(0x5410)) + p_store_words[i] = _prmt_b32( + own, peer, Int32(0x5410) + ) else: - p_store_words[i] = _prmt_b32(own, peer, Int32(0x3276)) + p_store_words[i] = _prmt_b32( + own, peer, Int32(0x3276) + ) destination_chunk = cute.make_tensor( - destination.iterator + chunk_idx * PAIR_P_CHUNK_PACKED_COLUMNS, + destination.iterator + + chunk_idx * PAIR_P_CHUNK_PACKED_COLUMNS, destination.layout, ) cute.copy(tiled_store, p_store_registers, destination_chunk) @@ -307,18 +320,32 @@ def _load_pack_v_half( @cute.struct class SharedStorage: q_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2] - pack_k_mbar_ptr: cute.struct.MemRange[cutlass.Int64, PAIR_STAGES * 2] - pack_v_mbar_ptr: cute.struct.MemRange[cutlass.Int64, PAIR_STAGES * 2] + pack_k_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, PAIR_STAGES * 2 + ] + pack_v_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, PAIR_STAGES * 2 + ] pair_score_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2] pair_o_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2] - final_stats: cute.struct.Align[cute.struct.MemRange[Float32, M * 2], 128] - route_partial: cute.struct.Align[cute.struct.MemRange[Float32, 4 * ROUTE_TILE_SIZE], 16] - column_masks: cute.struct.Align[cute.struct.MemRange[Float32, ROUTE_TILE_SIZE], 16] - route_packet: cute.struct.Align[cute.struct.MemRange[Int32, PACKET_WORDS], 16] + final_stats: cute.struct.Align[ + cute.struct.MemRange[Float32, M * 2], 128 + ] + route_partial: cute.struct.Align[ + cute.struct.MemRange[Float32, 4 * ROUTE_TILE_SIZE], 16 + ] + column_masks: cute.struct.Align[ + cute.struct.MemRange[Float32, ROUTE_TILE_SIZE], 16 + ] + route_packet: cute.struct.Align[ + cute.struct.MemRange[Int32, PACKET_WORDS], 16 + ] tmem_holding_buf: Int32 # Owner-warp 0 lane 0 appends both N128 route masks. The full-CTA # pre-exact join publishes the completed list to warp 5; no HBM indices. - route_indices: cute.struct.Align[cute.struct.MemRange[Int32, ROUTE_INDEX_CAPACITY], 16] + route_indices: cute.struct.Align[ + cute.struct.MemRange[Int32, ROUTE_INDEX_CAPACITY], 16 + ] @cute.kernel @@ -363,11 +390,21 @@ def _sol_attn_sm100_bf16_kernel( smem = utils.SmemAllocator() storage = smem.allocate(SharedStorage) - sFinalStats = storage.final_stats.get_tensor(cute.make_layout((M, 2))) - route_partial = storage.route_partial.get_tensor(cute.make_layout((4, ROUTE_TILE_SIZE))) - column_masks = storage.column_masks.get_tensor(cute.make_layout((ROUTE_TILE_SIZE,))) - route_packet = storage.route_packet.get_tensor(cute.make_layout((PACKET_WORDS,))) - route_indices = storage.route_indices.get_tensor(cute.make_layout((ROUTE_INDEX_CAPACITY,))) + sFinalStats = storage.final_stats.get_tensor( + cute.make_layout((M, 2)) + ) + route_partial = storage.route_partial.get_tensor( + cute.make_layout((4, ROUTE_TILE_SIZE)) + ) + column_masks = storage.column_masks.get_tensor( + cute.make_layout((ROUTE_TILE_SIZE,)) + ) + route_packet = storage.route_packet.get_tensor( + cute.make_layout((PACKET_WORDS,)) + ) + route_indices = storage.route_indices.get_tensor( + cute.make_layout((ROUTE_INDEX_CAPACITY,)) + ) sQ = smem.allocate_tensor( element_type=BFloat16, layout=q_layout.outer, @@ -389,11 +426,15 @@ def _sol_attn_sm100_bf16_kernel( # One independent physical N128 K stage and one N128 V stage. Every # runtime route/exact transaction stays in this completion domain. sPackKGather = cute.make_tensor( - cute.recast_ptr(sPackK.iterator, pack_k_gather_layout.inner, BFloat16), + cute.recast_ptr( + sPackK.iterator, pack_k_gather_layout.inner, BFloat16 + ), pack_k_gather_layout.outer, ) sPackVGather = cute.make_tensor( - cute.recast_ptr(sPackV.iterator, pack_v_gather_layout.inner, BFloat16), + cute.recast_ptr( + sPackV.iterator, pack_v_gather_layout.inner, BFloat16 + ), pack_v_gather_layout.outer, ) # KC/VC and exact K/V have disjoint lifetimes within each runtime group. @@ -409,11 +450,21 @@ def _sol_attn_sm100_bf16_kernel( ) tmem_barrier = pipeline.NamedBarrier(barrier_id=1, num_threads=THREADS) - score_loaded_barrier = pipeline.NamedBarrier(barrier_id=2, num_threads=4 * 32) - final_stats_ready_barrier = pipeline.NamedBarrier(barrier_id=3, num_threads=4 * 32) - pack_score_loaded_barrier = pipeline.NamedBarrier(barrier_id=4, num_threads=4 * 32) - route_packet_ready_barrier = pipeline.NamedBarrier(barrier_id=5, num_threads=5 * 32) - exact_pair_p_ready_barrier = pipeline.NamedBarrier(barrier_id=6, num_threads=5 * 32) + score_loaded_barrier = pipeline.NamedBarrier( + barrier_id=2, num_threads=4 * 32 + ) + final_stats_ready_barrier = pipeline.NamedBarrier( + barrier_id=3, num_threads=4 * 32 + ) + pack_score_loaded_barrier = pipeline.NamedBarrier( + barrier_id=4, num_threads=4 * 32 + ) + route_packet_ready_barrier = pipeline.NamedBarrier( + barrier_id=5, num_threads=5 * 32 + ) + exact_pair_p_ready_barrier = pipeline.NamedBarrier( + barrier_id=6, num_threads=5 * 32 + ) tmem = utils.TmemAllocator( storage.tmem_holding_buf.ptr, barrier_for_retrieve=tmem_barrier, @@ -421,12 +472,24 @@ def _sol_attn_sm100_bf16_kernel( tmem.allocate(TMEM_COLS) one_thread = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) - pack_owner_threads = pipeline.CooperativeGroup(pipeline.Agent.Thread, 4 * 32) - q_bytes = cute.size_in_bytes(BFloat16, cute.select(q_layout, mode=[0, 1, 2])) - route_k_bytes = cute.size_in_bytes(BFloat16, cute.select(route_k_layout, mode=[0, 1, 2])) - route_v_bytes = cute.size_in_bytes(BFloat16, cute.select(route_v_layout, mode=[0, 1, 2])) - pack_k_bytes = cute.size_in_bytes(BFloat16, cute.select(pack_k_layout, mode=[0, 1, 2])) - pack_v_bytes = cute.size_in_bytes(BFloat16, cute.select(pack_v_layout, mode=[0, 1, 2])) + pack_owner_threads = pipeline.CooperativeGroup( + pipeline.Agent.Thread, 4 * 32 + ) + q_bytes = cute.size_in_bytes( + BFloat16, cute.select(q_layout, mode=[0, 1, 2]) + ) + route_k_bytes = cute.size_in_bytes( + BFloat16, cute.select(route_k_layout, mode=[0, 1, 2]) + ) + route_v_bytes = cute.size_in_bytes( + BFloat16, cute.select(route_v_layout, mode=[0, 1, 2]) + ) + pack_k_bytes = cute.size_in_bytes( + BFloat16, cute.select(pack_k_layout, mode=[0, 1, 2]) + ) + pack_v_bytes = cute.size_in_bytes( + BFloat16, cute.select(pack_v_layout, mode=[0, 1, 2]) + ) assert route_k_bytes == pack_k_bytes assert route_v_bytes == pack_v_bytes q_pipe = fa_pipeline.PipelineTmaUmma.create( @@ -469,8 +532,12 @@ def _sol_attn_sm100_bf16_kernel( mKC_cur = mKC_nkl[None, None, head_idx, batch_idx] mVC_cur = mVC_nkl[None, None, head_idx, batch_idx] gQ = cute.local_tile(mQ_cur, (M, D), (None, 0)) - gPackK = cute.local_tile(mPackK_cur, (N_MEMBER, 64), (None, None)) - gPackV = cute.local_tile(mPackV_cur, (64, N_MEMBER), (None, None)) + gPackK = cute.local_tile( + mPackK_cur, (N_MEMBER, 64), (None, None) + ) + gPackV = cute.local_tile( + mPackV_cur, (64, N_MEMBER), (None, None) + ) gKC = cute.local_tile(mKC_cur, (N_PACK_HALF, D), (None, 0)) gVC = cute.local_tile(mVC_cur, (DV, N_PACK_HALF), (0, None)) thr_pack_qk = tiled_pack_qk.get_slice(0) @@ -520,7 +587,9 @@ def _sol_attn_sm100_bf16_kernel( cute.group_modes(tCgVC, 0, 3), ) - pack_score_shape = tiled_pack_qk.partition_shape_C(PACK_QK_TILE[:2]) + pack_score_shape = tiled_pack_qk.partition_shape_C( + PACK_QK_TILE[:2] + ) pack_score_template = tiled_pack_qk.make_fragment_C(pack_score_shape) pack_o_shape = tiled_pack_pv.partition_shape_C(PACK_PV_TILE[:2]) pack_o_template = tiled_pack_pv.make_fragment_C(pack_o_shape) @@ -552,27 +621,56 @@ def _sol_attn_sm100_bf16_kernel( # make_fragment_A drops the physical TMEM allocation base and addresses # packed BF16 columns in half-column units. Restore both facts so # 2*tmem_base + 2*PAIR_P_OFFSET names columns 64..127. - pair_tP_storage = cute.make_tensor(pair_tScore.iterator, pack_p_layout.outer) - pair_tP_base = tiled_pack_pv.make_fragment_A(pair_tP_storage)[None, None, None, 0] + pair_tP_storage = cute.make_tensor( + pair_tScore.iterator, pack_p_layout.outer + ) + pair_tP_base = tiled_pack_pv.make_fragment_A(pair_tP_storage)[ + None, None, None, 0 + ] pair_tP = cute.make_tensor( - pair_tP_base.iterator + tmem_base + tmem_base + Int32(PAIR_P_OFFSET * 2), + pair_tP_base.iterator + + tmem_base + + tmem_base + + Int32(PAIR_P_OFFSET * 2), pair_tP_base.layout, ) - q_producer = fa_pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) - q_consumer = fa_pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) - pack_k_producer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, PAIR_STAGES) - pack_k_consumer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, PAIR_STAGES) - pack_v_producer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, PAIR_STAGES) - pack_v_consumer = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, PAIR_STAGES) - pair_score_producer = fa_pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) - pair_score_consumer = fa_pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) - pair_o_producer = fa_pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) - pair_o_consumer = fa_pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) + q_producer = fa_pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, 1 + ) + q_consumer = fa_pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, 1 + ) + pack_k_producer = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, PAIR_STAGES + ) + pack_k_consumer = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, PAIR_STAGES + ) + pack_v_producer = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, PAIR_STAGES + ) + pack_v_consumer = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, PAIR_STAGES + ) + pair_score_producer = fa_pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, 1 + ) + pair_score_consumer = fa_pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, 1 + ) + pair_o_producer = fa_pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, 1 + ) + pair_o_consumer = fa_pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, 1 + ) route_start_base = Int32(0) q_len = token_count - q_block_idx * Int32(M) if q_len > Int32(M): q_len = Int32(M) - threshold = Float32(mThreshold_bnh[batch_idx, q_block_idx, head_idx]) + threshold = Float32( + mThreshold_bnh[batch_idx, q_block_idx, head_idx] + ) if warp_idx == Int32(5): cpasync.prefetch_descriptor(tma_atom_q) @@ -608,30 +706,44 @@ def _sol_attn_sm100_bf16_kernel( # The outer loop owns one logical G256 exact-index lifetime. The inner # loop consumes each physical score/PV half immediately; it appends only # integer indices, never a second score or probability fragment. - num_logical_groups = (num_route_tiles + Int32(ROUTE_HALVES_PER_GROUP - 1)) // Int32( - ROUTE_HALVES_PER_GROUP - ) + num_logical_groups = ( + num_route_tiles + Int32(ROUTE_HALVES_PER_GROUP - 1) + ) // Int32(ROUTE_HALVES_PER_GROUP) # BEGIN_G256_CURSOR_UNIFORM_INDUCTION # arch_make_warp_uniform is a lowering hint, not a value broadcast. Both # values are CTA-invariant integer scalars before the hint. logical_group_idx = cute.arch.make_warp_uniform(Int32(0)) remaining_group_tiles = cute.arch.make_warp_uniform(num_route_tiles) while logical_group_idx < num_logical_groups: - is_final_logical_group = logical_group_idx + Int32(1) == num_logical_groups - group_route_tile_base = logical_group_idx * Int32(ROUTE_HALVES_PER_GROUP) + is_final_logical_group = ( + logical_group_idx + Int32(1) == num_logical_groups + ) + group_route_tile_base = logical_group_idx * Int32( + ROUTE_HALVES_PER_GROUP + ) physical_halves_this_group = remaining_group_tiles if physical_halves_this_group > Int32(ROUTE_HALVES_PER_GROUP): physical_halves_this_group = Int32(ROUTE_HALVES_PER_GROUP) - for half_idx in cutlass.range(physical_halves_this_group, unroll=1): - route_tile_idx = cute.arch.make_warp_uniform(group_route_tile_base + half_idx) - is_final_route_tile = route_tile_idx + Int32(1) == num_route_tiles - is_logical_terminal_half = half_idx + Int32(1) == physical_halves_this_group + for half_idx in cutlass.range( + physical_halves_this_group, unroll=1 + ): + route_tile_idx = cute.arch.make_warp_uniform( + group_route_tile_base + half_idx + ) + is_final_route_tile = ( + route_tile_idx + Int32(1) == num_route_tiles + ) + is_logical_terminal_half = ( + half_idx + Int32(1) == physical_halves_this_group + ) route_start = cute.arch.make_warp_uniform( - route_start_base + route_tile_idx * Int32(ROUTE_TILE_SIZE) + route_start_base + + route_tile_idx * Int32(ROUTE_TILE_SIZE) ) remaining_route_count = cute.arch.make_warp_uniform( - route_valid_total - route_tile_idx * Int32(ROUTE_TILE_SIZE) + route_valid_total + - route_tile_idx * Int32(ROUTE_TILE_SIZE) ) valid_route_count = remaining_route_count if valid_route_count > Int32(ROUTE_TILE_SIZE): @@ -644,7 +756,9 @@ def _sol_attn_sm100_bf16_kernel( # a full-CTA phase boundary, so no K<->V alias handoff is required. if warp_idx == Int32(5): pack_k_pipe.producer_acquire(pack_k_producer) - route_k_barrier = pack_k_pipe.producer_get_barrier(pack_k_producer) + route_k_barrier = pack_k_pipe.producer_get_barrier( + pack_k_producer + ) cute.copy( tma_atom_kc, tKCgKC[(None, route_tile_idx)], @@ -654,7 +768,9 @@ def _sol_attn_sm100_bf16_kernel( pack_k_producer.advance() pack_v_pipe.producer_acquire(pack_v_producer) - route_v_barrier = pack_v_pipe.producer_get_barrier(pack_v_producer) + route_v_barrier = pack_v_pipe.producer_get_barrier( + pack_v_producer + ) cute.copy( tma_atom_vc, tVCgVC[(None, route_tile_idx)], @@ -697,7 +813,9 @@ def _sol_attn_sm100_bf16_kernel( pair_score_consumer.advance() owner_warp = owner_tidx // Int32(32) lane = owner_tidx % Int32(32) - semantic_row = (score_coords[0][0] + Int32(SEMANTIC_ROW_OFFSET)) & Int32(M - 1) + semantic_row = ( + score_coords[0][0] + Int32(SEMANTIC_ROW_OFFSET) + ) & Int32(M - 1) row_valid = semantic_row < q_len lane_col_parity = (lane // Int32(2)) % Int32(2) # Column-pair reduction: parity-0 lanes carry column 2*pair and @@ -708,7 +826,9 @@ def _sol_attn_sm100_bf16_kernel( # sees the same zero-padded operand streams, and the removed # chains only ever accumulated 0.0. Writer lanes 0 and 2 equal # 2*(col%2). - for pair_idx in cutlass.range_constexpr(0, ROUTE_TILE_SIZE // 2, 2): + for pair_idx in cutlass.range_constexpr( + 0, ROUTE_TILE_SIZE // 2, 2 + ): my_col0 = Int32(2 * pair_idx) + lane_col_parity partial0 = Float32(0.0) if row_valid and my_col0 < valid_route_count: @@ -724,34 +844,54 @@ def _sol_attn_sm100_bf16_kernel( (raw_partial0, raw_partial1), (softmax_scale_log2, softmax_scale_log2), ) - peer_scaled0 = cute.arch.shuffle_sync_bfly(scaled0, offset=1) - peer_scaled1 = cute.arch.shuffle_sync_bfly(scaled1, offset=1) + peer_scaled0 = cute.arch.shuffle_sync_bfly( + scaled0, offset=1 + ) + peer_scaled1 = cute.arch.shuffle_sync_bfly( + scaled1, offset=1 + ) partial0, partial1 = cute.arch.fma_packed_f32x2( (raw_partial0, raw_partial1), (softmax_scale_log2, softmax_scale_log2), (peer_scaled0, peer_scaled1), ) - peer0 = cute.arch.shuffle_sync_bfly(partial0, offset=16) - peer1 = cute.arch.shuffle_sync_bfly(partial1, offset=16) + peer0 = cute.arch.shuffle_sync_bfly( + partial0, offset=16 + ) + peer1 = cute.arch.shuffle_sync_bfly( + partial1, offset=16 + ) partial0, partial1 = cute.arch.add_packed_f32x2( (partial0, partial1), (peer0, peer1) ) - peer0 = cute.arch.shuffle_sync_bfly(partial0, offset=8) - peer1 = cute.arch.shuffle_sync_bfly(partial1, offset=8) + peer0 = cute.arch.shuffle_sync_bfly( + partial0, offset=8 + ) + peer1 = cute.arch.shuffle_sync_bfly( + partial1, offset=8 + ) partial0, partial1 = cute.arch.add_packed_f32x2( (partial0, partial1), (peer0, peer1) ) - peer0 = cute.arch.shuffle_sync_bfly(partial0, offset=4) - peer1 = cute.arch.shuffle_sync_bfly(partial1, offset=4) + peer0 = cute.arch.shuffle_sync_bfly( + partial0, offset=4 + ) + peer1 = cute.arch.shuffle_sync_bfly( + partial1, offset=4 + ) partial0, partial1 = cute.arch.add_packed_f32x2( (partial0, partial1), (peer0, peer1) ) if lane == Int32(0): route_partial[owner_warp, 2 * pair_idx] = partial0 - route_partial[owner_warp, 2 * (pair_idx + 1)] = partial1 + route_partial[owner_warp, 2 * (pair_idx + 1)] = ( + partial1 + ) if lane == Int32(2): route_partial[owner_warp, 2 * pair_idx + 1] = partial0 - route_partial[owner_warp, 2 * (pair_idx + 1) + 1] = partial1 + route_partial[ + owner_warp, 2 * (pair_idx + 1) + 1 + ] = partial1 cute.arch.fence_view_async_shared() score_loaded_barrier.arrive_and_wait() @@ -770,7 +910,9 @@ def _sol_attn_sm100_bf16_kernel( # A positive signed shift avoids materializing 1<<31: # lane 0 gets zero and lane 31 gets 0x7fffffff. - lane_mask_lt = Int32(0x7FFFFFFF) >> (Int32(31) - lane) + lane_mask_lt = Int32(0x7FFFFFFF) >> ( + Int32(31) - lane + ) preceding_word_count = Int32(0) for word in cutlass.range_constexpr(ROUTE_MASK_WORDS): off = Int32(word * 32) + lane @@ -793,11 +935,16 @@ def _sol_attn_sm100_bf16_kernel( ) # Sink is a KV-only contract. Text queries remain # a caller-side dense operation in MMDiT models. - exact_pred = exact_pred or ( - route_start + off >= sink_start_block - and route_start + off < sink_end_block + exact_pred = ( + exact_pred + or ( + route_start + off >= sink_start_block + and route_start + off < sink_end_block + ) ) - word_mask = Int32(cute.arch.vote_ballot_sync(exact_pred)) + word_mask = Int32( + cute.arch.vote_ballot_sync(exact_pred) + ) # Site 2: preserve the route decision and its four # ordered ballots, but materialize the resulting # approximate-column mask exactly once. Dedicated @@ -824,7 +971,10 @@ def _sol_attn_sm100_bf16_kernel( mask2 = word_mask else: mask3 = word_mask - preceding_word_count = preceding_word_count + sol_attn_popc_b32(word_mask) + preceding_word_count = ( + preceding_word_count + + sol_attn_popc_b32(word_mask) + ) # Every selected lane has a unique rank; lane 0 publishes # the packet after reconvergence. @@ -861,7 +1011,9 @@ def _sol_attn_sm100_bf16_kernel( # remains available for the later route-mass scratch. route_scores = score_raw assert cute.size(score_raw) % 2 == 0 - for i in cutlass.range_constexpr(0, cute.size(score_raw), 2): + for i in cutlass.range_constexpr( + 0, cute.size(score_raw), 2 + ): group_col0 = score_coords[i][1] group_col1 = score_coords[i + 1][1] mask0 = Float32(column_masks[group_col0]) @@ -879,7 +1031,9 @@ def _sol_attn_sm100_bf16_kernel( route_scores[i] = mask0 route_scores[i + 1] = mask1 - local_max = fa_utils.fmax_reduce(route_scores.load(), arch=100) + local_max = fa_utils.fmax_reduce( + route_scores.load(), arch=100 + ) local_max = Float32(local_max) * softmax_scale peer_max = cute.arch.shuffle_sync_bfly(local_max, offset=2) pair_max = local_max @@ -898,12 +1052,18 @@ def _sol_attn_sm100_bf16_kernel( fastmath=True, ) - route_probabilities = cute.make_rmem_tensor(route_scores.shape, Float32) + route_probabilities = cute.make_rmem_tensor( + route_scores.shape, Float32 + ) if new_max == -Float32.inf: - for i in cutlass.range(cute.size(route_scores), unroll_full=True): + for i in cutlass.range( + cute.size(route_scores), unroll_full=True + ): route_probabilities[i] = Float32(0.0) else: - for i in cutlass.range(cute.size(route_scores), unroll_full=True): + for i in cutlass.range( + cute.size(route_scores), unroll_full=True + ): route_probabilities[i] = cute.math.exp2( Float32(route_scores[i]) * softmax_scale_log2 - new_max * Float32(LOG2E), @@ -916,13 +1076,23 @@ def _sol_attn_sm100_bf16_kernel( # index order, and fadd_reduce preserves floating-point # reduction order and every phase edge. assert cute.size(route_probabilities) % 2 == 0 - for i in cutlass.range_constexpr(0, cute.size(route_probabilities), 2): + for i in cutlass.range_constexpr( + 0, cute.size(route_probabilities), 2 + ): block_idx0 = route_start + score_coords[i][1] - raw_length0 = token_count - block_idx0 * Int32(N_MEMBER) - block_length0 = max(Int32(0), min(raw_length0, Int32(N_MEMBER))) + raw_length0 = ( + token_count - block_idx0 * Int32(N_MEMBER) + ) + block_length0 = max( + Int32(0), min(raw_length0, Int32(N_MEMBER)) + ) block_idx1 = route_start + score_coords[i + 1][1] - raw_length1 = token_count - block_idx1 * Int32(N_MEMBER) - block_length1 = max(Int32(0), min(raw_length1, Int32(N_MEMBER))) + raw_length1 = ( + token_count - block_idx1 * Int32(N_MEMBER) + ) + block_length1 = max( + Int32(0), min(raw_length1, Int32(N_MEMBER)) + ) mass0, mass1 = cute.arch.mul_packed_f32x2( ( Float32(route_probabilities[i]), @@ -935,8 +1105,12 @@ def _sol_attn_sm100_bf16_kernel( ) route_scores[i] = mass0 route_scores[i + 1] = mass1 - current_sum = fa_utils.fadd_reduce(route_scores.load(), arch=100) - current_sum += cute.arch.shuffle_sync_bfly(current_sum, offset=2) + current_sum = fa_utils.fadd_reduce( + route_scores.load(), arch=100 + ) + current_sum += cute.arch.shuffle_sync_bfly( + current_sum, offset=2 + ) # KC is a block mean and VC a valid-token sum. Route mass uses # the true block length while PV still consumes p*VC once. running_sum = old_sum * row_alpha + current_sum @@ -980,14 +1154,20 @@ def _sol_attn_sm100_bf16_kernel( # G256 index stream is nonempty. Those score completions # prove this PV complete; only a final route-only CTA needs # an explicit O completion here. - if is_final_route_tile and Int32(route_packet[6]) == Int32(0): + if ( + is_final_route_tile + and Int32(route_packet[6]) == Int32(0) + ): pair_o_pipe.producer_commit(pair_o_producer) mma_o_initialized = Int32(1) pack_v_pipe.consumer_release(pack_v_consumer) pack_v_consumer.advance() if is_owner: cumulative_exact_count = Int32(route_packet[6]) - if is_final_route_tile and cumulative_exact_count == Int32(0): + if ( + is_final_route_tile + and cumulative_exact_count == Int32(0) + ): pair_o_pipe.consumer_wait(pair_o_consumer) # route_packet may be reused by the next physical half without a @@ -1023,7 +1203,9 @@ def _sol_attn_sm100_bf16_kernel( block1 = Int32(route_indices[ordinal0 + Int32(1)]) pack_k_pipe.producer_acquire(pack_k_producer) - pair_k_barrier = pack_k_pipe.producer_get_barrier(pack_k_producer) + pair_k_barrier = pack_k_pipe.producer_get_barrier( + pack_k_producer + ) _load_pack_k_half( tma_atom_pack_k, tPackKgK, @@ -1036,7 +1218,9 @@ def _sol_attn_sm100_bf16_kernel( pack_k_producer.advance() pack_v_pipe.producer_acquire(pack_v_producer) - pair_v_barrier = pack_v_pipe.producer_get_barrier(pack_v_producer) + pair_v_barrier = pack_v_pipe.producer_get_barrier( + pack_v_producer + ) _load_pack_v_half( tma_atom_pack_v, tPackVgV, @@ -1086,7 +1270,10 @@ def _sol_attn_sm100_bf16_kernel( # QK(i+1) completion dominates PV(i) completion for every # nonterminal transaction on this tcgen05 issuer. Commit one # explicit O-full generation only for the CTA's final PV. - if is_final_logical_group and pair_idx + Int32(1) == pair_count: + if ( + is_final_logical_group + and pair_idx + Int32(1) == pair_count + ): pair_o_pipe.producer_commit(pair_o_producer) mma_o_initialized = Int32(1) pack_v_pipe.consumer_release(pack_v_consumer) @@ -1099,7 +1286,9 @@ def _sol_attn_sm100_bf16_kernel( tiled_pack_qk, pair_tScore, tCrPackQ[None, None, None, q_consumer.index], - tCrPackK[None, None, None, pack_k_consumer.index], + tCrPackK[ + None, None, None, pack_k_consumer.index + ], zero_init=True, ) pair_score_pipe.producer_commit(pair_score_producer) @@ -1160,18 +1349,24 @@ def _sol_attn_sm100_bf16_kernel( pair_score_pipe.consumer_release(pair_score_consumer) pair_score_consumer.advance() - semantic_row = (pair_coords[0][0] + Int32(SEMANTIC_ROW_OFFSET)) & Int32(M - 1) + semantic_row = ( + pair_coords[0][0] + Int32(SEMANTIC_ROW_OFFSET) + ) & Int32(M - 1) row_valid = semantic_row < q_len row_mask = -Float32.inf if row_valid: row_mask = Float32(0.0) assert cute.size(pair_scores) % 2 == 0 - for i in cutlass.range_constexpr(0, cute.size(pair_scores), 2): + for i in cutlass.range_constexpr( + 0, cute.size(pair_scores), 2 + ): column0 = pair_coords[i][1] column1 = pair_coords[i + 1][1] mask0 = Float32(column_masks[column0]) mask1 = Float32(column_masks[column1]) - mask0, mask1 = cute.arch.add_packed_f32x2((mask0, mask1), (row_mask, row_mask)) + mask0, mask1 = cute.arch.add_packed_f32x2( + (mask0, mask1), (row_mask, row_mask) + ) mask0, mask1 = cute.arch.add_packed_f32x2( ( Float32(pair_scores[i]), @@ -1182,11 +1377,13 @@ def _sol_attn_sm100_bf16_kernel( pair_scores[i] = mask0 pair_scores[i + 1] = mask1 - probabilities, next_max, next_sum, row_alpha = _online_update_pair( - pair_scores, - running_max, - running_sum, - softmax_scale, + probabilities, next_max, next_sum, row_alpha = ( + _online_update_pair( + pair_scores, + running_max, + running_sum, + softmax_scale, + ) ) # For i>0, pair-score completion comes from QK(i), issued # after PV(i-1) on the same tcgen05 issuer. The score wait and @@ -1238,7 +1435,9 @@ def _sol_attn_sm100_bf16_kernel( # and the next owner score load cannot precede QK completion. # END_GENERAL_N128_PAIR - logical_group_idx = cute.arch.make_warp_uniform(logical_group_idx + Int32(1)) + logical_group_idx = cute.arch.make_warp_uniform( + logical_group_idx + Int32(1) + ) remaining_group_tiles = cute.arch.make_warp_uniform( remaining_group_tiles - Int32(ROUTE_HALVES_PER_GROUP) ) @@ -1280,46 +1479,69 @@ def _sol_attn_sm100_bf16_kernel( # semantic row and 4*w+2/4*w+3 to its row-plus-eight peer. Hoist # validity, final-sum LDS, reciprocal, and row base once per stratum. semantic_row0 = ( - owner_warp * Int32(16) + lane // Int32(4) + Int32(SEMANTIC_ROW_OFFSET) + owner_warp * Int32(16) + + lane // Int32(4) + + Int32(SEMANTIC_ROW_OFFSET) ) & Int32(M - 1) semantic_row1 = (semantic_row0 + Int32(8)) & Int32(M - 1) even_col_base = (lane % Int32(4)) * Int32(2) if semantic_row0 < q_len: - inv_sum0 = cute.arch.rcp_approx(Float32(sFinalStats[semantic_row0, 0])) + inv_sum0 = cute.arch.rcp_approx( + Float32(sFinalStats[semantic_row0, 0]) + ) query_idx0 = q_block_idx * Int32(M) + semantic_row0 destination_row0 = cute.domain_offset( (batch_idx, query_idx0, head_idx, Int32(0)), mO_bthd ) - for word_i in cutlass.range(O_PACKED_WORDS_PER_ROW_PER_THREAD, unroll_full=True): + for word_i in cutlass.range( + O_PACKED_WORDS_PER_ROW_PER_THREAD, unroll_full=True + ): even_i = word_i * 4 odd_i = even_i + 1 even_value = Float32(o_regs[even_i]) * inv_sum0 odd_value = Float32(o_regs[odd_i]) * inv_sum0 - packed_word = _cvt_bf16x2_f32(Float32(odd_value), Float32(even_value)) - even_col = even_col_base + word_i * O_PACKED_COLUMN_STRIDE - _store_global_u32_inline(destination_row0.iterator + even_col, packed_word) + packed_word = _cvt_bf16x2_f32( + Float32(odd_value), Float32(even_value) + ) + even_col = ( + even_col_base + word_i * O_PACKED_COLUMN_STRIDE + ) + _store_global_u32_inline( + destination_row0.iterator + even_col, packed_word + ) if semantic_row1 < q_len: - inv_sum1 = cute.arch.rcp_approx(Float32(sFinalStats[semantic_row1, 0])) + inv_sum1 = cute.arch.rcp_approx( + Float32(sFinalStats[semantic_row1, 0]) + ) query_idx1 = q_block_idx * Int32(M) + semantic_row1 destination_row1 = cute.domain_offset( (batch_idx, query_idx1, head_idx, Int32(0)), mO_bthd ) - for word_i in cutlass.range(O_PACKED_WORDS_PER_ROW_PER_THREAD, unroll_full=True): + for word_i in cutlass.range( + O_PACKED_WORDS_PER_ROW_PER_THREAD, unroll_full=True + ): even_i = word_i * 4 + 2 odd_i = even_i + 1 even_value = Float32(o_regs[even_i]) * inv_sum1 odd_value = Float32(o_regs[odd_i]) * inv_sum1 - packed_word = _cvt_bf16x2_f32(Float32(odd_value), Float32(even_value)) - even_col = even_col_base + word_i * O_PACKED_COLUMN_STRIDE - _store_global_u32_inline(destination_row1.iterator + even_col, packed_word) + packed_word = _cvt_bf16x2_f32( + Float32(odd_value), Float32(even_value) + ) + even_col = ( + even_col_base + word_i * O_PACKED_COLUMN_STRIDE + ) + _store_global_u32_inline( + destination_row1.iterator + even_col, packed_word + ) if (lane & Int32(2)) == Int32(0) and owner_row < q_len: query_idx = q_block_idx * Int32(M) + owner_row - mLSE_bth[batch_idx, query_idx, head_idx] = running_max + cute.math.log2( - running_sum, fastmath=True - ) * Float32(LN2) + mLSE_bth[batch_idx, query_idx, head_idx] = ( + running_max + + cute.math.log2(running_sum, fastmath=True) * Float32(LN2) + ) cute.arch.barrier() tmem.free(tmem_ptr) @@ -1340,9 +1562,15 @@ def _sol_attn_sm100_bf16_host( sink_end_block: Int32, stream: cuda.CUstream = None, ): - q, k, v, o, kc, vc = tuple(assume_tensor_aligned(t) for t in (q, k, v, o, kc, vc)) - q_mkl, k_nkl, kc_nkl = [layout_utils.select(t, [1, 3, 2, 0]) for t in (q, k, kc)] - v_nkl, vc_nkl = [layout_utils.select(t, [3, 1, 2, 0]) for t in (v, vc)] + q, k, v, o, kc, vc = tuple( + assume_tensor_aligned(t) for t in (q, k, v, o, kc, vc) + ) + q_mkl, k_nkl, kc_nkl = [ + layout_utils.select(t, [1, 3, 2, 0]) for t in (q, k, kc) + ] + v_nkl, vc_nkl = [ + layout_utils.select(t, [3, 1, 2, 0]) for t in (v, vc) + ] token_count = cute.size(q_mkl.shape[0]) num_blocks = cute.size(kc_nkl.shape[0]) num_heads = cute.size(q_mkl.shape[2]) @@ -1388,7 +1616,9 @@ def _sol_attn_sm100_bf16_host( cute.nvgpu.OperandMajorMode.MN, ) tiled_pack_pv_gather = cute.make_tiled_mma(pack_pv_quarter_op) - q_layout = sm100_utils.make_smem_layout_a(tiled_pack_qk, PACK_QK_TILE, BFloat16, 1) + q_layout = sm100_utils.make_smem_layout_a( + tiled_pack_qk, PACK_QK_TILE, BFloat16, 1 + ) pack_k_layout = sm100_utils.make_smem_layout_b( tiled_pack_qk, PACK_QK_TILE, BFloat16, PAIR_STAGES ) @@ -1407,7 +1637,9 @@ def _sol_attn_sm100_bf16_host( BFloat16, PAIR_STAGES * 4, ) - pack_p_layout = sm100_utils.make_smem_layout_a(tiled_pack_pv, PACK_PV_TILE, BFloat16, 1) + pack_p_layout = sm100_utils.make_smem_layout_a( + tiled_pack_pv, PACK_PV_TILE, BFloat16, 1 + ) route_k_layout = sm100_utils.make_smem_layout_b( tiled_pack_qk, PACK_QK_TILE, BFloat16, PAIR_STAGES ) diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/math.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/math.py index 65b0c821c1e6..e1120b09ce3c 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/math.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/math.py @@ -1,8 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# -# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see -# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. """Small tensor-core helpers used by the Blackwell mainloop.""" import cutlass diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/softmax.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/softmax.py index c6d21a0c8877..71b11f0e6ece 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/softmax.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/softmax.py @@ -1,12 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# -# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see -# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. -# -# Portions derive from the FlashAttention project -# (https://github.com/Dao-AILab/flash-attention), BSD-3-Clause; its license -# text is vendored at sol_attn/sm100/LICENSE.flash-attention. """Online-softmax helpers for the Blackwell mainloop.""" from __future__ import annotations @@ -15,9 +6,16 @@ import cutlass.cute as cute from cutlass import Float32, Int32 from cutlass.cute.nvgpu import tcgen05 + from flash_attn.cute import utils as fa_utils -from .tmem import _add_physical_tmem_base, _zero_based_tmem_tensor, tcgen05_wait_ld, tcgen05_wait_st +from .tmem import ( + _add_physical_tmem_base, + _zero_based_tmem_tensor, + tcgen05_wait_ld, + tcgen05_wait_st, +) + M = 64 N_HALF = 128 @@ -43,7 +41,9 @@ def _load_m64_n128_score( tiled_load = tcgen05.make_tmem_copy(load_atom, relative_score) thread_load = tiled_load.get_slice(owner_tidx) source_relative = thread_load.partition_S(relative_score) - source = _add_physical_tmem_base(source_relative, tmem_base + score_offset) + source = _add_physical_tmem_base( + source_relative, tmem_base + score_offset + ) coordinates = thread_load.partition_D( thr_mma_qk.partition_C(cute.make_identity_tensor((M, N_HALF))) ) @@ -67,11 +67,21 @@ def _rescale_m64_partial_o( relative_o = _zero_based_tmem_tensor(Float32, o_template.layout) correction_width = 16 - relative_fragment = cute.composition(relative_o, cute.make_layout((M, correction_width))) - load_atom = cute.make_copy_atom(tcgen05.copy.Ld16x64bOp(tcgen05.copy.Repetition(8)), Float32) - store_atom = cute.make_copy_atom(tcgen05.copy.St16x64bOp(tcgen05.copy.Repetition(8)), Float32) - thread_load = tcgen05.make_tmem_copy(load_atom, relative_fragment).get_slice(owner_tidx) - thread_store = tcgen05.make_tmem_copy(store_atom, relative_fragment).get_slice(owner_tidx) + relative_fragment = cute.composition( + relative_o, cute.make_layout((M, correction_width)) + ) + load_atom = cute.make_copy_atom( + tcgen05.copy.Ld16x64bOp(tcgen05.copy.Repetition(8)), Float32 + ) + store_atom = cute.make_copy_atom( + tcgen05.copy.St16x64bOp(tcgen05.copy.Repetition(8)), Float32 + ) + thread_load = tcgen05.make_tmem_copy( + load_atom, relative_fragment + ).get_slice(owner_tidx) + thread_store = tcgen05.make_tmem_copy( + store_atom, relative_fragment + ).get_slice(owner_tidx) source = _add_physical_tmem_base( thread_load.partition_S(relative_fragment), tmem_base + o_offset ) @@ -79,7 +89,9 @@ def _rescale_m64_partial_o( thread_store.partition_D(relative_fragment), tmem_base + o_offset ) for fragment_idx in cutlass.range_constexpr(DV // correction_width): - registers = cute.make_rmem_tensor(thread_load.partition_D(relative_fragment).shape, Float32) + registers = cute.make_rmem_tensor( + thread_load.partition_D(relative_fragment).shape, Float32 + ) source_i = cute.make_tensor( source.iterator + fragment_idx * correction_width, source.layout ) @@ -117,15 +129,22 @@ def _online_update_one_half( new_max = transaction_max alpha = Float32(0.0) if running_max != -Float32.inf: - alpha = cute.math.exp2((running_max - new_max) * Float32(LOG2E), fastmath=True) + alpha = cute.math.exp2( + (running_max - new_max) * Float32(LOG2E), fastmath=True + ) probabilities = cute.make_rmem_tensor(scores.shape, Float32) for i in cutlass.range(cute.size(scores), unroll_full=True): probabilities[i] = cute.math.exp2( - Float32(scores[i]) * softmax_scale * Float32(LOG2E) - new_max * Float32(LOG2E), + Float32(scores[i]) * softmax_scale * Float32(LOG2E) + - new_max * Float32(LOG2E), fastmath=True, ) - transaction_sum = fa_utils.fadd_reduce(probabilities.load(), arch=100) - transaction_sum += cute.arch.shuffle_sync_bfly(transaction_sum, offset=2) + transaction_sum = fa_utils.fadd_reduce( + probabilities.load(), arch=100 + ) + transaction_sum += cute.arch.shuffle_sync_bfly( + transaction_sum, offset=2 + ) new_sum = running_sum * alpha + transaction_sum return probabilities, new_max, new_sum, alpha diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/tmem.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/tmem.py index 8256aac562b8..aada77bea91d 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/tmem.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/tmem.py @@ -1,8 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# -# Vendored from https://github.com/NVlabs/Sana (Apache-2.0); see -# THIRD_PARTY_NOTICES.md in this directory for the pin and scope. """TMEM load helpers used by the SM100 mainloop.""" from __future__ import annotations @@ -12,6 +7,7 @@ from cutlass import Float32, Int32 from cutlass._mlir.dialects import llvm + M = 64 D = 128 O_OFFSET = 128 @@ -80,7 +76,9 @@ def _o_copy_views( assert o_template.element_type == Float32 assert cute.size(o_template) == M * D relative = _zero_based_tmem_tensor(Float32, o_template.layout) - coordinates = pv_thread.partition_C(cute.make_identity_tensor((M, D))) + coordinates = pv_thread.partition_C( + cute.make_identity_tensor((M, D)) + ) tiler = ( ( cute.size(relative, mode=[0, 0]), @@ -114,7 +112,9 @@ def load_m64_o_fp32_256b( thread_copy.partition_S(relative), physical_tmem_base + Int32(O_OFFSET), ) - register_coordinates = thread_copy.partition_D(coordinates)[None, None, Int32(0)] + register_coordinates = thread_copy.partition_D(coordinates)[ + None, None, Int32(0) + ] registers = cute.make_rmem_tensor( register_coordinates.shape, Float32, From 20b1424a5a0f1a3e2a9d78adc1bb44bc54aa6248 Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:24:00 -0700 Subject: [PATCH 20/22] [TRTLLM-15917][fix] Address review: typed dense_layers, TRTLLM_ strict 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 Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../features/visualgen-sparse-attention.md | 15 +- .../attention_backend/cute_dsl/sol_attn.py | 50 +++---- .../blackwell/sol_attn_backend.py | 129 +++++++----------- tensorrt_llm/visual_gen/sparse_attention.py | 56 ++------ .../test_attention_cute_dsl_sol_attn.py | 84 ++++-------- 5 files changed, 112 insertions(+), 222 deletions(-) diff --git a/docs/source/features/visualgen-sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md index 7e95663dfc76..cb734cb9c888 100644 --- a/docs/source/features/visualgen-sparse-attention.md +++ b/docs/source/features/visualgen-sparse-attention.md @@ -41,7 +41,7 @@ attention_config: tau: 2.0 # routing threshold; higher routes more blocks sparse thresh_type: diag # or "exact" disabled_until_timestep: 0.9090 # dense while normalized timestep >= cutoff - dense_layers: '0' # optional: layers forced dense + dense_layers: [0] # optional: layer indices forced dense ``` `disabled_until_timestep` has the same meaning as it does for Skip Softmax: @@ -49,12 +49,13 @@ attention runs dense while the normalized denoising timestep is at or above the cutoff, protecting the high-noise prefix, and switches to the sparse kernel below it. Use `None` rather than `0.0` to disable the prefix. -On an input the kernel cannot serve — an unsupported architecture, a -`head_dim` other than 128, a non-bfloat16 dtype — Sol-Attn falls back to dense -attention -- the configured backend's dense kernel where available, torch SDPA -otherwise -- logs the specific reason once, and counts the fallback. Set -`SOL_ATTN_STRICT=1` to raise instead of falling back, which is useful when -benchmarking to confirm the kernel actually ran. +On an input the kernel is known not to serve — an unsupported architecture, a +`head_dim` other than 128, a non-bfloat16 dtype — Sol-Attn runs dense +attention instead (the configured backend's dense kernel where available, +torch SDPA otherwise), logs the specific reason once, and counts the fallback. +Set `TRTLLM_SOL_ATTN_STRICT=1` to raise instead, which is useful when +benchmarking to confirm the kernel actually ran. Errors raised by the kernel +itself are not caught. ## Skip Softmax Attention diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py index fe5d7d1ad884..631a4a5756b6 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py @@ -84,27 +84,14 @@ def _cute_dense_available() -> bool: return True -def _parse_dense_layers(spec: Optional[str]) -> frozenset[int]: - layers: set[int] = set() - for item in str(spec or "").split(","): - item = item.strip() - if not item: - continue - if "-" in item: - start, end = item.split("-", 1) - layers.update(range(int(start), int(end) + 1)) - else: - layers.add(int(item)) - return frozenset(layers) - - class SolAttention(AttentionBackend): """Sol-Attn dynamic block-routing sparse attention (CuTeDSL, sm100/sm103). - The kernel wrapper already falls back to dense attention on any unsupported - shape/dtype/arch (see ``_run_sol_attn_bthd``); this class only adds the - ``dense_layers`` layer-skip guard (evaluated at construction time, no - external plumbing needed) and forwards the routing knobs from config. + Inputs the kernel cannot serve (shape, dtype, architecture) are delegated to + dense attention up front by ``_run_sol_attn_bthd``; a kernel error is not + recovered from and propagates. This class adds the ``dense_layers`` guard + (evaluated at construction time, no external plumbing needed) and forwards + the routing knobs from config. """ def __init__( @@ -140,7 +127,7 @@ def __init__( self.tau = getattr(cfg, "tau", 1.0) self.thresh_type = getattr(cfg, "thresh_type", "diag") self.disabled_until_timestep = getattr(cfg, "disabled_until_timestep", None) - self.dense_layers = _parse_dense_layers(getattr(cfg, "dense_layers", None)) + self.dense_layers = frozenset(getattr(cfg, "dense_layers", None) or ()) # Sol-Attn's dense steps must run the backend the user selected. Without # this they ran torch SDPA while a `backend: CUTEDSL` baseline ran @@ -167,13 +154,15 @@ def __init__( dtype=dtype, ) # Whether the CuTe DSL dense kernel can serve this device, decided once - # here. Doing it at construction (rather than lazily on the first call) - # keeps `_dense` free of attribute mutation, so it stays traceable and - # the dense step sits in the same place in the graph as the dense - # CUTEDSL baseline's does. Deciding it lazily and marking `_dense` - # `@torch.compiler.disable` instead moved the whole dense step out of - # the graph and reintroduced the very mismatch this is meant to remove: - # measured LPIPS 0.4044 compiled, against 0.2112 eager. + # here rather than lazily on the first call. Neither a kernel bug nor a + # torch.compile bug is involved; this is about graph structure. A lazy + # decision needs attribute mutation inside `_dense`, which forces + # `@torch.compiler.disable` on it, and that moves the whole dense step + # out of the compiled graph. The dense CUTEDSL baseline keeps its dense + # step inside the graph, so the two then differ on every dense step: + # measured LPIPS 0.4044 compiled against 0.2112 eager. Deciding here + # keeps `_dense` traceable and the dense step in the same place as the + # baseline's. self._cute_dense_ok = _cute_dense_available() if not self._cute_dense_ok: logger.warning_once( @@ -274,8 +263,13 @@ def _can_serve(self, q: torch.Tensor, k: torch.Tensor, **kwargs: Any) -> bool: Qwen-Image always, and WAN's ``attn1`` under async Ulysses -- silently costing those modules their configured backend. """ - # Cross-attention: K/V come from another sequence. Sol-Attn's routing - # assumes one self-attending sequence. + # Cross-attention: K/V come from another sequence, and Sol-Attn's + # routing assumes one self-attending sequence. Unequal Q/K lengths are + # a heuristic for that, not a definition: a cross-attention call whose + # context happens to match the query length is not caught here. The + # `Attention` module knows the answer (`encoder_hidden_states`), but the + # backend `forward` kwargs carry no such flag yet; TRTLLM-16475 tracks + # threading an explicit signal through and retiring this check. if k.shape[1] != q.shape[1]: return False # Masks: the sparse kernel is noncausal and takes no mask, so a masked diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py index 0dce5e086afb..829a3ce462f5 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py @@ -27,8 +27,11 @@ countable rather than only visible on stderr. * ``sol_attn_ineligible_reason()`` names the specific reason (architecture, head_dim, dtype) instead of returning one boolean. -* ``SOL_ATTN_STRICT=1`` also covers the eligibility path; upstream raises only - on kernel exceptions, so an ineligible run stayed silent. +* Kernel exceptions propagate. Upstream catches them and silently reruns the + call with torch SDPA; a failed CuTe launch can leave the device in a bad + state, so recovery is not attempted here. Only inputs known up front to be + unservable (shape, dtype, architecture) are routed to dense attention, and + ``TRTLLM_SOL_ATTN_STRICT=1`` turns even that into an error. * Dense paths route to ``cute_dsl_fmha_fwd`` through ``dense_fn``. Upstream falls back to torch SDPA; staying inside the configured backend is what lets a ``backend: CUTEDSL`` A/B isolate sparsity rather than also swapping the @@ -47,10 +50,11 @@ lives there too, keyed off the normalized timestep forward kwarg. CuTe DSL imports and compilation are deferred to first use. Calls the kernel -cannot serve -- wrong shape, dtype, or an architecture with no kernel -- -delegate to dense attention rather than failing, and increment -``_SOL_STATS["dense_fallback_calls"]`` so the degradation is countable. Set -``SOL_ATTN_STRICT=1`` to raise instead of falling back. +is known not to serve -- wrong shape, dtype, or an architecture with no kernel +-- are delegated to dense attention before any launch, logged once, and +counted in ``_SOL_STATS["dense_fallback_calls"]``. Set +``TRTLLM_SOL_ATTN_STRICT=1`` to raise instead. Errors raised by the kernel +itself are not caught. """ from __future__ import annotations @@ -61,6 +65,7 @@ import torch +from tensorrt_llm._utils import get_sm_version from tensorrt_llm.logger import logger HEAD_DIM = 128 @@ -83,10 +88,10 @@ def _load_sol_attn() -> Callable: return sol_attn -# Architectures with a Sol-Attn CuTe kernel. Kept in sync with -# ``sol_attn/interface.py::_CUTE_BACKENDS``; duplicated here so the eligibility -# check does not have to import the CuTe DSL. -SUPPORTED_ARCHS = frozenset({(10, 0), (10, 3)}) +# SM versions with a Sol-Attn CuTe kernel, in ``get_sm_version()`` form +# (major * 10 + minor). Kept in sync with ``sol_attn/interface.py::_CUTE_BACKENDS``; +# duplicated here so the eligibility check does not have to import the CuTe DSL. +SUPPORTED_ARCHS = frozenset({100, 103}) def sol_attn_ineligible_reason(q) -> Optional[str]: @@ -108,13 +113,10 @@ def sol_attn_ineligible_reason(q) -> Optional[str]: return f"head_dim must be {HEAD_DIM}, got {q.shape[-1]}" if q.dtype != torch.bfloat16: return f"dtype must be bfloat16, got {q.dtype}" - try: - arch = tuple(torch.cuda.get_device_capability(q.device)) - except Exception as exc: - return f"could not query device capability: {exc}" - if arch not in SUPPORTED_ARCHS: - return f"no Sol-Attn kernel for SM{arch[0]}{arch[1]}; supported: " + ", ".join( - f"SM{a}{b}" for a, b in sorted(SUPPORTED_ARCHS) + sm = get_sm_version() + if sm not in SUPPORTED_ARCHS: + return f"no Sol-Attn kernel for SM{sm}; supported: " + ", ".join( + f"SM{v}" for v in sorted(SUPPORTED_ARCHS) ) return None @@ -138,9 +140,9 @@ def _cute_runtime_available() -> bool: def _strict() -> bool: - """Whether SOL_ATTN_STRICT=1 asks us to raise instead of degrading.""" + """Whether TRTLLM_SOL_ATTN_STRICT=1 asks us to raise on an unservable input.""" - return os.environ.get("SOL_ATTN_STRICT", "0") == "1" + return os.environ.get("TRTLLM_SOL_ATTN_STRICT", "0") == "1" def _dense_bthd(q, k, v): @@ -151,36 +153,6 @@ def _dense_bthd(q, k, v): ).transpose(1, 2) -def _degradable_kernel_errors() -> tuple[type[BaseException], ...]: - """Exception types whose failure is safe to answer with dense attention. - - `CODING_GUIDELINES.md` asks for the smallest exception set. A bare - ``except Exception`` also swallowed ordinary programming and integration - errors -- ``TypeError``, ``NameError``, ``AttributeError`` -- turning a bug - into "Sol-Attn just didn't speed anything up". Those now propagate. - - ``DSLBaseError`` has to be named explicitly because CuTe DSL derives it - from ``Exception``, not ``RuntimeError``, so a narrower tuple would stop - catching real JIT and codegen failures that were previously degraded. It - is resolved lazily and cached: this module defers every CuTe DSL import to - first use, and the tuple is only needed once something has already raised. - """ - global _DEGRADABLE_KERNEL_ERRORS - if _DEGRADABLE_KERNEL_ERRORS is None: - types: list[type[BaseException]] = [ImportError, OSError, RuntimeError] - try: - from cutlass.base_dsl.common import DSLBaseError - except ImportError: - pass - else: - types.append(DSLBaseError) - _DEGRADABLE_KERNEL_ERRORS = tuple(types) - return _DEGRADABLE_KERNEL_ERRORS - - -_DEGRADABLE_KERNEL_ERRORS: tuple[type[BaseException], ...] | None = None - - # Opaque to Dynamo, like every other CuTe DSL launch boundary here (see # cute_dsl/fmha.py, video_sparse_attention/interface.py). Otherwise Dynamo # traces into the CuTe DSL JIT builder and retraces on every call: near two @@ -207,7 +179,11 @@ def _run_sol_attn_bthd( sink_tokens: int = 0, dense_fn: Callable | None = None, ): - """Run Sol-Attn on contiguous BTHD tensors, with a safe dense fallback.""" + """Run Sol-Attn on contiguous BTHD tensors. + + Inputs the kernel is known not to serve go to dense attention up front; + errors raised by the kernel itself propagate. + """ q0, k0, v0 = q.contiguous(), k.contiguous(), v.contiguous() @@ -223,52 +199,39 @@ def dense(): if reason is None and (k0.dtype != q0.dtype or v0.dtype != q0.dtype): reason = f"k/v dtype must match q {q0.dtype}" if reason is not None: - # Same strictness contract as the kernel-exception path below: this is - # the arm that silently turns Sol-Attn into a no-op for a whole run - # (wrong arch, head_dim, or dtype), so it must be visible. + # This is the arm that silently turns Sol-Attn into a no-op for a whole + # run (wrong arch, head_dim, or dtype), so it must be visible. if _strict(): raise RuntimeError(f"[sol-attn] cannot run the CuTe kernel: {reason}") logger.warning_once( f"[sol-attn] falling back to dense attention: {reason}. Sol-Attn will not " - "accelerate this run. Set SOL_ATTN_STRICT=1 to raise instead.", + "accelerate this run. Set TRTLLM_SOL_ATTN_STRICT=1 to raise instead.", key=("sol_attn_ineligible", reason), ) return dense() - try: - kernel = _load_sol_attn() - out = kernel( - q0, - k0, - v0, - tau=float(tau), - thresh_type=str(thresh_type), - # Only 1 split exists on the shipped sm100/sm103 kernels (2/4 was an - # SM90-only path), so this is not a user-facing knob. The kernel - # interface rejects anything else before the try below can swallow it. - kv_splits=int(kv_splits), - sink_start=sink_start, - sink_tokens=int(sink_tokens), - ) - except _degradable_kernel_errors() as exc: - if _strict(): - raise - logger.warning_once( - f"[sol-attn] kernel raised {type(exc).__name__}: {exc}; falling back to dense " - "attention for this call. Set SOL_ATTN_STRICT=1 to raise instead of silently falling " - "back.", - key=(type(exc).__name__, str(exc)), - ) - return dense() - + kernel = _load_sol_attn() + out = kernel( + q0, + k0, + v0, + tau=float(tau), + thresh_type=str(thresh_type), + # Only 1 split exists on the shipped sm100/sm103 kernels (2/4 was an + # SM90-only path), so this is not a user-facing knob; the kernel + # interface rejects anything else. + kv_splits=int(kv_splits), + sink_start=sink_start, + sink_tokens=int(sink_tokens), + ) _SOL_STATS["kernel_calls"] += 1 return out # Lightweight run-validation counters. `kernel_calls` is the census used to -# prove the CuTe kernel actually ran: because forward() falls back to dense -# SDPA on any kernel exception, a run that "works" but never increments this -# was silently dense. Set SOL_ATTN_STRICT=1 to raise instead of falling back. +# prove the CuTe kernel actually ran; `dense_fallback_calls` counts calls the +# kernel was known not to serve. Set TRTLLM_SOL_ATTN_STRICT=1 to raise instead +# of falling back on those. _SOL_STATS = {"kernel_calls": 0, "dense_fallback_calls": 0} diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 40716b120bef..bd4c22e22395 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -265,58 +265,24 @@ class SolAttentionConfig(BaseSparseAttentionConfig): "and runs sparse on every step (fail-open)." ), ) - dense_layers: Optional[str] = PydanticField( + dense_layers: Optional[list[int]] = PydanticField( None, description=( - "Comma-separated layer indices/ranges (e.g. '0,2-4') forced dense " - "regardless of the dense prefix. Evaluated per-layer at construction " - "time; no pipeline wiring required." + "Layer indices forced dense regardless of the dense prefix, e.g. " + "[0, 2, 3]. Evaluated per layer at construction time; no pipeline " + "wiring required." ), ) @field_validator("dense_layers") @classmethod - def _validate_dense_layers(cls, spec: Optional[str]) -> Optional[str]: - """Reject malformed specs here rather than deep in the backend. - - Without this a non-numeric token raises from ``_parse_dense_layers`` - during attention construction, far from the config that caused it, and - a descending range such as ``'4-2'`` raises nothing at all -- it yields - an empty set, so the layers the user asked to force dense silently stay - sparse. - """ - if spec is None: - return spec - for item in spec.split(","): - item = item.strip() - if not item: - # Not skipped: "," / "0,,2" / " " would otherwise be accepted - # and quietly force fewer layers dense than the user wrote. - raise ValueError( - f"dense_layers {spec!r} has an empty entry; expected a " - "comma-separated list such as '0,2-4'" - ) - try: - if "-" in item: - # A negative index cannot reach here: it also contains '-', - # so it takes this branch and the empty first part fails - # int() below. - start, end = (int(part) for part in item.split("-", 1)) - if start > end: - raise ValueError( - f"dense_layers range '{item}' is descending; " - f"write it as '{end}-{start}'" - ) - else: - int(item) - except ValueError as exc: - if "descending" in str(exc): - raise - raise ValueError( - f"dense_layers entry '{item}' is not a layer index or range; " - "expected a comma-separated list such as '0,2-4'" - ) from exc - return spec + def _validate_dense_layers(cls, layers: Optional[list[int]]) -> Optional[list[int]]: + if layers is None: + return None + for index in layers: + if index < 0: + raise ValueError(f"dense_layers contains a negative layer index: {index}") + return sorted(set(layers)) def to_sparse_params(self, **kwargs): # Sol-Attn's knobs are consumed directly by SolAttention.__init__ diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py index f328e14ae471..0da640d23d1e 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py @@ -32,10 +32,7 @@ from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxScheduler from tensorrt_llm._torch.visual_gen.attention_backend import CuTeDSLAttention -from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn import ( - SolAttention, - _parse_dense_layers, -) +from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn import SolAttention from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention from tensorrt_llm._torch.visual_gen.config import ( DiffusionModelConfig, @@ -196,23 +193,6 @@ def test_sol_attn_rejects_gqa_mqa(): SolAttention(layer_idx=0, num_heads=8, head_dim=128, num_kv_heads=2) -@pytest.mark.parametrize( - "spec,expected", - [ - (None, frozenset()), - ("", frozenset()), - ("0", frozenset({0})), - ("0,2,4", frozenset({0, 2, 4})), - ("0-3", frozenset({0, 1, 2, 3})), - ("0-1,5,7-8", frozenset({0, 1, 5, 7, 8})), - (" 0 , 2 ", frozenset({0, 2})), - ], - ids=["none", "empty", "single", "list", "range", "mixed", "whitespace"], -) -def test_parse_dense_layers(spec, expected): - assert _parse_dense_layers(spec) == expected - - @pytest.mark.parametrize( "timestep,expected", [ @@ -369,7 +349,7 @@ def test_cuda_graph_key_separates_dense_prefix_from_sparse_phase(): def test_cuda_graph_key_unregistered_without_prefix(): """dense_layers alone is fixed per layer, so it needs no graph key.""" - model = _make_solattn_model(disabled_until_timestep=None, dense_layers="0,2") + model = _make_solattn_model(disabled_until_timestep=None, dense_layers=[0, 2]) runner = _graph_runner() model.register_cuda_graph_extra_key_fns(runner) @@ -423,11 +403,11 @@ def test_ineligible_reason_is_reported(make, expect): def test_strict_raises_on_ineligible_input(monkeypatch): - """SOL_ATTN_STRICT=1 must cover the shape/dtype/arch path, not just kernel - exceptions. Without this, an unsupported arch degrades to dense silently - even under STRICT, and the counters the PR relies on cannot be trusted.""" + """TRTLLM_SOL_ATTN_STRICT=1 must turn an unservable input into an error. + Without it an unsupported arch degrades to dense silently and the counters + the PR relies on cannot be trusted.""" sab = _backend_mod() - monkeypatch.setenv("SOL_ATTN_STRICT", "1") + monkeypatch.setenv("TRTLLM_SOL_ATTN_STRICT", "1") q = k = v = torch.randn(1, 4, 2, 128) # CPU -> ineligible with pytest.raises(RuntimeError, match="cannot run the CuTe kernel"): sab._run_sol_attn_bthd(q, k, v) @@ -436,7 +416,7 @@ def test_strict_raises_on_ineligible_input(monkeypatch): def test_ineligible_falls_back_to_dense_and_counts(monkeypatch): """Without STRICT the same input degrades to dense and increments the counter.""" sab = _backend_mod() - monkeypatch.delenv("SOL_ATTN_STRICT", raising=False) + monkeypatch.delenv("TRTLLM_SOL_ATTN_STRICT", raising=False) sab.reset_sol_attn_stats() q = k = v = torch.randn(1, 4, 2, 128) out = sab._run_sol_attn_bthd(q, k, v) @@ -453,7 +433,9 @@ def test_supported_archs_matches_kernel_dispatch_map(): drift, eligibility silently rejects an arch the kernel actually supports.""" from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.sol_attn import interface - assert _backend_mod().SUPPORTED_ARCHS == frozenset(interface._CUTE_BACKENDS) + assert _backend_mod().SUPPORTED_ARCHS == frozenset( + major * 10 + minor for major, minor in interface._CUTE_BACKENDS + ) def test_datacenter_blackwell_archs_are_supported(): @@ -537,40 +519,24 @@ def test_zero_cutoff_rejected(): @pytest.mark.parametrize( - "spec,reason", - [ - ("4-2", "descending"), - ("abc", "not a layer index"), - ("0,x", "not a layer index"), - ("-1", "not a layer index"), - (",", "empty entry"), - ("0,,2", "empty entry"), - (" ", "empty entry"), - ], - ids=[ - "descending_range", - "non_numeric", - "non_numeric_in_list", - "negative", - "only_separator", - "empty_in_list", - "whitespace_only", - ], + "layers", + [[-1], [0, -2], ["x"], "0,2", [1.5]], + ids=["negative", "negative_in_list", "non_numeric", "string_spec", "float"], ) -def test_dense_layers_rejects_malformed_spec(spec, reason): - """Malformed dense_layers must fail at config time, not silently or late. - - A descending range is the dangerous one: `_parse_dense_layers('4-2')` - yields an empty set, so the layers the user asked to force dense would - quietly stay sparse with no error anywhere. - """ - with pytest.raises(ValueError, match=reason): - SolAttentionConfig(tau=2.0, dense_layers=spec) +def test_dense_layers_rejects_invalid_values(layers): + """dense_layers is a list of non-negative layer indices; anything else fails at + config time rather than at attention construction.""" + with pytest.raises(ValueError): + SolAttentionConfig(tau=2.0, dense_layers=layers) -@pytest.mark.parametrize("spec", [None, "0", "0,2-4", " 0 , 2 ", "0-0"]) -def test_dense_layers_accepts_valid_spec(spec): - assert SolAttentionConfig(tau=2.0, dense_layers=spec).dense_layers == spec +@pytest.mark.parametrize( + "layers,expected", + [(None, None), ([0], [0]), ([0, 2, 4], [0, 2, 4]), ([4, 2, 2, 0], [0, 2, 4])], + ids=["none", "single", "list", "unsorted_with_duplicates"], +) +def test_dense_layers_normalizes_valid_values(layers, expected): + assert SolAttentionConfig(tau=2.0, dense_layers=layers).dense_layers == expected @pytest.mark.skipif(not torch.cuda.is_available(), reason="Sol-Attn needs CUDA") From bf3652af719da162590f680e78b9903433f2b7e8 Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Wed, 16 Sep 2026 22:04:46 -0700 Subject: [PATCH 21/22] [TRTLLM-15917][chore] Clarify dense-backend initialization rationale 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> --- .../attention_backend/cute_dsl/sol_attn.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py index 631a4a5756b6..5d0fea9a1645 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py @@ -153,16 +153,8 @@ def __init__( num_kv_heads=self.num_kv_heads, dtype=dtype, ) - # Whether the CuTe DSL dense kernel can serve this device, decided once - # here rather than lazily on the first call. Neither a kernel bug nor a - # torch.compile bug is involved; this is about graph structure. A lazy - # decision needs attribute mutation inside `_dense`, which forces - # `@torch.compiler.disable` on it, and that moves the whole dense step - # out of the compiled graph. The dense CUTEDSL baseline keeps its dense - # step inside the graph, so the two then differ on every dense step: - # measured LPIPS 0.4044 compiled against 0.2112 eager. Deciding here - # keeps `_dense` traceable and the dense step in the same place as the - # baseline's. + # Resolve availability once at construction to keep capability checks + # and attribute initialization out of `_dense` during torch.compile tracing. self._cute_dense_ok = _cute_dense_available() if not self._cute_dense_ok: logger.warning_once( From b901f29f323b2aebbd3617347a8c9ed0a5970fcf Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:22:12 -0700 Subject: [PATCH 22/22] [TRTLLM-15917][chore] Align the SEPARATE_QKV fallback with main's cross-attention flag Rebase follow-up: main now gates the fallback on separate_qkv_cross_attention (the separate_qkv_is_self_attention flag from #18147); keep that condition and the Sol-Attn note together. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- tensorrt_llm/_torch/visual_gen/modules/attention.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index 2e995898c93c..196289d8c18f 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -104,8 +104,7 @@ def __init__( self.qkv_mode == QKVMode.SEPARATE_QKV and not separate_qkv_is_self_attention ) - - # SEPARATE_QKV fallback: TRTLLM and CUTEDSL VSA cannot serve it. + # Cross-attention fallback: TRTLLM and CUTEDSL VSA are self-attn only. # Sol-Attn is absent by design; see SolAttention._can_serve. if separate_qkv_cross_attention and (base_backend == "TRTLLM" or _is_vsa): backend_name = "VANILLA"