From 60ef12ad792f413e2a113c183524b488f5fcba60 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 1/9] [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> --- .../visual-gen/features/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 | 28 +- 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, 5363 insertions(+), 46 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/visual-gen/features/sparse-attention.md b/docs/source/visual-gen/features/sparse-attention.md index b4620cdc0618..4dd8c33e1026 100644 --- a/docs/source/visual-gen/features/sparse-attention.md +++ b/docs/source/visual-gen/features/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 @@ -90,7 +121,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. #### Python API diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py index 6b5d9b538d8b..72a3648fdd9a 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py @@ -23,6 +23,7 @@ from .cute_dsl import ( VSA_TILE_SIZE, CuTeDSLAttention, + SolAttnAttention, VSAAttention, VSAMetadata, VSAMetadataBuilder, @@ -44,6 +45,7 @@ "create_attention", "CuTeDSLAttention", "VSAAttention", + "SolAttnAttention", "FlashAttn4Attention", "TrtllmAttention", "TrtllmAttentionMetadata", 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 12108ad84ece..f28224d49a44 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py @@ -127,14 +127,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 d549cfce3d35..554e464294fd 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_backend.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 8c19484229c7..5d773d0f4ac4 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -97,16 +97,22 @@ 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" - # Cross-attention fallback: TRTLLM and CUTEDSL VSA are self-attn only. - if self.qkv_mode == QKVMode.SEPARATE_QKV and (base_backend == "TRTLLM" or _is_vsa): + # Cross-attention fallback: TRTLLM and CUTEDSL VSA/Sol-Attn are self-attn only. + if self.qkv_mode == QKVMode.SEPARATE_QKV 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( @@ -123,6 +129,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 71d13d91bb17..bda9ad2f894f 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, @@ -75,6 +76,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", @@ -130,6 +132,7 @@ def __dir__(): "SparseAttentionConfig", "SkipSoftmaxAttentionConfig", "VideoSparseAttentionConfig", + "SolAttnAttentionConfig", "CacheConfig", "TeaCacheConfig", "CacheDiTConfig", diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index f7d483389dd0..560e09dbe7e9 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -30,7 +30,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 @@ -89,7 +93,7 @@ class QuantAttentionConfig(StrictBaseModel): # Discriminated union of sparse attention configs. SparseAttentionConfig = Annotated[ - Union[SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig], + Union[SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig, SolAttnAttentionConfig], Field(discriminator="algorithm"), ] @@ -116,7 +120,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)." ), ) @@ -180,6 +185,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 @@ -195,19 +201,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 @@ -781,6 +791,7 @@ def from_yaml(cls, yaml_path: Union[str, Path], **overrides: Any) -> "VisualGenA "SparseAttentionConfig", "SkipSoftmaxAttentionConfig", "VideoSparseAttentionConfig", + "SolAttnAttentionConfig", "AttentionConfig", "ParallelConfig", "BaseCacheConfig", diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index d261c770ca9e..583fb38e4313 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 a3e8504ad21f..1cecdeaf47e0 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -252,6 +252,7 @@ l0_b200: - unittest/_torch/visual_gen/test_pertoken_adaln.py - unittest/_torch/visual_gen/test_attention_cute_dsl.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_trtllm_sage.py - unittest/_torch/visual_gen/test_attention_integration.py - unittest/_torch/visual_gen/test_attention_fa4.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 8bf2f15ab98661175e4b4f90cd9a40113fab324d 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 2/9] [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 01b82f054c00dd2deb0a4d23365b36364f2d279d 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 3/9] [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 4af47b65943cbcee8c14fc567543c7777d69a88f 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 4/9] [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> --- .../visual-gen/features/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/visual-gen/features/sparse-attention.md b/docs/source/visual-gen/features/sparse-attention.md index 4dd8c33e1026..7b9e1b8b1fae 100644 --- a/docs/source/visual-gen/features/sparse-attention.md +++ b/docs/source/visual-gen/features/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 583fb38e4313..d3d46356f307 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 f8d61561dc0e3aef58c6c6f95f7a643f08378ac0 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 5/9] [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 5d773d0f4ac4..d0ec54c8b2eb 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -102,10 +102,19 @@ def __init__( _is_sol_attn = base_backend == "CUTEDSL" and _sa_algo == "sol_attn" # 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 self.qkv_mode == QKVMode.SEPARATE_QKV 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 @@ -255,7 +264,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 117c2a9efc75526b6f1227db6320988ba0322f6c 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 6/9] [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> --- .../visual-gen/features/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/visual-gen/features/sparse-attention.md b/docs/source/visual-gen/features/sparse-attention.md index 7b9e1b8b1fae..71e715d190a0 100644 --- a/docs/source/visual-gen/features/sparse-attention.md +++ b/docs/source/visual-gen/features/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 554e464294fd..8be2fd007b99 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 d0ec54c8b2eb..b1a8acd33b86 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -101,16 +101,25 @@ def __init__( _is_vsa = base_backend == "CUTEDSL" and _sa_algo == "vsa" _is_sol_attn = base_backend == "CUTEDSL" and _sa_algo == "sol_attn" - # 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 self.qkv_mode == QKVMode.SEPARATE_QKV and ( base_backend == "TRTLLM" or _is_vsa or _is_sol_attn ): @@ -268,6 +277,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 d3d46356f307..4708a5eb2332 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 7efdc2856df92177a8d0a00ab74e4b9f4da05d3c 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 7/9] [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> --- .../visual-gen/features/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/visual-gen/features/sparse-attention.md b/docs/source/visual-gen/features/sparse-attention.md index 71e715d190a0..9ad32eefc6ec 100644 --- a/docs/source/visual-gen/features/sparse-attention.md +++ b/docs/source/visual-gen/features/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 72a3648fdd9a..31323f9687e4 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py @@ -23,7 +23,7 @@ from .cute_dsl import ( VSA_TILE_SIZE, CuTeDSLAttention, - SolAttnAttention, + SolAttention, VSAAttention, VSAMetadata, VSAMetadataBuilder, @@ -45,7 +45,7 @@ "create_attention", "CuTeDSLAttention", "VSAAttention", - "SolAttnAttention", + "SolAttention", "FlashAttn4Attention", "TrtllmAttention", "TrtllmAttentionMetadata", 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 f28224d49a44..9111b43815da 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py @@ -134,9 +134,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 8be2fd007b99..e1a884952ad8 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_backend.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 bda9ad2f894f..03e00bf8264f 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, @@ -76,7 +76,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", @@ -132,7 +132,7 @@ def __dir__(): "SparseAttentionConfig", "SkipSoftmaxAttentionConfig", "VideoSparseAttentionConfig", - "SolAttnAttentionConfig", + "SolAttentionConfig", "CacheConfig", "TeaCacheConfig", "CacheDiTConfig", diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index 560e09dbe7e9..58b9435fbc12 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -32,7 +32,7 @@ from .sparse_attention import ( SkipSoftmaxAttentionConfig, - SolAttnAttentionConfig, + SolAttentionConfig, VideoSparseAttentionConfig, ) @@ -93,7 +93,7 @@ class QuantAttentionConfig(StrictBaseModel): # Discriminated union of sparse attention configs. SparseAttentionConfig = Annotated[ - Union[SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig, SolAttnAttentionConfig], + Union[SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig, SolAttentionConfig], Field(discriminator="algorithm"), ] @@ -791,7 +791,7 @@ def from_yaml(cls, yaml_path: Union[str, Path], **overrides: Any) -> "VisualGenA "SparseAttentionConfig", "SkipSoftmaxAttentionConfig", "VideoSparseAttentionConfig", - "SolAttnAttentionConfig", + "SolAttentionConfig", "AttentionConfig", "ParallelConfig", "BaseCacheConfig", diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 4708a5eb2332..adba8a30dff5 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 b25ca70edf4d4267ff4916cad177897e4d8fb804 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 8/9] [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 b1a8acd33b86..01d73e74aefd 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -101,36 +101,18 @@ def __init__( _is_vsa = base_backend == "CUTEDSL" and _sa_algo == "vsa" _is_sol_attn = base_backend == "CUTEDSL" and _sa_algo == "sol_attn" - # 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 self.qkv_mode == QKVMode.SEPARATE_QKV 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 self.qkv_mode == QKVMode.SEPARATE_QKV 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( @@ -273,14 +255,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 dd788c75205e4d348db095adb56547323b51b11b 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 9/9] [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> --- .../visual-gen/features/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/visual-gen/features/sparse-attention.md b/docs/source/visual-gen/features/sparse-attention.md index 9ad32eefc6ec..85004bc6f3f8 100644 --- a/docs/source/visual-gen/features/sparse-attention.md +++ b/docs/source/visual-gen/features/sparse-attention.md @@ -6,6 +6,7 @@ This page is an unindexed draft until the VisualGen documentation hub is introdu - [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 adba8a30dff5..6bb650fe4615 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 "