diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 6ac5ee01c3b0..5dcceeea3d4c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1525,7 +1525,7 @@ legacy-files: &legacy_files | # list; the hook's own `files:` pattern only gates *when* the hook triggers. # Global exclude: vendored code + trtllm-gen FMHA artifacts (cubin pointers, export headers, cuda_ptx) -exclude: '(^cpp/tensorrt_llm/common/sha256/|^triton_kernels/|^tensorrt_llm/_torch/attention/backends/prims_ts/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$|trtllmGenKernels/gemm/trtllmGen_gemm_export/KernelMetaInfo\.h$|\.cubin\.tar\.zst$)' +exclude: '(^cpp/tensorrt_llm/common/sha256/|^triton_kernels/|^tensorrt_llm/_torch/attention/backends/prims_ts/|^tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$|trtllmGenKernels/gemm/trtllmGen_gemm_export/KernelMetaInfo\.h$|\.cubin\.tar\.zst$)' default_install_hook_types: [pre-commit, commit-msg] repos: diff --git a/3rdparty/vendor_patches/sana-sol-attn.patch b/3rdparty/vendor_patches/sana-sol-attn.patch new file mode 100644 index 000000000000..5751db4112ea --- /dev/null +++ b/3rdparty/vendor_patches/sana-sol-attn.patch @@ -0,0 +1,474 @@ +diff --git a/sol_attn/THIRD_PARTY_NOTICES.md b/sol_attn/THIRD_PARTY_NOTICES.md +index d3755083d1e7622ff509351594157dd879865e56..b226f77f96bc2316e902ddf91765b35c915a8a6c 100644 +--- a/sol_attn/THIRD_PARTY_NOTICES.md ++++ b/sol_attn/THIRD_PARTY_NOTICES.md +@@ -1,15 +1,49 @@ + # Third-party notices + +-The files under `sol_attn/_vendor/flash_attn/cute/` and portions of the SM89, +-SM90, and SM100 design scaffold derive from the FlashAttention project. Its +-BSD-3-Clause license is included at +-`sol_attn/sm100/LICENSE.flash-attention`. ++This package is vendored from ++[`github.com/NVlabs/Sana`](https://github.com/NVlabs/Sana), branch ++[`sol-engine`](https://github.com/NVlabs/Sana/tree/sol-engine), at commit ++[`5fe5feb`](https://github.com/NVlabs/Sana/commit/5fe5feb) (2026-08-17). ++Checked against the branch tip ++([`83e54df`](https://github.com/NVlabs/Sana/commit/83e54df), 2026-08-20) on ++2026-08-27. + +-The runtime also depends on NVIDIA CUTLASS / CuTe DSL, cuda-python, PyTorch, +-and Triton. Those dependencies are not redistributed by this repository and +-remain subject to their respective licenses. ++## Scope of the vendored subset ++ ++Only the pieces needed for the architectures TensorRT-LLM ships are carried: ++ ++| Carried | Not carried | ++|---|---| ++| `interface.py`, `preprocess.py`, `common/` | `sm89/`, `sm90/` (incl. `sm90/_compat/`) | ++| `sm100/` — serves SM100 (B200/GB200) and SM103 (B300/GB300) | `triton_ref/` Triton reference attention | ++| | `sm120/` (RTX Blackwell) | ++| | `_vendor/flash_attn/` (see below) | ++ ++`../sol_attn_backend.py` is not part of the vendored package but is a ++derivative work of upstream's ++`techniques/sparse_backends/sol_attn_backend.py` (same branch and commit). ++Only the kernel-wrapper subset is carried; upstream's model-integration half ++(diffusers dispatch hook, HunyuanVideo MMDiT padding, model-level Morton ++ordering) is not. ++ ++Implementation divergences from upstream are documented in the source itself, ++in the module docstrings of `../sol_attn_backend.py` and ++`attention_backend/cute_dsl/sol_attn.py`. + +-The SM120 warp-MMA/TMA execution skeleton and online-softmax helpers are +-adapted from NVIDIA cuDNN Frontend's block-sparse-attention reference at commit +-`74785165de2da954a2c879a5e3e6f95411c2292d`. That source is licensed under the +-Apache License 2.0; adapted files retain the corresponding SPDX header. ++## Licensing ++ ++The upstream package vendored a copy of FlashAttention's CuTe DSL helpers ++under `sol_attn/_vendor/flash_attn/cute/`. That copy is **not** carried here: ++TensorRT-LLM already depends on ++[`flash-attn-4`](https://github.com/Dao-AILab/flash-attention) (pinned in ++`requirements.txt`), which provides the same `flash_attn.cute` modules, and ++the SM100 kernels import them from that dependency directly. FlashAttention's ++BSD-3-Clause license is retained at `sol_attn/sm100/LICENSE.flash-attention` ++because portions of the SM100 design scaffold still derive from that project. ++ ++`preprocess.py` implements the routing/threshold stage in Triton, so Triton is ++a required runtime dependency on every Sol-Attn path, not only a fallback. ++ ++The runtime also depends on NVIDIA CUTLASS / CuTe DSL, cuda-python, and ++PyTorch. Those dependencies are not redistributed by this repository and ++remain subject to their respective licenses. +diff --git a/sol_attn/interface.py b/sol_attn/interface.py +index 43e9267b931c4f87620739182e596f1fe63b2475..2da4ea383947f5cd25bea439d7748d8ab5e4f28c 100644 +--- a/sol_attn/interface.py ++++ b/sol_attn/interface.py +@@ -7,11 +7,12 @@ import functools + import torch + + BLOCK_SIZE = 64 ++# TensorRT-LLM ships the SM100 kernel only. It serves both datacenter Blackwell ++# steppings: the CuTe DSL JIT targets the device it compiles on and the kernel ++# body uses no SM100-exclusive construct (measured identical on B200 and B300). + _CUTE_BACKENDS = { +- (8, 9): "cute_sm89", +- (9, 0): "cute_sm90", +- (10, 0): "cute_sm100", +- (12, 0): "cute_sm120", ++ (10, 0): "cute_sm100", # B200 / GB200 ++ (10, 3): "cute_sm100", # B300 / GB300 + } + _compiled = {} + +@@ -68,23 +69,30 @@ def _backend_for_arch( + *, + cute_available: bool | None = None, + ) -> str: +- """Select CuTe when specialized and available, otherwise Triton.""" ++ """Select the CuTe kernel for ``arch``, or raise if there is none. + +- if arch[0] < 8: ++ TensorRT-LLM raises instead of falling back to the Triton reference so a ++ missing kernel is visible; the caller decides what to do about it. ++ """ ++ ++ cute_backend = _CUTE_BACKENDS.get(arch) ++ if cute_backend is None: + raise RuntimeError( +- "Sol-Attn requires an NVIDIA GPU with compute capability >= 8.0; " +- f"got SM{arch[0]}{arch[1]}" ++ f"Sol-Attn has no kernel for SM{arch[0]}{arch[1]}; supported " ++ f"architectures are " ++ f"{', '.join(f'SM{a}{b}' for a, b in sorted(_CUTE_BACKENDS))}." + ) +- cute_backend = _CUTE_BACKENDS.get(arch) +- if cute_backend is not None: +- available = ( +- _cute_runtime_available() +- if cute_available is None +- else cute_available ++ available = ( ++ _cute_runtime_available() ++ if cute_available is None ++ else cute_available ++ ) ++ if not available: ++ raise RuntimeError( ++ "Sol-Attn requires the CuTe DSL runtime (cutlass.cute and " ++ "cuda.bindings.driver); neither could be imported." + ) +- if available: +- return cute_backend +- return "triton" ++ return cute_backend + + + def get_sol_attn_backend(device: torch.device | str | int | None = None) -> str: +@@ -96,8 +104,11 @@ def get_sol_attn_backend(device: torch.device | str | int | None = None) -> str: + + + def _validate_cute(arch, tokens, kv_splits): +- if arch != (9, 0) and kv_splits != 1: +- raise ValueError("kv_splits=2/4 is currently available on SM90 only") ++ if kv_splits != 1: ++ raise ValueError( ++ "kv_splits=2/4 was an SM90-only path; this build ships the SM100 " ++ "kernel only, so kv_splits must be 1." ++ ) + route_groups = ((tokens + 63) // 64 + 63) // 64 + if kv_splits > route_groups: + raise ValueError("each KV split must contain at least one N64 route group") +@@ -126,60 +137,6 @@ def _sink_block_range(tokens, sink_start, sink_tokens): + ) + + +-def _compile_sm90( +- key, +- tensors, +- scale, +- tokens, +- kv_splits, +- sink_range, +- stream, +-): +- import cutlass.cute as cute +- +- from .sm90 import make_kernel +- +- operator = make_kernel(tokens, kv_splits) +- args = _to_cute_tensors(tensors) +- compiled = cute.compile( +- operator, +- *args, +- scale, +- sink_range, +- stream=stream, +- options="--enable-tvm-ffi", +- ) +- _compiled[key] = compiled +- return compiled, args +- +- +-def _compile_sm89( +- key, +- tensors, +- scale, +- sink_start_block, +- sink_end_block, +- stream, +-): +- import cutlass.cute as cute +- +- from .sm89 import make_kernel +- +- operator = make_kernel() +- args = _to_cute_tensors(tensors) +- compiled = cute.compile( +- operator, +- *args, +- scale, +- sink_start_block, +- sink_end_block, +- stream=stream, +- options="--enable-tvm-ffi", +- ) +- _compiled[key] = compiled +- return compiled, args +- +- + def _compile_sm100( + key, + tensors, +@@ -206,33 +163,6 @@ def _compile_sm100( + return compiled, args + + +-def _compile_sm120( +- key, +- tensors, +- scale, +- sink_start_block, +- sink_end_block, +- stream, +-): +- import cutlass.cute as cute +- +- from .sm120 import make_kernel +- +- operator = make_kernel() +- args = _to_cute_tensors(tensors) +- compiled = cute.compile( +- operator, +- *args, +- scale, +- sink_start_block, +- sink_end_block, +- stream=stream, +- options="--enable-tvm-ffi", +- ) +- _compiled[key] = compiled +- return compiled, args +- +- + def _sol_attn_cute( + q, + k, +@@ -268,129 +198,36 @@ def _sol_attn_cute( + stream = _stream(q.device) + key = (q.device.index, arch, batch, tokens, heads, kv_splits) + +- if arch == (8, 9): +- sink_start_block, sink_end_block = _sink_block_range( +- tokens, +- sink_start, +- sink_tokens, +- ) +- tensors = [q, k, v, output, kc, vc, threshold, lse] +- compiled = _compiled.get(key) +- if compiled is None: +- compiled, args = _compile_sm89( +- key, +- tensors, +- scale, +- sink_start_block, +- sink_end_block, +- stream, +- ) +- else: +- args = _to_cute_tensors(tensors) +- compiled( +- *args, +- scale, +- sink_start_block, +- sink_end_block, +- stream=stream, +- ) +- elif arch == (9, 0): +- if sink_tokens: +- sink_start_block, sink_end_block = _sink_block_range( +- tokens, +- sink_start, +- sink_tokens, +- ) +- sink_range = sink_start_block | (sink_end_block << 16) +- else: +- sink_range = 0 +- tensors = [q, k, v, output, kc, vc, threshold, lse] +- if kv_splits > 1: +- tensors.extend( +- [ +- torch.empty( +- (batch, tokens, kv_splits * heads, 128), +- device=q.device, +- dtype=torch.bfloat16, +- ), +- torch.empty( +- (batch, tokens, kv_splits * heads), +- device=q.device, +- dtype=torch.float32, +- ), +- ] +- ) +- compiled = _compiled.get(key) +- if compiled is None: +- compiled, args = _compile_sm90( +- key, +- tensors, +- scale, +- tokens, +- kv_splits, +- sink_range, +- stream, +- ) +- else: +- args = _to_cute_tensors(tensors) +- compiled( +- *args, +- scale, +- sink_range, +- stream=stream, +- ) +- elif arch == (10, 0): +- sink_start_block, sink_end_block = _sink_block_range( +- tokens, +- sink_start, +- sink_tokens, +- ) +- tensors = [q, k, v, output, kc, vc, threshold, lse] +- compiled = _compiled.get(key) +- if compiled is None: +- compiled, args = _compile_sm100( +- key, +- tensors, +- scale, +- sink_start_block, +- sink_end_block, +- stream, +- ) +- else: +- args = _to_cute_tensors(tensors) +- compiled( +- *args, ++ if arch not in _CUTE_BACKENDS: ++ # Unreachable via sol_attn(): _backend_for_arch raises first. Kept ++ # explicit so a missed guard cannot return the uninitialised ++ # `output` buffer. Keyed off _CUTE_BACKENDS, not a literal. ++ raise ValueError(f"no Sol-Attn CuTe kernel for SM{arch[0]}{arch[1]}") ++ sink_start_block, sink_end_block = _sink_block_range( ++ tokens, ++ sink_start, ++ sink_tokens, ++ ) ++ tensors = [q, k, v, output, kc, vc, threshold, lse] ++ compiled = _compiled.get(key) ++ if compiled is None: ++ compiled, args = _compile_sm100( ++ key, ++ tensors, + scale, + sink_start_block, + sink_end_block, +- stream=stream, ++ stream, + ) + else: +- sink_start_block, sink_end_block = _sink_block_range( +- tokens, +- sink_start, +- sink_tokens, +- ) +- tensors = [q, k, v, output, kc, vc, threshold, lse] +- compiled = _compiled.get(key) +- if compiled is None: +- compiled, args = _compile_sm120( +- key, +- tensors, +- scale, +- sink_start_block, +- sink_end_block, +- stream, +- ) +- else: +- args = _to_cute_tensors(tensors) +- compiled( +- *args, +- scale, +- sink_start_block, +- sink_end_block, +- stream=stream, +- ) ++ args = _to_cute_tensors(tensors) ++ compiled( ++ *args, ++ scale, ++ sink_start_block, ++ sink_end_block, ++ stream=stream, ++ ) + return output + + +@@ -421,28 +258,15 @@ def sol_attn( + sink_tokens, + sink_start, + ) +- if kv_splits not in (1, 2, 4): +- raise ValueError("kv_splits must be 1, 2, or 4") +- backend = _backend_for_arch(arch) ++ if kv_splits != 1: ++ raise ValueError( ++ "kv_splits must be 1; the 2/4 path was SM90-only and this build " ++ "ships SM100 kernels only." ++ ) ++ _backend_for_arch(arch) # raises on an architecture with no kernel + scale = q.shape[-1] ** -0.5 if scale is None else float(scale) + tau = float(tau) + +- if backend == "triton": +- if kv_splits != 1: +- raise ValueError("kv_splits=2/4 is currently available on SM90 only") +- from .triton_ref import sol_attn as triton_sol_attn +- +- return triton_sol_attn( +- q, +- k, +- v, +- scale=scale, +- tau=tau, +- thresh_type=thresh_type, +- sink_tokens=sink_tokens, +- sink_start=sink_start, +- ) +- + _validate_cute(arch, q.shape[1], kv_splits) + return _sol_attn_cute( + q, +diff --git a/sol_attn/sm100/mainloop.py b/sol_attn/sm100/mainloop.py +index 36b3ad9d1450c836f7acdc80bbed1382eac4d3a9..24dee1cb0233ebc5ecaabed21eab3c7fe12e86ba 100644 +--- a/sol_attn/sm100/mainloop.py ++++ b/sol_attn/sm100/mainloop.py +@@ -12,13 +12,13 @@ import cutlass.cute as cute + import cutlass.pipeline as pipeline + import cutlass.utils as utils + import cutlass.utils.blackwell_helpers as sm100_utils +-import sol_attn._vendor.flash_attn.cute.pipeline as fa_pipeline +-import sol_attn._vendor.flash_attn.cute.utils as fa_utils ++import flash_attn.cute.pipeline as fa_pipeline ++import flash_attn.cute.utils as fa_utils + from cutlass import BFloat16, Float32, Int32 + from cutlass._mlir.dialects import llvm + from cutlass.cute.nvgpu import cpasync, tcgen05 + from cutlass.cutlass_dsl import T, dsl_user_op +-from sol_attn._vendor.flash_attn.cute.cute_dsl_utils import assume_tensor_aligned ++from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned + + from .softmax import ( + _load_m64_n128_score as _load_pair_score, +@@ -27,8 +27,8 @@ from .softmax import ( + ) + from . import math as mma_utils + +-from sol_attn.common import layout_utils +-from sol_attn.common.selector import ( ++from ..common import layout_utils ++from ..common.selector import ( + sol_attn_popc_b32, + sol_attn_route_is_exact, + ) +diff --git a/sol_attn/sm100/softmax.py b/sol_attn/sm100/softmax.py +index 876ef8007e12eb855cdaee1f1ee40c18775ae143..71b11f0e6ece8cabd7d6182bd1759e091c7c57db 100644 +--- a/sol_attn/sm100/softmax.py ++++ b/sol_attn/sm100/softmax.py +@@ -7,7 +7,7 @@ import cutlass.cute as cute + from cutlass import Float32, Int32 + from cutlass.cute.nvgpu import tcgen05 + +-from sol_attn._vendor.flash_attn.cute import utils as fa_utils ++from flash_attn.cute import utils as fa_utils + + from .tmem import ( + _add_physical_tmem_base, diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 5f5e8947b193..2adb1972965f 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -11,3 +11,19 @@ vendors: patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch patch_digest: sha256:ba917f330d8281781ce30ddb89f1772ee8f9017264d948b588b881a92f511e45 digest: sha256-tree-v1:fa65b674a9ec6121b4158069d5585cc56df28243d1ac1199eb97003d81135d76 + sana-sol-attn: + url: https://github.com/NVlabs/Sana.git + branch: sol-engine + commit: 5fe5febdf0f59fee1c0b44a5ce6665df0dabd247 + source: techniques/sparse_backends + destination: tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell + include: + - sol_attn/THIRD_PARTY_NOTICES.md + - sol_attn/__init__.py + - sol_attn/common/* + - sol_attn/interface.py + - sol_attn/preprocess.py + - sol_attn/sm100/* + patch: 3rdparty/vendor_patches/sana-sol-attn.patch + patch_digest: sha256:0d66669021683773aff34533abaf8f200abd09e3af93d19eed0e2a2e04f2f980 + digest: sha256-tree-v1:383f6085ca1bc2438e8d7d4796036af3a816dbfd396313704f7f4be2c06ac55b diff --git a/docs/source/features/visualgen-sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md index 38efb717765e..cb734cb9c888 100644 --- a/docs/source/features/visualgen-sparse-attention.md +++ b/docs/source/features/visualgen-sparse-attention.md @@ -6,6 +6,7 @@ This feature is in **beta** stage. APIs, supported models, and optimization opti - [Overview](#overview) - [Algorithms](#algorithms) + - [Sol-Attn](#sol-attn) - [Skip Softmax Attention](#skip-softmax-attention) - [Video Sparse Attention (VSA)](#video-sparse-attention-vsa) @@ -20,7 +21,41 @@ Sparse attention in VisualGen is configured through `VisualGenArgs.attention_con | `algorithm` | Config class | Status | |---|---|---| | `skip_softmax` | `SkipSoftmaxAttentionConfig` | Supported | -| VSA | TBD | TODO | +| `vsa` | `VideoSparseAttentionConfig` | Supported (CUTEDSL) | +| `sol_attn` | `SolAttentionConfig` | Supported (CUTEDSL, sm100/sm103) | + +### Sol-Attn + +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 datacenter +Blackwell -- sm100 (B200/GB200) and sm103 (B300/GB300) -- 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: layer indices 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 is known not to serve — an unsupported architecture, a +`head_dim` other than 128, a non-bfloat16 dtype — Sol-Attn runs dense +attention instead (the configured backend's dense kernel where available, +torch SDPA otherwise), logs the specific reason once, and counts the fallback. +Set `TRTLLM_SOL_ATTN_STRICT=1` to raise instead, which is useful when +benchmarking to confirm the kernel actually ran. Errors raised by the kernel +itself are not caught. ## Skip Softmax Attention @@ -92,7 +127,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 are currently the two supported sparse-attention algorithms, both available only through the CuTeDSL attention backend; quantized attention is not yet supported or enabled with either mode. #### Mapping `disabled_until_timestep` to Actual Denoising Steps diff --git a/pyproject.toml b/pyproject.toml index e198a15955fc..971c42858258 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,7 @@ force-exclude = true exclude = [ "**3rdparty/**", "tensorrt_llm/_torch/attention/backends/prims_ts/**", + "tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/**", "triton_kernels/**", ] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/skip_softmax/params.py b/tensorrt_llm/_torch/attention/backends/sparse/skip_softmax/params.py index 6b688806860e..a54636ea6d65 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/skip_softmax/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/skip_softmax/params.py @@ -382,17 +382,23 @@ def get_runtime_params( *, runtime_params: Optional[SparseRuntimeParams] = None, timestep: Any = None, + graph_phase: Optional[int] = None, ) -> SparseRuntimeParams: - """Return runtime parameters with skip-softmax thresholds.""" + """Return runtime parameters with skip-softmax thresholds. + + ``graph_phase`` lets a caller that already resolved the dense-prefix + phase host-side (the CUDA-graph runner does, to build its key) pass it + in, so this never has to read ``timestep`` -- a ``.item()`` on a CUDA + tensor -- while a graph is being captured. + """ if runtime_params is None: runtime_params = SparseRuntimeParams() - if ( - self.get_graph_phase_for_timestep( + if graph_phase is None: + graph_phase = self.get_graph_phase_for_timestep( timestep, disabled_until_timestep=self.disabled_until_timestep, ) - == 0 - ): + if graph_phase == 0: return replace( runtime_params, threshold_scale_factor_prefill=0.0, diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py index 12aa3ec4195c..bdc9fe9f33c1 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py @@ -24,6 +24,7 @@ from .cute_dsl import ( VSA_TILE_SIZE, CuTeDSLAttention, + SolAttention, VSAAttention, VSAMetadata, VSAMetadataBuilder, @@ -48,6 +49,7 @@ "FlashAttn4Attention", "FlashInferAttention", "RingAttention", + "SolAttention", "TrtllmAttention", "TrtllmAttentionMetadata", "UlyssesAttention", diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py index 9b70421c3b81..e5dccb71a7ba 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py @@ -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 from .vsa import ( VSA_KERNEL_MAX_CUBES, VSA_TILE_SIZE, @@ -42,4 +44,5 @@ "set_vsa_forward_context", "get_vsa_forward_context", "_cute_dsl_import_error", + "SolAttention", ] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py index e9638ef9075c..1d7edbbaa948 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py @@ -26,6 +26,7 @@ import torch from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxParams +from tensorrt_llm._torch.visual_gen.cuda_graph_runner import resolved_extra_key from tensorrt_llm.logger import logger from tensorrt_llm.visual_gen.args import QuantAttentionConfig @@ -86,7 +87,12 @@ def _resolve_skip_softmax_threshold_scale_factor( "through.", key="cute_dsl_skip_softmax_missing_timestep", ) - runtime_params = sparse_params.scheduler.get_runtime_params(timestep=timestep) + # Prefer the phase the CUDA-graph runner resolved host-side; the tensor + # read is a `.item()`, which is illegal while a graph is being captured. + runtime_params = sparse_params.scheduler.get_runtime_params( + timestep=timestep, + graph_phase=resolved_extra_key("skip_softmax_phase"), + ) threshold_scale_factor = runtime_params.threshold_scale_factor_prefill if threshold_scale_factor is None or threshold_scale_factor <= 0.0: return None diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py new file mode 100644 index 000000000000..5d0fea9a1645 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py @@ -0,0 +1,315 @@ +# 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._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxScheduler +from tensorrt_llm._torch.visual_gen.cuda_graph_runner import resolved_extra_key +from tensorrt_llm.logger import logger + +from ....attention.backends.interface import PredefinedAttentionMask +from ..interface import AttentionBackend, AttentionTensorLayout +from ..vanilla import VanillaAttention + +_sol_attn_import_error = None +try: + 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 _cute_dense_available() -> bool: + """Whether `cute_dsl_fmha_fwd` can run on the current device. + + Checked once at construction. Sol-Attn and the dense CuTe DSL kernel now + cover the same set (sm_100a/sm_103a), so in practice this is always true + wherever Sol-Attn runs; the negative branch exists so an unsupported device + degrades to SDPA instead of raising. + """ + try: + from .fmha import _check_cute_runtime_available, _get_gpu_arch + + _check_cute_runtime_available() + _get_gpu_arch() + except Exception: + return False + return True + + +class SolAttention(AttentionBackend): + """Sol-Attn dynamic block-routing sparse attention (CuTeDSL, sm100/sm103). + + Inputs the kernel cannot serve (shape, dtype, architecture) are delegated to + dense attention up front by ``_run_sol_attn_bthd``; a kernel error is not + recovered from and propagates. This class adds the ``dense_layers`` guard + (evaluated at construction time, no external plumbing needed) and forwards + the routing knobs from config. + """ + + def __init__( + 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.disabled_until_timestep = getattr(cfg, "disabled_until_timestep", None) + self.dense_layers = frozenset(getattr(cfg, "dense_layers", None) or ()) + + # Sol-Attn's dense steps must run the backend the user selected. Without + # this they ran torch SDPA while a `backend: CUTEDSL` baseline ran + # 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 + + # The only backend that consumes `key_padding_mask` (CuTeDSL's `_fwd` + # swallows it via **kwargs, silently). Masked self-attention is routed + # here rather than to the mask-blind sparse kernel or to `_inner`. + self._vanilla = VanillaAttention( + layer_idx=layer_idx, + num_heads=num_heads, + head_dim=head_dim, + num_kv_heads=self.num_kv_heads, + dtype=dtype, + ) + self._inner = CuTeDSLAttention( + layer_idx=layer_idx, + num_heads=num_heads, + head_dim=head_dim, + num_kv_heads=self.num_kv_heads, + dtype=dtype, + ) + # Resolve availability once at construction to keep capability checks + # and attribute initialization out of `_dense` during torch.compile tracing. + self._cute_dense_ok = _cute_dense_available() + if not self._cute_dense_ok: + logger.warning_once( + "[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: + # Under CUDA-graph capture the runner has already resolved the phase + # host-side (it is part of the graph key); reading the tensor here + # would `.item()` inside capture, which CUDA forbids. + phase = resolved_extra_key("sol_attn_phase") + if phase is None: + phase = SkipSoftmaxScheduler.get_graph_phase_for_timestep( + timestep, + disabled_until_timestep=self.disabled_until_timestep, + ) + if phase is None: + # Fail open, matching the CuTeDSL skip-softmax path: without a + # timestep we cannot tell which phase we are in, so run the + # 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, **kwargs: Any) -> torch.Tensor: + """Dense attention via torch SDPA, for devices CuTe DSL cannot serve. + + Honors the same two mask kwargs the backends accept: ``attention_mask`` + (``CAUSAL``) and ``key_padding_mask`` (``[B, S_kv]`` bool, True = valid). + """ + is_causal = kwargs.get("attention_mask", PredefinedAttentionMask.FULL) == ( + PredefinedAttentionMask.CAUSAL + ) + key_padding_mask = kwargs.get("key_padding_mask") + attn_mask = None + if key_padding_mask is not None: + # SDPA wants a broadcastable bool mask where True = attend. + attn_mask = key_padding_mask.to(torch.bool)[:, None, None, :] + return torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + attn_mask=attn_mask, + is_causal=is_causal and attn_mask is None, + ).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 kwargs.get("key_padding_mask") is not None: + # CuTeDSL does not consume `key_padding_mask`, so staying in-family + # here would silently attend to padded tokens; VANILLA honors it. + # VANILLA works in HND ([B, H, S, D]); this backend is NHD. + if self._vanilla.preferred_layout == AttentionTensorLayout.HND: + out = self._vanilla.forward( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), **kwargs + ) + return out.transpose(1, 2) + return self._vanilla.forward(q, k, v, **kwargs) + if self._cute_dense_ok and q.is_cuda: + # CAUSAL is fine here: CuTeDSL handles it via `_prepare_inputs`. + return self._inner.forward(q, k, v, **kwargs) + return self._sdpa(q, k, v, **kwargs) + + def _can_serve(self, q: torch.Tensor, k: torch.Tensor, **kwargs: Any) -> bool: + """Whether the sparse kernel applies to this particular call. + + Everything false here is delegated (see ``_delegate``). Deciding it from the + tensors, per call, is deliberate, and it is why ``modules/attention.py`` + has no ``SEPARATE_QKV`` rule for Sol-Attn: ``qkv_mode`` describes how + Q/K/V are *projected*, not whether K/V come from another sequence, so a + construction-time rule keyed on it mistakes self-attention for + cross-attention wherever that mode is chosen for other reasons -- + Qwen-Image always, and WAN's ``attn1`` under async Ulysses -- silently + costing those modules their configured backend. + """ + # Cross-attention: K/V come from another sequence, and Sol-Attn's + # routing assumes one self-attending sequence. Unequal Q/K lengths are + # a heuristic for that, not a definition: a cross-attention call whose + # context happens to match the query length is not caught here. The + # `Attention` module knows the answer (`encoder_hidden_states`), but the + # backend `forward` kwargs carry no such flag yet; TRTLLM-16475 tracks + # threading an explicit signal through and retiring this check. + if k.shape[1] != q.shape[1]: + return False + # Masks: the sparse kernel is noncausal and takes no mask, so a masked + # call -- HunyuanVideo1.5 / GLM-Image `key_padding_mask`, Cosmos3 + # `CAUSAL` -- must go to a backend that honors it. Equal Q/K lengths do + # not imply "unmasked self-attention". + if kwargs.get("key_padding_mask") is not None: + return False + if ( + kwargs.get("attention_mask", PredefinedAttentionMask.FULL) + != PredefinedAttentionMask.FULL + ): + return False + if self.layer_idx in self.dense_layers: + return False + if self.disabled_until_timestep is not None and self._dense_by_step(kwargs.get("timestep")): + 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, + # 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 4cdb62d28a94..80cedc79fa9c 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py @@ -145,14 +145,17 @@ def create_attention( ) kwargs["attention_metadata_state"] = attention_metadata_state if backend.upper() == "CUTEDSL" and attention_config is not None: - if ( - attention_config.sparse_attention_config is not None - and getattr(attention_config.sparse_attention_config, "algorithm", None) == "vsa" - ): + sparse_algo = getattr(attention_config.sparse_attention_config, "algorithm", None) + if sparse_algo == "vsa": from .cute_dsl.vsa import VSAAttention attn_cls = VSAAttention kwargs["sparse_attention_config"] = attention_config.sparse_attention_config + elif sparse_algo == "sol_attn": + from .cute_dsl.sol_attn import 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/cuda_graph_runner.py b/tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py index 3d79f73dde39..eebb6c6ad497 100644 --- a/tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py @@ -1,3 +1,5 @@ +import contextlib +import contextvars import functools import gc from dataclasses import dataclass @@ -9,6 +11,37 @@ from ..utils import make_weak_ref +# Extra-key values the runner resolved host-side for the call it is currently +# capturing. CUDA-graph capture forbids device-to-host syncs, so anything a +# forward would otherwise derive from a CUDA tensor with `.item()` -- the +# sparse-attention dense-prefix phase, for instance -- must be resolved before +# capture and read back from here inside it. The runner already computes these +# values to build the graph key; this just makes them visible to the callee. +_RESOLVED_EXTRA_KEYS: contextvars.ContextVar[Optional[Dict[str, Any]]] = contextvars.ContextVar( + "cuda_graph_resolved_extra_keys", default=None +) + + +def resolved_extra_key(name: str) -> Any: + """Return the extra-key value `name` resolved for the call being captured, else None. + + None means "not inside a runner-driven call" (or the key was omitted), and + callers must fall back to deriving the value themselves. + """ + resolved = _RESOLVED_EXTRA_KEYS.get() + return None if resolved is None else resolved.get(name) + + +@contextlib.contextmanager +def resolved_extra_keys_scope(values: Dict[str, Any]): + """Expose `values` through `resolved_extra_key` for the enclosed block.""" + token = _RESOLVED_EXTRA_KEYS.set(dict(values)) + try: + yield + finally: + _RESOLVED_EXTRA_KEYS.reset(token) + + # One named graph-key component, e.g. ("hidden_states", (1, 4096, 3072)). KeyPart: TypeAlias = Tuple[str, Hashable] # Full CUDA graph cache key, stored as a tuple so it can be used as a dict key. @@ -121,6 +154,9 @@ def _get_extra_key(self, *args, **kwargs) -> KeyType: value = fn(*args, **kwargs) if value is not None: parts.append((name, value)) + # Remembered so `capture()` can republish them to the callee; this runs + # outside capture, where the `.item()` these callbacks need is legal. + self._last_resolved_extra_keys = dict(parts) return tuple(parts) def get_graph_key(self, *args, **kwargs) -> KeyType: @@ -144,14 +180,18 @@ def capture( } graph = torch.cuda.CUDAGraph() - for _ in range(self.WARMUP_STEPS): - fn(*static_args, **static_kwargs) - torch.cuda.synchronize() - gc.collect() - torch.cuda.empty_cache() - - with torch.cuda.graph(graph, pool=self._get_pool()): - output = fn(*static_args, **static_kwargs) + # Warmup and capture both see the host-resolved extra keys, so a + # forward that needs the dense-prefix phase reads it from here instead + # of syncing a CUDA tensor -- which capture would reject. + with resolved_extra_keys_scope(getattr(self, "_last_resolved_extra_keys", {})): + for _ in range(self.WARMUP_STEPS): + fn(*static_args, **static_kwargs) + torch.cuda.synchronize() + gc.collect() + torch.cuda.empty_cache() + + with torch.cuda.graph(graph, pool=self._get_pool()): + output = fn(*static_args, **static_kwargs) self.graphs[key] = graph self.static_inputs[key] = (static_args, static_kwargs) diff --git a/tensorrt_llm/_torch/visual_gen/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..b226f77f96bc --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md @@ -0,0 +1,49 @@ +# 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). +Checked against the branch tip +([`83e54df`](https://github.com/NVlabs/Sana/commit/83e54df), 2026-08-20) on +2026-08-27. + +## Scope of the vendored subset + +Only the pieces needed for the architectures TensorRT-LLM ships are carried: + +| Carried | Not carried | +|---|---| +| `interface.py`, `preprocess.py`, `common/` | `sm89/`, `sm90/` (incl. `sm90/_compat/`) | +| `sm100/` — serves SM100 (B200/GB200) and SM103 (B300/GB300) | `triton_ref/` Triton reference attention | +| | `sm120/` (RTX Blackwell) | +| | `_vendor/flash_attn/` (see below) | + +`../sol_attn_backend.py` is not part of the vendored package but is a +derivative work of upstream's +`techniques/sparse_backends/sol_attn_backend.py` (same branch and commit). +Only the kernel-wrapper subset is carried; upstream's model-integration half +(diffusers dispatch hook, HunyuanVideo MMDiT padding, model-level Morton +ordering) is not. + +Implementation divergences from upstream are documented in the source itself, +in the module docstrings of `../sol_attn_backend.py` and +`attention_backend/cute_dsl/sol_attn.py`. + +## Licensing + +The upstream package vendored a copy of FlashAttention's CuTe DSL helpers +under `sol_attn/_vendor/flash_attn/cute/`. That copy is **not** carried here: +TensorRT-LLM already depends on +[`flash-attn-4`](https://github.com/Dao-AILab/flash-attention) (pinned in +`requirements.txt`), which provides the same `flash_attn.cute` modules, and +the SM100 kernels import them from that dependency directly. FlashAttention's +BSD-3-Clause license is retained at `sol_attn/sm100/LICENSE.flash-attention` +because portions of the SM100 design scaffold still derive from that project. + +`preprocess.py` implements the routing/threshold stage in Triton, so Triton is +a required runtime dependency on every Sol-Attn path, not only a fallback. + +The runtime also depends on NVIDIA CUTLASS / CuTe DSL, cuda-python, and +PyTorch. Those dependencies are not redistributed by this repository and +remain subject to their respective licenses. diff --git a/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..d37351dd8265 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/__init__.py @@ -0,0 +1,5 @@ +"""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..ec1d7ee9f173 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/__init__.py @@ -0,0 +1,5 @@ +"""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..67e4d0896a85 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/layout_utils.py @@ -0,0 +1,130 @@ +"""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..5182d0c49879 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/runtime.py @@ -0,0 +1,14 @@ +"""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..0321f0fe429c --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/common/selector.py @@ -0,0 +1,169 @@ +"""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..2da4ea383947 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/interface.py @@ -0,0 +1,285 @@ +"""Public Sol-Attn interface.""" + +from __future__ import annotations + +import functools + +import torch + +BLOCK_SIZE = 64 +# TensorRT-LLM ships the SM100 kernel only. It serves both datacenter Blackwell +# steppings: the CuTe DSL JIT targets the device it compiles on and the kernel +# body uses no SM100-exclusive construct (measured identical on B200 and B300). +_CUTE_BACKENDS = { + (10, 0): "cute_sm100", # B200 / GB200 + (10, 3): "cute_sm100", # B300 / GB300 +} +_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 is none. + + TensorRT-LLM raises instead of falling back to the Triton reference so a + missing kernel is visible; the caller decides what to do about it. + """ + + cute_backend = _CUTE_BACKENDS.get(arch) + if cute_backend is None: + raise RuntimeError( + 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 the SM100 " + "kernel only, so kv_splits must be 1." + ) + route_groups = ((tokens + 63) // 64 + 63) // 64 + if kv_splits > route_groups: + raise ValueError("each KV split must contain at least one N64 route group") + + +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 not in _CUTE_BACKENDS: + # Unreachable via sol_attn(): _backend_for_arch raises first. Kept + # explicit so a missed guard cannot return the uninitialised + # `output` buffer. Keyed off _CUTE_BACKENDS, not a literal. + raise ValueError(f"no Sol-Attn CuTe kernel for SM{arch[0]}{arch[1]}") + sink_start_block, sink_end_block = _sink_block_range( + tokens, + sink_start, + sink_tokens, + ) + tensors = [q, k, v, output, kc, vc, threshold, lse] + compiled = _compiled.get(key) + if compiled is None: + compiled, args = _compile_sm100( + key, + tensors, + scale, + sink_start_block, + sink_end_block, + stream, + ) + 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..77c4cccf1870 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/preprocess.py @@ -0,0 +1,463 @@ +"""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..7fb6167a6b14 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/__init__.py @@ -0,0 +1,5 @@ +"""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..d3edaff3ec14 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/kernel.py @@ -0,0 +1,5 @@ +"""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..24dee1cb0233 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/mainloop.py @@ -0,0 +1,1762 @@ +"""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 .softmax import ( + _load_m64_n128_score as _load_pair_score, + _online_update_one_half as _online_update_pair, + _rescale_m64_partial_o as _rescale_pair_o, +) +from . import math as mma_utils + +from ..common import layout_utils +from ..common.selector import ( + sol_attn_popc_b32, + sol_attn_route_is_exact, +) +from .tmem import ( + _add_physical_tmem_base, + _zero_based_tmem_tensor, + 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..e1120b09ce3c --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/math.py @@ -0,0 +1,29 @@ +"""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..71b11f0e6ece --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/softmax.py @@ -0,0 +1,156 @@ +"""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..aada77bea91d --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/sm100/tmem.py @@ -0,0 +1,138 @@ +"""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..829a3ce462f5 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py @@ -0,0 +1,248 @@ +# 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 which subset is +carried. Check that file before re-syncing against upstream. + +Deliberate divergences from upstream, which a re-sync must preserve rather +than overwrite: + +* ``logger.warning_once`` replaces ``print()``, so fallbacks are suppressible + and routed through this repository's logger. +* ``_SOL_STATS["dense_fallback_calls"]`` makes a silently-degraded run + countable rather than only visible on stderr. +* ``sol_attn_ineligible_reason()`` names the specific reason (architecture, + head_dim, dtype) instead of returning one boolean. +* Kernel exceptions propagate. Upstream catches them and silently reruns the + call with torch SDPA; a failed CuTe launch can leave the device in a bad + state, so recovery is not attempted here. Only inputs known up front to be + unservable (shape, dtype, architecture) are routed to dense attention, and + ``TRTLLM_SOL_ATTN_STRICT=1`` turns even that into an error. +* Dense paths route to ``cute_dsl_fmha_fwd`` through ``dense_fn``. Upstream + falls back to torch SDPA; staying inside the configured backend is what lets + a ``backend: CUTEDSL`` A/B isolate sparsity rather than also swapping the + dense kernel. +* ``@torch.compiler.disable`` guards the launch boundary; see the comment on + ``_run_sol_attn_bthd`` for why, and for the ``torch.library.custom_op`` + alternative upstream uses. + +The kernel-facing API accepts contiguous BF16 +``[batch, tokens, heads, 128]`` Q/K/V, ``tau``, ``thresh_type``, +``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 +is known not to serve -- wrong shape, dtype, or an architecture with no kernel +-- are delegated to dense attention before any launch, logged once, and +counted in ``_SOL_STATS["dense_fallback_calls"]``. Set +``TRTLLM_SOL_ATTN_STRICT=1`` to raise instead. Errors raised by the kernel +itself are not caught. +""" + +from __future__ import annotations + +import functools +import os +from typing import Callable, Optional + +import torch + +from tensorrt_llm._utils import get_sm_version +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 + + +# SM versions with a Sol-Attn CuTe kernel, in ``get_sm_version()`` form +# (major * 10 + minor). Kept in sync with ``sol_attn/interface.py::_CUTE_BACKENDS``; +# duplicated here so the eligibility check does not have to import the CuTe DSL. +SUPPORTED_ARCHS = frozenset({100, 103}) + + +def sol_attn_ineligible_reason(q) -> Optional[str]: + """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}" + sm = get_sm_version() + if sm not in SUPPORTED_ARCHS: + return f"no Sol-Attn kernel for SM{sm}; supported: " + ", ".join( + f"SM{v}" for v in sorted(SUPPORTED_ARCHS) + ) + return None + + +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 _strict() -> bool: + """Whether TRTLLM_SOL_ATTN_STRICT=1 asks us to raise on an unservable input.""" + + return os.environ.get("TRTLLM_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. +# +# Upstream instead wraps the same call in a `torch.library.custom_op` with a +# `register_fake`, which keeps the kernel in the compiled graph as an opaque +# node rather than breaking the graph at it. That is the better end state -- +# it removes the per-layer graph break -- and is tracked as a follow-up. It is +# not done here because `torch.compiler.disable` is the convention every other +# CuTe DSL entry point in this repository already follows, and is what the +# reported measurements were taken with. +@torch.compiler.disable +def _run_sol_attn_bthd( + q, + k, + v, + *, + tau: float = DEFAULT_TAU, + thresh_type: str = DEFAULT_THRESH_TYPE, + kv_splits: int = 1, + sink_start: int | None = None, + sink_tokens: int = 0, + dense_fn: Callable | None = None, +): + """Run Sol-Attn on contiguous BTHD tensors. + + Inputs the kernel is known not to serve go to dense attention up front; + errors raised by the kernel itself propagate. + """ + + q0, k0, v0 = q.contiguous(), k.contiguous(), v.contiguous() + + 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: + # 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 TRTLLM_SOL_ATTN_STRICT=1 to raise instead.", + key=("sol_attn_ineligible", reason), + ) + return dense() + + kernel = _load_sol_attn() + out = kernel( + q0, + k0, + v0, + tau=float(tau), + thresh_type=str(thresh_type), + # Only 1 split exists on the shipped sm100/sm103 kernels (2/4 was an + # SM90-only path), so this is not a user-facing knob; the kernel + # interface rejects anything else. + kv_splits=int(kv_splits), + sink_start=sink_start, + sink_tokens=int(sink_tokens), + ) + _SOL_STATS["kernel_calls"] += 1 + return out + + +# Lightweight run-validation counters. `kernel_calls` is the census used to +# prove the CuTe kernel actually ran; `dense_fallback_calls` counts calls the +# kernel was known not to serve. Set TRTLLM_SOL_ATTN_STRICT=1 to raise instead +# of falling back on those. +_SOL_STATS = {"kernel_calls": 0, "dense_fallback_calls": 0} + + +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/ltx2/pipeline_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py index f06c5d003b41..2b8d1685bc9b 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py @@ -20,7 +20,11 @@ from tensorrt_llm._torch.utils import make_weak_ref from tensorrt_llm._torch.visual_gen.cache.teacache import CacheContext, register_extractor from tensorrt_llm._torch.visual_gen.checkpoints.prefetch import prefetch_files_to_host_cache -from tensorrt_llm._torch.visual_gen.cuda_graph_runner import CUDAGraphRunner, CUDAGraphRunnerConfig +from tensorrt_llm._torch.visual_gen.cuda_graph_runner import ( + CUDAGraphRunner, + CUDAGraphRunnerConfig, + resolved_extra_keys_scope, +) from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline import ( BasePipeline, @@ -493,14 +497,19 @@ def capture(self, key, fn, args, kwargs): static_kwargs = {k: self._clone_value(v) for k, v in kwargs.items()} graph = torch.cuda.CUDAGraph() - for _ in range(self.WARMUP_STEPS): - fn(*static_args, **static_kwargs) - torch.cuda.synchronize() - gc.collect() - torch.cuda.empty_cache() - - with torch.cuda.graph(graph, pool=self._get_pool()): - output = fn(*static_args, **static_kwargs) + # Same contract as CUDAGraphRunner.capture: warmup and capture read the + # host-resolved extra keys (skip-softmax / Sol-Attn dense-prefix phase) + # from this scope instead of syncing the CUDA timestep, which capture + # would reject. + with resolved_extra_keys_scope(getattr(self, "_last_resolved_extra_keys", {})): + for _ in range(self.WARMUP_STEPS): + fn(*static_args, **static_kwargs) + torch.cuda.synchronize() + gc.collect() + torch.cuda.empty_cache() + + with torch.cuda.graph(graph, pool=self._get_pool()): + output = fn(*static_args, **static_kwargs) self.graphs[key] = graph self.static_inputs[key] = (static_args, static_kwargs) @@ -1714,7 +1723,6 @@ def forward( audio_scheduler = copy.deepcopy(self.scheduler) audio_scheduler.set_timesteps(num_inference_steps, latent=latents_5d) timesteps = self.scheduler.timesteps - num_steps = len(timesteps) # ---- 7. Build perturbation config for STG ----------------------- stg_perturbation: PerturbationConfig | None = None @@ -1842,7 +1850,11 @@ def _run_transformer( audio=audio_mod, perturbations=perturbations, text_cache=text_cache, - timestep=timestep_val.new_tensor(float(step_index) / num_steps), + # The scheduler sigma is already normalized to [0, 1] with the + # contract's sense (larger = noisier). The previous + # step_index / num_steps was ascending, which inverted the + # sparse-attention dense prefix on this model. + timestep=timestep_val, step_index=step_index, ) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py index 7c266fe0a2cd..e3529a7e819d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py @@ -1107,6 +1107,10 @@ def _setup_cuda_graphs(self): self._current_lora_cuda_graph_state, lambda: self.transformer.active_topology, ) + # Same registration the single-stage pipeline does: without it the + # dense-prefix and sparse phases of skip-softmax / Sol-Attn share one + # captured graph and the wrong kernel is replayed. + self.transformer.register_cuda_graph_extra_key_fns(runner) compile_note = " (with torch.compile)" if self.pipeline_config.torch_compile.enable else "" logger.info( "CUDA graph runner: wrapping LTX-2 two-stage transformer.forward " @@ -1793,6 +1797,7 @@ def _refinement_denoise( audio=audio_mod, text_cache=_s2_static, step_index=i, + timestep=timestep, ) # Video: velocity → x0 → post-process → Euler step diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py index 7cba785b1bf2..2ff493af9900 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py @@ -928,9 +928,18 @@ def forward( text_kv_video: tuple[torch.Tensor, torch.Tensor] | None = None, text_kv_audio: tuple[torch.Tensor, torch.Tensor] | None = None, step_index=None, + timestep: torch.Tensor | None = None, ) -> tuple[TransformerArgs | None, TransformerArgs | None]: """Forward with optional perturbation masking for STG. + ``timestep`` is the normalized scheduler time in ``[0, 1]`` (larger = + noisier) that the attention backends consume for timestep-dependent + decisions such as the sparse-attention dense prefix. It is deliberately + *not* ``video.timesteps`` / ``audio.timesteps``: those are the AdaLN + modulation outputs -- learned activations -- and were previously passed + to attention as if they were the scheduler time, so any + ``disabled_until_timestep`` cutoff was compared against noise. + Args: perturbations: Optional ``BatchedPerturbationConfig`` that masks attention outputs for selected blocks/modalities. @@ -991,9 +1000,7 @@ def forward( self._fuse_adaln, fp4_input_scale=get_nvfp4_self_attn_input_scale(self.attn1), ) - v_attn_raw = self.attn1( - norm_vx, pe=video.positional_embeddings, timestep=video.timesteps - ) + v_attn_raw = self.attn1(norm_vx, pe=video.positional_embeddings, timestep=timestep) if has_perturbations and perturbations.any_in_batch( PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx ): @@ -1018,7 +1025,7 @@ def forward( attn2_q_input, context=video.context, pre_projected_kv=text_kv_video, - timestep=video.timesteps, + timestep=timestep, ) # --- Audio self-attention + text cross-attention --- @@ -1049,7 +1056,7 @@ def forward( norm_ax, pe=audio.positional_embeddings, key_padding_mask=audio.audio_padding_mask, - timestep=audio.timesteps, + timestep=timestep, ) if has_perturbations and perturbations.any_in_batch( PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx @@ -1074,7 +1081,7 @@ def forward( audio_attn2_q_input, context=audio.context, pre_projected_kv=text_kv_audio, - timestep=audio.timesteps, + timestep=timestep, ) # --- Bidirectional audio ↔ video cross-attention --- @@ -1233,7 +1240,7 @@ def forward( pre_projected_kv=(k_a2v, v_a2v), pe=video.cross_positional_embeddings, key_padding_mask=audio.audio_padding_mask, - timestep=video.timesteps, + timestep=timestep, ) if has_perturbations and perturbations.any_in_batch( PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx @@ -1267,7 +1274,7 @@ def forward( freqs=a_cross_pe, kv_input=vx_scaled_v2a, kv_freqs=video.cross_positional_embeddings, - timestep=audio.timesteps, + timestep=timestep, ) else: k_v2a, v_v2a = self.video_to_audio_attn.project_kv( @@ -1286,7 +1293,7 @@ def forward( ax_v2a_local, pre_projected_kv=(k_v2a, v_v2a), pe=a_cross_pe, - timestep=audio.timesteps, + timestep=timestep, ) v2a_attn_raw = self._sp_all_gather(out_local, dim=1) elif self._async_ulysses and self.video_to_audio_attn.is_ulysses: @@ -1299,7 +1306,7 @@ def forward( freqs=audio.cross_positional_embeddings, kv_input=vx_scaled_v2a, kv_freqs=video.cross_positional_embeddings, - timestep=audio.timesteps, + timestep=timestep, ) else: # v2a sync: with a Ulysses wrapper, K/V (video) stay seq-sharded @@ -1320,7 +1327,7 @@ def forward( ax_scaled_v2a, pre_projected_kv=(k_v2a, v_v2a), pe=audio.cross_positional_embeddings, - timestep=audio.timesteps, + timestep=timestep, ) if has_perturbations and perturbations.any_in_batch( PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx @@ -2235,9 +2242,10 @@ def forward( text_cache: Pre-computed step-invariant outputs from ``prepare_text_cache()``. Always required — callers must invoke ``prepare_text_cache()`` first. timestep: Normalized denoising-time coordinate in ``[0, 1]``. - May be ``None`` for LTX-2 paths that rely only on per-modality - timestep values and do not need timestep-based CUDA graph - partitioning. + Threaded to every attention call; sparse-attention dense + prefixes (``disabled_until_timestep``) and the CUDA-graph phase + key both read it, so pipelines must pass the scheduler sigma + here. May be ``None`` only when neither is configured. LTX-2 also carries per-modality timestep values in ``video.timesteps`` / ``audio.timesteps`` for the reference time-embedding path. @@ -2333,6 +2341,7 @@ def forward( ax, perturbations=perturbations, step_index=step_index, + timestep=timestep, ) if video_args is not None and vx is not None: video_args = replace(video_args, x=vx) @@ -2347,6 +2356,7 @@ def forward( text_kv_video=v_kv[i] if v_kv else None, text_kv_audio=a_kv[i] if a_kv else None, step_index=step_index, + timestep=timestep, ) # Gather sequences back to full length for output processing. diff --git a/tensorrt_llm/_torch/visual_gen/models/modeling.py b/tensorrt_llm/_torch/visual_gen/models/modeling.py index 9fd8a02187ea..cd0f931eb3d6 100644 --- a/tensorrt_llm/_torch/visual_gen/models/modeling.py +++ b/tensorrt_llm/_torch/visual_gen/models/modeling.py @@ -21,7 +21,7 @@ from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxScheduler 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 +74,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: SkipSoftmaxScheduler.get_graph_phase_for_timestep( + kwargs.get("timestep"), + disabled_until_timestep=disabled_until_timestep, + ), + ) + return diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index be338b20564f..196289d8c18f 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -97,17 +97,15 @@ def __init__( cp_size = vgm.cp_size if vgm else 1 base_backend = config.attention.backend _sa_cfg = config.attention.sparse_attention_config - _is_vsa = ( - base_backend == "CUTEDSL" - and _sa_cfg is not None - and getattr(_sa_cfg, "algorithm", None) == "vsa" - ) - + _sa_algo = getattr(_sa_cfg, "algorithm", None) if _sa_cfg is not None else None + _is_vsa = base_backend == "CUTEDSL" and _sa_algo == "vsa" + _is_sol_attn = base_backend == "CUTEDSL" and _sa_algo == "sol_attn" separate_qkv_cross_attention = ( self.qkv_mode == QKVMode.SEPARATE_QKV and not separate_qkv_is_self_attention ) # Cross-attention fallback: TRTLLM and CUTEDSL VSA are self-attn only. + # Sol-Attn is absent by design; see SolAttention._can_serve. if separate_qkv_cross_attention and (base_backend == "TRTLLM" or _is_vsa): backend_name = "VANILLA" requested = f"{base_backend} (VSA)" if _is_vsa else base_backend @@ -121,9 +119,12 @@ def __init__( else: backend_name = base_backend - if _is_vsa and cp_size > 1: + # Every sparse algorithm here routes over the whole token sequence, so + # none of them can be split across context-parallel ranks. + if (_is_vsa or _is_sol_attn) and cp_size > 1: + _algo_name = "VSA" if _is_vsa else "Sol-Attn" raise ValueError( - f"VSA needs the full token sequence per rank, so it is incompatible " + f"{_algo_name} needs the full token sequence per rank, so it is incompatible " f"with context parallelism (Attention2D/Ring, cp_size={cp_size}). Use " f"ulysses or cfg parallelism instead." ) diff --git a/tensorrt_llm/visual_gen/__init__.py b/tensorrt_llm/visual_gen/__init__.py index ec714e3d2bfe..3b1ef26b8d98 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, @@ -76,6 +77,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", @@ -134,6 +136,7 @@ def __dir__(): "SparseAttentionConfig", "SkipSoftmaxAttentionConfig", "VideoSparseAttentionConfig", + "SolAttentionConfig", "VAEConfig", "CacheConfig", "TeaCacheConfig", diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index 8be9ea4357a9..1ccf02e1a90f 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -31,7 +31,11 @@ from tensorrt_llm.llmapi.utils import StrictBaseModel, set_api_status from tensorrt_llm.models.modeling_utils import QuantConfig -from .sparse_attention import SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig +from .sparse_attention import ( + SkipSoftmaxAttentionConfig, + SolAttentionConfig, + VideoSparseAttentionConfig, +) # ============================================================================= # Type aliases @@ -95,7 +99,7 @@ class QuantAttentionConfig(StrictBaseModel): # Discriminated union of sparse attention configs. SparseAttentionConfig = Annotated[ - Union[SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig], + Union[SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig, SolAttentionConfig], Field(discriminator="algorithm"), ] @@ -122,7 +126,8 @@ class AttentionConfig(StrictBaseModel): status="prototype", description=( "Sparse attention recipe. Discriminated by algorithm: " - "skip_softmax (TRTLLM / CUTEDSL backends) or VSA (CUTEDSL backend)." + "skip_softmax (TRTLLM / CUTEDSL backends), vsa (CUTEDSL backend), " + "or sol_attn (CUTEDSL backend)." ), ) @@ -220,6 +225,7 @@ def _validate_sparse_attention_config(self) -> "AttentionConfig": supported_backends = { "skip_softmax": ("TRTLLM", "CUTEDSL"), "vsa": ("CUTEDSL",), + "sol_attn": ("CUTEDSL",), }.get(algo) if supported_backends is None: return self @@ -235,19 +241,23 @@ def _validate_sparse_attention_config(self) -> "AttentionConfig": @model_validator(mode="after") def _validate_cutedsl_quant_sparse_mutex(self) -> "AttentionConfig": - # VSA replaces the dense CuTeDSL path and cannot compose with quantized - # attention. SkipSoftmax is part of that dense path and can compose. + # VSA and Sol-Attn each replace the dense CuTeDSL path and cannot + # compose with quantized attention: create_attention swaps in their own + # backend class, which never consumes quant_attention_config, so the + # request would be silently ignored. SkipSoftmax is part of the dense + # path itself and can compose. + _replaces_dense_path = ("vsa", "sol_attn") if ( self.backend == "CUTEDSL" and self.quant_attention_config is not None and self.sparse_attention_config is not None - and self.sparse_attention_config.algorithm == "vsa" + and self.sparse_attention_config.algorithm in _replaces_dense_path ): raise ValueError( - "CUTEDSL backend: quant_attention_config and VSA " - "sparse_attention_config are mutually exclusive (the " - "CuTeDSLAttention dispatcher selects either the dense path " - "or the sparse VSA path, not both)." + f"CUTEDSL backend: quant_attention_config and " + f"'{self.sparse_attention_config.algorithm}' sparse_attention_config " + "are mutually exclusive (the CuTeDSLAttention dispatcher selects " + "either the dense path or that sparse path, not both)." ) return self @@ -839,6 +849,7 @@ def from_yaml(cls, yaml_path: Union[str, Path], **overrides: Any) -> "VisualGenA "SparseAttentionConfig", "SkipSoftmaxAttentionConfig", "VideoSparseAttentionConfig", + "SolAttentionConfig", "AttentionConfig", "VAEConfig", "ParallelConfig", diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 2bccb10e2652..bd4c22e22395 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,76 @@ 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 on + datacenter Blackwell -- sm100 (B200/GB200) and sm103 (B300/GB300) -- + head_dim=128, bf16, MHA. + + On an unsupported *shape, dtype, or architecture* the kernel falls back to + dense attention -- the configured backend's dense kernel where available, + 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').", + ) + 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. " + "Read from the `timestep` forward kwarg, which must be the normalized " + "scheduler time (larger = noisier). WAN and LTX-2 pass it; a pipeline " + "that does not, or that passes something else, gets a one-time warning " + "and runs sparse on every step (fail-open)." + ), + ) + dense_layers: Optional[list[int]] = PydanticField( + None, + description=( + "Layer indices forced dense regardless of the dense prefix, e.g. " + "[0, 2, 3]. Evaluated per layer at construction time; no pipeline " + "wiring required." + ), + ) + + @field_validator("dense_layers") + @classmethod + def _validate_dense_layers(cls, layers: Optional[list[int]]) -> Optional[list[int]]: + if layers is None: + return None + for index in layers: + if index < 0: + raise ValueError(f"dense_layers contains a negative layer index: {index}") + return sorted(set(layers)) + + def to_sparse_params(self, **kwargs): + # Sol-Attn's knobs are consumed directly by SolAttention.__init__ + # (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 ffeb93d979ae..43f82850f760 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -292,6 +292,7 @@ l0_b200: - unittest/_torch/visual_gen/test_attention_cute_dsl.py - unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py - unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py + - unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py - unittest/_torch/visual_gen/test_attention_flashinfer.py - unittest/_torch/visual_gen/test_attention_trtllm_sage.py - unittest/_torch/visual_gen/test_attention_cudnn.py diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 4f67742c28a1..b8135b6efd0a 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -34,6 +34,8 @@ l0_b300: - unittest/_torch/thop/parallel TIMEOUT (90) - unittest/_torch/visual_gen/kernels/parallel - unittest/_torch/visual_gen/test_attention_flashinfer.py + # ------------- Visual Gen sparse attention (sm103) --------------- + - unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py - unittest/_torch/thop/serial - unittest/_torch/executor # 250s - unittest/_torch/disaggregation diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py new file mode 100644 index 000000000000..0da640d23d1e --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py @@ -0,0 +1,1032 @@ +# 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 is +covered by test_cute_kernel_matches_dense_on_a_single_block: Sol-Attn's +routing is score-derived, so no tau provably forces dense routing the way +VSA's full-top-k does, but a single KV block makes sparsity structurally +impossible and gives the same guarantee. +""" + +import math +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask +from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxScheduler +from tensorrt_llm._torch.visual_gen.attention_backend import CuTeDSLAttention +from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn import SolAttention +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.cuda_graph_runner import ( + CUDAGraphRunner, + CUDAGraphRunnerConfig, + resolved_extra_key, + resolved_extra_keys_scope, +) +from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode +from tensorrt_llm.visual_gen.args import AttentionConfig, SolAttentionConfig + + +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( + "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 ( + SkipSoftmaxScheduler.get_graph_phase_for_timestep(timestep, disabled_until_timestep=0.9545) + == expected + ) + + +def test_graph_phase_none_when_prefix_unset(): + assert ( + SkipSoftmaxScheduler.get_graph_phase_for_timestep(0.5, disabled_until_timestep=None) is None + ) + + +def test_graph_phase_accepts_tensor_timestep(): + """Pipelines pass a tensor; a 0-d or 1-element tensor must work.""" + assert ( + SkipSoftmaxScheduler.get_graph_phase_for_timestep( + torch.tensor(0.99), disabled_until_timestep=0.95 + ) + == 0 + ) + assert ( + SkipSoftmaxScheduler.get_graph_phase_for_timestep( + torch.tensor([0.10]), disabled_until_timestep=0.95 + ) + == 1 + ) + + +def test_dense_prefix_skips_kernel(monkeypatch): + """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 + + +def _fake_cuda_q(shape, dtype=torch.bfloat16): + """A tensor-like that reports `is_cuda=True` without needing a GPU. + + `sol_attn_ineligible_reason` checks `is_cuda` first and returns early, so a + CPU tensor can never reach the rank, head_dim, or dtype branches. These + stubs pass that first gate so each later reason is actually exercised; the + architecture check comes after them and is not reached. + """ + return SimpleNamespace(is_cuda=True, ndim=len(shape), shape=shape, dtype=dtype) + + +@pytest.mark.parametrize( + "make,expect", + [ + (lambda: torch.randn(1, 4, 2, 128), "not a CUDA tensor"), + (lambda: _fake_cuda_q((1, 4, 128)), "must be 4-D"), + (lambda: _fake_cuda_q((1, 4, 2, 64)), "head_dim must be 128"), + (lambda: _fake_cuda_q((1, 4, 2, 128), dtype=torch.float16), "dtype must be bfloat16"), + ], + ids=["cpu-tensor", "wrong-rank", "wrong-head-dim", "wrong-dtype"], +) +def test_ineligible_reason_is_reported(make, expect): + """Ineligibility must name the specific reason, never fail silently. + + Each case is built to fail exactly one check, in the order the function + applies them, so the id names the branch that actually fires. + """ + reason = _backend_mod().sol_attn_ineligible_reason(make()) + assert reason is not None and expect in reason + assert not _backend_mod().sol_attn_supported(make()) + + +def test_strict_raises_on_ineligible_input(monkeypatch): + """TRTLLM_SOL_ATTN_STRICT=1 must turn an unservable input into an error. + Without it an unsupported arch degrades to dense silently and the counters + the PR relies on cannot be trusted.""" + sab = _backend_mod() + monkeypatch.setenv("TRTLLM_SOL_ATTN_STRICT", "1") + q = k = v = torch.randn(1, 4, 2, 128) # CPU -> ineligible + with pytest.raises(RuntimeError, match="cannot run the CuTe kernel"): + sab._run_sol_attn_bthd(q, k, v) + + +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("TRTLLM_SOL_ATTN_STRICT", raising=False) + sab.reset_sol_attn_stats() + q = k = v = torch.randn(1, 4, 2, 128) + out = sab._run_sol_attn_bthd(q, k, v) + 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( + major * 10 + minor for major, minor in interface._CUTE_BACKENDS + ) + + +def test_datacenter_blackwell_archs_are_supported(): + """Both datacenter Blackwell steppings dispatch to the vendored kernel. + + sm103 (B300/GB300) runs the same `sm100/` kernel as sm100: the CuTe DSL JIT + targets the device it compiles on and the kernel body uses no sm100-only + construct. Verified on hardware against a B200 control; see the kernel's + THIRD_PARTY_NOTICES.md. + """ + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.sol_attn import interface + + # Resolve, don't just look up: a malformed value would satisfy a key check + # and then fail at dispatch. + for arch in ((10, 0), (10, 3)): + backend = interface._backend_for_arch(arch, cute_available=True) + assert backend in interface._CUTE_BACKENDS.values() + assert isinstance(backend, str) and backend + + +def test_no_arch_literal_outside_the_dispatch_map(): + """`_sol_attn_cute`'s guard must key off _CUTE_BACKENDS, not a literal. + + That guard is a second, deeper check than `_backend_for_arch`, and it is + the one no other test reaches: `test_supported_archs_matches_kernel_dispatch_map` + keeps SUPPORTED_ARCHS and _CUTE_BACKENDS in step, so widening both leaves a + hardcoded literal here as the only thing still rejecting the new arch -- + with every test green. Assert on the source so the coupling cannot regress. + """ + import ast + import inspect + import textwrap + + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.sol_attn import interface + + src = textwrap.dedent(inspect.getsource(interface._sol_attn_cute)) + assert "arch not in _CUTE_BACKENDS" in src, ( + "the guard in _sol_attn_cute must be keyed off _CUTE_BACKENDS" + ) + + # Reject *any* architecture tuple literal, not just (10, 0): pinning one + # value would let a future `arch != (10, 3)` reintroduce the same trap. + offenders = [] + for node in ast.walk(ast.parse(src)): + if not isinstance(node, ast.Compare): + continue + operands = [node.left, *node.comparators] + names = {n.id for n in operands if isinstance(n, ast.Name)} + if "arch" not in names: + continue + for operand in operands: + if isinstance(operand, ast.Tuple) and all( + isinstance(e, ast.Constant) and isinstance(e.value, int) for e in operand.elts + ): + offenders.append(ast.unparse(node)) + assert not offenders, ( + "hardcoded architecture literal(s) compared against `arch` in " + f"_sol_attn_cute: {offenders}; key the guard off _CUTE_BACKENDS instead" + ) + + +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( + "layers", + [[-1], [0, -2], ["x"], "0,2", [1.5]], + ids=["negative", "negative_in_list", "non_numeric", "string_spec", "float"], +) +def test_dense_layers_rejects_invalid_values(layers): + """dense_layers is a list of non-negative layer indices; anything else fails at + config time rather than at attention construction.""" + with pytest.raises(ValueError): + SolAttentionConfig(tau=2.0, dense_layers=layers) + + +@pytest.mark.parametrize( + "layers,expected", + [(None, None), ([0], [0]), ([0, 2, 4], [0, 2, 4]), ([4, 2, 2, 0], [0, 2, 4])], + ids=["none", "single", "list", "unsorted_with_duplicates"], +) +def test_dense_layers_normalizes_valid_values(layers, expected): + assert SolAttentionConfig(tau=2.0, dense_layers=layers).dense_layers == expected + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Sol-Attn needs CUDA") +def test_cute_kernel_matches_dense_on_a_single_block(): + """Kernel-level numerical check against dense attention. + + The obvious form of this test -- the analogue of VSA's + `test_cute_kernel_matches_dense_at_full_topk` -- needs a tau that provably + forces full non-sparse routing. Sol-Attn has no such value: its routing is + score-derived rather than a plain top-k, so tau=0 is not the dense limit. + + A single KV block sidesteps that entirely. With `tokens <= BLOCK_SIZE` + there is exactly one block, so there is nothing to route away and sparsity + is structurally impossible whatever tau says. The kernel must therefore + reproduce dense attention here, which makes this a genuine numerical test + of the kernel rather than of the routing policy. + + `kernel_calls` is asserted to have advanced: without that, a fallback to + dense would make this pass by comparing dense against itself. + """ + sab = _backend_mod() + if not sab.sol_attn_supported(torch.empty(1, 8, 8, 128, device="cuda", dtype=torch.bfloat16)): + pytest.skip("no Sol-Attn kernel for this device") + + tokens = 64 # == BLOCK_SIZE: exactly one KV block + torch.manual_seed(0) + shape = (1, tokens, 4, 128) + q, k, v = (torch.randn(shape, device="cuda", dtype=torch.bfloat16) for _ in range(3)) + + reference = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2).float(), k.transpose(1, 2).float(), v.transpose(1, 2).float() + ).transpose(1, 2) + + before = sab._SOL_STATS["kernel_calls"] + out = sab._run_sol_attn_bthd(q, k, v, tau=2.0, thresh_type="diag", kv_splits=1) + torch.cuda.synchronize() + assert sab._SOL_STATS["kernel_calls"] == before + 1, ( + "the kernel did not run; this comparison would be dense against dense" + ) + + torch.testing.assert_close(out.float(), reference, rtol=2e-2, atol=2e-2) + + +def test_kv_splits_is_not_a_config_knob(): + """Only one KV split exists on the shipped kernels, so the config exposes + no `kv_splits` and rejects it rather than accepting a no-op choice.""" + assert "kv_splits" not in SolAttentionConfig.model_fields + with pytest.raises(ValueError): + SolAttentionConfig(tau=2.0, kv_splits="1") + + +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 + + if not sol_attn_mod._cute_dense_available(): + # A CUDA device is not enough: the premise here is that the dense paths + # reach `cute_dsl_fmha_fwd`, which only exists on sm100/sm103. Without + # this the test fails rather than skips on any other CUDA runner. + pytest.skip("no CuTe DSL dense kernel for this device") + + device = torch.device("cuda") + q = k = v = torch.randn(1, 64, 2, 128, device=device, dtype=torch.bfloat16) + + 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" + + +# --------------------------------------------------------------------------- +# Dense-prefix phase under CUDA-graph capture +# --------------------------------------------------------------------------- + + +def test_dense_by_step_prefers_runner_resolved_phase(monkeypatch): + """Inside a runner-driven call the phase comes from the runner, not the tensor. + + CUDA-graph capture forbids device-to-host syncs, and reading the timestep + is a `.item()`. The runner resolves the phase host-side to build the graph + key and republishes it during capture; `_dense_by_step` must use that and + never touch the tensor. `torch.compiler.disable` does not help here -- it + only excludes Dynamo, not stream capture. + """ + reads = {"n": 0} + real = SkipSoftmaxScheduler._as_float + + def spy(value): + reads["n"] += 1 + return real(value) + + monkeypatch.setattr(SkipSoftmaxScheduler, "_as_float", staticmethod(spy)) + attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) + attn.disabled_until_timestep = 0.9 + + # Resolved phase wins even when the tensor says otherwise: phase 0 is the + # dense prefix, phase 1 the sparse phase. + with resolved_extra_keys_scope({"sol_attn_phase": 0}): + assert attn._dense_by_step(torch.tensor(0.1)) is True + with resolved_extra_keys_scope({"sol_attn_phase": 1}): + assert attn._dense_by_step(torch.tensor(0.99)) is False + assert reads["n"] == 0, "timestep tensor was read despite a runner-resolved phase" + + # Outside a runner-driven call the tensor is the only source of truth. + assert attn._dense_by_step(torch.tensor(0.99)) is True + assert reads["n"] == 1 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graphs need a GPU") +def _base_runner(): + return CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) + + +def _ltx2_runner(): + from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2 import _LTX2CUDAGraphRunner + + return _LTX2CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) + + +def _ltx2_two_stage_runner(): + from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2_two_stages import ( + _LTX2TwoStageCUDAGraphRunner, + ) + + return _LTX2TwoStageCUDAGraphRunner( + CUDAGraphRunnerConfig(use_cuda_graph=True), lambda: "original", lambda: "default" + ) + + +# Every runner that overrides ``capture`` must keep the phase-scope contract; +# LTX-2's runners re-implement capture for ``Modality`` inputs. +_RUNNER_FACTORIES = pytest.mark.parametrize( + "make_runner", + [_base_runner, _ltx2_runner, _ltx2_two_stage_runner], + ids=["base", "ltx2", "ltx2_two_stage"], +) + + +@_RUNNER_FACTORIES +def test_cuda_graph_runner_publishes_resolved_extra_keys_during_capture(make_runner): + """The runner exposes its host-resolved extra keys to the captured callee.""" + seen = [] + + def fn(x): + seen.append(resolved_extra_key("probe")) + return x * 2 + + runner = make_runner() + runner.register_extra_key_fn("probe", lambda *args, **kwargs: 7) + wrapped = runner.wrap(fn) + x = torch.ones(4, device="cuda") + out = wrapped(x) + torch.cuda.synchronize() + assert torch.equal(out, x * 2) + assert seen and all(v == 7 for v in seen), seen # warmup + capture passes + assert resolved_extra_key("probe") is None, "scope must close after capture" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Sol-Attn needs CUDA") +@_RUNNER_FACTORIES +def test_sol_attn_dense_prefix_survives_cuda_graph_capture(make_runner): + """Capture and replay a graph on each side of the cutoff without a sync. + + The dense-prefix decision must be baked into each captured graph via the + runner's phase key, and no `.item()` may run inside capture. `kernel_calls` + proves the phase-0 graph never launched the sparse kernel and the phase-1 + graph did. + """ + sab = _backend_mod() + if not sab.sol_attn_supported(torch.empty(1, 8, 8, 128, device="cuda", dtype=torch.bfloat16)): + pytest.skip("no Sol-Attn kernel for this device") + + cutoff = 0.9 + attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) + attn.disabled_until_timestep = cutoff + runner = make_runner() + runner.register_extra_key_fn( + "sol_attn_phase", + lambda *args, **kwargs: SkipSoftmaxScheduler.get_graph_phase_for_timestep( + kwargs.get("timestep"), disabled_until_timestep=cutoff + ), + ) + + def fwd(q, k, v, *, timestep): + return attn.forward(q, k, v, timestep=timestep) + + wrapped = runner.wrap(fwd) + torch.manual_seed(0) + q = k = v = torch.randn(1, 64, 2, 128, device="cuda", dtype=torch.bfloat16) + t_dense = torch.tensor(0.95, device="cuda") + t_sparse = torch.tensor(0.5, device="cuda") + + before = sab._SOL_STATS["kernel_calls"] + wrapped(q, k, v, timestep=t_dense) + torch.cuda.synchronize() + after_dense = sab._SOL_STATS["kernel_calls"] + wrapped(q, k, v, timestep=t_sparse) + torch.cuda.synchronize() + after_sparse = sab._SOL_STATS["kernel_calls"] + + assert len(runner.graphs) == 2, "one graph per phase" + assert after_dense == before, "dense prefix must not launch the sparse kernel" + assert after_sparse > after_dense, "sparse phase must launch the kernel" + # Replays must not error either. + wrapped(q, k, v, timestep=t_dense) + wrapped(q, k, v, timestep=t_sparse) + torch.cuda.synchronize() + + +# --------------------------------------------------------------------------- +# Kernel vs a PyTorch reference of the block-mean approximation +# --------------------------------------------------------------------------- + + +def _sol_attn_reference(q, k, v, *, tau, scale, block=64): + """PyTorch fp32 reference of Sol-Attn's routing + block-mean approximation. + + Mirrors the vendored kernel (`preprocess.py`, `common/selector.py`, + `sm100/mainloop.py`): + + * `kc` = per-block mean of K, `vc` = per-block *sum* of V. + * Per query block: threshold = mean + tau * std, in the log2 domain, where + mean/var come from the query centroid against the global K mean and + per-dim K variance (`_diag_threshold_kernel`). + * A KV block is *exact* for a query block if the column-mean route score + exceeds that threshold, or it is within one block of the diagonal + (`sol_attn_route_is_exact`). + * Non-exact blocks contribute `exp(q . kc) * vc` to the numerator and + `exp(q . kc) * block_len` to the denominator -- every key in the block + is treated as its mean. + + Returns the output and the exact-routing mask `[B, nb_q, nb_kv, H]`. + """ + q, k, v = (t.float() for t in (q, k, v)) + B, S, H, D = q.shape + nb = -(-S // block) + log2e = math.log2(math.e) + log2_scale = scale * log2e + dev = q.device + idx = torch.arange(S, device=dev) // block + blen = torch.bincount(idx, minlength=nb).float() + + def _block_sum(t): + return torch.zeros(B, nb, H, D, device=dev).index_add_(1, idx, t) + + kc = _block_sum(k) / blen[None, :, None, None] + vc = _block_sum(v) + qc = _block_sum(q) / blen[None, :, None, None] + kmean = k.mean(dim=1) + kvar = ((k * k).mean(dim=1) - kmean * kmean).clamp(min=0) + mean = (qc * kmean[:, None]).sum(-1) * log2_scale + var = (qc * qc * kvar[:, None]).sum(-1) * log2_scale**2 + thr = mean + tau * torch.sqrt(var.clamp(min=0) + 1e-6) # [B, nb_q, H] + + s2 = torch.einsum("bshd,bjhd->bsjh", q, kc) * log2_scale # [B, S, nb_kv, H] + col_mean = ( + torch.zeros(B, nb, nb, H, device=dev).index_add_(1, idx, s2) / blen[None, :, None, None] + ) + ar = torch.arange(nb, device=dev) + near_diag = (ar[:, None] - ar[None, :]).abs() <= 1 + exact = (col_mean > thr[:, :, None, :]) | near_diag[None, :, :, None] # [B, nb_q, nb_kv, H] + + exact_rows = exact[:, idx] # [B, S, nb_kv, H] + exact_pair = exact_rows[:, :, idx, :] # [B, S, T, H] + scores = torch.einsum("bshd,bthd->bsth", q, k) * scale + approx = s2 / log2e + neg = torch.finfo(torch.float32).min + m = torch.maximum( + scores.masked_fill(~exact_pair, neg).amax(2), + approx.masked_fill(exact_rows, neg).amax(2), + ) + pe = torch.exp(scores - m[:, :, None]) * exact_pair + pa = torch.exp(approx - m[:, :, None]) * (~exact_rows) + num = torch.einsum("bsth,bthd->bshd", pe, v) + torch.einsum("bsjh,bjhd->bshd", pa, vc) + den = pe.sum(2) + (pa * blen[None, None, :, None]).sum(2) + return num / den[..., None], exact + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Sol-Attn needs CUDA") +def test_cute_kernel_matches_reference_on_multi_block_with_approximation(): + """The real kernel agrees with the reference where blocks are approximated. + + Complements the single-block test, which can only reach the all-exact + path. Plain Gaussian inputs are the right stimulus: every KV block carries + comparable softmax mass and no block's column-mean score clears the + `mean + tau*std` threshold, so every block two or more away from the + diagonal is approximated -- and because keys vary within a block, treating + 64 keys as their mean is measurably different from dense attention (a + Jensen gap). Clustered inputs are the wrong stimulus: distant blocks get + routed to the approximation but carry no mass, so approximating them + changes nothing and the test would pass vacuously. The reference computes + the same routing, and the test asserts that some blocks were approximated + *and* that the result differs from dense, so neither an accidental + all-exact run nor a mass-free approximation can pass. It also asserts + `kernel_calls` advanced, so a dense fallback cannot pass either. + """ + sab = _backend_mod() + if not sab.sol_attn_supported(torch.empty(1, 8, 8, 128, device="cuda", dtype=torch.bfloat16)): + pytest.skip("no Sol-Attn kernel for this device") + + torch.manual_seed(0) + S, H, D, block = 256, 2, 128, 64 # 4 KV blocks; batch of 1 + dev = torch.device("cuda") + q, k, v = (torch.randn(1, S, H, D, device=dev, dtype=torch.bfloat16) for _ in range(3)) + scale = D**-0.5 + tau = 2.0 + + ref, exact = _sol_attn_reference(q, k, v, tau=tau, scale=scale, block=block) + assert (~exact).any(), "inputs did not force any block onto the approximation path" + dense = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2).float(), k.transpose(1, 2).float(), v.transpose(1, 2).float() + ).transpose(1, 2) + # Calibrated on B200: max |ref - dense| = 0.45 for this stimulus. A tight + # margin here would let a mass-free approximation pass vacuously. + assert (ref - dense).abs().max() > 1e-1, "approximation is not measurably different from dense" + + before = sab._SOL_STATS["kernel_calls"] + out = sab._run_sol_attn_bthd(q, k, v, tau=tau, thresh_type="diag", kv_splits=1) + torch.cuda.synchronize() + assert sab._SOL_STATS["kernel_calls"] == before + 1, "kernel did not run" + + # Calibrated on B200: max |kernel - ref| = 1.7e-3 (mean 2e-4); 5e-3 leaves + # ~3x headroom for bf16 rounding while still catching a routing or + # approximation-formula mismatch, which shows up at 1e-1 or worse. + torch.testing.assert_close(out.float(), ref, rtol=5e-3, atol=5e-3) + + +# --------------------------------------------------------------------------- +# Masks: routed to a mask-aware backend, and honored -- not just redirected +# --------------------------------------------------------------------------- + + +def _masked_sdpa_reference(q, k, v, *, key_padding_mask=None, is_causal=False): + attn_mask = ( + None if key_padding_mask is None else key_padding_mask.to(torch.bool)[:, None, None, :] + ) + return torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2).float(), + k.transpose(1, 2).float(), + v.transpose(1, 2).float(), + attn_mask=attn_mask, + is_causal=is_causal and attn_mask is None, + ).transpose(1, 2) + + +def test_key_padding_mask_routes_to_vanilla_and_is_honored(monkeypatch): + """Masked self-attention must never reach the mask-blind sparse kernel. + + Equal Q/K lengths do not mean "unmasked": HunyuanVideo1.5 and GLM-Image + pass a `[B, S]` `key_padding_mask` on self-attention. The sparse kernel + takes no mask and CuTeDSL's dense `_fwd(**kwargs)` swallows it, so the only + correct destination is VANILLA. The output is checked against a masked + SDPA reference so a mask that was routed but then dropped still fails. + """ + import tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn as sol_attn_mod + + monkeypatch.setattr( + sol_attn_mod, + "_sol_attn_run", + lambda *a, **k: pytest.fail("sparse kernel ran on a masked call"), + ) + attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) + attn.disabled_until_timestep = None + calls = {"vanilla": 0} + real = attn._vanilla.forward + + def spy(*a, **k): + calls["vanilla"] += 1 + return real(*a, **k) + + monkeypatch.setattr(attn._vanilla, "forward", spy) + + torch.manual_seed(0) + q, k, v = (torch.randn(2, 64, 2, 128, dtype=torch.float32) for _ in range(3)) + mask = torch.ones(2, 64, dtype=torch.bool) + mask[0, 48:] = False # pad the tail of sample 0 + mask[1, :16] = False # pad the head of sample 1 + + assert not attn._can_serve(q, k, key_padding_mask=mask) + out = attn.forward(q, k, v, key_padding_mask=mask) + assert calls["vanilla"] == 1 + ref = _masked_sdpa_reference(q, k, v, key_padding_mask=mask) + torch.testing.assert_close(out.float(), ref, rtol=1e-4, atol=1e-4) + # And the mask actually mattered: unmasked attention gives a different answer. + assert (ref - _masked_sdpa_reference(q, k, v)).abs().max() > 1e-3 + + +def test_causal_mask_routes_to_dense_and_is_honored(monkeypatch): + """CAUSAL disqualifies the noncausal sparse kernel and is honored downstream. + + On sm100 the delegate is dense CuTeDSL, which supports CAUSAL; on any other + device it is `_sdpa`, which previously dropped it. Either way the result + must match a causal reference. + """ + import tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn as sol_attn_mod + + monkeypatch.setattr( + sol_attn_mod, + "_sol_attn_run", + lambda *a, **k: pytest.fail("sparse kernel ran on a causal call"), + ) + attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) + attn.disabled_until_timestep = None + torch.manual_seed(0) + q, k, v = (torch.randn(1, 64, 2, 128, dtype=torch.float32) for _ in range(3)) + + assert not attn._can_serve(q, k, attention_mask=PredefinedAttentionMask.CAUSAL) + out = attn.forward(q, k, v, attention_mask=PredefinedAttentionMask.CAUSAL) + ref = _masked_sdpa_reference(q, k, v, is_causal=True) + torch.testing.assert_close(out.float(), ref, rtol=1e-4, atol=1e-4) + assert (ref - _masked_sdpa_reference(q, k, v)).abs().max() > 1e-3 + + +def test_unmasked_self_attention_is_still_served(): + """The routing change must not touch the measured path: no mask, sparse.""" + attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) + attn.disabled_until_timestep = None + q = k = torch.randn(1, 64, 2, 128, dtype=torch.bfloat16) + assert attn._can_serve(q, k) + assert attn._can_serve(q, k, attention_mask=PredefinedAttentionMask.FULL, key_padding_mask=None) diff --git a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py index 4a5c984f6e23..ce4614969841 100644 --- a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py @@ -1253,6 +1253,9 @@ def __init__(self): def forward(self, x): return self.lin(x) + def register_cuda_graph_extra_key_fns(self, runner): + self.registered_runner = runner + pipeline = object.__new__(ltx2_two_stages.LTX2TwoStagesPipeline) # __init__ is bypassed here; BasePipeline.device now reads self._device # (set in __init__), so initialize it explicitly for the warmup path. @@ -1310,9 +1313,14 @@ def test_two_stage_cuda_graph_setup_uses_pipeline_config(self): """CUDA graph setup runs before the two-stage model_config is assigned.""" class TinyTransformer: + registered_runner = None + def forward(self, *args, **kwargs): return args, kwargs + def register_cuda_graph_extra_key_fns(self, runner): + self.registered_runner = runner + pipeline = object.__new__(ltx2_two_stages.LTX2TwoStagesPipeline) pipeline.pipeline_config = DiffusionPipelineConfig( cuda_graph=CudaGraphConfig(enable=True), @@ -1328,6 +1336,9 @@ def forward(self, *args, **kwargs): runner = pipeline._cuda_graph_runners["transformer"] assert isinstance(runner, ltx2_two_stages._LTX2TwoStageCUDAGraphRunner) assert runner._lora_state_getter() == "original" + # Same contract as the single-stage pipeline: the transformer's extra + # CUDA-graph keys (skip-softmax / Sol-Attn phase) must be registered. + assert pipeline.transformer.registered_runner is runner assert pipeline.transformer.forward.__wrapped__.__self__ is pipeline.transformer def test_cuda_graph_rejects_nonpersistent_lora_bindings(self):