diff --git a/docs/source/visual-gen/features/sparse-attention.md b/docs/source/visual-gen/features/sparse-attention.md index b4620cdc0618..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) @@ -21,6 +22,38 @@ Sparse attention in VisualGen is configured through `VisualGenArgs.attention_con |---|---|---| | `skip_softmax` | `SkipSoftmaxAttentionConfig` | Supported | | VSA | TBD | TODO | +| `sol_attn` | `SolAttentionConfig` | Supported (CUTEDSL, sm100) | + +### Sol-Attn + +Sol-Attn ([arXiv:2607.24027](https://arxiv.org/abs/2607.24027)) folds dynamic block +routing, sparse computation, and an approximation-correction term into one +online-softmax pass. It runs on the **CUTEDSL** backend only, on sm100 +(B200/GB200), and requires `head_dim=128`, bfloat16, and MHA +(`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.9090 # 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 +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. ## Skip Softmax Attention @@ -90,7 +123,7 @@ User configuration is supplied through Python or YAML and controls how the check `threshold_scale_factor` and `target_sparsity` are alternatives: if both are present, `threshold_scale_factor` takes precedence and the calibration formula is not used. User-provided `target_sparsity` and `disabled_until_timestep` override checkpoint defaults. Checkpoint `ignore` patterns always disable Skip Softmax Attention for matching layers. -Skip Softmax Attention works with both the **TRTLLM** and **CUTEDSL** attention backends in VisualGen. Set `attention_config.backend` to either when enabling it. On CUTEDSL, Skip Softmax Attention can also be combined with `quant_attention_config`'s block-scaled Q/K recipes (MXFP8, NVFP4); VSA is the only CUTEDSL sparse-attention algorithm that is mutually exclusive with quantized attention. +Skip Softmax Attention works with both the **TRTLLM** and **CUTEDSL** attention backends in VisualGen. Set `attention_config.backend` to either when enabling it. On CUTEDSL, Skip Softmax Attention can also be combined with `quant_attention_config`'s block-scaled Q/K recipes (MXFP8, NVFP4); VSA and Sol-Attn each replace the dense CuTeDSL path and are therefore mutually exclusive with quantized attention. #### 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..31323f9687e4 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, + SolAttention, VSAAttention, VSAMetadata, VSAMetadataBuilder, @@ -44,6 +45,7 @@ "create_attention", "CuTeDSLAttention", "VSAAttention", + "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 9b70421c3b81..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 @@ -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 — SolAttention (Sol-Attn dynamic block routing, CuTe JIT + SDPA fallback) """ from .fmha import CuTeDSLAttention, _cute_dsl_import_error +from .sol_attn import SolAttention, 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", + "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 new file mode 100644 index 000000000000..048eb88fb24e --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py @@ -0,0 +1,304 @@ +# 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) +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 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 +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 _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[int]: + layers: set[int] = set() + for item in str(spec or "").split(","): + item = item.strip() + if not item: + continue + if "-" in item: + start, end = item.split("-", 1) + layers.update(range(int(start), int(end) + 1)) + else: + layers.add(int(item)) + return frozenset(layers) + + +class SolAttention(AttentionBackend): + """Sol-Attn dynamic block-routing sparse attention (CuTeDSL, sm100). + + 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. + """ + + 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( + "SolAttention 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 + 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) + 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)) + + # 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._inner = 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 + # phases still compile as separate graphs -- they run different kernels. + @torch.compiler.disable + def _dense_by_step(self, timestep: Any) -> bool: + phase = sol_attn_graph_phase( + 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( + "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 " + "unset disabled_until_timestep.", + key="sol_attn_missing_timestep", + ) + 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 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 _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 + 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._inner.forward(q, k, v, **kwargs) + return self._sdpa(q, k, v) + + def _can_serve(self, q: torch.Tensor, k: torch.Tensor, **kwargs: Any) -> bool: + """Whether the sparse kernel applies to this particular call. + + Everything false here is delegated to ``_inner``. Deciding it from the + tensors, per call, is deliberate: ``qkv_mode`` describes how Q/K/V are + *projected*, not whether K/V come from another sequence, so a + construction-time rule keyed on ``SEPARATE_QKV`` mistakes self-attention + for cross-attention wherever that mode is chosen for other reasons -- + Qwen-Image always, and WAN's ``attn1`` under async Ulysses. + """ + # 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, + k: torch.Tensor, + v: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + """q, k, v: [B, S, H, D] (NHD), same original token order in and out.""" + if not self._can_serve(q, k, **kwargs): + return self._delegate(q, k, v, **kwargs) + return _sol_attn_run( + q, + k, + v, + tau=self.tau, + thresh_type=self.thresh_type, + kv_splits=self.kv_splits, + # Shape/dtype/arch ineligibility is only detectable inside the + # wrapper, so that last delegation happens through this hook. + dense_fn=lambda a, b, c: self._delegate(a, b, c, **kwargs), + ) + + @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..9111b43815da 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 SolAttention + + attn_cls = SolAttention + 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..2b0ba8011fb4 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md @@ -0,0 +1,101 @@ +# 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 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. + +## A derived file outside this directory + +`../sol_attn_backend.py` is **not** part of the vendored package above, but it +is a derivative work and is recorded here because this file is where a future +currency check starts. It is adapted from upstream's +`techniques/sparse_backends/sol_attn_backend.py` (same branch and commit as the +package). Only the kernel-wrapper subset is carried -- the shape/dtype guard, +the dense fallback (routed to cute_dsl_fmha_fwd here, not torch SDPA), and the +call counters. Upstream's model-integration +half is not carried: the diffusers self-attention dispatch hook, HunyuanVideo's +padded `[video, text]` MMDiT handling, and model-level Morton ordering. + +Deliberate divergences from upstream in that file, all of which a re-sync must +preserve rather than overwrite: + +| Divergence | Why | +|---|---| +| `logger.warning_once` replaces `print()` | fallbacks must be suppressible and routed through the repo's logger | +| `dense_fallback_calls` counter added | makes a silently-degraded run countable, not just visible in stderr | +| `sol_attn_ineligible_reason()` added | names the specific reason (arch / head_dim / dtype) instead of one boolean | +| `SOL_ATTN_STRICT=1` also covers the eligibility path | upstream raises only on kernel exceptions, so an ineligible run stayed silent | +| `@torch.compiler.disable` on `_run_sol_attn_bthd` | see below | +| dense paths routed to `cute_dsl_fmha_fwd` via `dense_fn` | upstream's dense fallback is torch SDPA; staying in-backend is what makes a `backend: CUTEDSL` A/B isolate sparsity | + +**Upstream solves the `torch.compile` problem differently, and arguably +better.** Its `sol_attn_backend.py` wraps the same call in a +`torch.library.custom_op` (`sana_sol_attn::self_attention`) with a +`register_fake` returning `torch.empty_like(q)`, which keeps the kernel in the +compiled graph as an opaque node instead of breaking the graph at it; a second +consumer (`models/ltx2.5-refiner/GB200/sol_attention.py`) applies +`torch.compiler.disable` at the call site behind a flag. This repository uses +`@torch.compiler.disable` on the launch boundary instead, matching the +convention every other CuTe DSL entry point here already follows +(`attention_backend/cute_dsl/fmha.py`, +`cute_dsl_kernels/blackwell/video_sparse_attention/interface.py`). Without some +such guard Dynamo traces into the CuTe DSL JIT builder and retraces on every +call -- measured at near two orders of magnitude slower on B200. Migrating to +the `custom_op` form would +remove the per-layer graph break and is a reasonable follow-up; it was not done +here because the `torch.compiler.disable` form is what this repository's other +kernels use and what the measurements above were taken with. + +The runtime also depends on NVIDIA CUTLASS / CuTe DSL, cuda-python, and +PyTorch. Those dependencies are not redistributed by this repository and +remain subject to their respective licenses. + +SM120 (RTX Blackwell) was carried in an earlier revision of this port and has +been dropped: it had kernel-level evidence only, no end-to-end validation, and +no `cute_dsl_fmha_fwd` exists for that architecture, so Sol-Attn's dense +fallback could not match its own backend there. 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/__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..d9644c1fd8bb --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/interface.py @@ -0,0 +1,284 @@ +# 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 +} +_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 " + "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 _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): + # 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, + ) + 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 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_backend.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py new file mode 100644 index 000000000000..1f38ad39d356 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py @@ -0,0 +1,237 @@ +# 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 +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. + +TRT-LLM's dispatch path (``attention_backend/cute_dsl/sol_attn.py``, +``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. + +CuTe DSL imports and compilation are deferred to first use. Calls the kernel +cannot serve -- wrong shape, dtype, or an architecture with no kernel -- +delegate to dense attention rather than failing, and increment +``_SOL_STATS["dense_fallback_calls"]`` so the degradation is countable. Set +``SOL_ATTN_STRICT=1`` to raise instead of falling back. +""" + +from __future__ import annotations + +import functools +import os +from typing import Callable, Optional + +import torch + +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)}) + + +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 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): + return torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + ).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: 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, + 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 attention: {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 " + "attention 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..e1a884952ad8 100644 --- a/tensorrt_llm/_torch/visual_gen/models/modeling.py +++ b/tensorrt_llm/_torch/visual_gen/models/modeling.py @@ -20,8 +20,9 @@ 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, SolAttentionConfig if TYPE_CHECKING: from tensorrt_llm._torch.visual_gen.cuda_graph_runner import CUDAGraphRunner @@ -74,23 +75,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, 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 + # already baked into each captured graph and needs no key. + return + + # 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( + "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..01d73e74aefd 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -97,13 +97,19 @@ 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" - ) - - # Cross-attention fallback: TRTLLM and CUTEDSL VSA are self-attn only. + _sa_algo = getattr(_sa_cfg, "algorithm", None) if _sa_cfg is not None else None + _is_vsa = base_backend == "CUTEDSL" and _sa_algo == "vsa" + _is_sol_attn = base_backend == "CUTEDSL" and _sa_algo == "sol_attn" + + # SEPARATE_QKV fallback: TRTLLM and CUTEDSL VSA cannot serve it. + # + # Sol-Attn is deliberately absent: it decides per call (`_can_serve`) and + # delegates what it cannot serve to the dense backend of its own family. + # A rule here would have to guess from `qkv_mode`, which describes how + # Q/K/V are *projected* rather than whether K/V come from another + # sequence -- and that guess is wrong wherever SEPARATE_QKV is chosen for + # other reasons (Qwen-Image always; WAN's attn1 under async Ulysses), + # silently costing those modules their configured backend. 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 @@ -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..03e00bf8264f 100644 --- a/tensorrt_llm/visual_gen/__init__.py +++ b/tensorrt_llm/visual_gen/__init__.py @@ -49,6 +49,7 @@ QuantAttentionConfig, RuntimeLoRAConfig, SkipSoftmaxAttentionConfig, + SolAttentionConfig, 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", + "SolAttentionConfig": "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", + "SolAttentionConfig", "CacheConfig", "TeaCacheConfig", "CacheDiTConfig", diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index f7d483389dd0..58b9435fbc12 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, + SolAttentionConfig, + VideoSparseAttentionConfig, +) # ============================================================================= # Type aliases @@ -89,7 +93,7 @@ class QuantAttentionConfig(StrictBaseModel): # Discriminated union of sparse attention configs. SparseAttentionConfig = Annotated[ - Union[SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig], + Union[SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig, SolAttentionConfig], 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", + "SolAttentionConfig", "AttentionConfig", "ParallelConfig", "BaseCacheConfig", diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index d261c770ca9e..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 @@ -221,6 +222,112 @@ def _ckpt_sparse_attention_config_from_kwargs( return None +class SolAttentionConfig(BaseSparseAttentionConfig): + """Sol-Attn sparse attention configuration for visual generation. + + Dynamic block routing + sparse computation + approximation correction in + one online-softmax pass (arXiv:2607.24027). Kernel is CuTeDSL, sm100 + (B200/GB200) only, head_dim=128, bf16, MHA. + + On an unsupported *shape, dtype, or architecture* the kernel falls back to + 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. + """ + + 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 " + "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." + ), + ) + + @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 + # 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). 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/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..43799383c56a --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py @@ -0,0 +1,574 @@ +# 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 ( + SolAttention, + _parse_dense_layers, + sol_attn_graph_phase, +) +from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention +from tensorrt_llm._torch.visual_gen.config import ( + DiffusionModelConfig, + create_attention_metadata_state, +) +from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode +from tensorrt_llm.visual_gen.args import AttentionConfig, SolAttentionConfig + + +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 = 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", + layer_idx=0, + num_heads=8, + head_dim=128, + attention_config=sol_attn_config, + ) + + assert isinstance(dense_attention, CuTeDSLAttention) + assert isinstance(sol_attn_attention, SolAttention) + 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 = ( + SolAttentionConfig(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_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). + """ + 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 == "CUTEDSL", ( + f"expected CUTEDSL, got {cross_attn.attn_backend!r}" + ) + 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__}" + ) + # 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. + """ + attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) + attn.disabled_until_timestep = None + # 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" + + +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=SolAttentionConfig(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(ValueError, match="MHA-only"): + SolAttention(layer_idx=0, num_heads=8, head_dim=128, num_kv_heads=2) + + +@pytest.mark.parametrize( + "spec,expected", + [ + (None, frozenset()), + ("", frozenset()), + ("0", frozenset({0})), + ("0,2,4", frozenset({0, 2, 4})), + ("0-3", frozenset({0, 1, 2, 3})), + ("0-1,5,7-8", frozenset({0, 1, 5, 7, 8})), + (" 0 , 2 ", frozenset({0, 2})), + ], + ids=["none", "empty", "single", "list", "range", "mixed", "whitespace"], +) +def test_parse_dense_layers(spec, expected): + assert _parse_dense_layers(spec) == expected + + +@pytest.mark.parametrize( + "timestep,expected", + [ + (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_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): + raise AssertionError("kernel must not run inside the dense prefix") + + monkeypatch.setattr(sol_attn_mod, "_sol_attn_run", _fail_if_called) + + 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)) + 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 = 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 + assert called["n"] == 1, "expected the sparse kernel, not a silent dense fallback" + + +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 + + 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 = 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) + 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=SolAttentionConfig( + 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=SolAttentionConfig(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): + SolAttentionConfig(tau=2.0, disabled_until_timestep=0.0) + 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 " + "(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): + SolAttentionConfig(tau=2.0, kv_splits="4") + assert SolAttentionConfig(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: 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" + ) + + +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(SolAttention._dense_by_step), ( + "SolAttention._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): + """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 + 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 = SolAttention(layer_idx=0, num_heads=2, head_dim=128) + calls = {"n": 0} + real = a._inner.forward + + def _spy(*args, **kwargs): + calls["n"] += 1 + return real(*args, **kwargs) + + monkeypatch.setattr(a._inner, "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 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 was not delegated to 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"