diff --git a/docs/en/attention.md b/docs/en/attention.md index 08cad047..80ad5a72 100644 --- a/docs/en/attention.md +++ b/docs/en/attention.md @@ -31,6 +31,7 @@ class AttnImplType(Enum): # Sparse attention RADIAL_ATTN = auto() LOCAL_SPARSE_ATTN = auto() + SOL_ATTN = auto() ``` ### AttentionConfig @@ -51,6 +52,7 @@ Factory methods: - `AttentionConfig.dense_attention(attn_impl)` - Create dense attention config - `AttentionConfig.radial_attention(**kwargs)` - Create radial sparse attention config - `AttentionConfig.local_sparse_attention(**kwargs)` - Create local sparse attention config +- `AttentionConfig.sol_attention(**kwargs)` - Create dynamic Sol-Attn config ### SparseAttentionConfig @@ -59,13 +61,16 @@ Configuration for sparse attention: ```python @dataclass class SparseAttentionConfig: - sparse_impl: str | None = None # "radial", "local", etc. + sparse_impl: str | None = None # "radial", "local", "sol", etc. dense_timesteps: int = 40 # Use dense attention for initial timesteps dense_layers: int = 0 # Use dense attention for initial layers decay_factor: float = 1.0 # Decay factor for attention window local_window_size: int = 6 # Window size for local sparse attention block_size: int = 128 # Block size for sparse computation use_sage_attention: bool = False # Use sage attention backend + sol_tau: float = 1.0 # Sol-Attn routing threshold + sol_threshold_type: str = "diag" # "diag" or "exact" + sol_kv_splits: int | str = "auto" # "auto", 1, 2, or 4 ``` ## Calling Flow @@ -155,12 +160,12 @@ else: ## Pipeline Support Status -| Pipeline | Dense Attention | Sparse (Radial) | Notes | -|----------|-----------------|-----------------|-------| -| `Wan21VideoPipeline` | ✅ | ✅ | Full support for video generation | -| `Wan22VideoPipeline` | ✅ | ✅ | Full support for video generation | -| `QwenImagePipeline` | ✅ | ❌ | Image generation doesn't need temporal sparse attention | -| `ZImagePipeline` | ✅ | ❌ | Image generation doesn't need temporal sparse attention | +| Pipeline | Dense Attention | Radial | Sol-Attn | Notes | +|----------|-----------------|--------|----------|-------| +| `Wan21VideoPipeline` | Yes | Yes | Experimental | Sol-Attn covers eligible self-attention calls | +| `Wan22VideoPipeline` | Yes | Yes | No | Sol-Attn is not wired into Wan2.2 yet | +| `QwenImagePipeline` | Yes | No | No | Image generation doesn't need temporal sparse attention | +| `ZImagePipeline` | Yes | No | No | Image generation doesn't need temporal sparse attention | ### Wan21VideoPipeline / Wan22VideoPipeline @@ -185,6 +190,18 @@ When using radial attention: 3. Updates state per timestep/layer in denoising loop 4. Automatically falls back to dense for early timesteps/layers +Wan2.1 can select Sol-Attn through the same pipeline configuration surface: + +```python +config = AttentionConfig.sol_attention() +pipe_config.dit_config.attention_config = config +``` + +Sol-Attn is used only for contiguous, noncausal BF16 self-attention with equal Q/K/V +shapes and head dimension 128. Unsupported calls, dense warmup layers or timesteps, +and kernel runtime failures fall back to the existing dense attention path. Ring/USP +also remains dense because its online merge requires log-sum-exp output. + ### QwenImagePipeline / ZImagePipeline Supports only dense attention (image generation doesn't have temporal dimension): @@ -218,9 +235,14 @@ pipe_config.dit_config.attention_config = config |---------|-------------|--------------| | `RADIAL_ATTN` | Radial attention for video | `flashinfer` or `sageattention` (tf-kernel prioritized) | | `LOCAL_SPARSE_ATTN` | Local window sparse attention | `block_sparse_attn` | +| `SOL_ATTN` | Dynamic block-sparse video attention | Built in; BF16, head dimension 128, SM80+ | **Note on SageAttention Priority**: When `use_sage_attention=True` is set, the system will prioritize tf-kernel's sageattention implementation over the standalone `sageattention` package if both are available. This provides better performance and integration with the TeleFuser kernel library. +**Sol-Attn packaging**: Sol-Attn ships with TeleFuser under `telefuser.kernel.sol_attn`; it does not require +`tf-kernel`. The upstream runtime targets PyTorch 2.10+, CUDA 12.8+, and Triton 3.6+. Specialized CuTe DSL +kernels are selected when that optional runtime is available, otherwise SM80+ uses the Triton implementation. + ### Installing Sparge Attention To use `SPARGE_ATTN` backend or sparse sage attention in radial attention, you need to install `spas_sage_attn` from source: @@ -428,6 +450,7 @@ print(f"FlashInfer: {FLASHINFER_AVAILABLE}") | Flash Attention 4 | Build from source (cute interface) | SM90+ (H100, B100/B200) | | SageAttention | tf-kernel or [official source](https://github.com/thu-ml/SageAttention) | SM80+ | | Radial Attention | tf-kernel or [FlashInfer source](https://github.com/flashinfer-ai/flashinfer) | SM80+ | +| Sol-Attn | Built into TeleFuser | SM80+; CuTe on supported SM90/100, Triton fallback | | Block Sparse | tf-kernel or [official source](https://github.com/mit-han-lab/Block-Sparse-Attention) | SM80+ | | Sparge Attention | Install from source (see above) | SM80, SM86, SM89, SM90 | diff --git a/docs/zh/attention.md b/docs/zh/attention.md index abe0cb79..b44c0729 100644 --- a/docs/zh/attention.md +++ b/docs/zh/attention.md @@ -31,6 +31,7 @@ class AttnImplType(Enum): # 稀疏注意力 RADIAL_ATTN = auto() LOCAL_SPARSE_ATTN = auto() + SOL_ATTN = auto() ``` ### AttentionConfig @@ -51,6 +52,7 @@ class AttentionConfig: - `AttentionConfig.dense_attention(attn_impl)` - 创建密集注意力配置 - `AttentionConfig.radial_attention(**kwargs)` - 创建径向稀疏注意力配置 - `AttentionConfig.local_sparse_attention(**kwargs)` - 创建局部稀疏注意力配置 +- `AttentionConfig.sol_attention(**kwargs)` - 创建动态 Sol-Attn 配置 ### SparseAttentionConfig @@ -59,13 +61,16 @@ class AttentionConfig: ```python @dataclass class SparseAttentionConfig: - sparse_impl: str | None = None # "radial", "local" 等 + sparse_impl: str | None = None # "radial", "local", "sol" 等 dense_timesteps: int = 40 # 初始时间步使用密集注意力 dense_layers: int = 0 # 初始层使用密集注意力 decay_factor: float = 1.0 # 注意力窗口衰减因子 local_window_size: int = 6 # 局部稀疏注意力窗口大小 block_size: int = 128 # 稀疏计算块大小 use_sage_attention: bool = False # 使用 sage attention 后端 + sol_tau: float = 1.0 # Sol-Attn 路由阈值 + sol_threshold_type: str = "diag" # "diag" 或 "exact" + sol_kv_splits: int | str = "auto" # "auto"、1、2 或 4 ``` ## 调用流程 @@ -155,12 +160,12 @@ else: ## Pipeline 支持情况 -| Pipeline | 密集注意力 | 稀疏 (径向) | 说明 | -|----------|-----------|------------|------| -| `Wan21VideoPipeline` | ✅ | ✅ | 视频生成完整支持 | -| `Wan22VideoPipeline` | ✅ | ✅ | 视频生成完整支持 | -| `QwenImagePipeline` | ✅ | ❌ | 图像生成不需要时序稀疏注意力 | -| `ZImagePipeline` | ✅ | ❌ | 图像生成不需要时序稀疏注意力 | +| Pipeline | 密集注意力 | Radial | Sol-Attn | 说明 | +|----------|-----------|--------|----------|------| +| `Wan21VideoPipeline` | 支持 | 支持 | 实验性 | Sol-Attn 用于满足约束的 self-attention | +| `Wan22VideoPipeline` | 支持 | 支持 | 不支持 | 尚未接入 Wan2.2 | +| `QwenImagePipeline` | 支持 | 不支持 | 不支持 | 图像生成不需要时序稀疏注意力 | +| `ZImagePipeline` | 支持 | 不支持 | 不支持 | 图像生成不需要时序稀疏注意力 | ### Wan21VideoPipeline / Wan22VideoPipeline @@ -185,6 +190,17 @@ pipe_config.dit_config.attention_config = config 3. 在去噪循环中每时间步/层更新状态 4. 早期时间步/层自动回退到密集注意力 +Wan2.1 可以通过同一配置入口启用 Sol-Attn： + +```python +config = AttentionConfig.sol_attention() +pipe_config.dit_config.attention_config = config +``` + +Sol-Attn 仅用于连续、非因果、BF16、Q/K/V 形状相同且 head dimension 为 128 的 +self-attention。其他调用、dense 预热层/时间步以及内核运行失败都会回退到现有密集路径。 +Ring/USP 需要 LSE 做在线合并，因此仍使用支持 LSE 的密集后端。 + ### QwenImagePipeline / ZImagePipeline 仅支持密集注意力（图像生成没有时序维度）： @@ -218,9 +234,14 @@ pipe_config.dit_config.attention_config = config |------|------|------| | `RADIAL_ATTN` | 视频径向注意力 | `flashinfer` 或 `sageattention` (优先使用 tf-kernel) | | `LOCAL_SPARSE_ATTN` | 局部窗口稀疏注意力 | `block_sparse_attn` | +| `SOL_ATTN` | 动态块稀疏视频注意力 | 内置；BF16、head dimension 128、SM80+ | **SageAttention 优先级说明**: 当设置 `use_sage_attention=True` 时，如果 tf-kernel 和独立的 `sageattention` 包都可用，系统将优先使用 tf-kernel 的 sageattention 实现。这提供了更好的性能和与 TeleFuser 内核库的集成。 +**Sol-Attn 打包方式**：Sol-Attn 随 TeleFuser 发布，位于 `telefuser.kernel.sol_attn`，不依赖 +`tf-kernel`。上游运行时要求 PyTorch 2.10+、CUDA 12.8+ 和 Triton 3.6+。可选 CuTe DSL runtime +可用时选择专用内核，否则 SM80+ 使用 Triton 实现。 + ### 安装 Sparge Attention 要使用 `SPARGE_ATTN` 后端或径向注意力中的稀疏 sage attention，需要从源码安装 `spas_sage_attn`： @@ -426,6 +447,7 @@ print(f"FlashInfer: {FLASHINFER_AVAILABLE}") | Flash Attention 4 | 从源码编译（cute 接口） | SM90+ (H100, B100/B200) | | SageAttention | tf-kernel 或 [官方源码](https://github.com/thu-ml/SageAttention) | SM80+ | | Radial Attention | tf-kernel 或 [FlashInfer 源码](https://github.com/flashinfer-ai/flashinfer) | SM80+ | +| Sol-Attn | TeleFuser 内置 | SM80+；支持的 SM90/100 使用 CuTe，否则回退 Triton | | Block Sparse | tf-kernel 或 [官方源码](https://github.com/mit-han-lab/Block-Sparse-Attention) | SM80+ | | Sparge Attention | 从源码安装（见上文） | SM80, SM86, SM89, SM90 | diff --git a/examples/wan_video/README.md b/examples/wan_video/README.md index c97a1aa7..9a7e03fc 100644 --- a/examples/wan_video/README.md +++ b/examples/wan_video/README.md @@ -163,6 +163,17 @@ python examples/wan_video/wan21_1_3b_text_to_video_radial.py \ - Reduced memory usage for long videos - Requires flashinfer or sageattention backend +Wan2.1 also supports Sol-Attn through the same attention configuration: + +```python +from telefuser.core.config import AttentionConfig + +pipe_config.dit_config.attention_config = AttentionConfig.sol_attention() +``` + +Sol-Attn is built into TeleFuser. Eligible BF16 self-attention calls use the sparse kernel; unsupported calls +automatically use the existing dense fallback. The defaults follow the official Wan2.1 profile: Morton3D token ordering, dense layer 0, and 10 dense warm-up steps for the standard 50-step schedule. + #### wan21_1_3b_text_to_video_cache_calibrate.py Calibration tool for AdaTaylorCache. diff --git a/pyproject.toml b/pyproject.toml index c09b0782..45973c99 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -89,6 +89,7 @@ dev = [ "pytest-cov>=4.0.0", "pre-commit==4.0.1", "scikit-image>=0.19.0", + "tomli>=2.0.0; python_version < '3.11'", "uv" ] @@ -128,6 +129,12 @@ telefuser = "telefuser.entrypoints.cli.main:main" where = ["."] include = ["telefuser*"] +[tool.setuptools.package-data] +"telefuser.kernel.sol_attn" = [ + "THIRD_PARTY_NOTICES.md", + "sm100/LICENSE.flash-attention", +] + [tool.pytest.ini_options] minversion = "7.0" testpaths = ["tests"] @@ -204,7 +211,8 @@ exclude = [ "dist", "tf-kernel", "benchmarks", - "telefuser/_version.py" + "telefuser/_version.py", + "telefuser/kernel/sol_attn", ] [tool.ruff.lint] diff --git a/telefuser/core/config.py b/telefuser/core/config.py index 228af15c..505b481b 100644 --- a/telefuser/core/config.py +++ b/telefuser/core/config.py @@ -130,23 +130,35 @@ class AttnImplType(Enum): # Sparse attention implementations RADIAL_ATTN = auto() # Radial attention for video generation LOCAL_SPARSE_ATTN = auto() # Local window sparse attention + SOL_ATTN = auto() # Dynamic on-the-fly sparse attention for video generation @dataclass class SparseAttentionConfig: """Configuration for sparse attention implementations. - Used with radial or local sparse attention to reduce memory usage + Used with radial, local, or Sol sparse attention to reduce memory usage for long sequences like videos. """ - sparse_impl: str | None = None # "radial", "local", or None + sparse_impl: str | None = None # "radial", "local", "sol", or None dense_timesteps: int = 40 # Initial timesteps to use dense attention dense_layers: int = 0 # Initial layers to use dense attention decay_factor: float = 1.0 # Decay for radial attention window local_window_size: int = 6 # Window size for local attention block_size: int = 128 # Block size for sparse computation use_sage_attention: bool = False # Use sage attention backend + sol_tau: float = 1.0 # Sol-Attn routing threshold multiplier + sol_threshold_type: str = "diag" # Sol-Attn threshold estimator: "diag" or "exact" + sol_kv_splits: int | str = "auto" # Auto selects split 4 for long SM90 sequences + + def __post_init__(self) -> None: + if self.sparse_impl != "sol": + return + if self.sol_threshold_type not in ("diag", "exact"): + raise ValueError("Sol-Attn threshold type must be 'diag' or 'exact'") + if self.sol_kv_splits not in ("auto", 1, 2, 4): + raise ValueError("Sol-Attn KV splits must be 'auto', 1, 2, or 4") def should_use_dense(self, numeral_timestep: int, layer_idx: int) -> bool: """Check if dense attention should be used for current step/layer. @@ -209,14 +221,42 @@ def local_sparse_attention( **kwargs, ) + @classmethod + def sol_attention( + cls, + dense_timesteps: int = 10, + dense_layers: int = 1, + tau: float = 1.0, + threshold_type: str = "diag", + kv_splits: int | str = "auto", + **kwargs: any, + ) -> AttentionConfig: + """Create a Sol-Attn config for dynamic sparse video self-attention.""" + return cls( + attn_impl=AttnImplType.SOL_ATTN, + sparse_config=SparseAttentionConfig( + sparse_impl="sol", + dense_timesteps=dense_timesteps, + dense_layers=dense_layers, + sol_tau=tau, + sol_threshold_type=threshold_type, + sol_kv_splits=kv_splits, + ), + **kwargs, + ) + @classmethod def dense_attention(cls, attn_impl: AttnImplType = AttnImplType.FLASH_ATTN_2, **kwargs: any) -> AttentionConfig: """Create config for dense attention.""" return cls(attn_impl=attn_impl, sparse_config=None, **kwargs) def is_sparse(self) -> bool: - """Check if using sparse attention (radial or local).""" - return self.attn_impl in (AttnImplType.RADIAL_ATTN, AttnImplType.LOCAL_SPARSE_ATTN) + """Check if using a sparse attention implementation.""" + return self.attn_impl in ( + AttnImplType.RADIAL_ATTN, + AttnImplType.LOCAL_SPARSE_ATTN, + AttnImplType.SOL_ATTN, + ) def should_use_dense(self, numeral_timestep: int, layer_idx: int) -> bool: """Check if dense attention should be used for current step/layer.""" diff --git a/telefuser/kernel/sol_attn/THIRD_PARTY_NOTICES.md b/telefuser/kernel/sol_attn/THIRD_PARTY_NOTICES.md new file mode 100644 index 00000000..efe0a72c --- /dev/null +++ b/telefuser/kernel/sol_attn/THIRD_PARTY_NOTICES.md @@ -0,0 +1,23 @@ +# Third-party notices + +Sol-Attn was vendored from NVIDIA's `NVlabs/Sana` `sol-engine` branch at +commit `8a26fb0ec9e353125ead798cb2e312d5ce48cded`. The upstream repository +declares its code under the Apache License 2.0. Local changes move the +implementation under the internal `telefuser.kernel.sol_attn` namespace and +rewrite its absolute imports. + +Upstream source: https://github.com/NVlabs/Sana/tree/sol-engine/techniques/sparse_backends/sol_attn + +The files under `telefuser/kernel/sol_attn/_vendor/flash_attn/cute/` and portions of the SM90 +and SM100 design scaffold derive from the FlashAttention project. Its +BSD-3-Clause license is included at +`telefuser/kernel/sol_attn/sm100/LICENSE.flash-attention`. + +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. + +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. diff --git a/telefuser/kernel/sol_attn/__init__.py b/telefuser/kernel/sol_attn/__init__.py new file mode 100644 index 00000000..56e7bad5 --- /dev/null +++ b/telefuser/kernel/sol_attn/__init__.py @@ -0,0 +1,5 @@ +"""Sol-Attn.""" + +from .interface import sol_attn + +__all__ = ["sol_attn"] diff --git a/telefuser/kernel/sol_attn/_vendor/__init__.py b/telefuser/kernel/sol_attn/_vendor/__init__.py new file mode 100644 index 00000000..b5e34dbf --- /dev/null +++ b/telefuser/kernel/sol_attn/_vendor/__init__.py @@ -0,0 +1 @@ +"""Private source dependencies bundled with Sol-Attn.""" diff --git a/telefuser/kernel/sol_attn/_vendor/flash_attn/__init__.py b/telefuser/kernel/sol_attn/_vendor/flash_attn/__init__.py new file mode 100644 index 00000000..25565021 --- /dev/null +++ b/telefuser/kernel/sol_attn/_vendor/flash_attn/__init__.py @@ -0,0 +1 @@ +"""Local FlashAttention Cute shim for the SOL_ATTN SM90 release.""" diff --git a/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/__init__.py b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/__init__.py new file mode 100644 index 00000000..5db3a69e --- /dev/null +++ b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/__init__.py @@ -0,0 +1 @@ +"""Vendored FlashAttention Cute Python helpers used by SOL_ATTN SM90.""" diff --git a/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/ampere_helpers.py b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/ampere_helpers.py new file mode 100644 index 00000000..e3072d8c --- /dev/null +++ b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/ampere_helpers.py @@ -0,0 +1,103 @@ +# Copyright (c) 2025, Tri Dao. +from typing import Type, Callable, Optional + +import cutlass +import cutlass.cute as cute + + +def get_smem_layout_atom(dtype: Type[cutlass.Numeric], k_dim: int) -> cute.ComposedLayout: + dtype_byte = cutlass.const_expr(dtype.width // 8) + bytes_per_row = cutlass.const_expr(k_dim * dtype_byte) + smem_k_block_size = ( + cutlass.const_expr( + 128 + if bytes_per_row % 128 == 0 + else (64 if bytes_per_row % 64 == 0 else (32 if bytes_per_row % 32 == 0 else 16)) + ) + // dtype_byte + ) + swizzle_bits = ( + 4 + if smem_k_block_size == 128 + else (3 if smem_k_block_size == 64 else (2 if smem_k_block_size == 32 else 1)) + ) + swizzle_base = 2 if dtype_byte == 4 else (3 if dtype_byte == 2 else 4) + return cute.make_composed_layout( + cute.make_swizzle(swizzle_bits, swizzle_base, swizzle_base), + 0, + cute.make_ordered_layout( + (8 if cutlass.const_expr(k_dim % 32 == 0) else 16, smem_k_block_size), order=(1, 0) + ), + ) + + +@cute.jit +def gemm( + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + tCsA: cute.Tensor, + tCsB: cute.Tensor, + smem_thr_copy_A: cute.TiledCopy, + smem_thr_copy_B: cute.TiledCopy, + hook_fn: Optional[Callable] = None, + A_in_regs: cutlass.Constexpr[bool] = False, + B_in_regs: cutlass.Constexpr[bool] = False, + swap_AB: cutlass.Constexpr[bool] = False, +) -> None: + if cutlass.const_expr(swap_AB): + gemm( + tiled_mma, + acc, + tCrB, + tCrA, + tCsB, + tCsA, + smem_thr_copy_B, + smem_thr_copy_A, + hook_fn, + A_in_regs=B_in_regs, + B_in_regs=A_in_regs, + swap_AB=False, + ) + else: + tCrA_copy_view = smem_thr_copy_A.retile(tCrA) + tCrB_copy_view = smem_thr_copy_B.retile(tCrB) + if cutlass.const_expr(not A_in_regs): + cute.copy(smem_thr_copy_A, tCsA[None, None, 0], tCrA_copy_view[None, None, 0]) + if cutlass.const_expr(not B_in_regs): + cute.copy(smem_thr_copy_B, tCsB[None, None, 0], tCrB_copy_view[None, None, 0]) + for k in cutlass.range_constexpr(cute.size(tCsA.shape[2])): + if k < cute.size(tCsA.shape[2]) - 1: + if cutlass.const_expr(not A_in_regs): + cute.copy( + smem_thr_copy_A, tCsA[None, None, k + 1], tCrA_copy_view[None, None, k + 1] + ) + if cutlass.const_expr(not B_in_regs): + cute.copy( + smem_thr_copy_B, tCsB[None, None, k + 1], tCrB_copy_view[None, None, k + 1] + ) + cute.gemm(tiled_mma, acc, tCrA[None, None, k], tCrB[None, None, k], acc) + if cutlass.const_expr(k == 0 and hook_fn is not None): + hook_fn() + + +@cute.jit +def gemm_rs( + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + tCsB: cute.Tensor, + smem_thr_copy_B: cute.TiledCopy, + hook_fn: Optional[Callable] = None, +) -> None: + tCrB_copy_view = smem_thr_copy_B.retile(tCrB) + cute.copy(smem_thr_copy_B, tCsB[None, None, 0], tCrB_copy_view[None, None, 0]) + for k in cutlass.range_constexpr(cute.size(tCrA.shape[2])): + if cutlass.const_expr(k < cute.size(tCrA.shape[2]) - 1): + cute.copy(smem_thr_copy_B, tCsB[None, None, k + 1], tCrB_copy_view[None, None, k + 1]) + cute.gemm(tiled_mma, acc, tCrA[None, None, k], tCrB[None, None, k], acc) + if cutlass.const_expr(k == 0 and hook_fn is not None): + hook_fn() diff --git a/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/block_info.py b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/block_info.py new file mode 100644 index 00000000..149a23d3 --- /dev/null +++ b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/block_info.py @@ -0,0 +1,139 @@ +# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. +from typing import Tuple, Optional +from dataclasses import dataclass + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, const_expr + +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.seqlen_info import SeqlenInfoQK, SeqlenInfoQKNewK + + +@dataclass(frozen=True) +class BlockInfo: + tile_m: cutlass.Constexpr[int] + tile_n: cutlass.Constexpr[int] + is_causal: cutlass.Constexpr[bool] + is_local: cutlass.Constexpr[bool] = False + is_split_kv: cutlass.Constexpr[bool] = False + window_size_left: Optional[Int32] = None + window_size_right: Optional[Int32] = None + qhead_per_kvhead_packgqa: cutlass.Constexpr[int] = 1 + + @cute.jit + def get_n_block_min_max( + self, + seqlen_info: SeqlenInfoQK, + m_block: Int32, + split_idx: Int32 = 0, + num_splits: Int32 = 1, + ) -> Tuple[Int32, Int32]: + n_block_max = cute.ceil_div(seqlen_info.seqlen_k, self.tile_n) + if const_expr(self.is_causal or (self.is_local and self.window_size_right is not None)): + m_idx_max = (m_block + 1) * self.tile_m + if const_expr(self.qhead_per_kvhead_packgqa > 1): + m_idx_max = cute.ceil_div(m_idx_max, self.qhead_per_kvhead_packgqa) + n_idx = m_idx_max + seqlen_info.seqlen_k - seqlen_info.seqlen_q + n_idx_right = n_idx if const_expr(self.is_causal) else n_idx + self.window_size_right + n_block_max = min(n_block_max, cute.ceil_div(n_idx_right, self.tile_n)) + n_block_min = 0 + if const_expr(self.is_local and self.window_size_left is not None): + m_idx_min = m_block * self.tile_m + if const_expr(self.qhead_per_kvhead_packgqa > 1): + m_idx_min = m_idx_min // self.qhead_per_kvhead_packgqa + n_idx = m_idx_min + seqlen_info.seqlen_k - seqlen_info.seqlen_q + n_idx_left = n_idx - self.window_size_left + n_block_min = cutlass.max(n_idx_left // self.tile_n, 0) + if cutlass.const_expr(self.is_split_kv): + num_n_blocks_per_split = ( + Int32(0) + if n_block_max <= n_block_min + else (n_block_max - n_block_min + num_splits - 1) // num_splits + ) + n_block_min = n_block_min + split_idx * num_n_blocks_per_split + n_block_max = cutlass.min(n_block_min + num_n_blocks_per_split, n_block_max) + return n_block_min, n_block_max + + @cute.jit + def get_m_block_min_max(self, seqlen_info: SeqlenInfoQK, n_block: Int32) -> Tuple[Int32, Int32]: + m_block_max = cute.ceil_div(seqlen_info.seqlen_q, self.tile_m) + m_block_min = 0 + if const_expr(self.is_causal or (self.is_local and self.window_size_right is not None)): + n_idx_min = n_block * self.tile_n + m_idx = n_idx_min + seqlen_info.seqlen_q - seqlen_info.seqlen_k + m_idx_right = m_idx if const_expr(self.is_causal) else m_idx - self.window_size_right + m_block_min = max(m_block_min, m_idx_right // self.tile_m) + if const_expr(self.is_local and self.window_size_left is not None): + n_idx_max = (n_block + 1) * self.tile_n + m_idx = n_idx_max + seqlen_info.seqlen_q - seqlen_info.seqlen_k + m_idx_left = m_idx + self.window_size_left + m_block_max = min(m_block_max, cute.ceil_div(m_idx_left, self.tile_m)) + return m_block_min, m_block_max + + @cute.jit + def get_n_block_k_new_min_max( + self, + seqlen_info: SeqlenInfoQKNewK, + m_block: Int32, + split_idx: Int32 = 0, + num_splits: Int32 = 1, + ) -> Tuple[Int32, Int32]: + """Get the block range for new K tokens (append KV). + + First computes the full n_block range via get_n_block_min_max, then maps + those blocks into the new-K index space by subtracting seqlen_k_og. + """ + n_block_min, n_block_max = self.get_n_block_min_max( + seqlen_info, + m_block, + split_idx, + num_splits, + ) + idx_k_new_min = cutlass.max(n_block_min * self.tile_n - seqlen_info.seqlen_k_og, 0) + idx_k_new_max = cutlass.min( + n_block_max * self.tile_n - seqlen_info.seqlen_k_og, seqlen_info.seqlen_k_new + ) + n_block_new_min = idx_k_new_min // self.tile_n + n_block_new_max = ( + cute.ceil_div(idx_k_new_max, self.tile_n) + if idx_k_new_max > idx_k_new_min + else n_block_new_min + ) + return n_block_new_min, n_block_new_max + + @cute.jit + def get_n_block_min_causal_local_mask( + self, + seqlen_info: SeqlenInfoQK, + m_block: Int32, + n_block_min: Int32, + ) -> Int32: + """If we have separate iterations with causal or local masking at the start, where do we stop""" + m_idx_min = m_block * self.tile_m + if const_expr(self.qhead_per_kvhead_packgqa > 1): + m_idx_min = m_idx_min // self.qhead_per_kvhead_packgqa + n_idx = m_idx_min + seqlen_info.seqlen_k - seqlen_info.seqlen_q + n_idx_right = ( + n_idx + if const_expr(not self.is_local or self.window_size_right is None) + else n_idx + self.window_size_right + ) + return cutlass.max(n_block_min, n_idx_right // self.tile_n) + + @cute.jit + def get_n_block_min_before_local_mask( + self, + seqlen_info: SeqlenInfoQK, + m_block: Int32, + n_block_min: Int32, + ) -> Int32: + """If we have separate iterations with local masking at the end, where do we stop the non-masked iterations""" + if const_expr(not self.is_local or self.window_size_left is None): + return n_block_min + else: + m_idx_max = (m_block + 1) * self.tile_m + if const_expr(self.qhead_per_kvhead_packgqa > 1): + m_idx_max = cute.ceil_div(m_idx_max, self.qhead_per_kvhead_packgqa) + n_idx = m_idx_max + seqlen_info.seqlen_k - seqlen_info.seqlen_q + n_idx_left = n_idx - self.window_size_left + return cutlass.max(n_block_min, cute.ceil_div(n_idx_left, self.tile_n)) diff --git a/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/block_sparsity.py b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/block_sparsity.py new file mode 100644 index 00000000..8fc8c76c --- /dev/null +++ b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/block_sparsity.py @@ -0,0 +1,463 @@ +""" +Block-sparsity utilities for FlexAttention +""" + +from typing import Callable, NamedTuple, Tuple + +import cutlass.cute as cute +import torch + +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.cute_dsl_utils import get_broadcast_dims, to_cute_tensor + + +def ceildiv(a: int, b: int) -> int: + return (a + b - 1) // b + + +class BlockSparseTensors(NamedTuple): + mask_block_cnt: cute.Tensor + mask_block_idx: cute.Tensor + full_block_cnt: cute.Tensor | None + full_block_idx: cute.Tensor | None + + def __new_from_mlir_values__(self, values): + if len(values) == 2: + values = (*values, None, None) + return BlockSparseTensors(*values) + + +class BlockSparseTensorsTorch(NamedTuple): + mask_block_cnt: torch.Tensor + mask_block_idx: torch.Tensor + full_block_cnt: torch.Tensor | None = None + full_block_idx: torch.Tensor | None = None + block_size: tuple[int, int] | None = None + + +def get_sparse_q_block_size( + tensors: BlockSparseTensorsTorch | None, + seqlen_q: int, +) -> int | None: + """Return the Q sparse block size, or None when sparsity is unset or ambiguous.""" + if tensors is None: + return None + if tensors.block_size is not None: + return tensors.block_size[0] + num_m_blocks = tensors.mask_block_idx.shape[2] + min_block_size = ceildiv(seqlen_q, num_m_blocks) + max_block_size = seqlen_q if num_m_blocks == 1 else (seqlen_q - 1) // (num_m_blocks - 1) + if min_block_size != max_block_size: + return None + return min_block_size + + +def _expand_sparsity_tensor( + tensor: torch.Tensor, + expected_shape: Tuple[int, ...], + tensor_name: str, + context: str | None, + hint: str | Callable[[], str] | None, +) -> torch.Tensor: + """Check if we need to expand the tensor to expected shape, and do so if possible.""" + needs_expand = tensor.shape != expected_shape + if not needs_expand: + return tensor + can_expand = all(map(lambda cur, tgt: cur == tgt or cur == 1, tensor.shape, expected_shape)) + if not can_expand: + context_clause = f" ({context})" if context else "" + resolved_hint = hint() if callable(hint) else hint + hint_clause = f" Hint: {resolved_hint}" if resolved_hint else "" + raise ValueError( + f"{tensor_name}{context_clause} with shape {tensor.shape} cannot be expanded to expected shape {expected_shape}." + f"{hint_clause}" + ) + return tensor.expand(*expected_shape) + + +def _check_and_expand_block( + name: str, + cnt: torch.Tensor | None, + idx: torch.Tensor | None, + expected_count_shape: Tuple[int, int, int], + expected_index_shape: Tuple[int, int, int, int], + context: str | None, + hint: str | Callable[[], str] | None, +) -> Tuple[torch.Tensor | None, torch.Tensor | None]: + if (cnt is None) != (idx is None): + raise ValueError( + f"{name}_block_cnt and {name}_block_idx must both be provided or both be None" + ) + if cnt is None or idx is None: + return None, None + if cnt.dtype != torch.int32 or idx.dtype != torch.int32: + raise ValueError(f"{name}_block tensors must have dtype torch.int32") + if cnt.device != idx.device: + raise ValueError(f"{name}_block_cnt and {name}_block_idx must be on the same device") + if not cnt.is_cuda or not idx.is_cuda: + raise ValueError(f"{name}_block tensors must live on CUDA") + expanded_cnt = _expand_sparsity_tensor( + cnt, expected_count_shape, f"{name}_block_cnt", context, hint + ) + # [Note] Allow Compact block sparse indices + # Allow the last dimension (n_blocks) of idx to be <= expected, since + # FA4 only accesses indices 0..cnt-1 per query tile. This enables compact + # index tensors that avoid O(N^2) memory at long sequence lengths. + if idx.ndim == 4 and idx.shape[3] <= expected_index_shape[3]: + expected_index_shape = (*expected_index_shape[:3], idx.shape[3]) + expanded_idx = _expand_sparsity_tensor( + idx, expected_index_shape, f"{name}_block_idx", context, hint + ) + return expanded_cnt, expanded_idx + + +def get_block_sparse_expected_shapes( + batch_size: int, + num_head: int, + seqlen_q: int, + seqlen_k: int, + m_block_size: int, + n_block_size: int, + q_stage: int, +) -> Tuple[Tuple[int, int, int], Tuple[int, int, int, int]]: + """Return (expected_count_shape, expected_index_shape) for block sparse normalization.""" + m_block_size_effective = q_stage * m_block_size + expected_m_blocks = ceildiv(seqlen_q, m_block_size_effective) + expected_n_blocks = ceildiv(seqlen_k, n_block_size) + expected_count_shape = (batch_size, num_head, expected_m_blocks) + expected_index_shape = (batch_size, num_head, expected_m_blocks, expected_n_blocks) + return expected_count_shape, expected_index_shape + + +def infer_block_sparse_expected_shapes( + tensors: BlockSparseTensorsTorch, + *, + batch_size: int, + num_head: int, + seqlen_q: int, + seqlen_k: int, + m_block_size: int, + n_block_size: int, + q_stage: int, + context: str, + sparse_block_size_q: int | None = None, + sparse_block_size_kv: int | None = None, +) -> Tuple[Tuple[int, int, int], Tuple[int, int, int, int], int]: + """Infer shapes and scaling for block-sparse tensors. + + Expectations: + - mask_block_cnt is (B, H, M) and mask_block_idx is (B, H, M, N). + - Batch/head dims may be 1 for broadcast, or match the requested sizes. + - sparse_block_size_kv must match tile_n. + - sparse_block_size_q must be a multiple of q_stage * tile_m. + - If sparse_block_size_q is omitted and seqlen_q/num_m_blocks is ambiguous, + the caller must provide block_size to disambiguate. TODO will make this required in a future PR. + """ + base_m_block = q_stage * m_block_size + base_n_block = n_block_size + if sparse_block_size_kv is None: + sparse_block_size_kv = base_n_block + if sparse_block_size_kv != base_n_block: + raise ValueError(f"Block sparse tensors{context} require BLOCK_SIZE_KV={base_n_block}.") + if tensors.mask_block_idx is None: + raise ValueError("mask_block_cnt and mask_block_idx must be provided for block sparsity.") + num_m_blocks = tensors.mask_block_idx.shape[2] + + if sparse_block_size_q is None: + sparse_block_size_q = get_sparse_q_block_size(tensors, seqlen_q) + if sparse_block_size_q is None and base_m_block != 1: + raise ValueError( + f"Block sparse tensors{context} require explicit sparse_block_size[0] " + f"to disambiguate block size for seqlen_q={seqlen_q} and num_m_blocks={num_m_blocks}." + ) + if sparse_block_size_q is None: + sparse_block_size_q = ceildiv(seqlen_q, num_m_blocks) + + if sparse_block_size_q % base_m_block != 0: + raise ValueError( + f"Block sparse tensors{context} have block size {sparse_block_size_q}, " + f"which must be a multiple of {base_m_block}." + ) + + expected_m_blocks = ceildiv(seqlen_q, sparse_block_size_q) + expected_n_blocks = ceildiv(seqlen_k, sparse_block_size_kv) + q_subtile_factor = sparse_block_size_q // base_m_block + expected_count_shape = (batch_size, num_head, expected_m_blocks) + expected_index_shape = (batch_size, num_head, expected_m_blocks, expected_n_blocks) + + mask_block_cnt = tensors.mask_block_cnt + mask_block_idx = tensors.mask_block_idx + if mask_block_cnt is None or mask_block_idx is None: + raise ValueError("mask_block_cnt and mask_block_idx must be provided for block sparsity.") + if mask_block_cnt.ndim != 3 or mask_block_idx.ndim != 4: + raise ValueError( + f"Block sparse tensors{context} must have shapes (B, H, M) and (B, H, M, N)." + ) + for dim_name, cur, tgt in ( + ("batch", mask_block_cnt.shape[0], expected_count_shape[0]), + ("head", mask_block_cnt.shape[1], expected_count_shape[1]), + ): + if cur != tgt and cur != 1: + raise ValueError(f"Block sparse tensors{context} {dim_name} dim must be {tgt} or 1.") + for dim_name, cur, tgt in ( + ("batch", mask_block_idx.shape[0], expected_index_shape[0]), + ("head", mask_block_idx.shape[1], expected_index_shape[1]), + ): + if cur != tgt and cur != 1: + raise ValueError(f"Block sparse tensors{context} {dim_name} dim must be {tgt} or 1.") + if mask_block_cnt.shape[2] != mask_block_idx.shape[2]: + raise ValueError(f"Block sparse tensors{context} must share the same m-block dimension.") + # [Note] Allow Compact block sparse indices: FA4 only accesses indices 0..cnt-1 + # per query tile, so idx.shape[3] can be <= expected_n_blocks. + if mask_block_idx.shape[3] > expected_n_blocks: + raise ValueError( + f"Block sparse tensors{context} n-block dimension must be <= {expected_n_blocks}." + ) + if expected_m_blocks != num_m_blocks: + raise ValueError( + f"Block sparse tensors{context} m-block dimension {num_m_blocks} does not match " + f"sparse_block_size_q={sparse_block_size_q}. " + f"Set BlockSparseTensorsTorch.block_size to match the BlockMask BLOCK_SIZE." + ) + return expected_count_shape, expected_index_shape, q_subtile_factor + + +def get_block_sparse_expected_shapes_bwd( + batch_size: int, + num_head: int, + seqlen_q: int, + seqlen_k: int, + m_block_size: int, + n_block_size: int, + subtile_factor: int, +) -> Tuple[Tuple[int, int, int], Tuple[int, int, int, int]]: + """Return (expected_count_shape, expected_index_shape) for backward block sparse normalization. + + Backward uses Q-direction indexing (transposed from forward), where shapes are + indexed by N-blocks first, then M-blocks. The sparse_block_size_q is determined + by subtile_factor * m_block_size. + """ + sparse_block_size_q = subtile_factor * m_block_size + expected_m_blocks = ceildiv(seqlen_q, sparse_block_size_q) + expected_n_blocks = ceildiv(seqlen_k, n_block_size) + expected_count_shape = (batch_size, num_head, expected_n_blocks) + expected_index_shape = (batch_size, num_head, expected_n_blocks, expected_m_blocks) + return expected_count_shape, expected_index_shape + + +def normalize_block_sparse_tensors( + tensors: BlockSparseTensorsTorch, + *, + expected_count_shape: Tuple[int, int, int], + expected_index_shape: Tuple[int, int, int, int], + context: str | None = None, + hint: str | Callable[[], str] | None = None, +) -> BlockSparseTensorsTorch: + if tensors.mask_block_cnt is None or tensors.mask_block_idx is None: + raise ValueError("mask_block_cnt and mask_block_idx must be provided for block sparsity.") + + mask_cnt, mask_idx = _check_and_expand_block( + "mask", + tensors.mask_block_cnt, + tensors.mask_block_idx, + expected_count_shape, + expected_index_shape, + context, + hint, + ) + if mask_cnt is None or mask_idx is None: + raise ValueError("mask_block_cnt and mask_block_idx must be provided for block sparsity.") + + full_cnt, full_idx = _check_and_expand_block( + "full", + tensors.full_block_cnt, + tensors.full_block_idx, + expected_count_shape, + expected_index_shape, + context, + hint, + ) + if full_cnt is not None and mask_cnt.device != full_cnt.device: + raise ValueError("All block sparse tensors must be on the same device") + + return BlockSparseTensorsTorch( + mask_block_cnt=mask_cnt, + mask_block_idx=mask_idx, + full_block_cnt=full_cnt, + full_block_idx=full_idx, + block_size=tensors.block_size, + ) + + +def is_block_sparsity_enabled(tensors: BlockSparseTensorsTorch) -> bool: + return any(t is not None for t in (tensors.full_block_cnt, tensors.mask_block_cnt)) + + +def get_block_sparse_broadcast_pattern( + tensors: BlockSparseTensorsTorch, +) -> Tuple[Tuple[bool, ...], ...] | None: + """Return broadcast pattern for block sparse tensors by checking actual strides. + + Returns a tuple of broadcast patterns (one per tensor) where each pattern + is a tuple of bools indicating which dims have stride=0. + This is used in compile keys to ensure kernels are recompiled when + broadcast patterns change, since CuTe's mark_layout_dynamic() keeps + stride=0 as static. + + The tensors should already be expanded/normalized before calling this function. + + Returns None if block sparsity is not enabled. + """ + if not is_block_sparsity_enabled(tensors): + return None + + patterns = [] + for tensor in ( + tensors.mask_block_cnt, + tensors.mask_block_idx, + tensors.full_block_cnt, + tensors.full_block_idx, + ): + if tensor is not None: + patterns.append(get_broadcast_dims(tensor)) + else: + patterns.append(None) + return tuple(patterns) + + +def normalize_block_sparse_config( + tensors: BlockSparseTensorsTorch, + *, + batch_size: int, + num_head: int, + seqlen_q: int, + seqlen_k: int, + block_size: tuple[int, int], + q_stage: int, +) -> tuple[BlockSparseTensorsTorch, Tuple[Tuple[bool, ...], ...] | None, int]: + m_block_size, n_block_size = block_size + if tensors.block_size is None: + sparse_block_size_q, sparse_block_size_kv = None, n_block_size + else: + sparse_block_size_q, sparse_block_size_kv = tensors.block_size + if sparse_block_size_kv != n_block_size: + raise ValueError( + f"Block sparsity requires sparse_block_size[1]={n_block_size} to match tile_n." + ) + expected_count_shape, expected_index_shape, q_subtile_factor = ( + infer_block_sparse_expected_shapes( + tensors, + batch_size=batch_size, + num_head=num_head, + seqlen_q=seqlen_q, + seqlen_k=seqlen_k, + m_block_size=m_block_size, + n_block_size=n_block_size, + q_stage=q_stage, + context="forward", + sparse_block_size_q=sparse_block_size_q, + sparse_block_size_kv=sparse_block_size_kv, + ) + ) + normalized_tensors = normalize_block_sparse_tensors( + tensors, + expected_count_shape=expected_count_shape, + expected_index_shape=expected_index_shape, + ) + return ( + normalized_tensors, + get_block_sparse_broadcast_pattern(normalized_tensors), + q_subtile_factor, + ) + + +def normalize_block_sparse_config_bwd( + tensors: BlockSparseTensorsTorch, + *, + batch_size: int, + num_head: int, + seqlen_q: int, + seqlen_k: int, + block_size: tuple[int, int], + subtile_factor: int, +) -> tuple[BlockSparseTensorsTorch, Tuple[Tuple[bool, ...], ...] | None]: + m_block_size, n_block_size = block_size + if tensors.block_size is None: + sparse_block_size_q, sparse_block_size_kv = subtile_factor * m_block_size, n_block_size + else: + sparse_block_size_q, sparse_block_size_kv = tensors.block_size + if sparse_block_size_q != subtile_factor * m_block_size: + raise ValueError( + f"Block sparsity expects sparse_block_size_q={subtile_factor * m_block_size} " + f"for subtile_factor={subtile_factor}." + ) + if sparse_block_size_kv != n_block_size: + raise ValueError( + f"Block sparsity expects sparse_block_size[1]={n_block_size} to match tile_n." + ) + expected_count_shape, expected_index_shape = get_block_sparse_expected_shapes_bwd( + batch_size, + num_head, + seqlen_q, + seqlen_k, + m_block_size, + n_block_size, + subtile_factor, + ) + normalized_tensors = normalize_block_sparse_tensors( + tensors, + expected_count_shape=expected_count_shape, + expected_index_shape=expected_index_shape, + context="_flash_attn_bwd", + hint=lambda: ( + f"Backward expects Q-direction block-sparse tensors (q_mask_cnt/q_mask_idx, " + f"and optionally full_q_cnt/full_q_idx). Regenerate the backward BlockMask with " + f"BLOCK_SIZE=({subtile_factor * m_block_size}, {n_block_size})." + ), + ) + return normalized_tensors, get_block_sparse_broadcast_pattern(normalized_tensors) + + +def to_cute_block_sparse_tensors( + tensors: BlockSparseTensorsTorch, enable_tvm_ffi: bool = True +) -> BlockSparseTensors | None: + """Convert torch block sparsity tensors to CuTe tensors, optionally for tvm ffi""" + if not is_block_sparsity_enabled(tensors): + return None + + ( + mask_block_cnt, + mask_block_idx, + full_block_cnt, + full_block_idx, + *_, + ) = tensors + + ( + mask_block_cnt_tensor, + mask_block_idx_tensor, + ) = [ + to_cute_tensor(t, assumed_align=4, leading_dim=-1, enable_tvm_ffi=enable_tvm_ffi) + for t in (mask_block_cnt, mask_block_idx) + ] + ( + full_block_cnt_tensor, + full_block_idx_tensor, + ) = [ + to_cute_tensor(t, assumed_align=4, leading_dim=-1, enable_tvm_ffi=enable_tvm_ffi) + if t is not None + else None + for t in (full_block_cnt, full_block_idx) + ] + + return BlockSparseTensors( + mask_block_cnt_tensor, + mask_block_idx_tensor, + full_block_cnt_tensor, + full_block_idx_tensor, + ) + + +def fast_sampling(mask_mod): + """Convenience decorator to mark mask_mod as safe for 5-point fast sampling""" + mask_mod.use_fast_sampling = True + return mask_mod diff --git a/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/cute_dsl_utils.py b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/cute_dsl_utils.py new file mode 100644 index 00000000..79ebd9df --- /dev/null +++ b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/cute_dsl_utils.py @@ -0,0 +1,129 @@ +# Copyright (c) 2025, Tri Dao. + +import os +import pathlib +from typing import Tuple +from functools import partial, lru_cache + +import torch + +try: + from triton.tools.disasm import extract +except ImportError: + extract = None + +import cutlass +import cutlass.cute as cute +from cutlass.cutlass_dsl import NumericMeta +from cutlass.cute.runtime import from_dlpack + +StaticTypes = (cutlass.Constexpr, NumericMeta, int, bool, str, float, type(None)) + + +load_cubin_module_data_og = cutlass.base_dsl.runtime.cuda.load_cubin_module_data +cute_compile_og = cute.compile + + +torch2cute_dtype_map = { + torch.float16: cutlass.Float16, + torch.bfloat16: cutlass.BFloat16, + torch.float32: cutlass.Float32, +} + + +@lru_cache +def get_max_active_clusters(cluster_size): + return cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_size=cluster_size) + + +@lru_cache +def get_device_capacity(device: torch.device = None) -> Tuple[int, int]: + return torch.cuda.get_device_capability(device) + + +def load_cubin_module_data_patched(cubin_data, filepath): + pathlib.Path(filepath).write_bytes(cubin_data) + return load_cubin_module_data_og(cubin_data) + + +def cute_compile_patched(*args, **kwargs): + """A patched version of cute.compile that dump the SASS to a file if CUTE_CUBIN_PATH is set.""" + cubin_path = os.getenv("CUTE_CUBIN_PATH", None) + if cubin_path is not None: + cutlass.base_dsl.runtime.cuda.load_cubin_module_data = partial( + load_cubin_module_data_patched, filepath=cubin_path + ) + output = cute_compile_og(*args, **kwargs) + if cubin_path is not None: + cutlass.base_dsl.runtime.cuda.load_cubin_module_data = load_cubin_module_data_og + if extract is not None: + sass = extract(cubin_path, None) + pathlib.Path(cubin_path).with_suffix(".annotated.sass").write_text(sass) + return output + + +def assume_strides_aligned(t): + """Assume all strides except the last are divisible by 128 bits. + + Python int strides (e.g., stride=0 from GQA expand) are kept as-is + since they're static and don't need alignment assumptions. + """ + divby = 128 // t.element_type.width + strides = tuple(s if isinstance(s, int) else cute.assume(s, divby=divby) for s in t.stride[:-1]) + return (*strides, t.stride[-1]) + + +def assume_tensor_aligned(t): + """Rebuild a tensor with 128-bit aligned stride assumptions. Passes through None.""" + if t is None: + return None + return cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=assume_strides_aligned(t))) + + +def to_cute_tensor(t, assumed_align=16, leading_dim=-1, fully_dynamic=False, enable_tvm_ffi=True): + """Convert torch tensor to cute tensor for TVM FFI. leading_dim=-1 defaults to t.ndim-1.""" + tensor = from_dlpack(t.detach(), assumed_align=assumed_align, enable_tvm_ffi=enable_tvm_ffi) + if fully_dynamic: + return tensor.mark_layout_dynamic() + if leading_dim == -1: + leading_dim = t.ndim - 1 + return tensor.mark_layout_dynamic(leading_dim=leading_dim) + + +def to_cute_aux_tensor(t, enable_tvm_ffi=True): + """Convert torch tensor to cute tensor for TVM FFI, tailored to FlexAttention aux tensors. + This allows the user to specify alignment and leading dimension for aux tensors used in + custom score_mod callables. + """ + assumed_align: int = getattr(t, "__assumed_align__", None) + leading_dim: int = getattr(t, "__leading_dim__", None) + fully_dynamic: bool = leading_dim is None + + return to_cute_tensor( + t, + assumed_align=assumed_align, + leading_dim=leading_dim, + fully_dynamic=fully_dynamic, + enable_tvm_ffi=enable_tvm_ffi, + ) + + +def get_aux_tensor_metadata(aux_tensors): + return tuple( + ( + getattr(t, "__assumed_align__", 0), + getattr(t, "__leading_dim__", -1), + hasattr(t, "__leading_dim__"), + ) + for t in aux_tensors + ) + + +def get_broadcast_dims(tensor: torch.Tensor) -> Tuple[bool, ...]: + """Return tuple of bools indicating which dims have stride=0 (broadcast). + + This is useful for compile keys since CuTe's mark_layout_dynamic() keeps + stride=0 as static, meaning kernels compiled with different broadcast + patterns are not interchangeable. + """ + return tuple(s == 0 for s in tensor.stride()) diff --git a/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/fast_math.py b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/fast_math.py new file mode 100644 index 00000000..c56ea89e --- /dev/null +++ b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/fast_math.py @@ -0,0 +1,21 @@ +# Copyright (c) 2025, Tri Dao. + +import cutlass +import cutlass.cute as cute +from cutlass import Int32 + + +@cute.jit +def clz(x: Int32) -> Int32: + # for i in cutlass.range_constexpr(32): + # if (1 << (31 - i)) & x: + # return Int32(i) + # return Int32(32) + # Early exit is not supported yet + res = Int32(32) + done = False + for i in cutlass.range(32): + if ((1 << (31 - i)) & x) and not done: + res = Int32(i) + done = True + return res diff --git a/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/flash_fwd.py b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/flash_fwd.py new file mode 100644 index 00000000..9eb1e384 --- /dev/null +++ b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/flash_fwd.py @@ -0,0 +1,1218 @@ +# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. +# A reimplementation of +# https://github.com/Dao-AILab/flash-attention/blob/main/hopper/flash_fwd_kernel_sm80.h +# and https://github.com/Dao-AILab/flash-attention/blob/main/hopper/flash_fwd_kernel_sm90.h +# from Cutlass C++ to Cute-DSL. +# Built on Cute-DSL example: https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/ampere/flash_attention_v2.py + +import math +from types import SimpleNamespace +from typing import Type, Callable, Optional, List +from functools import partial + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass import Constexpr, Float32, Int32, const_expr, Boolean +from cutlass.cute.nvgpu import cpasync, warp +import cutlass.utils as utils_basic +from cutlass.base_dsl.arch import Arch +from cutlass.cutlass_dsl import BaseDSL + +from telefuser.kernel.sol_attn.sm90._compat import copy_utils +from telefuser.kernel.sol_attn.sm90._compat import layout_utils + +from telefuser.kernel.sol_attn._vendor.flash_attn.cute import ampere_helpers as sm80_utils +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.cute_dsl_utils import assume_tensor_aligned +from telefuser.kernel.sol_attn._vendor.flash_attn.cute import utils +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.mask import AttentionMask +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.softmax import Softmax +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.seqlen_info import SeqlenInfoQK +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.block_info import BlockInfo +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.pack_gqa import PackGQA +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.named_barrier import NamedBarrierFwd +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.block_sparsity import BlockSparseTensors +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.tile_scheduler import ( + SingleTileScheduler, + SingleTileVarlenScheduler, + TileSchedulerArguments, +) + + +class FlashAttentionForwardBase: + + def __init__( + self, + dtype: Type[cutlass.Numeric], + head_dim: int, + head_dim_v: Optional[int] = None, + qhead_per_kvhead: int = 1, + is_causal: bool = False, + is_local: bool = False, + pack_gqa: bool = True, + tile_m: int = 128, + tile_n: int = 128, + num_stages: int = 1, + num_threads: int = 128, + Q_in_regs: bool = False, + score_mod: Optional[cutlass.Constexpr] = None, + mask_mod: Optional[cutlass.Constexpr] = None, + has_aux_tensors: bool = False, + q_subtile_factor: int | None = None, + ): + """Initializes the configuration for a flash attention kernel. + + All contiguous dimensions must be at least 16 bytes aligned, which means that the head dimension + should be a multiple of 8. + + :param head_dim: head dimension + :type head_dim: int + :param tile_m: m block size + :type tile_m: int + :param tile_n: n block size + :type tile_n: int + :param num_threads: number of threads + :type num_threads: int + :param is_causal: is causal + :param score_mod: A callable that takes the attention scores and applies a modification. + Callable signature: ``score_mod(scores, batch_idx, head_idx, q_idx, kv_idx, aux_tensors) -> Any`` + :param mask_mod: A callable that takes the attention scores and returns a boolean representing whether that score should be masked. + Callable signature: ``mask_mod(batch_idx, head_idx, q_idx, kv_idx, aux_tensors) -> Boolean`` + """ + self.dtype = dtype + # padding head_dim to a multiple of 16 as k_block_size + hdim_multiple_of = 16 + self.tile_hdim = int(math.ceil(head_dim / hdim_multiple_of) * hdim_multiple_of) + head_dim_v = head_dim_v if head_dim_v is not None else head_dim + self.same_hdim_kv = head_dim == head_dim_v + self.tile_hdimv = int(math.ceil(head_dim_v / hdim_multiple_of) * hdim_multiple_of) + # Can save registers (and hence be faster) if we don't have to check hdim predication + self.check_hdim_oob = head_dim != self.tile_hdim + self.check_hdim_v_oob = head_dim_v != self.tile_hdimv + self.qhead_per_kvhead = qhead_per_kvhead + self.is_causal = is_causal + self.is_local = is_local + self.pack_gqa = pack_gqa + self.tile_m = tile_m + self.tile_n = tile_n + self.num_threads = num_threads + self.num_stages = num_stages + self.q_subtile_factor = q_subtile_factor + self.Q_in_regs = Q_in_regs + self.score_mod = score_mod + self.mask_mod = mask_mod + self.qk_acc_dtype = Float32 + self.vec_size: cutlass.Constexpr = getattr( + score_mod, "__vec_size__", 1 if cutlass.const_expr(has_aux_tensors) else 2 + ) + if self.vec_size > 2: + raise ValueError( + f"score_mod vec_size {self.vec_size} not supported on Sm80/90/120 " + "due to accumulator thread ownership pattern." + ) + self.arch = BaseDSL._get_dsl().get_arch_enum() + + @staticmethod + def can_implement( + dtype, + head_dim, + head_dim_v, + tile_m, + tile_n, + num_stages, + num_threads, + is_causal, + Q_in_regs=False, + ) -> bool: + """Check if the kernel can be implemented with the given parameters. + + :param dtype: data type + :type dtype: cutlass.Numeric + :param head_dim: head dimension + :type head_dim: int + :param tile_m: m block size + :type tile_m: int + :param tile_n: n block size + :type tile_n: int + :param num_threads: number of threads + :type num_threads: int + :param is_causal: is causal + :type is_causal: bool + + :return: True if the kernel can be implemented, False otherwise + :rtype: bool + """ + if dtype not in [cutlass.Float16, cutlass.BFloat16]: + return False + if head_dim % 8 != 0: + return False + if head_dim_v % 8 != 0: + return False + if tile_n % 16 != 0: + return False + if num_threads % 32 != 0: + return False + # Check if block size setting is out of shared memory capacity + # Shared memory usage: Q tile + (K tile + V tile) where K and V use the same tile size + smem_usage_Q = tile_m * head_dim * 2 + smem_usage_K = tile_n * head_dim * num_stages * 2 + smem_usage_V = tile_n * head_dim_v * num_stages * 2 + smem_usage_QV = ( + (smem_usage_Q + smem_usage_V) if not Q_in_regs else max(smem_usage_Q, smem_usage_V) + ) + smem_usage = smem_usage_QV + smem_usage_K + # TODO: sm86 and sm89 + smem_capacity = utils_basic.get_smem_capacity_in_bytes("sm_80") + if smem_usage > smem_capacity: + return False + # Check if twice the block size is divisible by the number of threads + if (tile_m * 2) % num_threads != 0: + return False + return True + + def _check_type( + self, + mQ_type: Type[cutlass.Numeric], + mK_type: Type[cutlass.Numeric], + mV_type: Type[cutlass.Numeric], + mO_type: Type[cutlass.Numeric], + mLSE_type: Type[cutlass.Numeric] | None, + mCuSeqlensQ_type: Type[cutlass.Numeric] | None, + mCuSeqlensK_type: Type[cutlass.Numeric] | None, + mSeqUsedQ_type: Type[cutlass.Numeric] | None, + mSeqUsedK_type: Type[cutlass.Numeric] | None, + ): + # Get the data type and check if it is fp16 or bf16 + if const_expr(not (mQ_type == mK_type == mV_type == mO_type)): + raise TypeError("All tensors must have the same data type") + if const_expr(mQ_type not in [cutlass.Float16, cutlass.BFloat16]): + raise TypeError("Only Float16 or BFloat16 is supported") + if const_expr(mLSE_type not in [None, Float32]): + raise TypeError("LSE tensor must be Float32") + if const_expr(mCuSeqlensQ_type not in [None, Int32]): + raise TypeError("cu_seqlens_q tensor must be Int32") + if const_expr(mCuSeqlensK_type not in [None, Int32]): + raise TypeError("cu_seqlens_k tensor must be Int32") + if const_expr(mSeqUsedQ_type not in [None, Int32]): + raise TypeError("seqused_q tensor must be Int32") + if const_expr(mSeqUsedK_type not in [None, Int32]): + raise TypeError("seqused_k tensor must be Int32") + assert mQ_type == self.dtype + + def _setup_attributes(self): + # /////////////////////////////////////////////////////////////////////////////// + # Shared memory layout: Q/K/V + # /////////////////////////////////////////////////////////////////////////////// + sQ_layout_atom, sK_layout_atom, sV_layout_atom, sO_layout_atom, sP_layout_atom = ( + self._get_smem_layout_atom() + ) + self.sQ_layout = cute.tile_to_shape( + sQ_layout_atom, + (self.tile_m, self.tile_hdim), + (0, 1), + ) + self.sK_layout = cute.tile_to_shape( + sK_layout_atom, + (self.tile_n, self.tile_hdim, self.num_stages), + (0, 1, 2), + ) + self.sV_layout = cute.tile_to_shape( + sV_layout_atom, + (self.tile_n, self.tile_hdimv, self.num_stages), + (0, 1, 2), + ) + self.sO_layout = cute.tile_to_shape( + sO_layout_atom, + (self.tile_m, self.tile_hdimv), + (0, 1), + ) + if const_expr(sP_layout_atom is not None): + self.sP_layout = cute.tile_to_shape( + sP_layout_atom, + (self.tile_m, self.tile_n), + (0, 1), + ) + else: + self.sP_layout = None + + # /////////////////////////////////////////////////////////////////////////////// + # GMEM Tiled copy: + # /////////////////////////////////////////////////////////////////////////////// + # Thread layouts for copies + universal_copy_bits = 128 + async_copy_elems = universal_copy_bits // self.dtype.width + # atom_async_copy: async copy atom for QKV load + atom_async_copy = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + self.dtype, + num_bits_per_copy=universal_copy_bits, + ) + # atom_universal_copy: universal copy atom for O store + atom_universal_copy = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.dtype, + num_bits_per_copy=universal_copy_bits, + ) + # tQ_layout and tK_layout: thread layout for QK load + tQK_shape_dim_1 = sQ_layout_atom.outer.shape[1] // async_copy_elems + assert self.num_Q_load_threads % tQK_shape_dim_1 == 0, ( + "num_threads must be divisible by tQK_shape_dim_1" + ) + assert self.num_producer_threads % tQK_shape_dim_1 == 0, ( + "num_threads must be divisible by tQK_shape_dim_1" + ) + tQ_layout = cute.make_ordered_layout( + (self.num_Q_load_threads // tQK_shape_dim_1, tQK_shape_dim_1), + order=(1, 0), + ) + tK_layout = cute.make_ordered_layout( + (self.num_producer_threads // tQK_shape_dim_1, tQK_shape_dim_1), + order=(1, 0), + ) + # So that we don't have to check if we overshoot kBlockM when we load Q + assert self.tile_m % tQ_layout.shape[0] == 0 + tV_shape_dim_1 = sV_layout_atom.outer.shape[1] // async_copy_elems + tV_layout = cute.make_ordered_layout( + (self.num_producer_threads // tV_shape_dim_1, tV_shape_dim_1), + order=(1, 0), + ) + # TODO: need a different layout for O if O dtype is not the same as V dtype + # tO_layout: thread layout for O store + tO_layout = cute.make_ordered_layout( + (self.num_epilogue_threads // tV_shape_dim_1, tV_shape_dim_1), + order=(1, 0), + ) + # So that we don't have to check if we overshoot kBlockM when we store O + assert self.tile_m % tO_layout.shape[0] == 0 + + # Value layouts for copies + vQKV_layout = cute.make_layout((1, async_copy_elems)) + vO_layout = vQKV_layout + + self.gmem_tiled_copy_Q = cute.make_tiled_copy_tv(atom_async_copy, tQ_layout, vQKV_layout) + self.gmem_tiled_copy_K = cute.make_tiled_copy_tv(atom_async_copy, tK_layout, vQKV_layout) + self.gmem_tiled_copy_V = cute.make_tiled_copy_tv(atom_async_copy, tV_layout, vQKV_layout) + # gmem_tiled_copy_O: tiled copy for O store + self.gmem_tiled_copy_O = cute.make_tiled_copy_tv(atom_universal_copy, tO_layout, vO_layout) + + def _get_smem_layout_atom(self): + raise NotImplementedError() + + def _get_tiled_mma(self): + raise NotImplementedError() + + def _get_shared_storage_cls(self): + raise NotImplementedError() + + @cute.jit + def __call__( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + softmax_scale: Float32, + # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI). + stream: cuda.CUstream = None, + ): + """Configures and launches the flash attention kernel. + + mQ/mK/mV/mO has same data types(supports fp16 and bf16) and same layout: + (batch_size, seqlen_q, num_head, head_dim):(_, _, _, 1) + """ + raise NotImplementedError() + + @cute.jit + def epilogue( + self, + acc_O: cute.Tensor, + lse: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + sO: cute.Tensor, + seqlen: SeqlenInfoQK, + gmem_tiled_copy_O: cute.TiledCopy, + tma_atom_O: Optional[cute.CopyAtom], + tiled_mma: cute.TiledMma, + tidx: Int32, + m_block: Int32, + head_idx: Int32, + batch_idx: Int32, + output_scale: Optional[cute.Tensor] = None, + ): + # store acc_O + rO = cute.make_fragment_like(acc_O, self.dtype) + if const_expr(output_scale is None): + rO.store(acc_O.load().to(self.dtype)) + else: + # Fuse the final row normalization with the FP32 -> output dtype + # conversion. This avoids a separate full traversal of acc_O. + acc_O_mn = layout_utils.reshape_acc_to_mn(acc_O) + rO_mn = layout_utils.reshape_acc_to_mn(rO) + assert cute.size(output_scale) == cute.size(acc_O_mn, mode=[0]) + for r in cutlass.range(cute.size(output_scale), unroll_full=True): + rO_mn[r, None].store( + (acc_O_mn[r, None].load() * output_scale[r]).to(self.dtype) + ) + # Make sure all threads have finished reading V + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.Epilogue), number_of_threads=self.num_epilogue_threads + ) + smem_copy_atom_O = utils.get_smem_store_atom(self.arch.major * 10 + self.arch.minor, self.dtype) + smem_thr_copy_O = cute.make_tiled_copy_C(smem_copy_atom_O, tiled_mma).get_slice(tidx) + taccOrO = smem_thr_copy_O.retile(rO) + taccOsO = smem_thr_copy_O.partition_D(sO) + # taccOsO = copy_utils.partition_D_position_independent(smem_thr_copy_O, sO) + # copy acc O from rmem to smem with the smem copy atom + cute.copy(smem_copy_atom_O, taccOrO, taccOsO) + + cO = cute.make_identity_tensor((self.tile_m, self.tile_hdimv)) + pack_gqa = PackGQA( + self.tile_m, self.tile_hdimv, self.check_hdim_v_oob, self.qhead_per_kvhead + ) + + # Write LSE from rmem -> gmem + if const_expr(mLSE is not None): + mLSE_cur = seqlen.offset_batch_Q(mLSE, batch_idx, dim=2)[None, head_idx] + if const_expr(not self.pack_gqa): + gLSE = cute.local_tile(mLSE_cur, (self.tile_m,), (m_block,)) + gLSE_expanded_layout = cute.append( + gLSE.layout, cute.make_layout((self.tile_hdimv,), stride=(0,)) + ) + gLSE_expanded = cute.make_tensor(gLSE.iterator, gLSE_expanded_layout) + thr_mma = tiled_mma.get_slice(tidx) + taccOgLSE = layout_utils.reshape_acc_to_mn(thr_mma.partition_C(gLSE_expanded)) + assert cute.size(taccOgLSE, mode=[0]) == cute.size(lse) + taccOcO = layout_utils.reshape_acc_to_mn(thr_mma.partition_C(cO)) + t0accOcO = layout_utils.reshape_acc_to_mn(thr_mma.get_slice(0).partition_C(cO)) + # Only the thread corresponding to column 0 writes out the lse to gmem + if taccOcO[0][1] == 0: + for m in cutlass.range(cute.size(taccOgLSE.shape[1]), unroll_full=True): + if ( + t0accOcO[m, 0][0] + < seqlen.seqlen_q - m_block * self.tile_m - taccOcO[0][0] + ): + taccOgLSE[m, 0] = lse[m] + else: + pack_gqa.store_LSE(mLSE_cur, lse, tiled_mma, tidx, m_block, seqlen.seqlen_q) + + ragged = self.use_tma_O and (seqlen.has_cu_seqlens_q or seqlen.has_seqused_q) + mO_cur = seqlen.offset_batch_Q(mO, batch_idx, dim=3, ragged=ragged)[None, None, head_idx] + # thr_mma = tiled_mma.get_slice(tidx) + # taccOgO = thr_mma.partition_C(gO) + # cute.autovec_copy(rO, taccOgO) + # sync to make sure all smem stores are done + if const_expr(self.use_tma_O): + # ensure smem writes are visible to TMA + cute.arch.fence_view_async_shared() + cute.arch.barrier_arrive( + barrier_id=int(NamedBarrierFwd.Epilogue), + number_of_threads=self.num_epilogue_threads + cute.arch.WARP_SIZE, + ) + gO = cute.local_tile(mO_cur, (self.tile_m, self.tile_hdimv), (m_block, 0)) + store_O, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_O, 0, cute.make_layout(1), sO, gO, single_stage=True + ) + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == 4: + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.Epilogue), + number_of_threads=self.num_epilogue_threads + cute.arch.WARP_SIZE, + ) + store_O() + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + else: + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.Epilogue), + number_of_threads=self.num_epilogue_threads, + ) + gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx) + tOsO = gmem_thr_copy_O.partition_S(sO) + tOrO = cute.make_fragment_like(tOsO, self.dtype) + # load acc O from smem to rmem for wider vectorization + cute.autovec_copy(tOsO, tOrO) + if const_expr(not self.pack_gqa): + gO = cute.local_tile(mO_cur, (self.tile_m, self.tile_hdimv), (m_block, 0)) + tOgO = gmem_thr_copy_O.partition_D(gO) + tOcO = gmem_thr_copy_O.partition_S(cO) + t0OcO = gmem_tiled_copy_O.get_slice(0).partition_S(cO) + tOpO = utils.predicate_k(tOcO, limit=mO.shape[1]) + # copy acc O from rmem to gmem + for rest_m in cutlass.range_constexpr(cute.size(tOrO.shape[1])): + if ( + t0OcO[0, rest_m, 0][0] + < seqlen.seqlen_q - m_block * self.tile_m - tOcO[0][0] + ): + cute.copy( + gmem_tiled_copy_O, + tOrO[None, rest_m, None], + tOgO[None, rest_m, None], + pred=tOpO[None, rest_m, None] + if const_expr(self.check_hdim_v_oob) + else None, + ) + else: + pack_gqa.store_O(mO_cur, tOrO, gmem_tiled_copy_O, tidx, m_block, seqlen.seqlen_q) + + @cute.jit + def advance_pipeline(self, pipeline_index): + return pipeline_index + 1 if pipeline_index < self.num_stages - 1 else 0 + + @cute.jit + def load_Q( + self, + gmem_thr_copy: cute.TiledCopy, + gQ: cute.Tensor, + sQ: cute.Tensor, + block: Int32, + seqlen: Int32, + headdim: Int32, + ): + tQsQ, tQgQ = gmem_thr_copy.partition_D(sQ), gmem_thr_copy.partition_S(gQ) + cQ = cute.make_identity_tensor((self.tile_m, self.tile_hdim)) + tQcQ = gmem_thr_copy.partition_S(cQ) + t0QcQ = gmem_thr_copy.get_slice(0).partition_S(cQ) + tQpQ = utils.predicate_k(tQcQ, limit=headdim) + for m in cutlass.range_constexpr(cute.size(tQsQ.shape[1])): + # Instead of using tQcQ, we using t0QcQ and subtract the offset from the limit + # (seqlen - block * kBlockM). This is because the entries of t0QcQ are known at compile time. + if t0QcQ[0, m, 0][0] < seqlen - block * self.tile_m - tQcQ[0][0]: + cute.copy( + gmem_thr_copy, + tQgQ[None, m, None], + tQsQ[None, m, None], + pred=tQpQ[None, m, None] if const_expr(self.check_hdim_oob) else None, + ) + # We don't need to clear the sQ smem tiles since we'll only write out the valid outputs + + @cute.jit + def load_K( + self, + gmem_tiled_copy: cute.TiledCopy, + tKgK: cute.Tensor, + tKsK: cute.Tensor, + tKcK: cute.Tensor, + t0KcK: cute.Tensor, + tKpK: cute.Tensor, + block: Int32, + smem_pipe_write: Int32, + seqlen: Int32, + need_predicates: cutlass.Constexpr, + ): + # Do we need to check if we overshoot kBlockN when we load K? + is_even_n_smem_k = self.tile_n % gmem_tiled_copy.tiler_mn[0].shape == 0 + if const_expr(need_predicates or not is_even_n_smem_k): + # Instead of using tKcK, we using t0KcK and subtract the offset from the limit + # (seqlen - block * kBlockN). This is because the entries of t0KcK are known at compile time. + if const_expr(is_even_n_smem_k): + seqlen_limit = seqlen - block * self.tile_n + else: + if const_expr(not need_predicates): + seqlen_limit = self.tile_n + else: + seqlen_limit = cutlass.min(seqlen - block * self.tile_n, self.tile_n) + seqlen_limit -= tKcK[0][0] + for n in cutlass.range_constexpr(cute.size(tKsK.shape[1])): + if t0KcK[0, n, 0][0] < seqlen_limit: + cute.copy( + gmem_tiled_copy, + tKgK[None, n, None, block], + tKsK[ + None, n, None, smem_pipe_write if const_expr(self.num_stages > 1) else 0 + ], + pred=tKpK[None, n, None] if const_expr(self.check_hdim_oob) else None, + ) + # We don't need to clear the sK smem tiles since we'll mask out the scores anyway. + else: + cute.copy( + gmem_tiled_copy, + tKgK[None, None, None, block], + tKsK[None, None, None, smem_pipe_write if const_expr(self.num_stages > 1) else 0], + pred=tKpK if const_expr(self.check_hdim_oob) else None, + ) + + @cute.jit + def load_V( + self, + gmem_tiled_copy: cute.TiledCopy, + tVgV: cute.Tensor, + tVsV: cute.Tensor, + tVcV: cute.Tensor, + t0VcV: cute.Tensor, + tVpV: cute.Tensor, + block: Int32, + smem_pipe_write: Int32, + seqlen: Int32, + need_predicates: cutlass.Constexpr, + ): + # Do we need to check if we overshoot kBlockN when we load V? + is_even_n_smem_v = self.tile_n % gmem_tiled_copy.tiler_mn[0].shape == 0 + if const_expr(need_predicates or not is_even_n_smem_v): + for n in cutlass.range_constexpr(cute.size(tVsV.shape[1])): + # If kBlockN doesn't evenly divide the tiled copy, only the last `n` needs to be checked + if ( + is_even_n_smem_v + or n < cute.size(tVsV.shape[1]) - 1 + or tVcV[0, n, 0][0] < self.tile_n + ): + predicate = tVpV[None, n, None] if const_expr(self.check_hdim_v_oob) else None + if const_expr(need_predicates): + seqlen_limit = seqlen - block * self.tile_n - tVcV[0][0] + predicate_n = t0VcV[0, n, 0][0] < seqlen_limit + predicate = cute.make_fragment_like(tVpV[None, 0, None]) + for k in cutlass.range_constexpr(cute.size(predicate.shape[1])): + for i in cutlass.range_constexpr(cute.size(predicate.shape[0])): + predicate[i, k] = ( + tVpV[i, n, k] if const_expr(self.check_hdim_v_oob) else True + ) and predicate_n + cute.copy( + gmem_tiled_copy, + tVgV[None, n, None, block], + tVsV[ + None, n, None, smem_pipe_write if const_expr(self.num_stages > 1) else 0 + ], + pred=predicate, + ) + else: + cute.copy( + gmem_tiled_copy, + tVgV[None, None, None, block], + tVsV[None, None, None, smem_pipe_write if const_expr(self.num_stages > 1) else 0], + pred=tVpV if const_expr(self.check_hdim_v_oob) else None, + ) + + +class FlashAttentionForwardSm80(FlashAttentionForwardBase): + def _get_smem_layout_atom(self): + sQ_layout_atom = sm80_utils.get_smem_layout_atom(self.dtype, self.tile_hdim) + sK_layout_atom = sQ_layout_atom + sV_layout_atom = sm80_utils.get_smem_layout_atom(self.dtype, self.tile_hdimv) + sO_layout_atom = sV_layout_atom + sP_layout_atom = None + return sQ_layout_atom, sK_layout_atom, sV_layout_atom, sO_layout_atom, sP_layout_atom + + def _get_tiled_mma(self): + tiled_mma_qk = cute.make_tiled_mma( + warp.MmaF16BF16Op(self.dtype, Float32, (16, 8, 16)), + (self.num_threads // 32, 1, 1), + permutation_mnk=(self.num_threads // 32 * 16, 16, 16), + ) + tiled_mma_pv = cute.make_tiled_mma( + warp.MmaF16BF16Op(self.dtype, Float32, (16, 8, 16)), + (self.num_threads // 32, 1, 1), + permutation_mnk=(self.num_threads // 32 * 16, 16, 16), + ) + return tiled_mma_qk, tiled_mma_pv + + def _get_shared_storage_cls(self): + sQ_struct, sK_struct, sV_struct = [ + cute.struct.Align[cute.struct.MemRange[self.dtype, cute.cosize(layout)], 1024] + for layout in (self.sQ_layout, self.sK_layout, self.sV_layout) + ] + cosize_sQV = max(cute.cosize(self.sQ_layout), cute.cosize(self.sV_layout)) + sQV_struct = cute.struct.Align[cute.struct.MemRange[self.dtype, cosize_sQV], 1024] + + @cute.struct + class SharedStorageQKV: + sV: sV_struct + sQ: sQ_struct + sK: sK_struct + + @cute.struct + class SharedStorageSharedQV: + sQ: sQV_struct + sK: sK_struct + + return SharedStorageQKV if const_expr(not self.Q_in_regs) else SharedStorageSharedQV + + @cute.jit + def __call__( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + softmax_scale: Float32, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + mSeqUsedQ: Optional[cute.Tensor] = None, + mSeqUsedK: Optional[cute.Tensor] = None, + mPageTable: Optional[cute.Tensor] = None, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + learnable_sink: Optional[cute.Tensor] = None, + blocksparse_tensors: Optional[BlockSparseTensors] = None, + aux_tensors=None, + # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI). + stream: cuda.CUstream = None, + ): + """Configures and launches the flash attention kernel. + + mQ/mK/mV/mO has same data types(supports fp16 and bf16) and same layout: + (batch_size, seqlen_q, num_head, head_dim):(_, _, _, 1) + """ + assert learnable_sink is None, "Learnable sink is not supported in this kernel" + self._check_type( + *(t.element_type if t is not None else None for t in (mQ, mK, mV, mO, mLSE, mCuSeqlensQ, mCuSeqlensK, mSeqUsedQ, mSeqUsedK)) + ) + tiled_mma_qk, tiled_mma_pv = self._get_tiled_mma() + self.num_mma_threads = tiled_mma_pv.size + self.num_producer_threads = self.num_threads + self.num_Q_load_threads = self.num_threads + self.num_epilogue_threads = self.num_threads + # This synchronous warp-MMA implementation only constructs the + # vectorized register/SMEM output copy. When compiled for SM100 as a + # SOL_ATTN scaffold, selecting TMA solely from the target arch leaves the + # TMA store atom unset and fails IR verification. + self.use_tma_O = False + self._setup_attributes() + SharedStorage = self._get_shared_storage_cls() + mQ, mK, mV, mO = [assume_tensor_aligned(t) for t in (mQ, mK, mV, mO)] + # Layout permutation: 4D non-varlen vs 3D varlen + QO_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1] + KV_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensK is None) else [0, 2, 1] + mQ, mO = [ + cute.make_tensor(t.iterator, cute.select(t.layout, mode=QO_layout_transpose)) + for t in (mQ, mO) + ] + mK, mV = [ + cute.make_tensor(t.iterator, cute.select(t.layout, mode=KV_layout_transpose)) + for t in (mK, mV) + ] + if const_expr(mLSE is not None): + LSE_layout_transpose = [2, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 0] + mLSE = cute.make_tensor(mLSE.iterator, cute.select(mLSE.layout, mode=LSE_layout_transpose)) + # TileScheduler for varlen, simple grid for non-varlen + if const_expr(mCuSeqlensQ is not None or mSeqUsedQ is not None): + TileScheduler = SingleTileVarlenScheduler + else: + TileScheduler = SingleTileScheduler + num_batch = ( + mCuSeqlensQ.shape[0] - 1 + if const_expr(mCuSeqlensQ is not None) + else mQ.shape[3] + ) + tile_sched_args = TileSchedulerArguments( + num_block=cute.ceil_div(mQ.shape[0], self.tile_m), + num_head=cute.size(mQ.shape[2]), + num_batch=num_batch, + num_splits=getattr(self, "sol_attn_v_splits", 1), + seqlen_k=0, + headdim=mQ.shape[1], + headdim_v=mV.shape[1], + total_q=cute.size(mQ.shape[0]) + if const_expr(mCuSeqlensQ is not None) + else cute.size(mQ.shape[0]) * cute.size(mQ.shape[3]), + tile_shape_mn=(self.tile_m, self.tile_n), + qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + mCuSeqlensQ=mCuSeqlensQ, + mSeqUsedQ=mSeqUsedQ, + is_split_kv=getattr(self, "sol_attn_v_split_d64", False), + ) + tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args) + grid_dim = TileScheduler.get_grid_shape(tile_sched_params) + softmax_scale_log2, softmax_scale = utils.compute_softmax_scale_log2(softmax_scale, self.score_mod) + fastdiv_mods = utils.compute_fastdiv_mods(mQ, mK, self.qhead_per_kvhead, self.pack_gqa, aux_tensors) + + kernel_args = ( + mQ, + mK, + mV, + mO, + mLSE, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + softmax_scale_log2, + softmax_scale, + window_size_left, + window_size_right, + self.sQ_layout, + self.sK_layout, + self.sV_layout, + self.sO_layout, + self.sP_layout, + self.gmem_tiled_copy_Q, + self.gmem_tiled_copy_K, + self.gmem_tiled_copy_V, + self.gmem_tiled_copy_O, + tiled_mma_qk, + tiled_mma_pv, + ) + kernel_tail = ( + SharedStorage, + tile_sched_params, + TileScheduler, + aux_tensors, + fastdiv_mods, + ) + kernel = self.kernel(*kernel_args, *kernel_tail) + kernel.launch( + grid=grid_dim, + block=[self.num_threads, 1, 1], + smem=SharedStorage.size_in_bytes(), + min_blocks_per_mp=getattr(self, "sol_attn_min_blocks_per_mp", 0), + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + mCuSeqlensQ: Optional[cute.Tensor], + mCuSeqlensK: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + mSeqUsedK: Optional[cute.Tensor], + softmax_scale_log2: Float32, + softmax_scale: Optional[Float32], + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + sQ_layout: cute.ComposedLayout, + sK_layout: cute.ComposedLayout, + sV_layout: cute.ComposedLayout, + sO_layout: cute.ComposedLayout, + sP_layout: cute.ComposedLayout | None, + gmem_tiled_copy_Q: cute.TiledCopy, + gmem_tiled_copy_K: cute.TiledCopy, + gmem_tiled_copy_V: cute.TiledCopy, + gmem_tiled_copy_O: cute.TiledCopy, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + SharedStorage: cutlass.Constexpr, + tile_sched_params, + TileScheduler: cutlass.Constexpr[Callable], + aux_tensors=None, + fastdiv_mods=None, + ): + # Thread index, block index + tidx, _, _ = cute.arch.thread_idx() + + tile_scheduler = TileScheduler.create(tile_sched_params) + work_tile = tile_scheduler.initial_work_tile_info() + m_block, num_head, batch_size, _ = work_tile.tile_idx + + block_info = BlockInfo( + self.tile_m, + self.tile_n, + self.is_causal, + self.is_local, + False, # is_split_kv + window_size_left, + window_size_right, + qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + seqlen = SeqlenInfoQK.create( + batch_idx=batch_size, + seqlen_q_static=mQ.shape[0], + seqlen_k_static=mK.shape[0], + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=mCuSeqlensK, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=mSeqUsedK, + ) + n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block) + # For varlen, wasted grid tiles (where batch_idx >= num_batch) will have + # seqlen_q=seqlen_k=0 and n_block_max=0. Clamp to 0 so we don't use a + # negative block index for K/V loads; the load/store predicates already + # guard all memory accesses when seqlen is 0. + n_block = cutlass.max(n_block_max - 1, 0) + + # /////////////////////////////////////////////////////////////////////////////// + # Get the appropriate tiles for this thread block. + # /////////////////////////////////////////////////////////////////////////////// + blkQ_shape = (self.tile_m, self.tile_hdim) + blkK_shape = (self.tile_n, self.tile_hdim) + blkV_shape = (self.tile_n, self.tile_hdimv) + num_head_kv = num_head // self.qhead_per_kvhead + if const_expr(not seqlen.has_cu_seqlens_q): + mQ_cur = mQ[None, None, num_head, batch_size] + else: + mQ_cur = cute.domain_offset((seqlen.offset_q, 0), mQ[None, None, num_head]) + if const_expr(not seqlen.has_cu_seqlens_k): + mK_cur = mK[None, None, num_head_kv, batch_size] + mV_cur = mV[None, None, num_head_kv, batch_size] + else: + mK_cur = cute.domain_offset((seqlen.offset_k, 0), mK[None, None, num_head_kv]) + mV_cur = cute.domain_offset((seqlen.offset_k, 0), mV[None, None, num_head_kv]) + gQ = cute.local_tile(mQ_cur, blkQ_shape, (m_block, 0)) + gK = cute.local_tile(mK_cur, blkK_shape, (None, 0)) + gV = cute.local_tile(mV_cur, blkV_shape, (None, 0)) + + # /////////////////////////////////////////////////////////////////////////////// + # Get shared memory buffer + # /////////////////////////////////////////////////////////////////////////////// + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + sQ = storage.sQ.get_tensor(sQ_layout) + sK = storage.sK.get_tensor(sK_layout) + if const_expr(not self.Q_in_regs): + sV = storage.sV.get_tensor(sV_layout) + else: + sV = cute.make_tensor(cute.recast_ptr(sQ.iterator, dtype=self.dtype), sV_layout) + # Transpose view of V to tensor with layout (head_dim_v, tile_n) for tiled mma + sVt = layout_utils.transpose_view(sV) + + gmem_thr_copy_K = gmem_tiled_copy_K.get_slice(tidx) + gmem_thr_copy_V = gmem_tiled_copy_V.get_slice(tidx) + # (CPY_Atom, CPY_N, CPY_K, n_block) + tKsK, tKgK = gmem_thr_copy_K.partition_D(sK), gmem_thr_copy_K.partition_S(gK) + # (CPY_Atom, CPY_N, CPY_K, n_block) + tVsV, tVgV = gmem_thr_copy_V.partition_D(sV), gmem_thr_copy_V.partition_S(gV) + + # /////////////////////////////////////////////////////////////////////////////// + # Tile MMA compute thread partitions and allocate accumulators + # /////////////////////////////////////////////////////////////////////////////// + thr_mma_qk = tiled_mma_qk.get_slice(tidx) + thr_mma_pv = tiled_mma_pv.get_slice(tidx) + tSrQ = thr_mma_qk.make_fragment_A(thr_mma_qk.partition_A(sQ)) + tSrK = thr_mma_qk.make_fragment_B(thr_mma_qk.partition_B(sK[None, None, 0])) + tOrVt = thr_mma_pv.make_fragment_B(thr_mma_pv.partition_B(sVt[None, None, 0])) + acc_shape_O = thr_mma_pv.partition_shape_C((self.tile_m, self.tile_hdimv)) + acc_O = cute.make_rmem_tensor(acc_shape_O, Float32) + acc_O.fill(0.0) + + # /////////////////////////////////////////////////////////////////////////////// + # Smem copy atom tiling + # /////////////////////////////////////////////////////////////////////////////// + smem_copy_atom_QK = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), + self.dtype, + ) + smem_copy_atom_V = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4), + self.dtype, + ) + smem_thr_copy_Q = utils.make_tiled_copy_A(smem_copy_atom_QK, tiled_mma_qk).get_slice(tidx) + smem_thr_copy_K = utils.make_tiled_copy_B(smem_copy_atom_QK, tiled_mma_qk).get_slice(tidx) + smem_thr_copy_V = utils.make_tiled_copy_B(smem_copy_atom_V, tiled_mma_pv).get_slice(tidx) + + tSsQ = smem_thr_copy_Q.partition_S(sQ) + tSsK = smem_thr_copy_K.partition_S(sK) + tOsVt = smem_thr_copy_V.partition_S(sVt) + + # /////////////////////////////////////////////////////////////////////////////// + # Predicate: Mark indices that need to copy when problem_shape isn't a multiple + # of tile_shape + # /////////////////////////////////////////////////////////////////////////////// + # Construct identity layout for KV + cK = cute.make_identity_tensor((self.tile_n, self.tile_hdim)) + tKcK = gmem_thr_copy_K.partition_S(cK) + t0KcK = gmem_thr_copy_K.get_slice(0).partition_S(cK) + if const_expr(self.tile_hdim == self.tile_hdimv): + tVcV = tKcK + t0VcV = t0KcK + else: + cV = cute.make_identity_tensor((self.tile_n, self.tile_hdimv)) + tVcV = gmem_thr_copy_V.partition_S(cV) + t0VcV = gmem_thr_copy_V.get_slice(0).partition_S(cV) + # Allocate predicate tensors for m and n, here we only allocate the tile of k, and + # use "if" on the mn dimension. + # This is to reduce register pressure and gets 2-3% performance gain. + tKpK = utils.predicate_k(tKcK, limit=mK.shape[1]) + if const_expr(self.same_hdim_kv): + tVpV = tKpK + else: + tVpV = utils.predicate_k(tVcV, limit=mV.shape[1]) + + # shape: (atom_v_m * rest_m) + softmax = Softmax.create( + softmax_scale_log2, + num_rows=acc_O.shape[0][0] * acc_O.shape[1], + softmax_scale=softmax_scale, + ) + softmax.reset() + + # group parameters for compute_one_n_block + mma_params = SimpleNamespace( + thr_mma_qk=thr_mma_qk, + thr_mma_pv=thr_mma_pv, + tSrQ=tSrQ, + tSrK=tSrK, + tOrVt=tOrVt, + acc_O=acc_O, + ) + smem_copy_params = SimpleNamespace( + smem_thr_copy_Q=smem_thr_copy_Q, + smem_thr_copy_K=smem_thr_copy_K, + smem_thr_copy_V=smem_thr_copy_V, + tSsQ=tSsQ, + tSsK=tSsK, + tOsVt=tOsVt, + ) + load_K = partial( + self.load_K, gmem_tiled_copy_K, tKgK, tKsK, tKcK, t0KcK, tKpK, seqlen=seqlen.seqlen_k + ) + load_V = partial( + self.load_V, gmem_tiled_copy_V, tVgV, tVsV, tVcV, t0VcV, tVpV, seqlen=seqlen.seqlen_k + ) + + compute_one_n_block = partial( + self.compute_one_n_block, + mma_params=mma_params, + smem_copy_params=smem_copy_params, + softmax=softmax, + load_K=load_K, + load_V=load_V, + score_mod=self.score_mod, + batch_idx=batch_size, + head_idx=num_head, + m_block=m_block, + aux_tensors=aux_tensors, + fastdiv_mods=fastdiv_mods, + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Prologue + # /////////////////////////////////////////////////////////////////////////////// + # Start async loads of the last mn-tile, where we take care of the mn residue + gmem_thr_copy_Q = gmem_tiled_copy_Q.get_slice(tidx) + self.load_Q(gmem_thr_copy_Q, gQ, sQ, m_block, seqlen=seqlen.seqlen_q, headdim=mQ.shape[1]) + cute.arch.cp_async_commit_group() + + def preprocess_Q(): + cute.arch.cp_async_wait_group(self.num_stages * 2 - 1) + if const_expr(self.Q_in_regs): + cute.arch.barrier() + tSrQ_copy_view = smem_thr_copy_Q.retile(tSrQ) + cute.copy(smem_thr_copy_Q, tSsQ, tSrQ_copy_view) + + # If Q_in_regs, we load Q, then load 1 stage of K, then (optionally) rotate Q and + # read from smem_q to registers, then load V. + # If !Q_in_regs, we load Q, load all stages of K & V, then (optionally) rotate Q. + if const_expr(self.Q_in_regs): + load_K(n_block, smem_pipe_write=0, need_predicates=True) + cute.arch.cp_async_commit_group() + preprocess_Q() + cute.arch.barrier() # Make sure all threads have read smem_q before loading V + + for stage in cutlass.range_constexpr(self.num_stages): + if const_expr(not self.Q_in_regs or stage > 0): + if stage == 0 or n_block - stage >= 0: + load_K(n_block - stage, smem_pipe_write=stage, need_predicates=stage == 0) + cute.arch.cp_async_commit_group() + if const_expr(stage < self.num_stages - 1): + if stage == 0 or n_block - stage >= 0: + load_V(n_block - stage, smem_pipe_write=stage, need_predicates=stage == 0) + cute.arch.cp_async_commit_group() + if const_expr(not self.Q_in_regs): + preprocess_Q() + + # /////////////////////////////////////////////////////////////////////////////// + # Mainloop + # /////////////////////////////////////////////////////////////////////////////// + # Start processing of the first n-block. + # For performance reason, we separate out two kinds of iterations: + # those that need masking on S, and those that don't. + # We need masking on S for the very last block when K and V has length not multiple of tile_n. + # We also need masking on S if it's causal, for the last several blocks. + mask = AttentionMask( + self.tile_m, + self.tile_n, + seqlen, + window_size_left, + window_size_right, + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + mask_fn = partial( + mask.apply_mask, + batch_idx=batch_size, + head_idx=num_head, + m_block=m_block, + thr_mma=thr_mma_qk, + mask_causal=self.is_causal, + mask_local=self.is_local, + aux_tensors=aux_tensors, + fastdiv_mods=fastdiv_mods if const_expr(self.mask_mod is not None) else None, + ) + + # First iteration with seqlen masking + smem_pipe_read = Int32(0) + smem_pipe_write = Int32(self.num_stages - 1) + compute_one_n_block( + n_block, + smem_pipe_read, + smem_pipe_write, + is_first_n_block=True, + seqlen=seqlen, + mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=True), + ) + smem_pipe_read = self.advance_pipeline(smem_pipe_read) + smem_pipe_write = self.advance_pipeline(smem_pipe_write) + # Next couple of iterations with causal masking + if const_expr(self.is_causal or self.is_local): + n_block_min_causal_local_mask = block_info.get_n_block_min_causal_local_mask( + seqlen, m_block, n_block_min + ) + for n_tile in cutlass.range(n_block_max - 1 - n_block_min_causal_local_mask, unroll=1): + n_block = n_block_max - 2 - n_tile + compute_one_n_block( + n_block, + smem_pipe_read, + smem_pipe_write, + seqlen=seqlen, + mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=True), + ) + smem_pipe_read = self.advance_pipeline(smem_pipe_read) + smem_pipe_write = self.advance_pipeline(smem_pipe_write) + # The remaining iterations have no masking + for n_tile in cutlass.range(n_block, unroll=1): + compute_one_n_block( + n_block - n_tile - 1, smem_pipe_read, smem_pipe_write, + seqlen=seqlen, is_first_n_block=False, + mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=False) + ) + smem_pipe_read = self.advance_pipeline(smem_pipe_read) + smem_pipe_write = self.advance_pipeline(smem_pipe_write) + # TODO: local + + # normalize acc_O by row_sum and calculate the lse + row_scale = softmax.finalize() + softmax.rescale_O(acc_O, row_scale) + + # /////////////////////////////////////////////////////////////////////////////// + # Epilogue + # /////////////////////////////////////////////////////////////////////////////// + # reuse sQ's data iterator + sO = cute.make_tensor(sQ.iterator, sO_layout) + self.epilogue( + acc_O, + softmax.row_sum, + mO, + mLSE, + sO, + seqlen, + gmem_tiled_copy_O, + None, + tiled_mma_pv, + tidx, + m_block, + num_head, + batch_size, + ) + + @cute.jit + def compute_one_n_block( + self, + n_block: Int32, + smem_pipe_read: Int32, + smem_pipe_write: Int32, + mma_params: SimpleNamespace, + smem_copy_params: SimpleNamespace, + softmax: Softmax, + load_K: Callable, + load_V: Callable, + score_mod: Callable | None, + batch_idx: cutlass.Int32, + head_idx: cutlass.Int32, + m_block: cutlass.Int32, + seqlen: SeqlenInfoQK, + aux_tensors=None, + fastdiv_mods=None, + mask_fn: Optional[Callable] = None, + is_first_n_block: cutlass.Constexpr = False, + check_inf: cutlass.Constexpr = True, + ): + """Compute one n_block of S/O. + + This function provides different variants for processing the first n block versus + subsequent blocks. + """ + + def sync(): + cute.arch.cp_async_wait_group(self.num_stages * 2 - 2) + cute.arch.barrier() + + acc_shape_S = mma_params.thr_mma_qk.partition_shape_C((self.tile_m, self.tile_n)) + acc_S = cute.make_rmem_tensor(acc_shape_S, Float32) + acc_S.fill(0.0) + # wait for smem tile QK before mma calculation for S + sync() + + # need predicates for the first tile + def load_V_next(): + if self.num_stages == 1 or n_block - self.num_stages + 1 >= 0: + load_V( + n_block - self.num_stages + 1, + smem_pipe_write, + need_predicates=is_first_n_block and self.num_stages == 1, + ) + cute.arch.cp_async_commit_group() + + load_V_next() + sm80_utils.gemm( + mma_params.thr_mma_qk, + acc_S, + mma_params.tSrQ, + mma_params.tSrK, + smem_copy_params.tSsQ, + smem_copy_params.tSsK[ + None, None, None, smem_pipe_read if const_expr(self.num_stages > 1) else 0 + ], + smem_copy_params.smem_thr_copy_Q, + smem_copy_params.smem_thr_copy_K, + # hook_fn=load_V_next, + A_in_regs=self.Q_in_regs, + ) + if const_expr(score_mod is not None): + self.apply_score_mod( + mma_params.thr_mma_qk, + batch_idx, + head_idx, + m_block, + acc_S, + n_block, + seqlen, + softmax_scale=softmax.softmax_scale, + aux_tensors=aux_tensors, + fastdiv_mods=fastdiv_mods, + ) + + smem_pipe_write = self.advance_pipeline(smem_pipe_write) + + def load_K_next(): + if n_block - self.num_stages >= 0: + load_K(n_block - self.num_stages, smem_pipe_write, need_predicates=False) + cute.arch.cp_async_commit_group() + + # wait for smem tile V for O + if const_expr(self.num_stages == 1): + sync() + load_K_next() + if const_expr(mask_fn is not None): + mask_fn(acc_S, n_block=n_block) + row_scale = softmax.online_softmax(acc_S, is_first=is_first_n_block, check_inf=check_inf) + softmax.rescale_O(mma_params.acc_O, row_scale) + rP = cute.make_fragment_like(acc_S, self.dtype) + rP.store(acc_S.load().to(self.dtype)) + tOrP = layout_utils.reshape_acc_to_frgA(rP) + if const_expr(self.num_stages > 1): + sync() + load_K_next() + sm80_utils.gemm_rs( + mma_params.thr_mma_pv, + mma_params.acc_O, + tOrP, + mma_params.tOrVt, + smem_copy_params.tOsVt[ + None, None, None, smem_pipe_read if const_expr(self.num_stages > 1) else 0 + ], + smem_copy_params.smem_thr_copy_V, + # hook_fn=load_K_next, + ) + # if const_expr(self.num_stages > 1): + # load_K_next() + + +def __getattr__(name): + if name == "FlashAttentionForwardSm90": + raise AttributeError("FlashAttentionForwardSm90 is not vendored in the SOL_ATTN release") + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/mask.py b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/mask.py new file mode 100644 index 00000000..0a6afacb --- /dev/null +++ b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/mask.py @@ -0,0 +1,712 @@ +# Copyright (c) 2025, Tri Dao. + +from typing import Optional, Callable, TypeAlias +from dataclasses import dataclass + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Uint32, const_expr + +from telefuser.kernel.sol_attn.sm90._compat import layout_utils +import telefuser.kernel.sol_attn._vendor.flash_attn.cute.utils as utils +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.seqlen_info import SeqlenInfoQK + +MaskGenFn: TypeAlias = Callable[[int], Uint32] +MASK_R2P_CHUNK_SIZE: int = 32 + + +@cute.jit +def r2p_bitmask_below(limit: Int32, s: int) -> Uint32: + """32-bit R2P bitmask keeping positions < limit (exclusive upper bound). + + Positions 0..limit-1 in chunk `s` get bit=1 (keep), the rest bit=0 (mask). + Uses inline PTX to avoid shift-by-type-width UB. + """ + m = max((s + 1) * MASK_R2P_CHUNK_SIZE - limit, 0) + return utils.shr_u32(Uint32(0xFFFFFFFF), Uint32(m)) + + +@cute.jit +def r2p_bitmask_above(limit: Int32, s: int) -> Uint32: + """32-bit R2P bitmask keeping positions >= limit (inclusive lower bound). + + Positions limit..31 in chunk `s` get bit=1 (keep), the rest bit=0 (mask). + Uses inline PTX to avoid shift-by-type-width UB. + """ + n = max(limit - s * MASK_R2P_CHUNK_SIZE, 0) + return utils.shl_u32(Uint32(0xFFFFFFFF), Uint32(n)) + + +@cute.jit +def mask_r2p_lambda( + X: cute.Tensor, + mask_gen_fn: cutlass.Constexpr[MaskGenFn], + rank1: bool = False, +) -> None: + """Apply R2P masking with a custom bitmask generator. + + mask_gen_fn(chunk_idx: constexpr int) -> Uint32: + Returns a 32-bit bitmask for the chunk. Bit i set means column + chunk_idx * chunk_size + i is KEPT; bit i clear means masked to -inf. + """ + ncol = const_expr(cute.size(X.shape[cute.rank(X) - 1]) if not rank1 else cute.size(X.shape)) + # 32-column chunks. The mask_gen_fn returns a Uint32 bitmask (1=keep). + CHUNK_SIZE = MASK_R2P_CHUNK_SIZE + for s in cutlass.range_constexpr(cute.ceil_div(ncol, CHUNK_SIZE)): + mask = mask_gen_fn(s) + # This needs to be range_constexpr, o/w the compiler can't generate the R2P instruction + for i in cutlass.range_constexpr(min(CHUNK_SIZE, ncol - s * CHUNK_SIZE)): + in_bound = cutlass.Boolean(mask & (Uint32(1) << i)) + c = s * CHUNK_SIZE + i + if const_expr(rank1): + X[c] = X[c] if in_bound else -Float32.inf + else: + for r in cutlass.range_constexpr(cute.size(X.shape[0])): + X[r, c] = X[r, c] if in_bound else -Float32.inf + + +@cute.jit +def sm90_col_to_r2p_idx(col_limit: Int32) -> Int32: + """Transform SM90 MMA column coordinate to R2P element index. + + SM90 MMA accumulator column indices are non-contiguous: 0, 1, 8, 9, 16, 17, ... + Element indices are contiguous: 0, 1, 2, 3, 4, 5, ... + This converts a column-space threshold to element-space for r2p_bitmask_below/above. + """ + return col_limit // 8 * 2 + min(col_limit % 8, 2) + + +@cute.jit +def row_to_r2p_idx(x: Int32, num_rep: int, num_wg: int) -> Int32: + """Convert a row coordinate to an R2P element index in the warp-group interleaved layout. + + In the SM100 backward pass, 2 warp groups share TMEM. The TMEM load atom + distributes rows in an interleaved pattern: elements 0..num_rep-1 map to + rows 0..num_rep-1 (warp group 0), elements num_rep..2*num_rep-1 map to + rows num_rep*num_wg..num_rep*num_wg+num_rep-1 (warp group 1), and so on. + Row-coordinate thresholds (causal limits, window bounds, uih_len) must be + converted to element indices before use with r2p_bitmask_above/below. + + Rows not owned by this thread (in the gap between warp groups) are clamped + to the boundary element index, which is safe because R2P thresholds are + monotonic. + + Example with num_rep=16, num_wg=2: + row 0 -> elem 0, row 15 -> elem 15, + row 16 -> elem 16 (clamped), row 31 -> elem 16 (clamped), + row 32 -> elem 16, row 33 -> elem 17, row 47 -> elem 31. + """ + return x // (num_rep * num_wg) * num_rep + min(x % (num_rep * num_wg), num_rep) + + +@dataclass(frozen=True) +class AttentionMask: + tile_m: cutlass.Constexpr[int] + tile_n: cutlass.Constexpr[int] + seqlen_info: SeqlenInfoQK + window_size_left: Optional[Int32] = None + window_size_right: Optional[Int32] = None + qhead_per_kvhead_packgqa: cutlass.Constexpr[int] = 1 # only pass in if we're doing PackGQA + swap_AB: cutlass.Constexpr[bool] = False + + @property + def seqlen_q(self) -> Int32: + return self.seqlen_info.seqlen_q + + @property + def seqlen_k(self) -> Int32: + return self.seqlen_info.seqlen_k + + @cute.jit + def apply_mask( + self, + acc_S: cute.Tensor, + batch_idx: cutlass.Int32, + head_idx: cutlass.Int32, + m_block: cutlass.Int32, + n_block: cutlass.Int32, + thr_mma: cute.TiledMma, + mask_seqlen: cutlass.Constexpr[bool], + mask_causal: cutlass.Constexpr[bool], + mask_local: cutlass.Constexpr[bool] = False, + mask_mod: cutlass.Constexpr[Optional[Callable]] = None, + aux_tensors: Optional[list] = None, + fastdiv_mods=(None, None), + ) -> None: + assert not (mask_causal and mask_local), "mask_causal and mask_local cannot be both True" + acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S, transpose=self.swap_AB) + acc_shape = (self.tile_m, self.tile_n) + cS = cute.make_identity_tensor(acc_shape if not self.swap_AB else acc_shape[::-1]) + tScS_mn = layout_utils.reshape_acc_to_mn(thr_mma.partition_C(cS), transpose=self.swap_AB) + # We use t0ScS as these indices are known at compile time. We then must subtract the + # column limit by the thread column offset. + t0ScS_mn = layout_utils.reshape_acc_to_mn( + thr_mma.get_slice(0).partition_C(cS), transpose=self.swap_AB + ) + ROW = 0 if const_expr(not self.swap_AB) else 1 + COL = 1 if const_expr(not self.swap_AB) else 0 + thr_col_offset = tScS_mn[0][COL] + # To handle edge cases of completely masked out rows where n_block_max = 0, + # we treat negative n_blocks as 0th n_block + # TODO: find more transparent solution + if n_block < 0: + n_block = 0 + seqlenk_col_limit = self.seqlen_k - n_block * self.tile_n - thr_col_offset + if const_expr(not mask_causal and not mask_local and mask_mod is None): + if const_expr(mask_seqlen): + r2p = const_expr(not self.swap_AB) + if const_expr(not r2p): + # traverse column index. + for c in cutlass.range(cute.size(tScS_mn.shape[1]), unroll_full=True): + oob = t0ScS_mn[0, c][COL] >= seqlenk_col_limit + for r in cutlass.range(cute.size(tScS_mn.shape[0]), unroll_full=True): + acc_S_mn[r, c] = -Float32.inf if oob else acc_S_mn[r, c] + else: + seqlenk_col_limit_r2p = sm90_col_to_r2p_idx(seqlenk_col_limit) + mask_r2p_lambda(acc_S_mn, lambda s: r2p_bitmask_below(seqlenk_col_limit_r2p, s)) + + elif const_expr( + not mask_causal and not mask_local and mask_mod is not None + ): # FlexAttention mask mod + nrow = const_expr(cute.size(tScS_mn.shape[0])) + ncol = const_expr(cute.size(tScS_mn.shape[1])) + has_fastdiv = const_expr( + fastdiv_mods is not None + and fastdiv_mods[0] is not None + and fastdiv_mods[1] is not None + ) + wrap_aux_indices = const_expr( + has_fastdiv and mask_seqlen and const_expr(aux_tensors is not None) + ) + + for r in cutlass.range_constexpr(nrow): + # Respect swap_AB: ROW/COL determine which coordinate component corresponds to Q/KV. + local_row = tScS_mn[r, 0][ROW] + global_row_idx = local_row + m_block * self.tile_m + row_for_mod = global_row_idx + head_idx_for_mod = head_idx + if const_expr(self.qhead_per_kvhead_packgqa != 1): + head_offset = global_row_idx % self.qhead_per_kvhead_packgqa + head_idx_for_mod = head_idx * self.qhead_per_kvhead_packgqa + head_offset + row_for_mod = global_row_idx // self.qhead_per_kvhead_packgqa + row_for_seqlen = row_for_mod + if const_expr(wrap_aux_indices): + _, row_for_mod = divmod(row_for_mod, fastdiv_mods[0]) + + for col in cutlass.range_constexpr(ncol): + col_idx_local = t0ScS_mn[0, col][COL] + # Convert to absolute column index + global_col_idx = thr_col_offset + col_idx_local + n_block * self.tile_n + col_for_mod = global_col_idx + if const_expr(wrap_aux_indices): + _, col_for_mod = divmod(global_col_idx, fastdiv_mods[1]) + + batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32) + head_idx_ssa = utils.scalar_to_ssa(head_idx_for_mod, cutlass.Int32) + q_idx_ssa = utils.scalar_to_ssa(row_for_mod, cutlass.Int32) + kv_idx_ssa = utils.scalar_to_ssa(col_for_mod, cutlass.Int32) + mask_value = mask_mod( + batch_idx_ssa, + head_idx_ssa, + q_idx_ssa, + kv_idx_ssa, + self.seqlen_info, + aux_tensors, + ) + cond = cutlass.Boolean(utils.ssa_to_scalar(mask_value)) + if const_expr(mask_seqlen): + out_of_bounds = (row_for_seqlen >= self.seqlen_q) or ( + global_col_idx >= self.seqlen_k + ) + if out_of_bounds: + acc_S_mn[r, col] = -cutlass.Float32.inf + else: + acc_S_mn[r, col] = acc_S_mn[r, col] if cond else -cutlass.Float32.inf + else: + acc_S_mn[r, col] = acc_S_mn[r, col] if cond else -cutlass.Float32.inf + + else: # Causal or local + if const_expr(not self.swap_AB): + # If PackGQA, we split the work of compute divmod among threads in the same row + threads_per_row = thr_mma.tv_layout_C.shape[0][0] + mma_m_idx = None + if const_expr(self.qhead_per_kvhead_packgqa != 1): + assert not self.swap_AB, "swap_AB with PackGQA not supported yet" + assert cute.arch.WARP_SIZE % threads_per_row == 0, ( + "threads_per_row must divide WARP_SIZE" + ) + assert cute.size(acc_S_mn.shape[0]) <= threads_per_row + tidx = thr_mma.thr_idx + mma_m_idx = ( + m_block * self.tile_m + tScS_mn[tidx % threads_per_row, 0][0] + ) // self.qhead_per_kvhead_packgqa + causal_row_offset = ( + 1 + self.seqlen_k - n_block * self.tile_n - self.seqlen_q - thr_col_offset + ) + if const_expr(mask_causal): + r2p = const_expr(not self.swap_AB) # R2P trick, see apply_mask_sm100 + for r in cutlass.range(cute.size(tScS_mn.shape[0]), unroll_full=True): + # get the column index limit based on current row. Only consider the row index, so the column index sets to 0. + if const_expr(self.qhead_per_kvhead_packgqa == 1): + row_idx = tScS_mn[r, 0][0] + m_block * self.tile_m + else: + row_idx = utils.shuffle_sync( + mma_m_idx, r % threads_per_row, width=threads_per_row + ) + col_limit_right = row_idx + causal_row_offset + if const_expr(mask_seqlen): + col_limit_right = cutlass.min(col_limit_right, seqlenk_col_limit) + if const_expr(not r2p): + # traverse column index. + for c in cutlass.range(cute.size(tScS_mn.shape[1]), unroll_full=True): + acc_S_mn[r, c] = ( + -Float32.inf + if t0ScS_mn[0, c][1] >= col_limit_right + else acc_S_mn[r, c] + ) + else: + col_limit_r2p = sm90_col_to_r2p_idx(col_limit_right) + mask_r2p_lambda( + acc_S_mn[r, None], + lambda s: r2p_bitmask_below(col_limit_r2p, s), + rank1=True, + ) + else: # Local + local_row_offset_right = ( + causal_row_offset + self.window_size_right + if const_expr(self.window_size_right is not None) + else None + ) + local_row_offset_left = ( + causal_row_offset - 1 - self.window_size_left + if const_expr(self.window_size_left is not None) + else None + ) + r2p_local = const_expr(not self.swap_AB) + for r in cutlass.range(cute.size(tScS_mn.shape[0]), unroll_full=True): + if const_expr(self.qhead_per_kvhead_packgqa == 1): + row_idx = tScS_mn[r, 0][0] + m_block * self.tile_m + else: + row_idx = utils.shuffle_sync( + mma_m_idx, r % threads_per_row, width=threads_per_row + ) + if const_expr(self.window_size_right is not None): + col_limit_right = row_idx + local_row_offset_right + else: + col_limit_right = self.tile_n + if const_expr(mask_seqlen): + col_limit_right = cutlass.min(col_limit_right, seqlenk_col_limit) + col_limit_left = ( + row_idx + local_row_offset_left + if const_expr(self.window_size_left is not None) + else 0 + ) + if const_expr(not r2p_local): + # traverse column index. + for c in cutlass.range(cute.size(tScS_mn.shape[1]), unroll_full=True): + col_idx = t0ScS_mn[0, c][1] + if col_idx >= col_limit_right or col_idx < col_limit_left: + acc_S_mn[r, c] = -Float32.inf + else: + col_limit_right_r2p = sm90_col_to_r2p_idx(col_limit_right) + col_limit_left_r2p = sm90_col_to_r2p_idx(col_limit_left) + + def mask_gen_fn(s: int) -> Uint32: + return r2p_bitmask_below( + col_limit_right_r2p, s + ) & r2p_bitmask_above(col_limit_left_r2p, s) + + mask_r2p_lambda(acc_S_mn[r, None], mask_gen_fn, rank1=True) + else: # swap_AB + assert self.qhead_per_kvhead_packgqa == 1 + thr_row_offset = tScS_mn[0][ROW] + causal_row_offset = ( + seqlenk_col_limit - self.seqlen_q + m_block * self.tile_m + thr_row_offset + ) + if const_expr(mask_causal): + for c in cutlass.range(cute.size(tScS_mn.shape[1]), unroll_full=True): + col0 = t0ScS_mn[0, c][COL] + # If col0 is beyond the column limit, we want to mask out the entire + # column, by setting row limit to be self.tile_m. + row_limit_top = ( + self.tile_m + if col0 >= seqlenk_col_limit and mask_seqlen + else col0 - causal_row_offset + ) + for r in cutlass.range(cute.size(tScS_mn.shape[0]), unroll_full=True): + acc_S_mn[r, c] = ( + -Float32.inf + if t0ScS_mn[r, 0][ROW] < row_limit_top + else acc_S_mn[r, c] + ) + else: + for c in cutlass.range(cute.size(tScS_mn.shape[1]), unroll_full=True): + col0 = t0ScS_mn[0, c][COL] + # If col0 is beyond the column limit, we want to mask out the entire + # column, by setting row limit to be self.tile_m. + row_limit_top = ( + self.tile_m + if col0 >= seqlenk_col_limit and mask_seqlen + else ( + col0 - causal_row_offset - self.window_size_right + if const_expr(self.window_size_right is not None) + else 0 + ) + ) + row_limit_bot = ( + col0 - causal_row_offset + self.window_size_left + if const_expr(self.window_size_left is not None) + else self.tile_m + ) + for r in cutlass.range(cute.size(tScS_mn.shape[0]), unroll_full=True): + row_idx = t0ScS_mn[r, 0][ROW] + acc_S_mn[r, c] = ( + -Float32.inf + if row_idx < row_limit_top or row_idx > row_limit_bot + else acc_S_mn[r, c] + ) + + @cute.jit + def apply_mask_sm100( + self, + acc_S: cute.Tensor, + m_block: Int32, + n_block: Int32, + thr_mma: cute.TiledMma, + thr_tmem_load: cute.TiledCopy, + mask_seqlen: cutlass.Constexpr[bool], + mask_causal: cutlass.Constexpr[bool], + mask_local: cutlass.Constexpr[bool] = False, + mask_mod: cutlass.Constexpr[Optional[Callable]] = None, + batch_idx: Int32 = None, + head_idx: Int32 = None, + aux_tensors: Optional[list] = None, + fastdiv_mods=(None, None), + head_divmod=None, + check_q_boundary: bool = False, + vec_size: cutlass.Constexpr[int] = 1, + ) -> None: + assert not (mask_causal and mask_local), "mask_causal and mask_local cannot be both True" + acc_shape = (self.tile_m, self.tile_n) + cS = cute.make_identity_tensor(acc_shape if not self.swap_AB else acc_shape[::-1]) + tScS = thr_mma.partition_C(cS) + tScS = tScS[(None, None), 0, 0] + tScS_t2r = thr_tmem_load.partition_D(tScS) + # To handle edge cases of completely masked out rows where n_block_max = 0, + # we treat negative n_blocks as 0th n_block + # TODO: find more transparent solution + if n_block < 0: + n_block = 0 + seqlenk_col_limit = self.seqlen_k - n_block * self.tile_n + r2p = True + if const_expr(not mask_causal and not mask_local and mask_mod is None): + if const_expr(mask_seqlen): + if const_expr(not r2p): + for i in cutlass.range(cute.size(tScS_t2r.shape), unroll_full=True): + # if tScS_t2r[i][1] >= seqlenk_col_limit: + # acc_S[i] = -Float32.inf + # For some reason the 2 lines above generate really bad SASS + acc_S[i] = -Float32.inf if tScS_t2r[i][1] >= seqlenk_col_limit else acc_S[i] + else: + mask_r2p_lambda( + acc_S, + lambda s: r2p_bitmask_below(seqlenk_col_limit, s), + rank1=True, + ) + + elif const_expr(not mask_causal and not mask_local and mask_mod is not None): + # Block sparse case w/ mask_mod + has_fastdiv = const_expr( + fastdiv_mods is not None + and fastdiv_mods[0] is not None + and fastdiv_mods[1] is not None + ) + batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32) + + ncol = const_expr(cute.size(tScS_t2r.shape)) + for i in cutlass.range_constexpr(ncol): + row_coord = tScS_t2r[i][0] if not self.swap_AB else tScS_t2r[i][1] + col_coord = tScS_t2r[i][1] if not self.swap_AB else tScS_t2r[i][0] + global_row = row_coord + m_block * self.tile_m + global_col = col_coord + n_block * self.tile_n + + if const_expr(self.qhead_per_kvhead_packgqa != 1): + assert head_divmod is not None + mask_row, head_offset = divmod(global_row, head_divmod) + head_idx_for_mod = head_idx * self.qhead_per_kvhead_packgqa + head_offset + else: + head_idx_for_mod = head_idx + mask_row = global_row + + mask_row_for_mod = mask_row + if const_expr(has_fastdiv and aux_tensors is not None): + if check_q_boundary: + _, mask_row_for_mod = divmod(mask_row, fastdiv_mods[0]) + global_col_for_mod = global_col + if const_expr(has_fastdiv and mask_seqlen and aux_tensors is not None): + _, global_col_for_mod = divmod(global_col, fastdiv_mods[1]) + + head_idx_ssa = utils.scalar_to_ssa(head_idx_for_mod, cutlass.Int32) + mask_row_ssa = utils.scalar_to_ssa(mask_row_for_mod, cutlass.Int32) + kv_idx_ssa = utils.scalar_to_ssa(global_col_for_mod, cutlass.Int32) + mask_value = mask_mod( + batch_idx_ssa, + head_idx_ssa, + mask_row_ssa, + kv_idx_ssa, + self.seqlen_info, + aux_tensors, + ) + cond = cutlass.Boolean(utils.ssa_to_scalar(mask_value)) + acc_S[i] = acc_S[i] if cond else -Float32.inf + if const_expr(mask_seqlen): + acc_S[i] = -Float32.inf if global_col >= self.seqlen_k else acc_S[i] + if check_q_boundary: + acc_S[i] = -Float32.inf if mask_row >= self.seqlen_q else acc_S[i] + + else: # Causal or local + causal_row_offset = self.seqlen_k - n_block * self.tile_n - self.seqlen_q + row_idx = tScS_t2r[0][0] + m_block * self.tile_m + if const_expr(self.qhead_per_kvhead_packgqa != 1): + row_idx = row_idx // self.qhead_per_kvhead_packgqa + if const_expr(mask_causal): + col_limit_right = row_idx + causal_row_offset + 1 + if const_expr(mask_seqlen): + col_limit_right = cutlass.min(col_limit_right, seqlenk_col_limit) + # if cute.arch.thread_idx()[0] % 32 == 0: + # cute.printf("tidx = %d, tidx tmem = %d, row_idx = %d, col_limit_right = %d, causal_row_offset = %d\n", cute.arch.thread_idx()[0], thr_tmem_load.thr_idx, row_idx, col_limit_right, causal_row_offset) + ncol = const_expr(cute.size(tScS_t2r.shape)) + if const_expr(not r2p): + for i in cutlass.range(ncol, unroll_full=True): + acc_S[i] = -Float32.inf if tScS_t2r[i][1] >= col_limit_right else acc_S[i] + else: + mask_r2p_lambda( + acc_S, + lambda s: r2p_bitmask_below(col_limit_right, s), + rank1=True, + ) + else: + local_row_offset_right = ( + causal_row_offset + 1 + self.window_size_right + if const_expr(self.window_size_right is not None) + else None + ) + local_row_offset_left = ( + causal_row_offset - self.window_size_left + if const_expr(self.window_size_left is not None) + else None + ) + if const_expr(self.window_size_right is not None): + col_limit_right = row_idx + local_row_offset_right + else: + col_limit_right = self.tile_n + if const_expr(mask_seqlen): + col_limit_right = cutlass.min(col_limit_right, seqlenk_col_limit) + col_limit_left = ( + row_idx + local_row_offset_left + if const_expr(self.window_size_left is not None) + else 0 + ) + if const_expr(not r2p): + # if cute.arch.thread_idx()[0] == 0 or cute.arch.thread_idx()[0] == 128: cute.printf("m_block = {}, n_block = {}, row_idx = {}, causal_row_offset = {}, col_limit_right = {}, col_limit_left = {}", m_block, n_block, row_idx, causal_row_offset, col_limit_right, col_limit_left) + for i in cutlass.range(cute.size(tScS_t2r.shape), unroll_full=True): + col_idx = tScS_t2r[i][1] + acc_S[i] = ( + -Float32.inf + if col_idx >= col_limit_right or col_idx < col_limit_left + else acc_S[i] + ) + else: + # Dual-bound R2P masking for SM100. + # Masks elements where: NOT (col_limit_left <= col < col_limit_right) + + def mask_gen_fn(s: int) -> Uint32: + return r2p_bitmask_below(col_limit_right, s) & r2p_bitmask_above( + col_limit_left, s + ) + + mask_r2p_lambda(acc_S, mask_gen_fn, rank1=True) + + @cute.jit + def apply_mask_sm100_transposed( + self, + acc_S: cute.Tensor, + tScS_t2r: cute.Tensor, + t0ScS_t2r: cute.Tensor, + m_block: cutlass.Int32, + n_block: cutlass.Int32, + mask_seqlen: cutlass.Constexpr, + mask_causal: cutlass.Constexpr, + mask_local: cutlass.Constexpr, + mask_mod: cutlass.Constexpr[Optional[Callable]] = None, + batch_idx: Int32 = None, + head_idx: Int32 = None, + aux_tensors: Optional[list] = None, + fastdiv_mods=(None, None), + is_full_block: bool = False, + check_m_boundary: bool = True, + ) -> None: + """ + Backward pass: mask S = K @ Q.T where n_block tiles seqlen_k and m_block tiles seqlen_q. + + Coordinate conventio: + - ROW corresponds to Q (m_block) + - COL corresponds to KV (n_block) + + is_full_block: If True, skip mask_mod (all elements valid). Only apply seqlen masking. + check_m_boundary: If False, skip seqlen_q boundary check (optimization for non-boundary m_blocks). + When iterating m_blocks in forward order, only the last m_block may be partial. + """ + assert not (mask_causal and mask_local), "mask_causal and mask_local cannot be both True" + ROW = 0 if const_expr(not self.swap_AB) else 1 + COL = 1 if const_expr(not self.swap_AB) else 0 + # assert t0ScS_t2r[0][COL] == 0, "col0 == 0" # tmp comment for 2-cta bwd + thr_col_offset = tScS_t2r[0][COL] + seqlenk_col_limit = self.seqlen_k - n_block * self.tile_n - thr_col_offset + + if const_expr(not mask_causal and not mask_local and mask_mod is not None): + # Block sparse case with mask_mod (backward) + # + # Coordinate convention: ROW → Q (m_block), COL → KV (n_block). + # These already account for swap_AB. + # + # FULL blocks: mask_mod returns True for all elements, so skip it. + # Still need seqlen bounds check (elements may be OOB on last m_block). + # PARTIAL blocks: apply mask_mod element-wise, then seqlen bounds. + if is_full_block: + if const_expr(mask_seqlen): + if seqlenk_col_limit <= 0: + # Entire tile is OOB for K + for i in cutlass.range(cute.size(acc_S.shape), unroll_full=True): + acc_S[i] = -cutlass.Float32.inf + elif check_m_boundary: + # Last m_block: check Q and K boundaries + ncol = const_expr(cute.size(tScS_t2r.shape)) + for i in cutlass.range_constexpr(ncol): + row_coord = tScS_t2r[i][ROW] + col_coord = tScS_t2r[i][COL] + global_q = row_coord + m_block * self.tile_m + global_kv = col_coord + n_block * self.tile_n + q_out_of_bounds = global_q >= self.seqlen_q + kv_out_of_bounds = global_kv >= self.seqlen_k + out_of_bounds = q_out_of_bounds or kv_out_of_bounds + acc_S[i] = -cutlass.Float32.inf if out_of_bounds else acc_S[i] + else: + # Partial block + has_fastdiv = const_expr( + fastdiv_mods is not None + and fastdiv_mods[0] is not None + and fastdiv_mods[1] is not None + ) + wrap_aux_indices = const_expr( + has_fastdiv and mask_seqlen and const_expr(aux_tensors is not None) + ) + batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32) + head_idx_ssa = utils.scalar_to_ssa(head_idx, cutlass.Int32) + + ncol = const_expr(cute.size(tScS_t2r.shape)) + for i in cutlass.range_constexpr(ncol): + row_coord = tScS_t2r[i][ROW] + col_coord = tScS_t2r[i][COL] + global_q = row_coord + m_block * self.tile_m + global_kv = col_coord + n_block * self.tile_n + + q_idx_for_mod = global_q + kv_idx_for_mod = global_kv + if const_expr(wrap_aux_indices): + _, q_idx_for_mod = divmod(global_q, fastdiv_mods[0]) + _, kv_idx_for_mod = divmod(global_kv, fastdiv_mods[1]) + + q_idx_ssa = utils.scalar_to_ssa(q_idx_for_mod, cutlass.Int32) + kv_idx_ssa = utils.scalar_to_ssa(kv_idx_for_mod, cutlass.Int32) + + mask_value = mask_mod( + batch_idx_ssa, + head_idx_ssa, + q_idx_ssa, + kv_idx_ssa, + self.seqlen_info, + aux_tensors, + ) + cond = cutlass.Boolean(utils.ssa_to_scalar(mask_value)) + acc_S[i] = acc_S[i] if cond else -cutlass.Float32.inf + + if const_expr(mask_seqlen): + # check_m_boundary=False skips q check for non-boundary m_blocks + q_out_of_bounds = check_m_boundary and (global_q >= self.seqlen_q) + kv_out_of_bounds = global_kv >= self.seqlen_k + out_of_bounds = q_out_of_bounds or kv_out_of_bounds + acc_S[i] = -cutlass.Float32.inf if out_of_bounds else acc_S[i] + + elif const_expr(not mask_causal and not mask_local): + if const_expr(mask_seqlen): + if seqlenk_col_limit <= 0: + for i in cutlass.range(cute.size(acc_S.shape), unroll_full=True): + acc_S[i] = -cutlass.Float32.inf + else: # Causal or local + thr_row_offset = tScS_t2r[0][ROW] + seqlenq_row_limit = self.seqlen_q - m_block * self.tile_m - thr_row_offset + causal_offset = seqlenq_row_limit - seqlenk_col_limit + if const_expr(mask_causal): + # tidx = cute.arch.thread_idx()[0] % 256 + # if tidx < 32: + # cute.printf("tidx = {}, {} {}, {} {}", tidx, tScS_t2r[0][0], tScS_t2r[0][1], tScS_t2r[1][0], tScS_t2r[1][1]) + row_limit_top = causal_offset + if const_expr(mask_seqlen): + # If col is beyond the column limit, we want to mask out the entire + # column, by setting row limit to be self.tile_m. + if seqlenk_col_limit <= 0: + row_limit_top = self.tile_m + r2p = True + if const_expr(not r2p): + for i in cutlass.range(cute.size(acc_S.shape), unroll_full=True): + acc_S[i] = ( + -cutlass.Float32.inf if t0ScS_t2r[i][ROW] < row_limit_top else acc_S[i] + ) + else: + num_rep = cute.size(tScS_t2r, mode=[0]) # 16 or 32 + num_wg = 2 + row_limit = row_to_r2p_idx(row_limit_top, num_rep, num_wg) + mask_r2p_lambda( + acc_S, + lambda s: r2p_bitmask_above(row_limit, s), + rank1=True, + ) + else: + if const_expr(self.window_size_right is not None): + row_limit_top = causal_offset - self.window_size_right + else: + row_limit_top = 0 + if const_expr(self.window_size_left is not None): + row_limit_bot = causal_offset + self.window_size_left + if const_expr(mask_seqlen): + if seqlenk_col_limit <= 0: + row_limit_top = self.tile_m + r2p = True + if const_expr(not r2p): + for i in cutlass.range(cute.size(acc_S.shape), unroll_full=True): + row_idx = t0ScS_t2r[i][ROW] + local_mask = row_idx < row_limit_top + if const_expr(self.window_size_left is not None): + local_mask |= row_idx > row_limit_bot + acc_S[i] = -cutlass.Float32.inf if local_mask else acc_S[i] + else: + + def mask_gen_fn(s: int) -> Uint32: + num_rep = cute.size(tScS_t2r, mode=[0]) + num_wg = 2 + + row_limit = row_to_r2p_idx(row_limit_top, num_rep, num_wg) + mask = r2p_bitmask_above(row_limit, s) + + if const_expr(self.window_size_left is not None): + row_limit_bottom = row_to_r2p_idx(row_limit_bot + 1, num_rep, num_wg) + mask = mask & r2p_bitmask_below(row_limit_bottom, s) + + return mask + + mask_r2p_lambda( + acc_S, + mask_gen_fn, + rank1=True, + ) diff --git a/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/named_barrier.py b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/named_barrier.py new file mode 100644 index 00000000..dd0d1988 --- /dev/null +++ b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/named_barrier.py @@ -0,0 +1,47 @@ +# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. + +import enum + + +class NamedBarrierFwd(enum.IntEnum): + Epilogue = enum.auto() # starts from 1 as barrier 0 is reserved for sync_threads() + WarpSchedulerWG1 = enum.auto() + WarpSchedulerWG2 = enum.auto() + WarpSchedulerWG3 = enum.auto() + PFull = enum.auto() + PEmpty = enum.auto() + + +class NamedBarrierFwdSm100(enum.IntEnum): + Epilogue = enum.auto() # starts from 1 as barrier 0 is reserved for sync_threads() + TmemPtr = enum.auto() + SoftmaxStatsW0 = enum.auto() + SoftmaxStatsW1 = enum.auto() + SoftmaxStatsW2 = enum.auto() + SoftmaxStatsW3 = enum.auto() + SoftmaxStatsW4 = enum.auto() + SoftmaxStatsW5 = enum.auto() + SoftmaxStatsW6 = enum.auto() + SoftmaxStatsW7 = enum.auto() + + +class NamedBarrierBwd(enum.IntEnum): + Epilogue = enum.auto() + WarpSchedulerWG1 = enum.auto() + WarpSchedulerWG2 = enum.auto() + WarpSchedulerWG3 = enum.auto() + PdS = enum.auto() + dQFullWG0 = enum.auto() + dQFullWG1 = enum.auto() + dQFullWG2 = enum.auto() + dQEmptyWG0 = enum.auto() + dQEmptyWG1 = enum.auto() + dQEmptyWG2 = enum.auto() + + +class NamedBarrierBwdSm100(enum.IntEnum): + EpilogueWG1 = enum.auto() + EpilogueWG2 = enum.auto() + Compute = enum.auto() + dQaccReduce = enum.auto() + TmemPtr = enum.auto() diff --git a/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/pack_gqa.py b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/pack_gqa.py new file mode 100644 index 00000000..d09778a4 --- /dev/null +++ b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/pack_gqa.py @@ -0,0 +1,263 @@ +# Copyright (c) 2025, Tri Dao. + +from dataclasses import dataclass +from typing import Union, Tuple + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync + + +from telefuser.kernel.sol_attn.sm90._compat import layout_utils +import telefuser.kernel.sol_attn._vendor.flash_attn.cute.utils as utils + + +def pack_gqa_layout(T, qhead_per_kvhead, nheads_kv, head_idx): + """Reshape a tensor to fold qhead_per_kvhead into the seqlen dimension (mode 0). + + The head dimension is at mode ``head_idx``. Modes before it (1..head_idx-1) + are kept as-is (e.g. headdim for Q/O tensors), and modes after it are kept + as-is (e.g. batch). + + For Q/O tensors (head_idx=2): + (seqlen_q, headdim, nheads, batch, ...) -> ((qhead_per_kvhead, seqlen_q), headdim, nheads_kv, batch, ...) + For LSE tensors (head_idx=1): + (seqlen_q, nheads, batch, ...) -> ((qhead_per_kvhead, seqlen_q), nheads_kv, batch, ...) + """ + head_stride = T.stride[head_idx] + shape_packed = ( + (qhead_per_kvhead, T.shape[0]), + *[T.shape[i] for i in range(1, head_idx)], + nheads_kv, + *[T.shape[i] for i in range(head_idx + 1, len(T.shape))], + ) + stride_packed = ( + (head_stride, T.stride[0]), + *[T.stride[i] for i in range(1, head_idx)], + head_stride * qhead_per_kvhead, + *[T.stride[i] for i in range(head_idx + 1, len(T.shape))], + ) + return cute.make_tensor(T.iterator, cute.make_layout(shape_packed, stride=stride_packed)) + + +def make_packgqa_tiled_tma_atom( + op: cute.atom.CopyOp, + gmem_tensor: cute.Tensor, + smem_layout: Union[cute.Layout, cute.ComposedLayout], + cta_tiler: Tuple[int, int], + qhead_per_kvhead: int, + head_idx: int, +): + # This packing and unpacking of the layout is so that we keep the same TMA dimension as usual. + # e.g. for (seqlen, d, nheads, b) layout, we still have 4D TMA after packing to + # ((nheads, seqlen), d, b). + # If we instead pack directly to ((qhead_per_kvhead, seqlen), d, nheads_kv, b) we'd have 5D TMA. + # Pack headdim and seqlen dim into 1: (seqlen, d, nheads, b) -> ((nheads, seqlen), d, b) + gmem_tensor = layout_utils.select( + gmem_tensor, [head_idx, *range(head_idx), *range(head_idx + 1, cute.rank(gmem_tensor))] + ) + gmem_tensor = cute.group_modes(gmem_tensor, 0, 2) + assert cta_tiler[0] % qhead_per_kvhead == 0, ( + "CTA tile size in the seqlen dimension must be divisible by qhead_per_kvhead" + ) + tma_atom, tma_tensor = cpasync.make_tiled_tma_atom( + op, + gmem_tensor, + smem_layout, + ((qhead_per_kvhead, cta_tiler[0] // qhead_per_kvhead), cta_tiler[1]), # No mcast + ) + # Unpack from ((nheads, seqlen), d, b) -> ((qhead_per_kvhead, seqlen), d, nheads_kv, b) + T = tma_tensor + shape_packed = ( + (qhead_per_kvhead, T.shape[0][1]), + *[T.shape[i] for i in range(1, head_idx)], + T.shape[0][0] // qhead_per_kvhead, + *[T.shape[i] for i in range(head_idx, len(T.shape))], + ) + stride_packed = ( + *[T.stride[i] for i in range(head_idx)], + T.stride[0][0] * qhead_per_kvhead, + *[T.stride[i] for i in range(head_idx, len(T.shape))], + ) + tma_tensor = cute.make_tensor(T.iterator, cute.make_layout(shape_packed, stride=stride_packed)) + return tma_atom, tma_tensor + + +def unpack_gqa_layout(T, qhead_per_kvhead, head_idx): + """Reverse of pack_gqa_layout: unfold qhead_per_kvhead from the seqlen dimension (mode 0). + + The head dimension is at mode ``head_idx``. Modes before it (1..head_idx-1) + are kept as-is (e.g. headdim for Q/O tensors), and modes after it are kept + as-is (e.g. batch). + + For Q/O tensors (head_idx=2): + ((qhead_per_kvhead, seqlen_q), headdim, nheads_kv, batch, ...) -> (seqlen_q, headdim, nheads, batch, ...) + For LSE tensors (head_idx=1): + ((qhead_per_kvhead, seqlen_q), nheads_kv, batch, ...) -> (seqlen_q, nheads, batch, ...) + """ + seqlen_stride = T.stride[0][1] + head_stride = T.stride[0][0] + shape_unpacked = ( + T.shape[0][1], + *[T.shape[i] for i in range(1, head_idx)], + T.shape[head_idx] * qhead_per_kvhead, + *[T.shape[i] for i in range(head_idx + 1, len(T.shape))], + ) + stride_unpacked = ( + seqlen_stride, + *[T.stride[i] for i in range(1, head_idx)], + head_stride, + *[T.stride[i] for i in range(head_idx + 1, len(T.shape))], + ) + return cute.make_tensor(T.iterator, cute.make_layout(shape_unpacked, stride=stride_unpacked)) + + +@dataclass +class PackGQA: + m_block_size: cutlass.Constexpr[int] + head_dim_padded: cutlass.Constexpr[int] + check_hdim_oob: cutlass.Constexpr[bool] + qhead_per_kvhead: cutlass.Constexpr[bool] + + @cute.jit + def compute_ptr( + self, + tensor: cute.Tensor, + cRows: cute.Tensor, + tidx: cutlass.Int32, + block: cutlass.Int32, + threads_per_row: cutlass.Constexpr[int], + num_threads: cutlass.Constexpr[int], + ): + num_ptr_per_thread = cute.ceil_div(cute.size(cRows), threads_per_row) + tPrPtr = cute.make_rmem_tensor(num_ptr_per_thread, cutlass.Int64) + for i in cutlass.range_constexpr(num_ptr_per_thread): + row = i * num_threads + cRows[tidx % threads_per_row][0] + idx = block * self.m_block_size + row + m_idx = idx // self.qhead_per_kvhead + h_idx = idx - m_idx * self.qhead_per_kvhead + tPrPtr[i] = utils.elem_pointer(tensor, ((h_idx, m_idx),)).toint() + return tPrPtr + + @cute.jit + def load_Q( + self, + mQ: cute.Tensor, # ((qhead_per_kvhead, seqlen_q), headdim) + sQ: cute.Tensor, # (m_block_size, head_dim_padded) + gmem_tiled_copy: cute.TiledCopy, + tidx: cutlass.Int32, + block: cutlass.Int32, + seqlen: cutlass.Int32, + ): + gmem_thr_copy = gmem_tiled_copy.get_slice(tidx) + cQ = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded)) + tQsQ = gmem_thr_copy.partition_D(sQ) + tQcQ = gmem_thr_copy.partition_S(cQ) + t0QcQ = gmem_thr_copy.get_slice(0).partition_S(cQ) + tQpQ = utils.predicate_k(tQcQ, limit=mQ.shape[1]) + tQcQ_row = tQcQ[0, None, 0] + threads_per_row = gmem_tiled_copy.layout_tv_tiled.shape[0][0] + assert cute.arch.WARP_SIZE % threads_per_row == 0, "threads_per_row must divide WARP_SIZE" + num_threads = gmem_tiled_copy.size + tPrQPtr = self.compute_ptr(mQ[None, 0], tQcQ_row, tidx, block, threads_per_row, num_threads) + for m in cutlass.range_constexpr(cute.size(tQsQ.shape[1])): + q_ptr_i64 = utils.shuffle_sync( + tPrQPtr[m // threads_per_row], m % threads_per_row, width=threads_per_row + ) + q_gmem_ptr = cute.make_ptr( + mQ.element_type, q_ptr_i64, cute.AddressSpace.gmem, assumed_align=16 + ) + if ( + t0QcQ[0, m, 0][0] + < seqlen * self.qhead_per_kvhead - block * self.m_block_size - tQcQ_row[0][0] + ): + mQ_cur = cute.make_tensor(q_gmem_ptr, (self.head_dim_padded,)) + elems_per_load = cute.size(tQsQ.shape[0][0]) + mQ_cur_copy = cute.tiled_divide(mQ_cur, (elems_per_load,)) + for k in cutlass.range_constexpr(cute.size(tQsQ.shape[2])): + ki = tQcQ[0, 0, k][1] // elems_per_load + cute.copy( + gmem_thr_copy, + mQ_cur_copy[None, ki], + tQsQ[None, m, k], + pred=tQpQ[None, m, k] if cutlass.const_expr(self.check_hdim_oob) else None, + ) + # We don't need to clear the sQ smem tiles since we'll only write out the valid outputs + + @cute.jit + def store_LSE( + self, + mLSE: cute.Tensor, # (qhead_per_kvhead, seqlen_q) + tLSErLSE: cute.Tensor, # (m_block_size, head_dim_padded) + tiled_mma: cute.TiledMma, + tidx: cutlass.Int32, + block: cutlass.Int32, + seqlen: cutlass.Int32, + ): + thr_mma = tiled_mma.get_slice(tidx) + caccO = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded)) + taccOcO = thr_mma.partition_C(caccO) + taccOcO_row = layout_utils.reshape_acc_to_mn(taccOcO)[None, 0] + assert cute.size(tLSErLSE) == cute.size(taccOcO_row) + threads_per_row = tiled_mma.tv_layout_C.shape[0][0] + assert cute.arch.WARP_SIZE % threads_per_row == 0, "threads_per_row must divide WARP_SIZE" + assert cute.size(tLSErLSE) <= threads_per_row + num_threads = tiled_mma.size + tPrLSEPtr = self.compute_ptr(mLSE, taccOcO_row, tidx, block, threads_per_row, num_threads) + for m in cutlass.range_constexpr(cute.size(tLSErLSE)): + lse_ptr_i64 = utils.shuffle_sync( + tPrLSEPtr[m // threads_per_row], + m % threads_per_row, + width=threads_per_row, + ) + lse_gmem_ptr = cute.make_ptr( + mLSE.element_type, lse_ptr_i64, cute.AddressSpace.gmem, assumed_align=4 + ) + row = block * self.m_block_size + taccOcO_row[m][0] + # Only the thread corresponding to column 0 writes out the lse to gmem + if taccOcO[0][1] == 0 and row < seqlen * self.qhead_per_kvhead: + mLSE_copy = cute.make_tensor(lse_gmem_ptr, (1,)) + mLSE_copy[0] = tLSErLSE[m] + + @cute.jit + def store_O( + self, + mO: cute.Tensor, # ((qhead_per_kvhead, seqlen_q), headdim) + tOrO: cute.Tensor, # (m_block_size, head_dim_padded) split across threads according to gmem_tiled_copy + gmem_tiled_copy: cute.TiledCopy, + tidx: cutlass.Int32, + block: cutlass.Int32, + seqlen: cutlass.Int32, + ): + gmem_thr_copy = gmem_tiled_copy.get_slice(tidx) + cO = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded)) + tOcO = gmem_thr_copy.partition_S(cO) + t0OcO = gmem_thr_copy.get_slice(0).partition_S(cO) + tOpO = utils.predicate_k(tOcO, limit=mO.shape[1]) + tOcO_row = tOcO[0, None, 0] + threads_per_row = gmem_tiled_copy.layout_tv_tiled.shape[0][0] + assert cute.arch.WARP_SIZE % threads_per_row == 0, "threads_per_row must divide WARP_SIZE" + num_threads = gmem_tiled_copy.size + tPrOPtr = self.compute_ptr(mO[None, 0], tOcO_row, tidx, block, threads_per_row, num_threads) + for m in cutlass.range_constexpr(cute.size(tOrO.shape[1])): + o_ptr_i64 = utils.shuffle_sync( + tPrOPtr[m // threads_per_row], m % threads_per_row, width=threads_per_row + ) + o_gmem_ptr = cute.make_ptr( + mO.element_type, o_ptr_i64, cute.AddressSpace.gmem, assumed_align=16 + ) + if ( + t0OcO[0, m, 0][0] + < seqlen * self.qhead_per_kvhead - block * self.m_block_size - tOcO_row[0][0] + ): + mO_cur = cute.make_tensor(o_gmem_ptr, (self.head_dim_padded,)) + elems_per_load = cute.size(tOrO.shape[0][0]) + mO_cur_copy = cute.tiled_divide(mO_cur, (elems_per_load,)) + for k in cutlass.range_constexpr(cute.size(tOrO.shape[2])): + ki = tOcO[0, 0, k][1] // elems_per_load + cute.copy( + gmem_thr_copy, + tOrO[None, m, k], + mO_cur_copy[None, ki], + pred=tOpO[None, m, k] if cutlass.const_expr(self.check_hdim_oob) else None, + ) diff --git a/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/pipeline.py b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/pipeline.py new file mode 100644 index 00000000..f8fdc1e8 --- /dev/null +++ b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/pipeline.py @@ -0,0 +1,402 @@ +# Copyright (c) 2025, Tri Dao. + +from typing import Optional +from dataclasses import dataclass + +import cutlass.cute as cute +from cutlass import Boolean, Int32, const_expr +from cutlass.cutlass_dsl import if_generate, dsl_user_op +from cutlass.pipeline import PipelineState +from cutlass.pipeline import PipelineUserType +from cutlass.pipeline import NamedBarrier as NamedBarrierOg +from cutlass.pipeline import PipelineAsync as PipelineAsyncOg +from cutlass.pipeline import PipelineCpAsync as PipelineCpAsyncOg +from cutlass.pipeline import PipelineTmaAsync as PipelineTmaAsyncOg +from cutlass.pipeline import PipelineTmaUmma as PipelineTmaUmmaOg +from cutlass.pipeline import PipelineUmmaAsync as PipelineUmmaAsyncOg +from cutlass.pipeline import PipelineAsyncUmma as PipelineAsyncUmmaOg + + +def _override_create(parent_cls, child_cls): + """Create a static factory that constructs parent_cls then re-classes to child_cls.""" + + @staticmethod + def create(*args, **kwargs): + obj = parent_cls.create(*args, **kwargs) + # Can't assign to __class__ directly since the dataclass is frozen + object.__setattr__(obj, "__class__", child_cls) + return obj + + return create + + +def _make_state(index: Int32, phase: Int32) -> PipelineState: + """Construct a PipelineState from index and phase (count/stages unused by callers).""" + return PipelineState(stages=0, count=Int32(0), index=index, phase=phase) + + +class PipelineStateSimple: + """ + Pipeline state contains an index and phase bit corresponding to the current position in the circular buffer. + Use a single Int32 to store both the index and phase bit, then we use divmod to get the + index and phase. If stages is a power of 2, divmod turns into bit twiddling. + """ + + def __init__(self, stages: int, phase_index: Int32): + self._stages = stages + self._phase_index = phase_index + + def clone(self) -> "PipelineStateSimple": + return PipelineStateSimple(self.stages, self._phase_index) + + @property + def stages(self) -> int: + return self._stages + + @property + def index(self) -> Int32: + if const_expr(self._stages == 1): + return Int32(0) + else: + return self._phase_index % self._stages + + @property + def phase(self) -> Int32: + # PTX docs say that the phase parity needs to be 0 or 1, so by right we need to + # take modulo 2. But in practice just passing the phase in without modulo works fine. + if const_expr(self._stages == 1): + return self._phase_index + else: + return self._phase_index // self._stages + + def advance(self): + if const_expr(self._stages == 1): + self._phase_index ^= 1 + else: + self._phase_index += 1 + + def __extract_mlir_values__(self): + phase_index = self._phase_index + return [phase_index.ir_value()] + + def __new_from_mlir_values__(self, values): + return PipelineStateSimple(self.stages, Int32(values[0])) + + +def make_pipeline_state(type: PipelineUserType, stages: int): + """ + Creates a pipeline state. Producers are assumed to start with an empty buffer and have a flipped phase bit of 1. + """ + if type is PipelineUserType.Producer: + return PipelineStateSimple(stages, Int32(stages)) + elif type is PipelineUserType.Consumer: + return PipelineStateSimple(stages, Int32(0)) + else: + assert False, "Error: invalid PipelineUserType specified for make_pipeline_state." + + +# ── Shared helpers ─────────────────────────────────────────────────────────── + + +def _call_with_elect_one(parent_method, self, state, elect_one, syncwarp, loc, ip): + """Optionally wrap a parent pipeline method call in sync_warp + elect_one.""" + if const_expr(elect_one): + if const_expr(syncwarp): + cute.arch.sync_warp() + with cute.arch.elect_one(): + parent_method(self, state, loc=loc, ip=ip) + else: + parent_method(self, state, loc=loc, ip=ip) + + +# ── Mixin: _w_index / _w_index_phase variants that delegate to parent ─────── +# Each parent class has PipelineState-based methods (producer_acquire, producer_commit, +# consumer_wait, consumer_release). The _w_index_phase variants just construct a +# PipelineState from (index, phase) and delegate. + + +class _PipelineIndexPhaseMixin: + """Mixin providing _w_index_phase / _w_index methods that delegate to PipelineState-based parents.""" + + @dsl_user_op + def producer_acquire_w_index_phase( + self, + index: Int32, + phase: Int32, + try_acquire_token: Optional[Boolean] = None, + *, + loc=None, + ip=None, + ): + state = _make_state(index, phase) + # Call the parent's producer_acquire (which takes PipelineState) + self.producer_acquire(state, try_acquire_token, loc=loc, ip=ip) + + @dsl_user_op + def producer_commit_w_index(self, index: Int32, *, loc=None, ip=None): + state = _make_state(index, Int32(0)) + self.producer_commit(state, loc=loc, ip=ip) + + @dsl_user_op + def consumer_wait_w_index_phase( + self, + index: Int32, + phase: Int32, + try_wait_token: Optional[Boolean] = None, + *, + loc=None, + ip=None, + ): + state = _make_state(index, phase) + self.consumer_wait(state, try_wait_token, loc=loc, ip=ip) + + @dsl_user_op + def consumer_release_w_index(self, index: Int32, *, loc=None, ip=None): + state = _make_state(index, Int32(0)) + self.consumer_release(state, loc=loc, ip=ip) + + +# ── NamedBarrier ───────────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class NamedBarrier(NamedBarrierOg): + create = _override_create(NamedBarrierOg, None) # patched below + + @dsl_user_op + def arrive_w_index(self, index: Int32, *, loc=None, ip=None) -> None: + """ + The aligned flavor of arrive is used when all threads in the CTA will execute the + same instruction. See PTX documentation. + """ + cute.arch.barrier_arrive( + barrier_id=self.barrier_id + index, + number_of_threads=self.num_threads, + loc=loc, + ip=ip, + ) + + @dsl_user_op + def arrive_and_wait_w_index(self, index: Int32, *, loc=None, ip=None) -> None: + cute.arch.barrier( + barrier_id=self.barrier_id + index, + number_of_threads=self.num_threads, + loc=loc, + ip=ip, + ) + + +NamedBarrier.create = _override_create(NamedBarrierOg, NamedBarrier) + + +# ── PipelineAsync ──────────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class PipelineAsync(_PipelineIndexPhaseMixin, PipelineAsyncOg): + """ + PipelineAsync with optional elect_one for producer_commit and consumer_release. + + When elect_one_*=True (set at create time), only one elected thread per warp + signals the barrier arrive. This is useful when the mask count is set to 1 per warp. + + Args (to create): + elect_one_commit: If True, only elected thread signals producer_commit. + syncwarp_before_commit: If True (default), issue syncwarp before elect_one. + elect_one_release: If True, only elected thread signals consumer_release. + syncwarp_before_release: If True (default), issue syncwarp before elect_one. + Set syncwarp to False when threads are already converged (e.g. after wgmma wait_group). + """ + + _elect_one_commit: bool = False + _syncwarp_before_commit: bool = True + _elect_one_release: bool = False + _syncwarp_before_release: bool = True + + @staticmethod + def create( + *args, + elect_one_commit: bool = False, + syncwarp_before_commit: bool = True, + elect_one_release: bool = False, + syncwarp_before_release: bool = True, + **kwargs, + ): + obj = PipelineAsyncOg.create(*args, **kwargs) + object.__setattr__(obj, "__class__", PipelineAsync) + object.__setattr__(obj, "_elect_one_commit", elect_one_commit) + object.__setattr__(obj, "_syncwarp_before_commit", syncwarp_before_commit) + object.__setattr__(obj, "_elect_one_release", elect_one_release) + object.__setattr__(obj, "_syncwarp_before_release", syncwarp_before_release) + return obj + + @dsl_user_op + def producer_commit(self, state: PipelineState, *, loc=None, ip=None): + _call_with_elect_one( + PipelineAsyncOg.producer_commit, + self, + state, + self._elect_one_commit, + self._syncwarp_before_commit, + loc, + ip, + ) + + @dsl_user_op + def consumer_release(self, state: PipelineState, *, loc=None, ip=None): + _call_with_elect_one( + PipelineAsyncOg.consumer_release, + self, + state, + self._elect_one_release, + self._syncwarp_before_release, + loc, + ip, + ) + + # _w_index variants inherited from _PipelineIndexPhaseMixin, which delegate + # to producer_commit / consumer_release above. + + +# ── PipelineCpAsync ────────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class PipelineCpAsync(_PipelineIndexPhaseMixin, PipelineCpAsyncOg): + _elect_one_release: bool = False + _syncwarp_before_release: bool = True + + @staticmethod + def create( + *args, + elect_one_release: bool = False, + syncwarp_before_release: bool = True, + **kwargs, + ): + obj = PipelineCpAsyncOg.create(*args, **kwargs) + object.__setattr__(obj, "__class__", PipelineCpAsync) + object.__setattr__(obj, "_elect_one_release", elect_one_release) + object.__setattr__(obj, "_syncwarp_before_release", syncwarp_before_release) + return obj + + @dsl_user_op + def consumer_release(self, state: PipelineState, *, loc=None, ip=None): + _call_with_elect_one( + PipelineCpAsyncOg.consumer_release, + self, + state, + self._elect_one_release, + self._syncwarp_before_release, + loc, + ip, + ) + + # _w_index variants inherited from _PipelineIndexPhaseMixin. + + +# ── PipelineTmaAsync ──────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class PipelineTmaAsync(_PipelineIndexPhaseMixin, PipelineTmaAsyncOg): + """Override producer_acquire to take in extra_tx_count parameter.""" + + @dsl_user_op + def producer_acquire( + self, + state: PipelineState, + try_acquire_token: Optional[Boolean] = None, + extra_tx_count: int = 0, + *, + loc=None, + ip=None, + ): + """ + TMA producer commit conditionally waits on buffer empty and sets the transaction barrier for leader threadblocks. + """ + if_generate( + try_acquire_token is None or try_acquire_token == 0, + lambda: self.sync_object_empty.wait(state.index, state.phase, loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + if const_expr(extra_tx_count == 0): + self.sync_object_full.arrive(state.index, self.producer_mask, loc=loc, ip=ip) + else: + tx_count = self.sync_object_full.tx_count + extra_tx_count + self.sync_object_full.arrive_and_expect_tx(state.index, tx_count, loc=loc, ip=ip) + + +PipelineTmaAsync.create = _override_create(PipelineTmaAsyncOg, PipelineTmaAsync) + + +# ── PipelineTmaUmma ───────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class PipelineTmaUmma(_PipelineIndexPhaseMixin, PipelineTmaUmmaOg): + """Override producer_acquire to take in extra_tx_count parameter.""" + + @dsl_user_op + def producer_acquire( + self, + state: PipelineState, + try_acquire_token: Optional[Boolean] = None, + extra_tx_count: int = 0, + *, + loc=None, + ip=None, + ): + """ + TMA producer commit conditionally waits on buffer empty and sets the transaction barrier for leader threadblocks. + """ + if_generate( + try_acquire_token is None or try_acquire_token == 0, + lambda: self.sync_object_empty.wait(state.index, state.phase, loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + if const_expr(extra_tx_count == 0): + if_generate( + self.is_leader_cta, + lambda: self.sync_object_full.arrive( + state.index, self.producer_mask, loc=loc, ip=ip + ), + loc=loc, + ip=ip, + ) + else: + tx_count = self.sync_object_full.tx_count + extra_tx_count + if_generate( + self.is_leader_cta, + lambda: self.sync_object_full.arrive_and_expect_tx( + state.index, tx_count, loc=loc, ip=ip + ), + loc=loc, + ip=ip, + ) + + +PipelineTmaUmma.create = _override_create(PipelineTmaUmmaOg, PipelineTmaUmma) + + +# ── PipelineUmmaAsync ─────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class PipelineUmmaAsync(_PipelineIndexPhaseMixin, PipelineUmmaAsyncOg): + pass + + +PipelineUmmaAsync.create = _override_create(PipelineUmmaAsyncOg, PipelineUmmaAsync) + + +# ── PipelineAsyncUmma ─────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class PipelineAsyncUmma(_PipelineIndexPhaseMixin, PipelineAsyncUmmaOg): + pass + + +PipelineAsyncUmma.create = _override_create(PipelineAsyncUmmaOg, PipelineAsyncUmma) diff --git a/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/seqlen_info.py b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/seqlen_info.py new file mode 100644 index 00000000..d495b89b --- /dev/null +++ b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/seqlen_info.py @@ -0,0 +1,290 @@ +from typing import Optional +from dataclasses import dataclass + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, const_expr + +from telefuser.kernel.sol_attn.sm90._compat import copy_utils + +""" +This consolidates all the info related to sequence length. This is so that we can do all +the gmem reads once at the beginning of each tile, rather than having to repeat these reads +to compute various things like n_block_min, n_block_max, etc. +""" + + +@dataclass(frozen=True) +class SeqlenInfo: + offset: Int32 + offset_padded: Int32 + seqlen: Int32 + has_cu_seqlens: cutlass.Constexpr[bool] = False + + @staticmethod + def create( + batch_idx: Int32, + seqlen_static: Int32, + cu_seqlens: Optional[cute.Tensor] = None, + seqused: Optional[cute.Tensor] = None, + tile: cutlass.Constexpr[int] = 128, + ): + offset = 0 if const_expr(cu_seqlens is None) else cu_seqlens[batch_idx] + offset_padded = ( + 0 + if const_expr(cu_seqlens is None) + # Add divby so that the compiler knows the alignment when moving by offset_padded + else cute.assume((offset + batch_idx * tile) // tile * tile, divby=tile) + ) + if const_expr(seqused is not None): + seqlen = seqused[batch_idx] + elif const_expr(cu_seqlens is not None): + seqlen = cu_seqlens[batch_idx + 1] - cu_seqlens[batch_idx] + else: + seqlen = seqlen_static + return SeqlenInfo(offset, offset_padded, seqlen, has_cu_seqlens=cu_seqlens is not None) + + def offset_batch( + self, + mT: cute.Tensor, + batch_idx: Int32, + dim: int, + padded: cutlass.Constexpr[bool] = False, + multiple: int = 1, + ) -> cute.Tensor: + """Offset a tensor by batch index. batch dim is at position `dim`, seqlen is at dim=0.""" + if const_expr(not self.has_cu_seqlens): + idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mT) - 1 - dim) + return mT[idx] + else: + off = multiple * (self.offset if const_expr(not padded) else self.offset_padded) + offset = off if const_expr(cute.rank(mT.shape[0]) == 1) else (0, off) + idx = (offset,) + (None,) * (cute.rank(mT) - 1) + return cute.domain_offset(idx, mT) + + +@dataclass(frozen=True) +class SeqlenInfoQK: + offset_q: Int32 + offset_k: Int32 + padded_offset_q: Int32 + padded_offset_k: Int32 + seqlen_q: Int32 + seqlen_k: Int32 + has_cu_seqlens_q: cutlass.Constexpr[bool] + has_cu_seqlens_k: cutlass.Constexpr[bool] + has_seqused_q: cutlass.Constexpr[bool] + has_seqused_k: cutlass.Constexpr[bool] + + @staticmethod + def create( + batch_idx: Int32, + seqlen_q_static: Int32, + seqlen_k_static: Int32, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + mSeqUsedQ: Optional[cute.Tensor] = None, + mSeqUsedK: Optional[cute.Tensor] = None, + mCuTotalMBlocks: Optional[cute.Tensor] = None, + mCuBlockIdxOffsets: Optional[cute.Tensor] = None, + tile_m: cutlass.Constexpr[Int32] = 128, + tile_n: cutlass.Constexpr[Int32] = 128, + ): + del mCuTotalMBlocks, mCuBlockIdxOffsets + offset_q = 0 if const_expr(mCuSeqlensQ is None) else mCuSeqlensQ[batch_idx] + offset_k = 0 if const_expr(mCuSeqlensK is None) else mCuSeqlensK[batch_idx] + padded_offset_q = ( + 0 + if const_expr(mCuSeqlensQ is None) + else cute.assume((offset_q + batch_idx * tile_m) // tile_m * tile_m, divby=tile_m) + ) + padded_offset_k = ( + 0 + if const_expr(mCuSeqlensK is None) + else cute.assume((offset_k + batch_idx * tile_n) // tile_n * tile_n, divby=tile_n) + ) + if const_expr(mSeqUsedQ is not None): + seqlen_q = mSeqUsedQ[batch_idx] + else: + seqlen_q = ( + seqlen_q_static + if const_expr(mCuSeqlensQ is None) + else mCuSeqlensQ[batch_idx + 1] - offset_q + ) + if const_expr(mSeqUsedK is not None): + seqlen_k = mSeqUsedK[batch_idx] + else: + seqlen_k = ( + seqlen_k_static + if const_expr(mCuSeqlensK is None) + else mCuSeqlensK[batch_idx + 1] - offset_k + ) + return SeqlenInfoQK( + offset_q, + offset_k, + padded_offset_q, + padded_offset_k, + seqlen_q, + seqlen_k, + has_cu_seqlens_q=mCuSeqlensQ is not None, + has_cu_seqlens_k=mCuSeqlensK is not None, + has_seqused_q=mSeqUsedQ is not None, + has_seqused_k=mSeqUsedK is not None, + ) + + def offset_batch_Q( + self, + mQ: cute.Tensor, + batch_idx: Int32, + dim: int, + padded: cutlass.Constexpr[bool] = False, + ragged: cutlass.Constexpr[bool] = False, + ) -> cute.Tensor: + """Seqlen must be the first dimension of mQ""" + if const_expr(not ragged): + if const_expr(not self.has_cu_seqlens_q): + idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mQ) - 1 - dim) + return mQ[idx] + else: + offset_q = self.offset_q if const_expr(not padded) else self.padded_offset_q + offset_q = offset_q if const_expr(cute.rank(mQ.shape[0]) == 1) else (None, offset_q) + idx = (offset_q,) + (None,) * (cute.rank(mQ) - 1) + return cute.domain_offset(idx, mQ) + else: + if const_expr(not self.has_cu_seqlens_q): + offset_q = 0 + idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mQ) - 1 - dim) + mQ = mQ[idx] + else: + offset_q = self.offset_q if const_expr(not padded) else self.padded_offset_q + if const_expr(cute.rank(mQ.shape[0]) == 1): + return copy_utils.offset_ragged_tensor( + mQ, offset_q, self.seqlen_q, ragged_dim=0, ptr_shift=True + ) + else: # PackGQA + assert cute.rank(mQ.shape[0]) == 2 + # Unpack before calling offset_ragged_tensor, then pack + idx = ((None, None),) + (None,) * (cute.rank(mQ) - 1) + mQ = mQ[idx] + mQ = copy_utils.offset_ragged_tensor( + mQ, offset_q, self.seqlen_q, ragged_dim=1, ptr_shift=True + ) + return cute.group_modes(mQ, 0, 2) + + def offset_batch_K( + self, + mK: cute.Tensor, + batch_idx: Int32, + dim: int, + padded: cutlass.Constexpr[bool] = False, + ragged: cutlass.Constexpr[bool] = False, + multiple: int = 1, + ) -> cute.Tensor: + """Seqlen must be the first dimension of mK""" + if const_expr(not ragged): + if const_expr(not self.has_cu_seqlens_k): + idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mK) - 1 - dim) + return mK[idx] + else: + offset_k = self.offset_k if const_expr(not padded) else self.padded_offset_k + offset_k *= multiple + idx = (offset_k,) + (None,) * (cute.rank(mK) - 1) + return cute.domain_offset(idx, mK) + else: + if const_expr(not self.has_cu_seqlens_k): + offset_k = 0 + idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mK) - 1 - dim) + mK = mK[idx] + else: + offset_k = self.offset_k if const_expr(not padded) else self.padded_offset_k + offset_k *= multiple + return copy_utils.offset_ragged_tensor( + mK, offset_k, self.seqlen_k, ragged_dim=0, ptr_shift=True + ) + + +@dataclass(frozen=True) +class SeqlenInfoQKNewK: + """Sequence length info for append-KV with left-padding and new K support. + + Extends SeqlenInfoQK with: + - leftpad_k: left padding for K (tokens to skip at the start of the KV cache) + - offset_k_new: offset into the new K tensor + - seqlen_k_og: original K length (before appending new K), excluding leftpad + - seqlen_k_new: length of new K to append + - seqlen_k: total K length (seqlen_k_og + seqlen_k_new) + - seqlen_rotary: position for rotary embedding computation + """ + + leftpad_k: Int32 + offset_q: Int32 + offset_k: Int32 + offset_k_new: Int32 + seqlen_q: Int32 + seqlen_k_og: Int32 + seqlen_k_new: Int32 + seqlen_k: Int32 + seqlen_rotary: Int32 + + @staticmethod + def create( + batch_idx: Int32, + seqlen_q_static: Int32, + seqlen_k_static: Int32, + shape_K_new_0: Int32, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + mCuSeqlensKNew: Optional[cute.Tensor] = None, + mSeqUsedQ: Optional[cute.Tensor] = None, + mSeqUsedK: Optional[cute.Tensor] = None, + mLeftpadK: Optional[cute.Tensor] = None, + mSeqlensRotary: Optional[cute.Tensor] = None, + ): + leftpad_k = 0 if const_expr(mLeftpadK is None) else mLeftpadK[batch_idx] + offset_q = 0 if const_expr(mCuSeqlensQ is None) else mCuSeqlensQ[batch_idx] + if const_expr(mCuSeqlensK is not None): + offset_k = mCuSeqlensK[batch_idx] + leftpad_k + else: + offset_k = leftpad_k if const_expr(mCuSeqlensQ is not None) else 0 + offset_k_new = 0 if const_expr(mCuSeqlensKNew is None) else mCuSeqlensKNew[batch_idx] + # seqlen_q + if const_expr(mSeqUsedQ is not None): + seqlen_q = mSeqUsedQ[batch_idx] + elif const_expr(mCuSeqlensQ is not None): + seqlen_q = mCuSeqlensQ[batch_idx + 1] - mCuSeqlensQ[batch_idx] + else: + seqlen_q = seqlen_q_static + # seqlen_k_og: original K length (excluding leftpad) + if const_expr(mSeqUsedK is not None): + seqlen_k_og = mSeqUsedK[batch_idx] - leftpad_k + elif const_expr(mCuSeqlensK is not None): + seqlen_k_og = mCuSeqlensK[batch_idx + 1] - mCuSeqlensK[batch_idx] - leftpad_k + else: + seqlen_k_og = ( + seqlen_k_static - leftpad_k + if const_expr(mCuSeqlensQ is not None) + else seqlen_k_static + ) + # seqlen_k_new + if const_expr(mCuSeqlensKNew is None): + seqlen_k_new = 0 if const_expr(mCuSeqlensQ is None) else shape_K_new_0 + else: + seqlen_k_new = mCuSeqlensKNew[batch_idx + 1] - mCuSeqlensKNew[batch_idx] + seqlen_k = seqlen_k_og if const_expr(mCuSeqlensQ is None) else seqlen_k_og + seqlen_k_new + + # seqlen_rotary: defaults to seqlen_k_og + leftpad_k unless explicitly provided + if const_expr(mSeqlensRotary is not None): + seqlen_rotary = mSeqlensRotary[batch_idx] + else: + seqlen_rotary = seqlen_k_og + leftpad_k + return SeqlenInfoQKNewK( + leftpad_k, + offset_q, + offset_k, + offset_k_new, + seqlen_q, + seqlen_k_og, + seqlen_k_new, + seqlen_k, + seqlen_rotary, + ) diff --git a/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/softmax.py b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/softmax.py new file mode 100644 index 00000000..90ef5891 --- /dev/null +++ b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/softmax.py @@ -0,0 +1,639 @@ +# Copyright (c) 2025, Tri Dao. + +import math +import operator +from typing import Tuple +from dataclasses import dataclass + +import cutlass +import cutlass.cute as cute +from cutlass import Float32 + +from telefuser.kernel.sol_attn.sm90._compat import layout_utils +import telefuser.kernel.sol_attn._vendor.flash_attn.cute.utils as utils +from telefuser.kernel.sol_attn.sm90._compat.cute_dsl_utils import ParamsBase +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.seqlen_info import SeqlenInfoQK + + +@dataclass +class Softmax(ParamsBase): + scale_log2: Float32 + num_rows: cutlass.Constexpr[int] + row_max: cute.Tensor + row_sum: cute.Tensor + row_ref_max: cute.Tensor + arch: cutlass.Constexpr[int] = 80 + softmax_scale: Float32 | None = None + + @staticmethod + def create( + scale_log2: Float32, + num_rows: cutlass.Constexpr[int], + arch: cutlass.Constexpr[int] = 80, + softmax_scale: Float32 | None = None, + ): + row_max = cute.make_rmem_tensor(num_rows, Float32) + row_sum = cute.make_rmem_tensor(num_rows, Float32) + row_ref_max = cute.make_rmem_tensor(num_rows, Float32) + return Softmax( + scale_log2, + num_rows, + row_max, + row_sum, + row_ref_max, + arch, + softmax_scale, + ) + + def reset(self) -> None: + self.row_max.fill(-Float32.inf) + self.row_sum.fill(0.0) + self.row_ref_max.fill(-Float32.inf) + + def _compute_row_max( + self, acc_S_row: cute.TensorSSA, init_val: float | Float32 | None = None + ) -> Float32: + return utils.fmax_reduce(acc_S_row, init_val, arch=self.arch) + + def _compute_row_sum( + self, acc_S_row_exp: cute.TensorSSA, init_val: float | Float32 | None = None + ) -> Float32: + return utils.fadd_reduce(acc_S_row_exp, init_val, arch=self.arch) + + @cute.jit + def online_softmax( + self, + acc_S: cute.Tensor, + is_first: cutlass.Constexpr[bool] = False, + check_inf: cutlass.Constexpr[bool] = True, + ) -> cute.Tensor: + """Apply online softmax and return the row_scale to rescale O. + + :param acc_S: acc_S tensor + :type acc_S: cute.Tensor + :param is_first: is first n_block + :type is_first: cutlass.Constexpr + """ + # Change acc_S to M,N layout view. + acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S) + row_scale = cute.make_fragment_like(self.row_max, Float32) + + row_max = self.row_max + row_sum = self.row_sum + scale_log2 = self.scale_log2 + arch = self.arch + + # Each iteration processes one row of acc_S + for r in cutlass.range(cute.size(row_max), unroll_full=True): + acc_S_row = acc_S_mn[r, None].load() # (n_block_size) + + row_max_cur = utils.fmax_reduce( + acc_S_row, + init_val=row_max[r] if cutlass.const_expr(not is_first) else None, + arch=arch, + ) + + row_max_cur = cute.arch.warp_reduction_max(row_max_cur, threads_in_group=4) + # Update row_max before changing row_max_cur to safe value for -inf + row_max_prev = row_max[r] + row_max[r] = row_max_cur + + if cutlass.const_expr(check_inf): + row_max_cur = 0.0 if row_max_cur == -Float32.inf else row_max_cur + + if cutlass.const_expr(is_first): + row_max_cur_scaled = row_max_cur * scale_log2 + acc_S_row_exp = cute.math.exp2( + acc_S_row * scale_log2 - row_max_cur_scaled, fastmath=True + ) + acc_S_row_sum = utils.fadd_reduce(acc_S_row_exp, init_val=None, arch=arch) + row_scale[r] = 1.0 + else: + row_max_cur_scaled = row_max_cur * scale_log2 + acc_S_row_exp = cute.math.exp2( + acc_S_row * scale_log2 - row_max_cur_scaled, fastmath=True + ) + # row_scale[r] = cute.math.exp2(row_max_prev * self.scale_log2 - row_max_cur_scaled) + row_scale[r] = cute.math.exp2( + (row_max_prev - row_max_cur) * scale_log2, fastmath=True + ) + acc_S_row_sum = utils.fadd_reduce( + acc_S_row_exp, init_val=row_sum[r] * row_scale[r], arch=arch + ) + + row_sum[r] = acc_S_row_sum + acc_S_mn[r, None].store(acc_S_row_exp) + + return row_scale + + @cute.jit + def finalize( + self, final_scale: Float32 = 1.0, sink_val: Float32 | cute.Tensor | None = None + ) -> cute.Tensor: + """Finalize the online softmax by computing the scale and logsumexp.""" + if cutlass.const_expr(sink_val is not None and isinstance(sink_val, cute.Tensor)): + assert cute.size(sink_val) == cute.size(self.row_sum) + row_sum = self.row_sum + row_max = self.row_max + scale_log2 = self.scale_log2 + + # quad reduction for row_sum as we didn't do it during each iteration of online softmax + row_sum.store(utils.warp_reduce(row_sum.load(), operator.add, width=4)) + row_scale = cute.make_fragment_like(row_max, Float32) + + for r in cutlass.range(cute.size(row_sum), unroll_full=True): + if cutlass.const_expr(sink_val is not None): + sink_val_cur = sink_val if not isinstance(sink_val, cute.Tensor) else sink_val[r] + LOG2_E = math.log2(math.e) + row_sum[r] += cute.math.exp2( + sink_val_cur * LOG2_E - row_max[r] * scale_log2, fastmath=True + ) + + # if row_sum is zero or nan, set acc_O_mn_row to 1.0 + acc_O_mn_row_is_zero_or_nan = row_sum[r] == 0.0 or row_sum[r] != row_sum[r] + row_scale[r] = ( + cute.arch.rcp_approx(row_sum[r] if not acc_O_mn_row_is_zero_or_nan else 1.0) + ) * final_scale + row_sum_cur = row_sum[r] + LN2 = math.log(2.0) + row_sum[r] = ( + (row_max[r] * scale_log2 + cute.math.log2(row_sum_cur, fastmath=True)) * LN2 + if not acc_O_mn_row_is_zero_or_nan + else -Float32.inf + ) + return row_scale + + @cute.jit + def rescale_O(self, acc_O: cute.Tensor, row_scale: cute.Tensor) -> None: + """Scale each row of acc_O by the given scale tensor. + :param acc_O: input tensor + :type acc_O: cute.Tensor + :param row_scale: row_scale tensor + :type row_scale: cute.Tensor + """ + acc_O_mn = layout_utils.reshape_acc_to_mn(acc_O) + assert cute.size(row_scale) == cute.size(acc_O_mn, mode=[0]) + for r in cutlass.range(cute.size(row_scale), unroll_full=True): + acc_O_mn[r, None].store(acc_O_mn[r, None].load() * row_scale[r]) + + +@dataclass +class SoftmaxSm100(Softmax): + rescale_threshold: cutlass.Constexpr[float] = 0.0 + max_offset: cutlass.Constexpr[int] = 0 + + @staticmethod + def create( + scale_log2: Float32, + rescale_threshold: cutlass.Constexpr[float] = 0.0, + softmax_scale: Float32 | None = None, + max_offset: cutlass.Constexpr[int] = 0, + ): + num_rows = 1 + arch = 100 + row_max = cute.make_rmem_tensor(num_rows, Float32) + row_sum = cute.make_rmem_tensor(num_rows, Float32) + row_ref_max = cute.make_rmem_tensor(num_rows, Float32) + return SoftmaxSm100( + scale_log2, + num_rows, + row_max, + row_sum, + row_ref_max, + arch, + softmax_scale, + rescale_threshold=rescale_threshold, + max_offset=max_offset, + ) + + @cute.jit + def compute_row_max_local(self, acc_S_row: cute.TensorSSA, is_first: int) -> Float32: + if cutlass.const_expr(is_first): + row_max_new = self._compute_row_max(acc_S_row) + else: + row_max_old = self.row_max[0] + row_max_new = self._compute_row_max(acc_S_row, init_val=row_max_old) + return row_max_new + + @cute.jit + def update_row_max_from_local( + self, + row_max_new: Float32, + is_first: int, + ) -> Tuple[Float32, Float32]: + if cutlass.const_expr(is_first): + row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0 + acc_scale = 0.0 + else: + row_max_old = self.row_max[0] + row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0 + acc_scale_ = (row_max_old - row_max_safe) * self.scale_log2 + acc_scale = cute.math.exp2(acc_scale_, fastmath=True) + if cutlass.const_expr(self.rescale_threshold > 0.0): + if acc_scale_ >= -self.rescale_threshold: + row_max_new = row_max_old + row_max_safe = row_max_old + acc_scale = 1.0 + self.row_max[0] = row_max_new + return row_max_safe, acc_scale + + @cute.jit + def update_row_max(self, acc_S_row: cute.TensorSSA, is_first: int) -> Tuple[Float32, Float32]: + if cutlass.const_expr(is_first): + row_max_new = self._compute_row_max(acc_S_row) + row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0 + acc_scale = 0.0 + else: + row_max_old = self.row_max[0] + row_max_new = self._compute_row_max(acc_S_row, init_val=row_max_old) + row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0 + acc_scale_ = (row_max_old - row_max_safe) * self.scale_log2 + acc_scale = cute.math.exp2(acc_scale_, fastmath=True) + if cutlass.const_expr(self.rescale_threshold > 0.0): + if acc_scale_ >= -self.rescale_threshold: + row_max_new = row_max_old + row_max_safe = row_max_old + acc_scale = 1.0 + self.row_max[0] = row_max_new + return row_max_safe, acc_scale + + def update_row_sum( + self, acc_S_row_exp: cute.TensorSSA, row_scale: Float32, is_first: int = False + ) -> None: + init_val = self.row_sum[0] * row_scale if cutlass.const_expr(not is_first) else None + # self.row_sum[0] = self._compute_row_sum(acc_S_row_exp, init_val=self.row_sum[0] * row_scale) + self.row_sum[0] = self._compute_row_sum(acc_S_row_exp, init_val=init_val) + # tmp = self._compute_row_sum(acc_S_row_exp) + # self.row_sum[0] = self.row_sum[0] * row_scale + tmp + + @cute.jit + def scale_subtract_rowmax( + self, + acc_S_row: cute.Tensor, + row_max: Float32, + ): + assert cute.size(acc_S_row.shape) % 2 == 0, "acc_S_row must have an even number of elements" + row_max_scaled = row_max * self.scale_log2 + for i in cutlass.range(0, cute.size(acc_S_row.shape), 2, unroll_full=True): + acc_S_row[i], acc_S_row[i + 1] = cute.arch.fma_packed_f32x2( + (acc_S_row[i], acc_S_row[i + 1]), + (self.scale_log2, self.scale_log2), + (-row_max_scaled, -row_max_scaled), + ) + + @cute.jit + def apply_exp2_convert( + self, + acc_S_row: cute.Tensor, + acc_S_row_converted: cute.Tensor, + ex2_emu_freq: cutlass.Constexpr[int] = 0, + ex2_emu_res: cutlass.Constexpr[int] = 4, + ex2_emu_start_frg: cutlass.Constexpr[int] = 0, + ): + assert cute.size(acc_S_row.shape) % 2 == 0, "acc_S_row must have an even number of elements" + frg_tile = 32 + assert frg_tile % 2 == 0 + frg_cnt = cute.size(acc_S_row) // frg_tile + assert cute.size(acc_S_row) % frg_tile == 0 + acc_S_row_frg = cute.logical_divide(acc_S_row, cute.make_layout(frg_tile)) + acc_S_row_converted_frg = cute.logical_divide( + acc_S_row_converted, cute.make_layout(frg_tile) + ) + for j in cutlass.range_constexpr(frg_cnt): + for k in cutlass.range_constexpr(0, cute.size(acc_S_row_frg, mode=[0]), 2): + # acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True) + # acc_S_row_frg[k + 1, j] = cute.math.exp2(acc_S_row_frg[k + 1, j], fastmath=True) + if cutlass.const_expr(ex2_emu_freq == 0): + acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True) + acc_S_row_frg[k + 1, j] = cute.math.exp2(acc_S_row_frg[k + 1, j], fastmath=True) + else: + if cutlass.const_expr( + k % ex2_emu_freq < ex2_emu_freq - ex2_emu_res + or j >= frg_cnt - 1 + or j < ex2_emu_start_frg + ): + acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True) + acc_S_row_frg[k + 1, j] = cute.math.exp2( + acc_S_row_frg[k + 1, j], fastmath=True + ) + else: + # acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j] = utils.e2e_asm2(acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j]) + acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j] = utils.ex2_emulation_2( + acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j] + ) + acc_S_row_converted_frg[None, j].store( + acc_S_row_frg[None, j].load().to(acc_S_row_converted.element_type) + ) + + @cute.jit + def scale_apply_exp2_convert( + self, + acc_S_row: cute.Tensor, + row_max: Float32, + acc_S_row_converted: cute.Tensor, + ): + assert cute.size(acc_S_row.shape) % 2 == 0, "acc_S_row must have an even number of elements" + minus_row_max_scaled = -row_max * self.scale_log2 + for i in cutlass.range_constexpr(0, cute.size(acc_S_row.shape), 2): + acc_S_row[i], acc_S_row[i + 1] = cute.arch.fma_packed_f32x2( + (acc_S_row[i], acc_S_row[i + 1]), + (self.scale_log2, self.scale_log2), + (minus_row_max_scaled, minus_row_max_scaled), + ) + + # for i in cutlass.range_constexpr(0, cute.size(acc_S_row.shape), 2): + # acc_S_row[i], acc_S_row[i + 1] = cute.arch.fma_packed_f32x2( + # (acc_S_row[i], acc_S_row[i + 1]), + # (self.scale_log2, self.scale_log2), + # (minus_row_max_scaled, minus_row_max_scaled), + # ) + # acc_S_row[i] = cute.math.exp2(acc_S_row[i], fastmath=True) + # acc_S_row[i + 1] = cute.math.exp2(acc_S_row[i + 1], fastmath=True) + + frg_tile = 32 + assert frg_tile % 2 == 0 + frg_cnt = cute.size(acc_S_row) // frg_tile + assert cute.size(acc_S_row) % frg_tile == 0 + acc_S_row_frg = cute.logical_divide(acc_S_row, cute.make_layout(frg_tile)) + acc_S_row_converted_frg = cute.logical_divide( + acc_S_row_converted, cute.make_layout(frg_tile) + ) + for j in cutlass.range_constexpr(frg_cnt): + for k in cutlass.range_constexpr(0, cute.size(acc_S_row_frg, mode=[0]), 2): + # acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j] = ( + # cute.arch.fma_packed_f32x2( + # (acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j]), + # (self.scale_log2, self.scale_log2), + # (minus_row_max_scaled, minus_row_max_scaled), + # ) + # ) + # acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True) + # acc_S_row_frg[k + 1, j] = cute.math.exp2(acc_S_row_frg[k + 1, j], fastmath=True) + acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True) + acc_S_row_frg[k + 1, j] = cute.math.exp2(acc_S_row_frg[k + 1, j], fastmath=True) + acc_S_row_converted_frg[None, j].store( + acc_S_row_frg[None, j].load().to(acc_S_row_converted.element_type) + ) + + +@cute.jit +def floor_if_packed( + q_idx, + qhead_per_kvhead: cutlass.Constexpr[int], +) -> cute.Tensor: + """Convert q_idx to packed format for Pack-GQA.""" + if cutlass.const_expr(qhead_per_kvhead == 1): + return q_idx + return q_idx // qhead_per_kvhead + + +@cute.jit +def apply_score_mod_inner( + score_tensor, + index_tensor, + score_mod: cutlass.Constexpr, + batch_idx, + head_idx, + softmax_scale, + vec_size: cutlass.Constexpr, + qk_acc_dtype: cutlass.Constexpr, + aux_tensors, + fastdiv_mods, + seqlen_info: SeqlenInfoQK, + constant_q_idx: cutlass.Constexpr, + qhead_per_kvhead: cutlass.Constexpr[int] = 1, + transpose_indices: cutlass.Constexpr[bool] = False, +): + """Shared implementation for applying score modification. + + Args: + score_tensor: The scores to modify (acc_S for flash_fwd, tSrS_t2r for sm100) + index_tensor: Index positions (tScS for flash_fwd, tScS_t2r for sm100) + score_mod: The score modification function to apply + batch_idx: Batch index + head_idx: Head index + softmax_scale: Scale to apply + vec_size: Vector size for processing elements + qk_acc_dtype: Data type for accumulator + aux_tensors: Optional aux_tensors for FlexAttention + fastdiv_mods: Tuple of (seqlen_q_divmod, seqlen_k_divmod) for wrapping + seqlen_info: Sequence length info + constant_q_idx: If provided, use this constant for all q_idx values + If None, compute q_idx per-element + qhead_per_kvhead_packgqa: Pack-GQA replication factor. Divide q_idx by this + when greater than 1 so score mods see logical heads. + transpose_indices: If True, swap q_idx/kv_idx in index_tensor (for bwd kernel where S is transposed) + """ + # Index positions in the index_tensor tuple + # Forward: index_tensor[...][0] = q_idx, index_tensor[...][1] = kv_idx + # Backward (transposed): index_tensor[...][0] = kv_idx, index_tensor[...][1] = q_idx + if cutlass.const_expr(transpose_indices): + q_idx_pos = cutlass.const_expr(1) + kv_idx_pos = cutlass.const_expr(0) + else: + q_idx_pos = cutlass.const_expr(0) + kv_idx_pos = cutlass.const_expr(1) + + n_vals = cutlass.const_expr(cute.size(score_tensor.shape)) + score_vec = cute.make_rmem_tensor(vec_size, qk_acc_dtype) + kv_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32) + + # SSA values for batch (constant across all elements) + batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32).broadcast_to((vec_size,)) + + # Handle q_idx based on whether it's constant + q_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32) + + # For Pack-GQA with non-constant q_idx, we need per-element head indices + # since a thread my process multiple query head indices + if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None): + head_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32) + + for i in cutlass.range(0, n_vals, vec_size, unroll_full=True): + for j in cutlass.range(vec_size, unroll_full=True): + score_vec[j] = score_tensor[i + j] * softmax_scale + + # Extract head offset from packed q_idx for Pack-GQA + if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None): + q_idx_packed = index_tensor[i + j][q_idx_pos] + # Building up the logical q_head idx: final_q_head = kv_head * qhead_per_kvhead + (q_physical % qhead_per_kvhead) + q_idx_logical = q_idx_packed // qhead_per_kvhead + head_offset = q_idx_packed - q_idx_logical * qhead_per_kvhead + head_idx_vec[j] = head_idx * qhead_per_kvhead + head_offset + + # If we will do loads we mod, in order to not read OOB + if cutlass.const_expr(aux_tensors is not None and fastdiv_mods is not None): + if cutlass.const_expr(constant_q_idx is None): + seqlen_q_divmod, seqlen_k_divmod = fastdiv_mods + q_idx_floored = floor_if_packed( + index_tensor[i + j][q_idx_pos], qhead_per_kvhead + ) + _, q_idx_wrapped = divmod(q_idx_floored, seqlen_q_divmod) + q_idx_vec[j] = q_idx_wrapped + else: + _, seqlen_k_divmod = fastdiv_mods + + _, kv_idx_wrapped = divmod(index_tensor[i + j][kv_idx_pos], seqlen_k_divmod) + kv_idx_vec[j] = kv_idx_wrapped + else: + # No bounds checking - direct indexing + if constant_q_idx is None: + q_idx_vec[j] = floor_if_packed(index_tensor[i + j][q_idx_pos], qhead_per_kvhead) + kv_idx_vec[j] = index_tensor[i + j][kv_idx_pos] + + # Convert to SSA for score_mod call + score_ssa = score_vec.load() + kv_idx_ssa = kv_idx_vec.load() + if cutlass.const_expr(constant_q_idx is None): + q_idx_ssa = q_idx_vec.load() + else: + # NB we do not apply Pack-GQA division here, as constant_q_idx is assumed to already be logical + q_idx_const = constant_q_idx + q_idx_ssa = utils.scalar_to_ssa(q_idx_const, cutlass.Int32).broadcast_to((vec_size,)) + + # Compute head_idx_ssa: per-element for Pack-GQA with non-constant q_idx, constant otherwise + if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None): + head_idx_ssa = head_idx_vec.load() + else: + head_idx_ssa = utils.scalar_to_ssa(head_idx, cutlass.Int32).broadcast_to((vec_size,)) + + aux_args = [] + if cutlass.const_expr(aux_tensors is not None): + aux_args = aux_tensors + + post_mod_scores = score_mod( + score_ssa, + batch_idx_ssa, + head_idx_ssa, + q_idx=q_idx_ssa, + kv_idx=kv_idx_ssa, + seqlen_info=seqlen_info, + aux_tensors=aux_args, + ) + + # Write back modified scores + score_vec.store(post_mod_scores) + for j in cutlass.range(vec_size, unroll_full=True): + score_tensor[i + j] = score_vec[j] + + +@cute.jit +def apply_score_mod_bwd_inner( + grad_tensor, + score_tensor, + index_tensor, + score_mod_bwd: cutlass.Constexpr, + batch_idx, + head_idx, + softmax_scale, + vec_size: cutlass.Constexpr, + qk_acc_dtype: cutlass.Constexpr, + aux_tensors, + fastdiv_mods, + seqlen_info, + constant_q_idx: cutlass.Constexpr, + qhead_per_kvhead: cutlass.Constexpr[int] = 1, + transpose_indices: cutlass.Constexpr[bool] = False, +): + """Apply backward score modification (joint graph). + + Args: + grad_tensor: in/out: dlogits rewritten in-place with d(scaled_scores) + score_tensor: pre-mod scores (unscaled QK tile), scaled by softmax_scale internally + index_tensor: Index positions (same as forward) + score_mod_bwd: The backward score modification function (joint graph) + batch_idx: Batch index + head_idx: Head index + softmax_scale: Scale to apply to score_tensor + vec_size: Vector size for processing elements + qk_acc_dtype: Data type for accumulator + aux_tensors: Optional aux_tensors for FlexAttention + fastdiv_mods: Tuple of (seqlen_q_divmod, seqlen_k_divmod) for wrapping + seqlen_info: Sequence length info + constant_q_idx: If provided, use this constant for all q_idx values + qhead_per_kvhead: Pack-GQA replication factor + transpose_indices: If True, swap q_idx/kv_idx in index_tensor + """ + # Index positions in the index_tensor tuple + # Forward: index_tensor[...][0] = q_idx, index_tensor[...][1] = kv_idx + # Backward (transposed): index_tensor[...][0] = kv_idx, index_tensor[...][1] = q_idx + if cutlass.const_expr(transpose_indices): + q_idx_pos = cutlass.const_expr(1) + kv_idx_pos = cutlass.const_expr(0) + else: + q_idx_pos = cutlass.const_expr(0) + kv_idx_pos = cutlass.const_expr(1) + n_vals = cutlass.const_expr(cute.size(grad_tensor.shape)) + grad_vec = cute.make_rmem_tensor(vec_size, qk_acc_dtype) + score_vec = cute.make_rmem_tensor(vec_size, qk_acc_dtype) + kv_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32) + batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32).broadcast_to((vec_size,)) + q_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32) + + # For Pack-GQA with non-constant q_idx, we need per-element head indices + if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None): + head_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32) + + for i in cutlass.range(0, n_vals, vec_size, unroll_full=True): + for j in cutlass.range(vec_size, unroll_full=True): + grad_vec[j] = grad_tensor[i + j] + # Scale score so joint graph sees same value as forward score_mod + score_vec[j] = score_tensor[i + j] * softmax_scale + + if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None): + q_idx_packed = index_tensor[i + j][q_idx_pos] + q_idx_logical = q_idx_packed // qhead_per_kvhead + head_offset = q_idx_packed - q_idx_logical * qhead_per_kvhead + head_idx_vec[j] = head_idx * qhead_per_kvhead + head_offset + + if cutlass.const_expr(aux_tensors is not None and fastdiv_mods is not None): + if cutlass.const_expr(constant_q_idx is None): + seqlen_q_divmod, seqlen_k_divmod = fastdiv_mods + q_idx_floored = floor_if_packed( + index_tensor[i + j][q_idx_pos], qhead_per_kvhead + ) + _, q_idx_wrapped = divmod(q_idx_floored, seqlen_q_divmod) + q_idx_vec[j] = q_idx_wrapped + else: + _, seqlen_k_divmod = fastdiv_mods + + _, kv_idx_wrapped = divmod(index_tensor[i + j][kv_idx_pos], seqlen_k_divmod) + kv_idx_vec[j] = kv_idx_wrapped + else: + # No bounds checking - direct indexing + if constant_q_idx is None: + q_idx_vec[j] = floor_if_packed(index_tensor[i + j][q_idx_pos], qhead_per_kvhead) + kv_idx_vec[j] = index_tensor[i + j][kv_idx_pos] + + grad_ssa = grad_vec.load() + score_ssa = score_vec.load() + kv_idx_ssa = kv_idx_vec.load() + + if cutlass.const_expr(constant_q_idx is None): + q_idx_ssa = q_idx_vec.load() + else: + q_idx_ssa = utils.scalar_to_ssa(constant_q_idx, cutlass.Int32).broadcast_to((vec_size,)) + + if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None): + head_idx_ssa = head_idx_vec.load() + else: + head_idx_ssa = utils.scalar_to_ssa(head_idx, cutlass.Int32).broadcast_to((vec_size,)) + + aux_args = [] + if cutlass.const_expr(aux_tensors is not None): + aux_args = aux_tensors + + grad_out_ssa = score_mod_bwd( + grad_ssa, + score_ssa, + batch_idx_ssa, + head_idx_ssa, + q_idx=q_idx_ssa, + kv_idx=kv_idx_ssa, + seqlen_info=seqlen_info, + aux_tensors=aux_args, + ) + + grad_vec.store(grad_out_ssa) + for j in cutlass.range(vec_size, unroll_full=True): + grad_tensor[i + j] = grad_vec[j] diff --git a/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/tile_scheduler.py b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/tile_scheduler.py new file mode 100644 index 00000000..28403790 --- /dev/null +++ b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/tile_scheduler.py @@ -0,0 +1,1087 @@ +# Copyright (c) 2025, Tri Dao. + +from enum import IntEnum, auto +from typing import Optional, Tuple, Protocol, runtime_checkable +from dataclasses import dataclass + +try: + from typing import override +except ImportError: # Python < 3.12 + from typing_extensions import override + +import cutlass +from cutlass.pipeline import PipelineClcFetchAsync, PipelineState +from cutlass._mlir import ir +import cutlass.cute as cute +from cutlass import Int32, const_expr +from cutlass.cute import FastDivmodDivisor +from cutlass.utils import ClcDynamicPersistentTileScheduler, ClcDynamicPersistentTileSchedulerParams + +from telefuser.kernel.sol_attn.sm90._compat.cute_dsl_utils import ParamsBase + +import telefuser.kernel.sol_attn._vendor.flash_attn.cute.utils as utils +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.fast_math import clz + + +class SchedulingMode(IntEnum): + NONE = auto() + STATIC = auto() + DYNAMIC = auto() + CLC = auto() + + +@dataclass +class ClcState(ParamsBase): + """Owns the runtime state shared by CLC-capable tile schedulers. + + `FlashAttentionForwardSm100` constructs this state because it owns the CLC + response buffer, mbarrier storage, and launch geometry needed to initialize + the hardware scheduler and async pipeline. Individual tile schedulers then + consume this state and map the returned hardware work tiles into their own + logical `WorkTileInfo` coordinates. + + To add CLC support to a scheduler: + - implement `clc_problem_shape(params)` so the kernel can create the hardware scheduler + - accept `clc: ClcState | None` in `create(...)` / `__init__` + - map `clc.initial_work_tile_info()` and `clc.get_current_work()` into scheduler coordinates + """ + + _hw_scheduler: ClcDynamicPersistentTileScheduler + _pipeline: PipelineClcFetchAsync + _consumer_state: PipelineState + _producer_state: PipelineState + + @staticmethod + def create( + *, + hw_scheduler: ClcDynamicPersistentTileScheduler, + pipeline: PipelineClcFetchAsync, + consumer_state: PipelineState, + producer_state: PipelineState, + ) -> "ClcState": + return ClcState(hw_scheduler, pipeline, consumer_state, producer_state) + + def initial_work_tile_info(self): + return self._hw_scheduler.initial_work_tile_info() + + def get_current_work(self): + return self._hw_scheduler.get_current_work() + + def prefetch_next_work(self, *, loc=None, ip=None): + self._pipeline.producer_acquire(self._producer_state, loc=loc, ip=ip) + mbarrier_addr = self._pipeline.producer_get_barrier(self._producer_state, loc=loc, ip=ip) + self._hw_scheduler.advance_to_next_work(mbarrier_addr, loc=loc, ip=ip) + self._producer_state.advance(loc=loc, ip=ip) + + def consumer_wait(self, *, loc=None, ip=None): + self._pipeline.consumer_wait(self._consumer_state, loc=loc, ip=ip) + + def consumer_release(self, *, loc=None, ip=None): + self._pipeline.consumer_release(self._consumer_state, loc=loc, ip=ip) + self._consumer_state.advance(loc=loc, ip=ip) + + def producer_tail(self, *, loc=None, ip=None): + self._pipeline.producer_tail(self._producer_state, loc=loc, ip=ip) + + +class WorkTileInfo(cutlass.utils.WorkTileInfo): + """Altered WorkTileInfo which includes four axes: (block, head, batch, split)""" + + @override + def __new_from_mlir_values__(self, values: list[ir.Value]) -> "WorkTileInfo": + assert len(values) == 5 + new_tile_idx = cutlass.new_from_mlir_values(self._tile_idx, values[:-1]) + new_is_valid_tile = cutlass.new_from_mlir_values(self._is_valid_tile, [values[-1]]) + return WorkTileInfo(new_tile_idx, new_is_valid_tile) + + +@runtime_checkable +class TileSchedulerProtocol(Protocol): + """Protocol defining the interface all tile schedulers must implement. + + Schedulers are responsible for: + 1. Coordinate mapping: linear tile index -> (m_block, head, batch, split) + 2. Work distribution: how to get the next tile (static grid-stride vs CLC dynamic) + """ + + def get_current_work(self) -> WorkTileInfo: + """Get the current work tile coordinates.""" + ... + + def initial_work_tile_info(self) -> WorkTileInfo: + """Get the initial work tile for this CTA.""" + ... + + def advance_to_next_work(self, *, loc=None, ip=None): + """Consumer-side advance: move to next tile and return it. + + For static schedulers: grid-stride increment + get_current_work. + For CLC schedulers: consumer wait + get_current_work + consumer release + state advance. + """ + ... + + def prefetch_next_work(self, *, loc=None, ip=None) -> None: + """Producer-side prefetch of next work tile (no-op for static schedulers). + + For CLC schedulers: producer acquire + issue CLC query + producer state advance. + Only called by the scheduler warp. + """ + ... + + def producer_tail(self, *, loc=None, ip=None) -> None: + """Producer-side cleanup after the last tile. + + No-op for static schedulers. For CLC schedulers: pipeline producer_tail. + """ + ... + + +@dataclass +class TileSchedulerArguments(ParamsBase): + num_block: Int32 + num_head: Int32 + num_batch: Int32 + num_splits: Int32 + seqlen_k: Int32 + headdim: Int32 + headdim_v: Int32 + total_q: Int32 + tile_shape_mn: cutlass.Constexpr[Tuple[int, int]] + cluster_shape_mn: cutlass.Constexpr[Tuple[int, int]] = (1, 1) + mCuSeqlensQ: Optional[cute.Tensor] = None + mSeqUsedQ: Optional[cute.Tensor] = None + qhead_per_kvhead_packgqa: cutlass.Constexpr[int] = 1 + element_size: cutlass.Constexpr[int] = 2 + is_persistent: cutlass.Constexpr[bool] = False + lpt: cutlass.Constexpr[bool] = False + is_split_kv: cutlass.Constexpr[bool] = False + head_swizzle: cutlass.Constexpr[bool] = False + use_cluster_idx: cutlass.Constexpr[bool] = False + + +class SingleTileScheduler: + @dataclass + class Params(ParamsBase): + num_block: Int32 + num_head: Int32 + num_batch: Int32 + num_splits: Int32 + num_splits_divmod: FastDivmodDivisor + is_split_kv: cutlass.Constexpr[bool] = False + cluster_shape_mn: cutlass.Constexpr[Tuple[int, int]] = (1, 1) + use_cluster_idx: cutlass.Constexpr[bool] = False + + @staticmethod + def create( + args: TileSchedulerArguments, *, loc=None, ip=None + ) -> "SingleTileScheduler.Params": + return SingleTileScheduler.Params( + args.num_block, + args.num_head, + args.num_batch, + args.num_splits, + FastDivmodDivisor(args.num_splits), + args.is_split_kv, + args.cluster_shape_mn, + args.use_cluster_idx, + ) + + def __init__(self, params: Params, blk_coord: cute.Coord, *, loc=None, ip=None): + self.params = params + self._blk_coord = blk_coord + self._is_first_block = True + self._loc = loc + self._ip = ip + + @staticmethod + def to_underlying_arguments( + args: TileSchedulerArguments, + *, + scheduling_mode: SchedulingMode = SchedulingMode.STATIC, + loc=None, + ip=None, + ) -> Params: + assert scheduling_mode == SchedulingMode.STATIC, ( + f"SingleTileScheduler only supports STATIC, got {scheduling_mode!r}" + ) + return SingleTileScheduler.Params.create(args, loc=loc, ip=ip) + + @staticmethod + def create( + params: Params, clc: ClcState | None = None, *, loc=None, ip=None + ) -> "SingleTileScheduler": + if const_expr(cute.size(params.cluster_shape_mn) == 1 or not params.use_cluster_idx): + blk_coord = cute.arch.block_idx() + else: + blk_coord = cute.arch.cluster_idx() + return SingleTileScheduler(params, blk_coord, loc=loc, ip=ip) + + # called by host + @staticmethod + def get_grid_shape( + params: Params, + *, + loc=None, + ip=None, + ) -> Tuple[Int32, Int32, Int32]: + # TODO: this hard-codes the fact that we only use cluster = (1, 1) or (2, 1) + assert params.cluster_shape_mn[1] == 1, "Only cluster_shape_mn[1] == 1 is supported" + if const_expr(params.use_cluster_idx): + # Grid must have num_block * cluster_m physical blocks so that there are num_block clusters + grid_x = params.num_block * params.cluster_shape_mn[0] + else: + grid_x = cute.round_up(params.num_block, params.cluster_shape_mn[0]) + return ( + grid_x, + params.num_head * params.num_splits, + params.num_batch, + ) + + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + block_idx, head_idx, batch_idx = self._blk_coord + if const_expr(self.params.is_split_kv): + head_idx, split_idx = divmod(head_idx, self.params.num_splits_divmod) + else: + split_idx = Int32(0) + return WorkTileInfo( + (block_idx, head_idx, batch_idx, split_idx), + self._is_first_block, + ) + + def initial_work_tile_info(self, *, loc=None, ip=None): + return self.get_current_work(loc=loc, ip=ip) + + def prefetch_next_work(self, *, loc=None, ip=None): + pass + + def advance_to_next_work(self, *, loc=None, ip=None): + self._is_first_block = False + return self.get_current_work() + + def producer_tail(self, *, loc=None, ip=None): + pass + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + for obj in [self.params, self._blk_coord]: + obj_values = cutlass.extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + return values + + def __new_from_mlir_values__(self, values): + obj_list = [] + for obj, n_items in zip([self.params, self._blk_coord], self._values_pos): + obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items])) + values = values[n_items:] + return SingleTileScheduler(*(tuple(obj_list)), loc=self._loc) + + +class StaticPersistentTileScheduler: + @dataclass + class Params(ParamsBase): + num_block_cluster_divmod: FastDivmodDivisor + num_head_divmod: FastDivmodDivisor + total_blocks_cluster: Int32 + cluster_shape_m: cutlass.Constexpr[int] = 1 + + @staticmethod + def create( + args: TileSchedulerArguments, *, loc=None, ip=None + ) -> "StaticPersistentTileScheduler.Params": + num_block_cluster = cute.ceil_div(args.num_block, cute.size(args.cluster_shape_mn)) + total_blocks_cluster = num_block_cluster * args.num_head * args.num_batch + return StaticPersistentTileScheduler.Params( + FastDivmodDivisor(num_block_cluster), + FastDivmodDivisor(args.num_head), + total_blocks_cluster, + cluster_shape_m=args.cluster_shape_mn[0], + ) + + def __init__(self, params: Params, tile_idx: Int32, *, loc=None, ip=None): + self.params = params + self._tile_idx = tile_idx + self._loc = loc + self._ip = ip + + @staticmethod + def to_underlying_arguments( + args: TileSchedulerArguments, + *, + scheduling_mode: SchedulingMode = SchedulingMode.STATIC, + loc=None, + ip=None, + ) -> Params: + assert scheduling_mode == SchedulingMode.STATIC, ( + f"StaticPersistentTileScheduler only supports STATIC, got {scheduling_mode!r}" + ) + return StaticPersistentTileScheduler.Params.create(args, loc=loc, ip=ip) + + @staticmethod + def create( + params: Params, clc: ClcState | None = None, *, loc=None, ip=None + ) -> "StaticPersistentTileScheduler": + if const_expr(cute.size(params.cluster_shape_m) == 1): + tile_idx = cute.arch.block_idx()[0] + else: + tile_idx = cute.arch.cluster_idx()[0] + return StaticPersistentTileScheduler(params, tile_idx, loc=loc, ip=ip) + + @staticmethod + def get_grid_shape( + params: Params, + *, + loc=None, + ip=None, + ) -> Tuple[Int32, Int32, Int32]: + hardware_info = cutlass.utils.HardwareInfo() + sm_count = hardware_info.get_device_multiprocessor_count() + max_ctas = (sm_count // params.cluster_shape_m) * params.cluster_shape_m + grid_x = cutlass.min(max_ctas, params.total_blocks_cluster * params.cluster_shape_m) + return (grid_x, Int32(1), Int32(1)) + + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + hn_idx, block_idx = divmod(self._tile_idx, self.params.num_block_cluster_divmod) + batch_idx, head_idx = divmod(hn_idx, self.params.num_head_divmod) + is_valid = self._tile_idx < self.params.total_blocks_cluster + return WorkTileInfo( + (Int32(block_idx), Int32(head_idx), Int32(batch_idx), Int32(0)), is_valid + ) + + def initial_work_tile_info(self, *, loc=None, ip=None): + return self.get_current_work(loc=loc, ip=ip) + + def prefetch_next_work(self, *, loc=None, ip=None): + pass + + def advance_to_next_work(self, *, loc=None, ip=None): + if const_expr(self.params.cluster_shape_m == 1): + self._tile_idx += cute.arch.grid_dim()[0] + else: + self._tile_idx += cute.arch.cluster_dim()[0] + return self.get_current_work() + + def producer_tail(self, *, loc=None, ip=None): + pass + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + for obj in [self.params, self._tile_idx]: + obj_values = cutlass.extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + return values + + def __new_from_mlir_values__(self, values): + obj_list = [] + for obj, n_items in zip( + [self.params, self._tile_idx], + self._values_pos, + ): + obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items])) + values = values[n_items:] + return StaticPersistentTileScheduler(*(tuple(obj_list)), loc=self._loc) + + +class SingleTileLPTScheduler: + @dataclass + class Params(ParamsBase): + total_blocks: Int32 + num_splits: Int32 + num_block: Int32 + num_head: Int32 + num_batch: Int32 + l2_minor: Int32 + num_head_divmod: FastDivmodDivisor + l2_minor_divmod: FastDivmodDivisor + l2_major_divmod: FastDivmodDivisor + l2_minor_residual_divmod: FastDivmodDivisor + num_hb_quotient: Int32 + num_splits_divmod: FastDivmodDivisor + is_split_kv: cutlass.Constexpr[bool] = False + cluster_shape_m: cutlass.Constexpr[int] = 1 + scheduling_mode: cutlass.Constexpr[SchedulingMode] = SchedulingMode.STATIC + lpt: cutlass.Constexpr[bool] = True + + @staticmethod + @cute.jit + def create( + args: TileSchedulerArguments, + *, + scheduling_mode: SchedulingMode = SchedulingMode.STATIC, + loc=None, + ip=None, + ) -> "SingleTileLPTScheduler.Params": + assert scheduling_mode in (SchedulingMode.STATIC, SchedulingMode.CLC), ( + f"Only STATIC and CLC are supported, got {scheduling_mode!r}" + ) + size_one_kv_head = args.seqlen_k * (args.headdim + args.headdim_v) * args.element_size + size_one_head = size_one_kv_head + size_l2 = 50 * 1024 * 1024 # 40 MB for K & V + # Swizzle is the size of each "section". Round swizzle to a power of 2 + # Need to be careful about the case where only one head will fit + # swizzle is how many heads can fit in L2 + # Seems faster if swizzle is a power of 2 + log2_floor = lambda n: 31 - clz(n) + swizzle = 1 if size_l2 < size_one_head else (1 << log2_floor(size_l2 // size_one_head)) + # If we're in the last section (called residual), we don't want to divide by + # swizzle. Instead we want to divide by the remainder. + num_hb_quotient = (args.num_head * args.num_batch) // swizzle + num_hb_remainder = (args.num_head * args.num_batch) % swizzle + return SingleTileLPTScheduler.Params( + total_blocks=args.num_block * args.num_head * args.num_batch, + num_block=args.num_block, + num_head=args.num_head, + num_batch=args.num_batch, + l2_minor=Int32(swizzle), + num_head_divmod=FastDivmodDivisor(args.num_head), + l2_minor_divmod=FastDivmodDivisor(swizzle), + l2_major_divmod=FastDivmodDivisor(swizzle * args.num_block), + l2_minor_residual_divmod=FastDivmodDivisor(max(num_hb_remainder, 1)), + num_hb_quotient=Int32(num_hb_quotient), + num_splits=args.num_splits, + num_splits_divmod=FastDivmodDivisor(args.num_splits), + is_split_kv=args.is_split_kv, + cluster_shape_m=args.cluster_shape_mn[0], + scheduling_mode=scheduling_mode, + lpt=args.lpt, + ) + + def __init__( + self, + params: Params, + tile_idx: Int32, + split_idx: Int32, + clc: ClcState | None = None, + *, + loc=None, + ip=None, + ): + self.params = params + self._tile_idx = tile_idx + self._split_idx = split_idx + self.clc = clc + self._loc = loc + self._ip = ip + + @staticmethod + def to_underlying_arguments( + args: TileSchedulerArguments, + *, + scheduling_mode: SchedulingMode = SchedulingMode.STATIC, + loc=None, + ip=None, + ) -> Params: + return SingleTileLPTScheduler.Params.create( + args, scheduling_mode=scheduling_mode, loc=loc, ip=ip + ) + + @staticmethod + def _clc_grid_shape(params: Params): + num_batch_splits = ( + params.num_batch * params.num_splits + if const_expr(params.is_split_kv) + else params.num_batch + ) + return ( + cute.round_up(params.num_block, params.cluster_shape_m), + params.num_head, + num_batch_splits, + ) + + @staticmethod + @cute.jit + def clc_problem_shape(params: Params): + return ClcDynamicPersistentTileSchedulerParams( + problem_shape_ntile_mnl=SingleTileLPTScheduler._clc_grid_shape(params), + cluster_shape_mnk=(params.cluster_shape_m, 1, 1), + ) + + @staticmethod + @cute.jit + def create( + params: Params, clc: ClcState | None = None, *, loc=None, ip=None + ) -> "SingleTileLPTScheduler": + if const_expr(params.scheduling_mode == SchedulingMode.CLC): + return SingleTileLPTScheduler( + params, cute.arch.block_idx()[0], Int32(0), clc, loc=loc, ip=ip + ) + tile_idx, split_idx, _ = cute.arch.block_idx() + return SingleTileLPTScheduler(params, tile_idx, split_idx, loc=loc, ip=ip) + + @staticmethod + def get_grid_shape( + params: Params, + *, + loc=None, + ip=None, + ) -> Tuple[Int32, Int32, Int32]: + if const_expr(params.scheduling_mode == SchedulingMode.CLC): + return SingleTileLPTScheduler._clc_grid_shape(params) + return (params.total_blocks, params.num_splits, Int32(1)) + + @cute.jit + def clc_work_to_coords(self, work) -> WorkTileInfo: + """Convert CLC response (block, head, batch_split) to WorkTileInfo. + + CLC returns raw grid coordinates — no L2 swizzle (hardware decides order). + We only apply cluster division, optional LPT block reversal, and split_kv unpacking. + """ + block_idx = work.tile_idx[0] + if const_expr(self.params.cluster_shape_m > 1): + block_idx = block_idx // self.params.cluster_shape_m + if const_expr(self.params.lpt): + # Longest-processing-time-first: reverse block order + block_idx = self.params.num_block - 1 - block_idx + split_idx = Int32(0) + if const_expr(self.params.is_split_kv): + batch_idx, split_idx = divmod(work.tile_idx[2], self.params.num_splits_divmod) + else: + batch_idx = work.tile_idx[2] + return WorkTileInfo( + (Int32(block_idx), Int32(work.tile_idx[1]), Int32(batch_idx), Int32(split_idx)), + work.is_valid_tile, + ) + + @cute.jit + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + work = self.clc.get_current_work() + self._tile_idx = work.tile_idx[0] + return self.clc_work_to_coords(work) + # Static path: L2-swizzled coordinate mapping + params = self.params + # Implement LPT scheduling coordinate calculation + bidhb, l2_mod = divmod(self._tile_idx, params.l2_major_divmod) + # If we're in the last section (called residual), we don't want to divide by + # swizzle. Instead we want to divide by the remainder. + block, bidhb_residual = 0, 0 + if bidhb < params.num_hb_quotient: + block, bidhb_residual = divmod(l2_mod, params.l2_minor_divmod) + else: + block, bidhb_residual = divmod(l2_mod, params.l2_minor_residual_divmod) + bidhb_actual = bidhb * params.l2_minor + bidhb_residual + batch_idx, head_idx = divmod(bidhb_actual, params.num_head_divmod) + # Longest-processing-time-first + if const_expr(params.lpt): + block = params.num_block - 1 - block + is_valid = self._tile_idx < params.total_blocks + return WorkTileInfo( + (Int32(block), Int32(head_idx), Int32(batch_idx), Int32(self._split_idx)), is_valid + ) + + @cute.jit + def initial_work_tile_info(self, *, loc=None, ip=None): + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + work = self.clc.initial_work_tile_info() + self._tile_idx = work.tile_idx[0] + return self.clc_work_to_coords(work) + return self.get_current_work(loc=loc, ip=ip) + + def prefetch_next_work(self, *, loc=None, ip=None): + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + self.clc.prefetch_next_work(loc=loc, ip=ip) + + def advance_to_next_work(self, *, loc=None, ip=None): + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + self.clc.consumer_wait(loc=loc, ip=ip) + work = self.get_current_work() + self.clc.consumer_release(loc=loc, ip=ip) + return work + # Single tile scheduler - set to invalid tile_idx to indicate no more work + self._tile_idx = self.params.total_blocks + return self.get_current_work() + + def producer_tail(self, *, loc=None, ip=None): + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + self.clc.producer_tail(loc=loc, ip=ip) + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + objs = [self.params, self._tile_idx, self._split_idx] + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + objs += [self.clc] + for obj in objs: + obj_values = cutlass.extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + return values + + def __new_from_mlir_values__(self, values): + obj_list = [] + objs = [self.params, self._tile_idx, self._split_idx] + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + objs += [self.clc] + for obj, n_items in zip(objs, self._values_pos): + obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items])) + values = values[n_items:] + return self.__class__(*obj_list, loc=self._loc) + + +class SingleTileLPTBwdScheduler: + @dataclass + class Params(ParamsBase): + total_blocks: Int32 + num_block: Int32 + l2_minor: Int32 + num_head_divmod: FastDivmodDivisor + l2_minor_divmod: FastDivmodDivisor + l2_major_divmod: FastDivmodDivisor + l2_minor_residual_divmod: FastDivmodDivisor + num_hb_quotient: Int32 + cluster_shape_mn: cutlass.Constexpr[Tuple[int, int]] = (1, 1) + spt: cutlass.Constexpr[bool] = True + + @staticmethod + @cute.jit + def create( + args: TileSchedulerArguments, *, loc=None, ip=None + ) -> "SingleTileLPTBwdScheduler.Params": + size_l2 = 50 * 1024 * 1024 + size_one_qdo_head = args.seqlen_k * (args.headdim + args.headdim_v) * args.element_size + size_one_dqaccum_head = args.seqlen_k * (args.headdim) * 4 + # size_one_dqaccum_head = 0 + size_one_head = size_one_qdo_head + size_one_dqaccum_head + log2_floor = lambda n: 31 - clz(n) + swizzle = 1 if size_l2 < size_one_head else (1 << log2_floor(size_l2 // size_one_head)) + # swizzle = 8 + # If we're in the last section (called residual), we don't want to divide by + # swizzle. Instead we want to divide by the remainder. + num_hb_quotient = (args.num_head * args.num_batch) // swizzle + num_hb_remainder = (args.num_head * args.num_batch) % swizzle + num_block = cute.ceil_div(args.num_block, args.cluster_shape_mn[0]) + return SingleTileLPTBwdScheduler.Params( + total_blocks=(num_block * args.cluster_shape_mn[0]) + * args.num_head + * args.num_batch, + num_block=num_block, + l2_minor=Int32(swizzle), + num_head_divmod=FastDivmodDivisor(args.num_head), + l2_minor_divmod=FastDivmodDivisor(swizzle), + l2_major_divmod=FastDivmodDivisor(swizzle * num_block), + l2_minor_residual_divmod=FastDivmodDivisor( + max(num_hb_remainder, 1) + ), # don't divide by 0 + num_hb_quotient=Int32(num_hb_quotient), + cluster_shape_mn=args.cluster_shape_mn, + spt=args.lpt, + ) + + def __init__(self, params: Params, tile_idx: Int32, *, loc=None, ip=None): + self.params = params + self._tile_idx = tile_idx + self._loc = loc + self._ip = ip + + @staticmethod + def to_underlying_arguments( + args: TileSchedulerArguments, + *, + scheduling_mode: SchedulingMode = SchedulingMode.STATIC, + loc=None, + ip=None, + ) -> Params: + assert scheduling_mode == SchedulingMode.STATIC, ( + f"SingleTileLPTBwdScheduler only supports STATIC, got {scheduling_mode!r}" + ) + return SingleTileLPTBwdScheduler.Params.create(args, loc=loc, ip=ip) + + @staticmethod + @cute.jit + def create(params: Params, *, loc=None, ip=None) -> "SingleTileLPTBwdScheduler": + tile_idx = cute.arch.block_idx()[0] + return SingleTileLPTBwdScheduler(params, tile_idx, loc=loc, ip=ip) + + # called by host + @staticmethod + def get_grid_shape( + params: Params, + *, + loc=None, + ip=None, + ) -> Tuple[Int32, Int32, Int32]: + return (params.total_blocks, Int32(1), Int32(1)) + + @cute.jit + def get_current_work(self, *, loc=None, ip=None) -> cutlass.utils.WorkTileInfo: + cluster_idx = self._tile_idx // self.params.cluster_shape_mn[0] + params = self.params + # Implement LPT scheduling coordinate calculation + bidhb, l2_mod = divmod(cluster_idx, params.l2_major_divmod) + # If we're in the last section (called residual), we don't want to divide by + # swizzle. Instead we want to divide by the remainder. + block, bidhb_residual = 0, 0 + if bidhb < params.num_hb_quotient: + block, bidhb_residual = divmod(l2_mod, params.l2_minor_divmod) + else: + block, bidhb_residual = divmod(l2_mod, params.l2_minor_residual_divmod) + bidhb_actual = bidhb * params.l2_minor + bidhb_residual + batch_idx, head_idx = divmod(bidhb_actual, params.num_head_divmod) + if cutlass.const_expr(params.spt): + block = params.num_block - 1 - block + if cutlass.const_expr(params.cluster_shape_mn[0] > 1): + bidx_in_cluster = cute.arch.block_in_cluster_idx() + block = block * params.cluster_shape_mn[0] + bidx_in_cluster[0] + is_valid = self._tile_idx < params.total_blocks + return WorkTileInfo((Int32(block), Int32(head_idx), Int32(batch_idx), Int32(0)), is_valid) + + def initial_work_tile_info(self, *, loc=None, ip=None): + return self.get_current_work(loc=loc, ip=ip) + + def prefetch_next_work(self, *, loc=None, ip=None): + pass + + def advance_to_next_work(self, *, loc=None, ip=None): + # Single tile scheduler - set to invalid tile_idx to indicate no more work + self._tile_idx = self.params.total_blocks + return self.get_current_work() + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + for obj in [self.params, self._tile_idx]: + obj_values = cutlass.extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + return values + + def __new_from_mlir_values__(self, values): + obj_list = [] + for obj, n_items in zip([self.params, self._tile_idx], self._values_pos): + obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items])) + values = values[n_items:] + return self.__class__(*(tuple(obj_list)), loc=self._loc) + + +class SingleTileVarlenScheduler: + @dataclass + class Params(ParamsBase): + num_head: Int32 + num_batch: Int32 + total_q: Int32 + num_splits: Int32 + max_kvblock_in_l2: Int32 + tile_shape_mn: cutlass.Constexpr[Tuple[int, int]] + mCuSeqlensQ: Optional[cute.Tensor] = None + mSeqUsedQ: Optional[cute.Tensor] = None + qhead_per_kvhead_packgqa: cutlass.Constexpr[int] = 1 + lpt: cutlass.Constexpr[bool] = False + is_split_kv: cutlass.Constexpr[bool] = False + head_swizzle: cutlass.Constexpr[bool] = False + cluster_shape_m: cutlass.Constexpr[int] = 1 + scheduling_mode: cutlass.Constexpr[SchedulingMode] = SchedulingMode.STATIC + + @staticmethod + @cute.jit + def create( + args: TileSchedulerArguments, + *, + scheduling_mode: SchedulingMode = SchedulingMode.STATIC, + loc=None, + ip=None, + ) -> "SingleTileVarlenScheduler.Params": + assert scheduling_mode in (SchedulingMode.STATIC, SchedulingMode.CLC), ( + f"Only STATIC and CLC are supported, got {scheduling_mode!r}" + ) + size_l2 = 50 * 1024 * 1024 # 50 MB for K & V + # if backward, this is qdo block size + kv_block_size = ( + (args.headdim + args.headdim_v) * args.element_size * args.tile_shape_mn[1] + ) + # if backward, add dqaccum block size to calculate swizzle + if args.head_swizzle: + kv_block_size += args.headdim * 4 * args.tile_shape_mn[1] + max_kvblock_in_l2 = size_l2 // kv_block_size + assert args.mCuSeqlensQ is not None or args.mSeqUsedQ is not None, ( + "At least one of mCuSeqlensQ or mSeqUsedQ must be provided" + ) + assert args.cluster_shape_mn[1] == 1, "Only cluster_shape_mn[1] == 1 is supported" + # TODO: Support varlen CLC with cluster_shape_m > 1 by refactoring the + # flattened-tile decode so cluster unpacking semantics are explicit. + assert scheduling_mode != SchedulingMode.CLC or args.cluster_shape_mn[0] == 1, ( + "Varlen CLC currently requires cluster_shape_mn[0] == 1" + ) + return SingleTileVarlenScheduler.Params( + num_head=args.num_head, + num_batch=args.num_batch, + total_q=args.total_q, + num_splits=args.num_splits, + max_kvblock_in_l2=max_kvblock_in_l2, + tile_shape_mn=args.tile_shape_mn, + mCuSeqlensQ=args.mCuSeqlensQ, + mSeqUsedQ=args.mSeqUsedQ, + qhead_per_kvhead_packgqa=args.qhead_per_kvhead_packgqa, + lpt=args.lpt, + is_split_kv=args.is_split_kv, + head_swizzle=args.head_swizzle, + cluster_shape_m=args.cluster_shape_mn[0], + scheduling_mode=scheduling_mode, + ) + + def __init__( + self, + params: Params, + tile_idx: Int32, + split_idx: Int32, + clc: ClcState | None = None, + *, + loc=None, + ip=None, + ): + self.params = params + self._tile_idx = tile_idx + self._split_idx = split_idx + self._is_first_block = True + self.clc = clc + self._loc = loc + self._ip = ip + + @staticmethod + def to_underlying_arguments( + args: TileSchedulerArguments, + *, + scheduling_mode: SchedulingMode = SchedulingMode.STATIC, + loc=None, + ip=None, + ) -> Params: + return SingleTileVarlenScheduler.Params.create( + args, scheduling_mode=scheduling_mode, loc=loc, ip=ip + ) + + @staticmethod + @cute.jit + def clc_problem_shape(params: Params): + return ClcDynamicPersistentTileSchedulerParams( + problem_shape_ntile_mnl=SingleTileVarlenScheduler.get_grid_shape(params), + cluster_shape_mnk=(1, 1, 1), + ) + + @staticmethod + @cute.jit + def create( + params: Params, clc: ClcState | None = None, *, loc=None, ip=None + ) -> "SingleTileVarlenScheduler": + if const_expr(params.scheduling_mode == SchedulingMode.CLC): + block_idx = cute.arch.block_idx() + split_idx = Int32(0) + if const_expr(params.is_split_kv): + split_idx = block_idx[1] + return SingleTileVarlenScheduler( + params, + block_idx[0], + split_idx, + clc, + loc=loc, + ip=ip, + ) + tile_idx, split_idx, _ = cute.arch.block_idx() + return SingleTileVarlenScheduler(params, tile_idx, split_idx, loc=loc, ip=ip) + + # called by host + @staticmethod + def get_grid_shape( + params: Params, + *, + loc=None, + ip=None, + ) -> Tuple[Int32, Int32, Int32]: + total_blocks_max = ( + params.total_q + + params.num_batch * (params.cluster_shape_m * params.tile_shape_mn[0] - 1) + ) // params.tile_shape_mn[0] + # Round down to nearest multiple of cluster since odd excess is always padding. + total_blocks_max = total_blocks_max // params.cluster_shape_m * params.cluster_shape_m + return (total_blocks_max * params.num_head, params.num_splits, Int32(1)) + + @cute.jit + def _get_num_m_blocks(self, lane: Int32, bidb_start: Int32) -> Int32: + params = self.params + batch_idx = lane + bidb_start + if cutlass.const_expr(params.mSeqUsedQ is not None): + seqlen = Int32(0) + if batch_idx < params.num_batch: + seqlen = params.mSeqUsedQ[batch_idx] + else: + assert params.mCuSeqlensQ is not None + cur_cu_seqlen = Int32(0) + if batch_idx <= params.num_batch: + cur_cu_seqlen = params.mCuSeqlensQ[batch_idx] + next_cu_seqlen = cute.arch.shuffle_sync_down(cur_cu_seqlen, offset=1) + seqlen = next_cu_seqlen - cur_cu_seqlen + if cutlass.const_expr(params.qhead_per_kvhead_packgqa > 1): + seqlen *= params.qhead_per_kvhead_packgqa + return ( + cute.ceil_div(cute.ceil_div(seqlen, params.tile_shape_mn[0]), params.cluster_shape_m) + if batch_idx < params.num_batch and lane < cute.arch.WARP_SIZE - 1 + else Int32(0) + ) + + @cute.jit + def _varlen_coord_map(self) -> WorkTileInfo: + """Map self._tile_idx to (block, head, batch) via warp-level prefix sums.""" + params = self.params + lane_idx = cute.arch.lane_idx() + num_m_blocks = self._get_num_m_blocks(lane_idx, bidb_start=0) + num_m_blocks_cumulative = utils.warp_prefix_sum(num_m_blocks, lane_idx) + # Total number of blocks for the next 31 batches + m_blocks_in_group = cute.arch.shuffle_sync(num_m_blocks_cumulative, cute.arch.WARP_SIZE - 1) + # Same for all lanes + group_end_tile = m_blocks_in_group * params.num_head + # if cute.arch.thread_idx()[0] == 128 + 31: cute.printf("SingleTileVarlenScheduler: tile_idx=%d, group_end_tile = %d, num_m_blocks=%d, num_m_blocks_cumulative = %d, m_blocks_in_group = %d", self._tile_idx, group_end_tile, num_m_blocks, num_m_blocks_cumulative, m_blocks_in_group) + block, head_idx, batch_idx = Int32(0), Int32(0), Int32(0) + next_tile_idx = self._tile_idx // params.cluster_shape_m + while group_end_tile <= next_tile_idx: + batch_idx += cute.arch.WARP_SIZE - 1 + if batch_idx >= params.num_batch: + batch_idx = Int32(params.num_batch) + group_end_tile = next_tile_idx + 1 + else: + num_m_blocks = self._get_num_m_blocks(lane_idx, bidb_start=batch_idx) + num_m_blocks_cumulative = utils.warp_prefix_sum(num_m_blocks, lane_idx) + m_blocks_in_group = cute.arch.shuffle_sync( + num_m_blocks_cumulative, cute.arch.WARP_SIZE - 1 + ) + group_end_tile += m_blocks_in_group * params.num_head + is_valid = False + if batch_idx >= params.num_batch: + block, head_idx, batch_idx = Int32(0), Int32(0), Int32(params.num_batch) + else: + group_start_tile = group_end_tile - m_blocks_in_group * params.num_head + # if cute.arch.thread_idx()[0] == 128 + 31: cute.printf("SingleTileVarlenScheduler: tile_idx=%d, group_end_tile = %d, num_m_blocks=%d, batch_idx = %d", self._tile_idx, group_end_tile, num_m_blocks, batch_idx) + # The next problem to process is the first one that does not have ending tile position + # that is greater than or equal to tile index. + batch_idx_in_group = cute.arch.popc( + cute.arch.vote_ballot_sync( + group_start_tile + num_m_blocks_cumulative * params.num_head <= next_tile_idx + ) + ) + batch_idx += batch_idx_in_group + num_m_blocks_prev_lane = ( + 0 + if batch_idx_in_group == 0 + else cute.arch.shuffle_sync(num_m_blocks_cumulative, batch_idx_in_group - 1) + ) + num_m_blocks = cute.arch.shuffle_sync(num_m_blocks, batch_idx_in_group) + mh_block = next_tile_idx - group_start_tile - num_m_blocks_prev_lane * params.num_head + if cutlass.const_expr(params.lpt or params.head_swizzle): + # This is a version of the SingleTileLPTScheduler, complicated by the fact that + # the seqlen can vary per batch. + # TODO: is there any case where num_m_blocks is 0? + # TODO: by right we should read the seqlen_kv but we're assuming seqlen_q == seqlen_k here + num_n_blocks = ( + num_m_blocks + * params.tile_shape_mn[0] + * params.cluster_shape_m + // params.qhead_per_kvhead_packgqa + // params.tile_shape_mn[1] + ) + # nheads_in_l2 = min(max(self.max_kvblock_in_l2 // num_n_blocks, 1), self.num_head) + # Seems faster to have this be a power of 2 + nheads_in_l2 = ( + 16 + if num_n_blocks * 16 <= params.max_kvblock_in_l2 + else ( + 8 + if num_n_blocks * 8 <= params.max_kvblock_in_l2 + else ( + 4 + if num_n_blocks * 4 <= params.max_kvblock_in_l2 + else (2 if num_n_blocks * 2 <= params.max_kvblock_in_l2 else 1) + ) + ) + ) + nheads_in_l2 = min(nheads_in_l2, params.num_head) + mh_in_l2 = nheads_in_l2 * num_m_blocks + section_idx = mh_block // mh_in_l2 + l2_mod = mh_block - section_idx * mh_in_l2 + # Deal with tail section + nheads_in_this_section = ( + nheads_in_l2 + if nheads_in_l2 * (section_idx + 1) <= params.num_head + else params.num_head - section_idx * nheads_in_l2 + ) + block = l2_mod // nheads_in_this_section + head_idx_residual = l2_mod - block * nheads_in_this_section + head_idx = section_idx * nheads_in_l2 + head_idx_residual + if cutlass.const_expr(params.lpt): + block = num_m_blocks - 1 - block + else: + head_idx = mh_block // num_m_blocks + block = mh_block - head_idx * num_m_blocks + is_valid = self._is_first_block and batch_idx < params.num_batch + if cutlass.const_expr(params.cluster_shape_m > 1): + bidx_in_cluster = cute.arch.block_in_cluster_idx() + block = block * params.cluster_shape_m + bidx_in_cluster[0] + # if cute.arch.thread_idx()[0] == 128: cute.printf("SingleTileVarlenScheduler: tile_idx=%d, batch_idx=%d, head_idx=%d, block=%d, is_valid = %d", self._tile_idx, batch_idx, head_idx, block, is_valid) + split_idx = self._split_idx if const_expr(params.is_split_kv) else Int32(0) + return WorkTileInfo((Int32(block), Int32(head_idx), Int32(batch_idx), split_idx), is_valid) + + @cute.jit + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + clc_work = self.clc.get_current_work() + # Default to grid_dim (one past last valid flat index) so _varlen_coord_map + # returns is_valid=False when CLC is exhausted. CLC tile_idx is garbage when + # invalid, so we can't trust it. Local-then-assign avoids CuTe DSL structural + # mismatch on self inside the runtime if. + new_tile_idx = cute.arch.grid_dim()[0] + new_split_idx = Int32(0) + if clc_work.is_valid_tile: + new_tile_idx = clc_work.tile_idx[0] + if const_expr(self.params.is_split_kv): + new_split_idx = clc_work.tile_idx[1] + self._tile_idx = new_tile_idx + self._split_idx = new_split_idx + return self._varlen_coord_map() + + @cute.jit + def initial_work_tile_info(self, *, loc=None, ip=None): + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + clc_work = self.clc.initial_work_tile_info() + # See get_current_work for why grid_dim and local-then-assign. + new_tile_idx = cute.arch.grid_dim()[0] + new_split_idx = Int32(0) + if clc_work.is_valid_tile: + new_tile_idx = clc_work.tile_idx[0] + if const_expr(self.params.is_split_kv): + new_split_idx = clc_work.tile_idx[1] + self._tile_idx = new_tile_idx + self._split_idx = new_split_idx + return self._varlen_coord_map() + + def prefetch_next_work(self, *, loc=None, ip=None): + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + self.clc.prefetch_next_work(loc=loc, ip=ip) + + def advance_to_next_work(self, *, loc=None, ip=None): + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + self.clc.consumer_wait(loc=loc, ip=ip) + work = self.get_current_work() + self.clc.consumer_release(loc=loc, ip=ip) + return work + self._is_first_block = False + return self.get_current_work() + + def producer_tail(self, *, loc=None, ip=None): + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + self.clc.producer_tail(loc=loc, ip=ip) + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + objs = [self.params, self._tile_idx, self._split_idx] + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + objs += [self.clc] + for obj in objs: + obj_values = cutlass.extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + return values + + def __new_from_mlir_values__(self, values): + obj_list = [] + objs = [self.params, self._tile_idx, self._split_idx] + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + objs += [self.clc] + for obj, n_items in zip(objs, self._values_pos): + obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items])) + values = values[n_items:] + return self.__class__(*obj_list, loc=self._loc) diff --git a/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/utils.py b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/utils.py new file mode 100644 index 00000000..6c98becd --- /dev/null +++ b/telefuser/kernel/sol_attn/_vendor/flash_attn/cute/utils.py @@ -0,0 +1,800 @@ +# Copyright (c) 2025, Tri Dao. + +import math +import hashlib +import inspect +import os +from typing import Type, Callable, Optional, Tuple, overload + +import cutlass +import cutlass.cute as cute + +from cutlass import Float32, Int32, const_expr +from cutlass.cute import FastDivmodDivisor +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass._mlir.dialects import nvvm, llvm +from cutlass.cute.runtime import from_dlpack + + +from telefuser.kernel.sol_attn.sm90._compat import activation + +_MIXER_ATTRS = ("__vec_size__",) + +# Obtained from sollya: +# fpminimax(exp(x * log(2.0)), 1, [|1,24...|],[0;1],relative); +POLY_EX2 = { + 0: (1.0), + 1: ( + 1.0, + 0.922497093677520751953125, + ), + 2: ( + 1.0, + 0.6657850742340087890625, + 0.330107033252716064453125, + ), + 3: ( + 1.0, + 0.695146143436431884765625, + 0.227564394474029541015625, + 0.077119089663028717041015625, + ), + 4: ( + 1.0, + 0.693042695522308349609375, + 0.2412912547588348388671875, + 5.2225358784198760986328125e-2, + 1.3434938155114650726318359375e-2, + ), + 5: ( + 1.0, + 0.693151414394378662109375, + 0.24016360938549041748046875, + 5.5802188813686370849609375e-2, + 9.01452265679836273193359375e-3, + 1.86810153536498546600341796875e-3, + ), +} + +_fa_clc_enabled: bool = os.environ.get("FA_CLC", "0") == "1" +_fa_disable_2cta_enabled: bool = os.environ.get("FA_DISABLE_2CTA", "0") == "1" + + +def _get_use_clc_scheduler_default() -> bool: + return _fa_clc_enabled + + +def _get_disable_2cta_default() -> bool: + return _fa_disable_2cta_enabled + + +def _compute_base_hash(func: Callable) -> str: + """Compute hash from source code or bytecode and closure values.""" + try: + data = inspect.getsource(func).encode() + except (OSError, TypeError): + if hasattr(func, "__code__") and func.__code__ is not None: + data = func.__code__.co_code + else: + data = repr(func).encode() + + hasher = hashlib.sha256(data) + + if hasattr(func, "__closure__") and func.__closure__ is not None: + for cell in func.__closure__: + hasher.update(repr(cell.cell_contents).encode()) + + return hasher.hexdigest() + + +def hash_callable( + func: Callable, mixer_attrs: Tuple[str] = _MIXER_ATTRS, set_cute_hash: bool = True +) -> str: + """Hash a callable based on the source code or bytecode and closure values. + Fast-path: if the callable (or its __wrapped__ base) has a ``__cute_hash__`` + attribute, that value is returned immediately as the base hash, then + metadata dunders are mixed in to produce the final dict-key hash. + set_cute_hash: whether or not to set func.__cute_hash__ + """ + # Resolve base hash + if hasattr(func, "__cute_hash__"): + base_hash = func.__cute_hash__ + else: + # Unwrap decorated functions (e.g., cute.jit wrappers). + base_func = getattr(func, "__wrapped__", func) + + if hasattr(base_func, "__cute_hash__"): + base_hash = base_func.__cute_hash__ + else: + base_hash = _compute_base_hash(base_func) + + if set_cute_hash: + base_func.__cute_hash__ = base_hash + + # Mix in mutable metadata dunders + mixer_values = tuple(getattr(func, attr, None) for attr in mixer_attrs) + + if all(v is None for v in mixer_values): + return base_hash + + hasher = hashlib.sha256(base_hash.encode()) + + for attr, val in zip(_MIXER_ATTRS, mixer_values): + hasher.update(f"{attr}={val!r}".encode()) + + return hasher.hexdigest() + + +def create_softcap_scoremod(softcap_val): + inv_softcap = 1.0 / softcap_val + + @cute.jit + def scoremod_premask_fn(acc_S_SSA, batch_idx, head_idx, q_idx, kv_idx, aux_tensors): + scores = acc_S_SSA * inv_softcap + return scores * cute.math.tanh(scores, fastmath=True) + + return scoremod_premask_fn + + +LOG2_E = math.log2(math.e) + + +def compute_softmax_scale_log2(softmax_scale, score_mod): + """Compute softmax_scale_log2 and adjusted softmax_scale based on whether score_mod is used. + + When score_mod is None, fold the log2(e) factor into softmax_scale_log2 and set softmax_scale + to None. When score_mod is present, keep softmax_scale separate so it can be applied before + the score_mod, and set softmax_scale_log2 to just the change-of-base constant. + + Returns (softmax_scale_log2, softmax_scale). + """ + if const_expr(score_mod is None): + return softmax_scale * LOG2_E, None + else: + return LOG2_E, softmax_scale + + +def compute_fastdiv_mods(mQ, mK, qhead_per_kvhead, pack_gqa, aux_tensors, mPageTable=None): + """Compute FastDivmodDivisor pairs for aux_tensors index computation. + + Returns a (seqlen_q_divmod, seqlen_k_divmod) tuple, or None if aux_tensors is None. + """ + if const_expr(aux_tensors is None): + return None + seqlen_q = cute.size(mQ.shape[0]) // (qhead_per_kvhead if const_expr(pack_gqa) else 1) + seqlen_k = ( + cute.size(mK.shape[0]) + if const_expr(mPageTable is None) + else mK.shape[0] * mPageTable.shape[1] + ) + return (FastDivmodDivisor(seqlen_q), FastDivmodDivisor(seqlen_k)) + + +def convert_from_dlpack(x, leading_dim, alignment=16, divisibility=1) -> cute.Tensor: + return ( + from_dlpack(x, assumed_align=alignment) + .mark_layout_dynamic(leading_dim=leading_dim) + .mark_compact_shape_dynamic( + mode=leading_dim, stride_order=x.dim_order(), divisibility=divisibility + ) + ) + + +def convert_from_dlpack_leading_static( + x, leading_dim, alignment=16, static_modes=None, stride_order=None +) -> cute.Tensor: + if stride_order is None: + stride_order = x.dim_order() + x_ = from_dlpack(x, assumed_align=alignment) + for i in range(x.ndim): + if i != leading_dim and (static_modes is None or i not in static_modes): + x_ = x_.mark_compact_shape_dynamic(mode=i, stride_order=stride_order) + return x_ + + +def make_tiled_copy_A( + copy_atom: cute.CopyAtom, tiled_mma: cute.TiledMma, swapAB: cutlass.Constexpr[bool] = False +) -> cute.TiledCopy: + if const_expr(swapAB): + return cute.make_tiled_copy_B(copy_atom, tiled_mma) + else: + return cute.make_tiled_copy_A(copy_atom, tiled_mma) + + +def make_tiled_copy_B( + copy_atom: cute.CopyAtom, tiled_mma: cute.TiledMma, swapAB: cutlass.Constexpr[bool] = False +) -> cute.TiledCopy: + if const_expr(swapAB): + return cute.make_tiled_copy_A(copy_atom, tiled_mma) + else: + return cute.make_tiled_copy_B(copy_atom, tiled_mma) + + +def mma_make_fragment_A( + smem: cute.Tensor, thr_mma: cute.ThrMma, swapAB: cutlass.Constexpr[bool] = False +) -> cute.Tensor: + if const_expr(swapAB): + return mma_make_fragment_B(smem, thr_mma) + else: + return thr_mma.make_fragment_A(thr_mma.partition_A(smem)) + + +def mma_make_fragment_B( + smem: cute.Tensor, thr_mma: cute.ThrMma, swapAB: cutlass.Constexpr[bool] = False +) -> cute.Tensor: + if const_expr(swapAB): + return mma_make_fragment_A(smem, thr_mma) + else: + return thr_mma.make_fragment_B(thr_mma.partition_B(smem)) + + +def get_smem_store_atom( + arch: cutlass.Constexpr[int], element_type: Type[cute.Numeric], transpose: bool = False +) -> cute.CopyAtom: + if const_expr(arch < 90 or element_type.width != 16): + return cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + element_type, + num_bits_per_copy=2 * element_type.width, + ) + else: + return cute.make_copy_atom( + cute.nvgpu.warp.StMatrix8x8x16bOp(transpose=transpose, num_matrices=4), + element_type, + ) + + +@cute.jit +def warp_reduce( + val: cute.TensorSSA | cute.Numeric, + op: Callable, + width: cutlass.Constexpr[int] = cute.arch.WARP_SIZE, +) -> cute.TensorSSA | cute.Numeric: + if const_expr(isinstance(val, cute.TensorSSA)): + res = cute.make_rmem_tensor(val.shape, val.dtype) + res.store(val) + for i in cutlass.range_constexpr(cute.size(val.shape)): + res[i] = warp_reduce(res[i], op, width) + return res.load() + else: + for i in cutlass.range_constexpr(int(math.log2(width))): + val = op(val, cute.arch.shuffle_sync_bfly(val, offset=1 << i)) + return val + + +@dsl_user_op +def smid(*, loc=None, ip=None) -> Int32: + return Int32( + llvm.inline_asm( + T.i32(), + [], + "mov.u32 $0, %smid;", + "=r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def fmax( + a: float | Float32, b: float | Float32, c: float | Float32 | None = None, *, loc=None, ip=None +) -> Float32: + from cutlass import CUDA_VERSION + + # * NVVM call based on nvvm version + if CUDA_VERSION.major == 12 and CUDA_VERSION.minor == 9: + # Old API: requires explicit result type as first positional argument + return Float32( + nvvm.fmax( + T.f32(), + Float32(a).ir_value(loc=loc, ip=ip), + Float32(b).ir_value(loc=loc, ip=ip), + c=Float32(c).ir_value(loc=loc, ip=ip) if c is not None else None, + loc=loc, + ip=ip, + ) + ) + else: + # New API: infers result type automatically + return Float32( + nvvm.fmax( + Float32(a).ir_value(loc=loc, ip=ip), + Float32(b).ir_value(loc=loc, ip=ip), + c=Float32(c).ir_value(loc=loc, ip=ip) if c is not None else None, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def fmax_reduce( + x: cute.TensorSSA, init_val: float | Float32 | None = None, arch: cutlass.Constexpr[int] = 80 +) -> Float32: + if const_expr(arch < 100 or cute.size(x.shape) % 8 != 0): + # if const_expr(init_val is None): + # init_val = -cutlass.Float32.if + # return x.reduce(cute.ReductionOp.MAX, init_val, 0) + res = cute.make_rmem_tensor(x.shape, Float32) + res.store(x) + # local_max = [res[0], res[1]] + # for i in cutlass.range_constexpr(2, cute.size(x.shape), 2): + # local_max[0] = fmax(local_max[0], res[i + 0]) + # local_max[1] = fmax(local_max[1], res[i + 1]) + # local_max[0] = fmax(local_max[0], local_max[1]) + # return local_max[0] if const_expr(init_val is None) else fmax(local_max[0], init_val) + local_max = [res[0], res[1], res[2], res[3]] + for i in cutlass.range_constexpr(4, cute.size(x.shape), 4): + local_max[0] = fmax(local_max[0], res[i + 0]) + local_max[1] = fmax(local_max[1], res[i + 1]) + local_max[2] = fmax(local_max[2], res[i + 2]) + local_max[3] = fmax(local_max[3], res[i + 3]) + local_max[0] = fmax(local_max[0], local_max[1]) + local_max[2] = fmax(local_max[2], local_max[3]) + local_max[0] = fmax(local_max[0], local_max[2]) + return local_max[0] if const_expr(init_val is None) else fmax(local_max[0], init_val) + else: + # [2025-06-15] x.reduce only seems to use 50% 3-input max and 50% 2-input max + # We instead force the 3-input max. + res = cute.make_rmem_tensor(x.shape, Float32) + res.store(x) + local_max_0 = ( + fmax(init_val, res[0], res[1]) + if const_expr(init_val is not None) + else fmax(res[0], res[1]) + ) + local_max = [ + local_max_0, + fmax(res[2], res[3]), + fmax(res[4], res[5]), + fmax(res[6], res[7]), + ] + for i in cutlass.range_constexpr(8, cute.size(x.shape), 8): + local_max[0] = fmax(local_max[0], res[i], res[i + 1]) + local_max[1] = fmax(local_max[1], res[i + 2], res[i + 3]) + local_max[2] = fmax(local_max[2], res[i + 4], res[i + 5]) + local_max[3] = fmax(local_max[3], res[i + 6], res[i + 7]) + local_max[0] = fmax(local_max[0], local_max[1]) + return fmax(local_max[0], local_max[2], local_max[3]) + + +@cute.jit +def fadd_reduce( + x: cute.TensorSSA, init_val: float | Float32 | None = None, arch: cutlass.Constexpr[int] = 80 +) -> Float32: + if const_expr(arch < 100 or cute.size(x.shape) % 8 != 0): + if const_expr(init_val is None): + init_val = Float32.zero + return x.reduce(cute.ReductionOp.ADD, init_val, 0) + # res = cute.make_rmem_tensor(x.shape, Float32) + # res.store(x) + # local_sum = [res[0], res[1], res[2], res[3]] + # for i in cutlass.range_constexpr(4, cute.size(x.shape), 4): + # local_sum[0] += res[i + 0] + # local_sum[1] += res[i + 1] + # local_sum[2] += res[i + 2] + # local_sum[3] += res[i + 3] + # local_sum[0] += local_sum[1] + # local_sum[2] += local_sum[3] + # local_sum[0] += local_sum[2] + # return local_sum[0] if const_expr(init_val is None) else local_sum[0] + init_val + else: + res = cute.make_rmem_tensor(x.shape, Float32) + res.store(x) + local_sum_0 = ( + cute.arch.add_packed_f32x2((init_val, 0.0), (res[0], res[1])) + # cute.arch.add_packed_f32x2((init_val / 2, init_val / 2), (res[0], res[1])) + if const_expr(init_val is not None) + else (res[0], res[1]) + ) + local_sum = [local_sum_0, (res[2], res[3]), (res[4], res[5]), (res[6], res[7])] + for i in cutlass.range_constexpr(8, cute.size(x.shape), 8): + local_sum[0] = cute.arch.add_packed_f32x2(local_sum[0], (res[i + 0], res[i + 1])) + local_sum[1] = cute.arch.add_packed_f32x2(local_sum[1], (res[i + 2], res[i + 3])) + local_sum[2] = cute.arch.add_packed_f32x2(local_sum[2], (res[i + 4], res[i + 5])) + local_sum[3] = cute.arch.add_packed_f32x2(local_sum[3], (res[i + 6], res[i + 7])) + local_sum[0] = cute.arch.add_packed_f32x2(local_sum[0], local_sum[1]) + local_sum[2] = cute.arch.add_packed_f32x2(local_sum[2], local_sum[3]) + local_sum[0] = cute.arch.add_packed_f32x2(local_sum[0], local_sum[2]) + return local_sum[0][0] + local_sum[0][1] + + +@dsl_user_op +def atomic_add_fp32(a: float | Float32, gmem_ptr: cute.Pointer, *, loc=None, ip=None) -> None: + # gmem_ptr_i64 = gmem_ptr.toint(loc=loc, ip=ip).ir_value() + # # cache_hint = cutlass.Int64(0x12F0000000000000) + # llvm.inline_asm( + # None, + # [gmem_ptr_i64, Float32(a).ir_value(loc=loc, ip=ip)], + # # [gmem_ptr_i64, Float32(a).ir_value(loc=loc, ip=ip), cache_hint.ir_value()], + # "red.global.add.f32 [$0], $1;", + # # "red.global.add.L2::cache_hint.f32 [$0], $1, 0x12F0000000000000;", + # # "red.global.add.L2::cache_hint.f32 [$0], $1, $2;", + # "l,f", + # # "l,f,l", + # has_side_effects=True, + # is_align_stack=False, + # asm_dialect=llvm.AsmDialect.AD_ATT, + # ) + nvvm.atomicrmw( + res=T.f32(), op=nvvm.AtomicOpKind.FADD, ptr=gmem_ptr.llvm_ptr, a=Float32(a).ir_value() + ) + + +@dsl_user_op +def elem_pointer(x: cute.Tensor, coord: cute.Coord, *, loc=None, ip=None) -> cute.Pointer: + return x.iterator + cute.crd2idx(coord, x.layout, loc=loc, ip=ip) + + +@cute.jit +def predicate_k(tAcA: cute.Tensor, limit: cutlass.Int32) -> cute.Tensor: + # Only compute predicates for the "k" dimension. For the mn dimension, we will use "if" + tApA = cute.make_rmem_tensor( + cute.make_layout( + (cute.size(tAcA, mode=[0, 1]), cute.size(tAcA, mode=[1]), cute.size(tAcA, mode=[2])), + stride=(cute.size(tAcA, mode=[2]), 0, 1), + ), + cutlass.Boolean, + ) + for rest_v in cutlass.range_constexpr(tApA.shape[0]): + for rest_k in cutlass.range_constexpr(tApA.shape[2]): + tApA[rest_v, 0, rest_k] = cute.elem_less(tAcA[(0, rest_v), 0, rest_k][1], limit) + return tApA + + +def canonical_warp_group_idx(sync: bool = True) -> cutlass.Int32: + warp_group_idx = cute.arch.thread_idx()[0] // 128 + if const_expr(sync): + warp_group_idx = cute.arch.make_warp_uniform(warp_group_idx) + return warp_group_idx + + +# @dsl_user_op +# def warp_vote_any_lt(a: float | Float32, b: float | Float32, *, loc=None, ip=None) -> cutlass.Boolean: +# mask = cutlass.Int32(-1) +# return cutlass.Boolean( +# llvm.inline_asm( +# T.i32(), +# [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip), mask.ir_value(loc=loc, ip=ip)], +# ".pred p1, p2;\n" +# "setp.lt.f32 p1, $1, $2;\n" +# "vote.sync.any.pred p2, p1, $3;\n" +# "selp.u32 $0, 1, 0, p2;", +# # "selp.u32 $0, 1, 0, p1;", +# "=r,f,f,r", +# has_side_effects=False, +# is_align_stack=False, +# asm_dialect=llvm.AsmDialect.AD_ATT, +# ) +# ) + + +@cute.jit +def shuffle_sync( + value: cute.Numeric, + offset: cute.typing.Int, + width: cutlass.Constexpr[int] = cute.arch.WARP_SIZE, +) -> cute.Numeric: + assert value.width % 32 == 0, "value type must be a multiple of 32 bits" + # 1 -> 0b11111, 2 -> 0b11110, 4 -> 0b11100, 8 -> 0b11000, 16 -> 0b10000, 32 -> 0b00000 + mask = cute.arch.WARP_SIZE - width + clamp = cute.arch.WARP_SIZE - 1 + mask_and_clamp = mask << 8 | clamp + # important: need stride 1 and not 0 for recast_tensor to work + val = cute.make_rmem_tensor(cute.make_layout((1,), stride=(1,)), type(value)) + val[0] = value + val_i32 = cute.recast_tensor(val, cutlass.Int32) + for i in cutlass.range_constexpr(cute.size(val_i32)): + val_i32[i] = cute.arch.shuffle_sync(val_i32[i], offset, mask_and_clamp=mask_and_clamp) + return val[0] + + +@dsl_user_op +def shl_u32(val: cutlass.Uint32, shift: cutlass.Uint32, *, loc=None, ip=None) -> cutlass.Uint32: + """ + Left-shift val by shift bits using PTX shl.b32 (sign-agnostic). + + Named ``shl_u32`` (not ``shl_b32``) because python type annotations + distinguish signed/unsigned. + + PTX semantics (§9.7.8.8): "Shift amounts greater than the register width N + are clamped to N." So ``shl.b32 d, a, 32`` is well-defined and yields 0. + + This differs from C/C++ and LLVM IR, where shifting by >= the type width is + undefined behavior. CuTeDSL compiles through MLIR -> LLVM IR, so a plain + Python-level ``Uint32(x) << Uint32(n)`` inherits LLVM's UB: the optimizer + may treat the result as poison and eliminate dependent code. Inline PTX + bypasses the LLVM IR shift entirely — the instruction is emitted verbatim + into PTX where clamping makes it safe for all shift amounts. + """ + return cutlass.Uint32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Uint32(val).ir_value(loc=loc, ip=ip), + cutlass.Uint32(shift).ir_value(loc=loc, ip=ip), + ], + "shl.b32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def shr_u32(val: cutlass.Uint32, shift: cutlass.Uint32, *, loc=None, ip=None) -> cutlass.Uint32: + """ + Unsigned right-shift val by shift bits using PTX shr.u32 (zero-fills). + + See ``shl_u32`` docstring for why inline PTX is used instead of plain + CuTeDSL shift operators (LLVM shift-by-type-width UB). + """ + return cutlass.Uint32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Uint32(val).ir_value(loc=loc, ip=ip), + cutlass.Uint32(shift).ir_value(loc=loc, ip=ip), + ], + "shr.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@cute.jit +def warp_prefix_sum(val: cutlass.Int32, lane: Optional[cutlass.Int32] = None) -> cutlass.Int32: + if const_expr(lane is None): + lane = cute.arch.lane_idx() + # if cute.arch.thread_idx()[0] >= 128 and cute.arch.thread_idx()[0] < 128 + 32 and cute.arch.block_idx()[0] == 0: cute.printf("tidx = %d, val = %d", cute.arch.thread_idx()[0] % 32, val) + for i in cutlass.range_constexpr(int(math.log2(cute.arch.WARP_SIZE))): + offset = 1 << i + # Very important that we set mask_and_clamp to 0 + partial_sum = cute.arch.shuffle_sync_up(val, offset=offset, mask_and_clamp=0) + if lane >= offset: + val += partial_sum + # if cute.arch.thread_idx()[0] >= 128 and cute.arch.thread_idx()[0] < 128 + 32 and cute.arch.block_idx()[0] == 0: cute.printf("tidx = %d, partial_sum = %d, val = %d", cute.arch.thread_idx()[0] % 32, partial_sum, val) + return val + + +@dsl_user_op +def cvt_f16x2_f32( + a: float | Float32, b: float | Float32, to_dtype: Type, *, loc=None, ip=None +) -> cutlass.Int32: + assert to_dtype in [cutlass.BFloat16, cutlass.Float16], "to_dtype must be BFloat16 or Float16" + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip)], + f"cvt.rn.{'bf16x2' if to_dtype is cutlass.BFloat16 else 'f16x2'}.f32 $0, $2, $1;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@overload +def cvt_f16(src: cute.Tensor, dst: cute.Tensor) -> None: ... + + +@overload +def cvt_f16(src: cute.Tensor, dtype: Type[cute.Numeric]) -> cute.Tensor: ... + + +@cute.jit +def cvt_f16(src: cute.Tensor, dst_or_dtype): + """Convert Float32 tensor to Float16/BFloat16. + + Args: + src: Source tensor with Float32 element type + dst_or_dtype: Either a destination tensor or a dtype (Float16/BFloat16) + + Returns: + None if dst is a tensor, or a new tensor if dtype is provided + """ + if const_expr(isinstance(dst_or_dtype, type)): + # dtype variant: create new tensor and call the tensor variant + dtype = dst_or_dtype + dst = cute.make_rmem_tensor(src.shape, dtype) + cvt_f16(src, dst) + return dst + else: + # tensor variant: write to dst + dst = dst_or_dtype + assert cute.size(dst.shape) == cute.size(src.shape), "dst and src must have the same size" + assert cute.size(src.shape) % 2 == 0, "src must have an even number of elements" + assert dst.element_type in [cutlass.BFloat16, cutlass.Float16], ( + "dst must be BFloat16 or Float16" + ) + assert src.element_type is Float32, "src must be Float32" + dst_i32 = cute.recast_tensor(dst, cutlass.Int32) + assert cute.size(dst_i32.shape) * 2 == cute.size(src.shape) + for i in cutlass.range_constexpr(cute.size(dst_i32)): + dst_i32[i] = cvt_f16x2_f32(src[2 * i], src[2 * i + 1], dst.element_type) + + +@dsl_user_op +@cute.jit +def evaluate_polynomial(x: Float32, poly: Tuple[Float32, ...], *, loc=None, ip=None) -> Float32: + deg = len(poly) - 1 + out = poly[deg] + for i in cutlass.range_constexpr(deg - 1, -1, -1): + out = out * x + poly[i] + return out + + +@dsl_user_op +@cute.jit +def evaluate_polynomial_2( + x: Float32, y: Float32, poly: Tuple[Float32, ...], *, loc=None, ip=None +) -> Tuple[Float32, Float32]: + deg = len(poly) - 1 + out = (poly[deg], poly[deg]) + for i in cutlass.range_constexpr(deg - 1, -1, -1): + out = cute.arch.fma_packed_f32x2(out, (x, y), (poly[i], poly[i])) + return out + + +@dsl_user_op +def add_round_down(x: float | Float32, y: float | Float32, *, loc=None, ip=None) -> Float32: + # There's probably a way to call llvm or nvvm to do this instead of ptx + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [Float32(x).ir_value(loc=loc, ip=ip), Float32(y).ir_value(loc=loc, ip=ip)], + "add.rm.ftz.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def combine_int_frac_ex2(x_rounded: Float32, frac_ex2: Float32, *, loc=None, ip=None) -> Float32: + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [ + Float32(x_rounded).ir_value(loc=loc, ip=ip), + Float32(frac_ex2).ir_value(loc=loc, ip=ip), + ], + "{\n\t" + ".reg .s32 x_rounded_i, frac_ex_i, x_rounded_e, out_i;\n\t" + "mov.b32 x_rounded_i, $1;\n\t" + "mov.b32 frac_ex_i, $2;\n\t" + "shl.b32 x_rounded_e, x_rounded_i, 23;\n\t" + # add.u32 generates IMAD instruction and add.s32 generates LEA instruction + # IMAD uses the FMA pipeline and LEA uses the ALU pipeline, afaik + "add.s32 out_i, x_rounded_e, frac_ex_i;\n\t" + "mov.b32 $0, out_i;\n\t" + "}\n", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def ex2_emulation(x: Float32, *, poly_degree: int = 3, loc=None, ip=None) -> Float32: + assert poly_degree in POLY_EX2, f"Polynomial degree {poly_degree} not supported" + # We assume x <= 127.0 + fp32_round_int = float(2**23 + 2**22) + x_clamped = cute.arch.fmax(x, -127.0) + # We want to round down here, so that the fractional part is in [0, 1) + x_rounded = add_round_down(x_clamped, fp32_round_int, loc=loc, ip=ip) + # The integer floor of x is now in the last 8 bits of x_rounded + # We assume the next 2 ops round to nearest even. The rounding mode is important. + x_rounded_back = x_rounded - fp32_round_int + x_frac = x_clamped - x_rounded_back + x_frac_ex2 = evaluate_polynomial(x_frac, POLY_EX2[poly_degree], loc=loc, ip=ip) + return combine_int_frac_ex2(x_rounded, x_frac_ex2, loc=loc, ip=ip) + + +# TODO: check that the ex2_emulation_2 produces the same SASS as the ptx version +@dsl_user_op +def ex2_emulation_2( + x: Float32, y: Float32, *, poly_degree: int = 3, loc=None, ip=None +) -> Tuple[Float32, Float32]: + # We assume x <= 127.0 and y <= 127.0 + fp32_round_int = float(2**23 + 2**22) + xy_clamped = (cute.arch.fmax(x, -127.0), cute.arch.fmax(y, -127.0)) + # We want to round down here, so that the fractional part is in [0, 1) + xy_rounded = cute.arch.add_packed_f32x2(xy_clamped, (fp32_round_int, fp32_round_int), rnd="rm") + # The integer floor of x & y are now in the last 8 bits of xy_rounded + # We want the next 2 ops to round to nearest even. The rounding mode is important. + xy_rounded_back = activation.sub_packed_f32x2( + xy_rounded, (fp32_round_int, fp32_round_int) + ) + xy_frac = activation.sub_packed_f32x2(xy_clamped, xy_rounded_back) + xy_frac_ex2 = evaluate_polynomial_2(*xy_frac, POLY_EX2[poly_degree], loc=loc, ip=ip) + x_out = combine_int_frac_ex2(xy_rounded[0], xy_frac_ex2[0], loc=loc, ip=ip) + y_out = combine_int_frac_ex2(xy_rounded[1], xy_frac_ex2[1], loc=loc, ip=ip) + return x_out, y_out + + +@dsl_user_op +def e2e_asm2(x: Float32, y: Float32, *, loc=None, ip=None) -> Tuple[Float32, Float32]: + out_f32x2 = llvm.inline_asm( + llvm.StructType.get_literal([T.f32(), T.f32()]), + [Float32(x).ir_value(loc=loc, ip=ip), Float32(y, loc=loc, ip=ip).ir_value()], + "{\n\t" + ".reg .f32 f1, f2, f3, f4, f5, f6, f7;\n\t" + ".reg .b64 l1, l2, l3, l4, l5, l6, l7, l8, l9, l10;\n\t" + ".reg .s32 r1, r2, r3, r4, r5, r6, r7, r8;\n\t" + "max.ftz.f32 f1, $2, 0fC2FE0000;\n\t" + "max.ftz.f32 f2, $3, 0fC2FE0000;\n\t" + "mov.b64 l1, {f1, f2};\n\t" + "mov.f32 f3, 0f4B400000;\n\t" + "mov.b64 l2, {f3, f3};\n\t" + "add.rm.ftz.f32x2 l7, l1, l2;\n\t" + "sub.rn.ftz.f32x2 l8, l7, l2;\n\t" + "sub.rn.ftz.f32x2 l9, l1, l8;\n\t" + "mov.f32 f7, 0f3D9DF09D;\n\t" + "mov.b64 l6, {f7, f7};\n\t" + "mov.f32 f6, 0f3E6906A4;\n\t" + "mov.b64 l5, {f6, f6};\n\t" + "mov.f32 f5, 0f3F31F519;\n\t" + "mov.b64 l4, {f5, f5};\n\t" + "mov.f32 f4, 0f3F800000;\n\t" + "mov.b64 l3, {f4, f4};\n\t" + "fma.rn.ftz.f32x2 l10, l9, l6, l5;\n\t" + "fma.rn.ftz.f32x2 l10, l10, l9, l4;\n\t" + "fma.rn.ftz.f32x2 l10, l10, l9, l3;\n\t" + "mov.b64 {r1, r2}, l7;\n\t" + "mov.b64 {r3, r4}, l10;\n\t" + "shl.b32 r5, r1, 23;\n\t" + "add.s32 r7, r5, r3;\n\t" + "shl.b32 r6, r2, 23;\n\t" + "add.s32 r8, r6, r4;\n\t" + "mov.b32 $0, r7;\n\t" + "mov.b32 $1, r8;\n\t" + "}\n", + "=r,=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + out0 = Float32(llvm.extractvalue(T.f32(), out_f32x2, [0], loc=loc, ip=ip)) + out1 = Float32(llvm.extractvalue(T.f32(), out_f32x2, [1], loc=loc, ip=ip)) + return out0, out1 + + +@dsl_user_op +def domain_offset_aligned( + coord: cute.Coord, tensor: cute.Tensor, *, loc=None, ip=None +) -> cute.Tensor: + assert isinstance(tensor.iterator, cute.Pointer) + # We assume that applying the offset does not change the pointer alignment + new_ptr = cute.make_ptr( + tensor.element_type, + elem_pointer(tensor, coord).toint(), + tensor.memspace, + assumed_align=tensor.iterator.alignment, + ) + return cute.make_tensor(new_ptr, tensor.layout) + + +@cute.jit +def scalar_to_ssa(a: cute.Numeric, dtype) -> cute.TensorSSA: + """Convert a scalar to a cute TensorSSA of shape (1,) and given dtype""" + vec = cute.make_rmem_tensor(1, dtype) + vec[0] = a + return vec.load() + + +def ssa_to_scalar(val): + """Could inline but nice for reflecting the above api""" + return val[0] diff --git a/telefuser/kernel/sol_attn/common/__init__.py b/telefuser/kernel/sol_attn/common/__init__.py new file mode 100644 index 00000000..ec1d7ee9 --- /dev/null +++ b/telefuser/kernel/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/telefuser/kernel/sol_attn/common/layout_utils.py b/telefuser/kernel/sol_attn/common/layout_utils.py new file mode 100644 index 00000000..67e4d089 --- /dev/null +++ b/telefuser/kernel/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/telefuser/kernel/sol_attn/common/runtime.py b/telefuser/kernel/sol_attn/common/runtime.py new file mode 100644 index 00000000..5182d0c4 --- /dev/null +++ b/telefuser/kernel/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/telefuser/kernel/sol_attn/common/selector.py b/telefuser/kernel/sol_attn/common/selector.py new file mode 100644 index 00000000..0321f0fe --- /dev/null +++ b/telefuser/kernel/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/telefuser/kernel/sol_attn/interface.py b/telefuser/kernel/sol_attn/interface.py new file mode 100644 index 00000000..933513ab --- /dev/null +++ b/telefuser/kernel/sol_attn/interface.py @@ -0,0 +1,399 @@ +"""Public Sol-Attn interface.""" + +from __future__ import annotations + +import functools + +import torch + +BLOCK_SIZE = 64 +_CUTE_BACKENDS = { + (9, 0): "cute_sm90", + (10, 0): "cute_sm100", + (12, 0): "cute_sm120", +} +_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 CuTe when specialized and available, otherwise Triton.""" + + if arch[0] < 8: + raise RuntimeError( + "Sol-Attn requires an NVIDIA GPU with compute capability >= 8.0; " + f"got SM{arch[0]}{arch[1]}" + ) + cute_backend = _CUTE_BACKENDS.get(arch) + if cute_backend is not None: + available = ( + _cute_runtime_available() + if cute_available is None + else cute_available + ) + if available: + return cute_backend + return "triton" + + +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") + 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_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_sm100( + key, + tensors, + scale, + sink_start_block, + sink_end_block, + stream, +): + import cutlass.cute as cute + + from .sm100 import forward + + args = _to_cute_tensors(tensors) + compiled = cute.compile( + forward, + *args, + scale, + sink_start_block, + sink_end_block, + stream=stream, + options="--enable-tvm-ffi", + ) + _compiled[key] = compiled + return compiled, args + + +def _compile_sm120( + key, + tensors, + scale, + sink_start_block, + sink_end_block, + stream, +): + import cutlass.cute as cute + + from .sm120 import make_kernel + + operator = make_kernel() + args = _to_cute_tensors(tensors) + compiled = cute.compile( + operator, + *args, + scale, + sink_start_block, + sink_end_block, + stream=stream, + options="--enable-tvm-ffi", + ) + _compiled[key] = compiled + return compiled, args + + +def _sol_attn_cute( + q, + k, + v, + *, + arch, + scale, + tau, + thresh_type, + kv_splits, + sink_tokens, + sink_start, +): + from .preprocess import prepare + + batch, tokens, heads, _ = q.shape + + with torch.cuda.device(q.device): + kc, vc, threshold = prepare( + q, + k, + v, + scale=scale, + tau=tau, + thresh_type=thresh_type, + ) + output = torch.empty_like(v) + lse = torch.empty( + (batch, tokens, heads), + device=q.device, + dtype=torch.float32, + ) + stream = _stream(q.device) + key = (q.device.index, arch, batch, tokens, heads, kv_splits) + + if arch == (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, + scale, + sink_start_block, + sink_end_block, + stream=stream, + ) + else: + sink_start_block, sink_end_block = _sink_block_range( + tokens, + sink_start, + sink_tokens, + ) + tensors = [q, k, v, output, kc, vc, threshold, lse] + compiled = _compiled.get(key) + if compiled is None: + compiled, args = _compile_sm120( + key, + tensors, + scale, + sink_start_block, + sink_end_block, + stream, + ) + else: + args = _to_cute_tensors(tensors) + compiled( + *args, + scale, + sink_start_block, + sink_end_block, + stream=stream, + ) + return output + + +def sol_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + scale: float | None = None, + tau: float = 1.0, + thresh_type: str = "diag", + kv_splits: int = 1, + sink_tokens: int = 0, + sink_start: int | None = None, +) -> torch.Tensor: + """Compute noncausal Sol-Attn for contiguous BF16 BTHD tensors. + + ``sink_start`` and ``sink_tokens`` keep every KV block overlapping the + corresponding contiguous token range exact for all queries. Omitting + ``sink_start`` places the range at the token suffix. + """ + + arch = _validate_inputs( + q, + k, + v, + thresh_type, + sink_tokens, + sink_start, + ) + if kv_splits not in (1, 2, 4): + raise ValueError("kv_splits must be 1, 2, or 4") + backend = _backend_for_arch(arch) + 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, + 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__ = ["sol_attn"] diff --git a/telefuser/kernel/sol_attn/preprocess.py b/telefuser/kernel/sol_attn/preprocess.py new file mode 100644 index 00000000..77c4cccf --- /dev/null +++ b/telefuser/kernel/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/telefuser/kernel/sol_attn/sm100/LICENSE.flash-attention b/telefuser/kernel/sol_attn/sm100/LICENSE.flash-attention new file mode 100644 index 00000000..5860e4b3 --- /dev/null +++ b/telefuser/kernel/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/telefuser/kernel/sol_attn/sm100/__init__.py b/telefuser/kernel/sol_attn/sm100/__init__.py new file mode 100644 index 00000000..7fb6167a --- /dev/null +++ b/telefuser/kernel/sol_attn/sm100/__init__.py @@ -0,0 +1,5 @@ +"""Blackwell backend.""" + +from .kernel import forward + +__all__ = ["forward"] diff --git a/telefuser/kernel/sol_attn/sm100/kernel.py b/telefuser/kernel/sol_attn/sm100/kernel.py new file mode 100644 index 00000000..d3edaff3 --- /dev/null +++ b/telefuser/kernel/sol_attn/sm100/kernel.py @@ -0,0 +1,5 @@ +"""Blackwell kernel entry.""" + +from .mainloop import forward + +__all__ = ["forward"] diff --git a/telefuser/kernel/sol_attn/sm100/mainloop.py b/telefuser/kernel/sol_attn/sm100/mainloop.py new file mode 100644 index 00000000..353868a5 --- /dev/null +++ b/telefuser/kernel/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 telefuser.kernel.sol_attn._vendor.flash_attn.cute.pipeline as fa_pipeline +import telefuser.kernel.sol_attn._vendor.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 telefuser.kernel.sol_attn._vendor.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 telefuser.kernel.sol_attn.common import layout_utils +from telefuser.kernel.sol_attn.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/telefuser/kernel/sol_attn/sm100/math.py b/telefuser/kernel/sol_attn/sm100/math.py new file mode 100644 index 00000000..e1120b09 --- /dev/null +++ b/telefuser/kernel/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/telefuser/kernel/sol_attn/sm100/softmax.py b/telefuser/kernel/sol_attn/sm100/softmax.py new file mode 100644 index 00000000..129199af --- /dev/null +++ b/telefuser/kernel/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 telefuser.kernel.sol_attn._vendor.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/telefuser/kernel/sol_attn/sm100/tmem.py b/telefuser/kernel/sol_attn/sm100/tmem.py new file mode 100644 index 00000000..aada77be --- /dev/null +++ b/telefuser/kernel/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/telefuser/kernel/sol_attn/sm120/__init__.py b/telefuser/kernel/sol_attn/sm120/__init__.py new file mode 100644 index 00000000..72060720 --- /dev/null +++ b/telefuser/kernel/sol_attn/sm120/__init__.py @@ -0,0 +1,5 @@ +"""GeForce Blackwell (SM120) backend.""" + +from .kernel import make_kernel + +__all__ = ["make_kernel"] diff --git a/telefuser/kernel/sol_attn/sm120/kernel.py b/telefuser/kernel/sol_attn/sm120/kernel.py new file mode 100644 index 00000000..1af031bd --- /dev/null +++ b/telefuser/kernel/sol_attn/sm120/kernel.py @@ -0,0 +1,19 @@ +"""SM120 kernel recipe.""" + +from .mainloop import SolAttnForwardSm120 + + +def make_kernel( + *, + debug_route_trace: bool = False, + prefetch_first_exact_k: bool = True, + prefetch_next_route_k: bool = True, +): + return SolAttnForwardSm120( + debug_route_trace=debug_route_trace, + prefetch_first_exact_k=prefetch_first_exact_k, + prefetch_next_route_k=prefetch_next_route_k, + ) + + +__all__ = ["make_kernel"] diff --git a/telefuser/kernel/sol_attn/sm120/mainloop.py b/telefuser/kernel/sol_attn/sm120/mainloop.py new file mode 100644 index 00000000..aa493073 --- /dev/null +++ b/telefuser/kernel/sol_attn/sm120/mainloop.py @@ -0,0 +1,1172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Fused Sol-Attn forward kernel for GeForce Blackwell SM120. + +The warp-MMA/TMA execution skeleton and online-softmax helpers are adapted +from NVIDIA cuDNN Frontend's SM120 block-sparse-attention kernel. Sol-specific +routing, CTA-local exact-index compaction, approximate block mass, and the +mixed approximate/exact mainloop are implemented here. +""" + +from __future__ import annotations + +import operator + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.hopper_helpers as sm90_utils + +from telefuser.kernel.sol_attn._vendor.flash_attn.cute import utils as kernel_utils +from telefuser.kernel.sol_attn.common import layout_utils +from telefuser.kernel.sol_attn.common.selector import ( + sol_attn_popc_b32, + sol_attn_route_is_exact, +) + + +M = 64 +N = 64 +D = 128 +DV = 128 +THREADS = 128 +STAGES = 1 + + +class SolAttnForwardSm120: + """M64/N64 warp-MMA Sol-Attn kernel for BF16 D128 inputs.""" + + def __init__( + self, + *, + debug_route_trace: bool = False, + prefetch_first_exact_k: bool = True, + prefetch_next_route_k: bool = True, + ): + self.dtype = cutlass.BFloat16 + self.acc_dtype = cutlass.Float32 + self.tile_shape_qk = (M, N, D) + self.tile_shape_pv = (M, DV, N) + self.num_threads = THREADS + self.q_stage = 1 + self.kv_stage = STAGES + self.debug_route_trace = debug_route_trace + self.prefetch_first_exact_k = prefetch_first_exact_k + self.prefetch_next_route_k = prefetch_next_route_k + + @cute.kernel + def kernel( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mO: cute.Tensor, + mKC: cute.Tensor, + mVC: cute.Tensor, + mThreshold: cute.Tensor, + mLSE: cute.Tensor, + tma_atom_Q: cute.CopyAtom, + tma_atom_K: cute.CopyAtom, + tma_atom_V: cute.CopyAtom, + tma_atom_KC: cute.CopyAtom, + tma_atom_VC: cute.CopyAtom, + tma_atom_O: cute.CopyAtom, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + Q_smem_layout: cute.ComposedLayout, + K_smem_layout: cute.ComposedLayout, + V_smem_layout: cute.ComposedLayout, + O_smem_layout: cute.ComposedLayout, + scale_softmax_log2e: cutlass.Float32, + sink_start_block: cutlass.Int32, + sink_end_block: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + lane = cute.arch.lane_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + q_tile_idx, head_idx, batch_idx = cute.arch.block_idx() + q_tile_idx = cute.arch.make_warp_uniform(q_tile_idx) + head_idx = cute.arch.make_warp_uniform(head_idx) + batch_idx = cute.arch.make_warp_uniform(batch_idx) + + token_count = mK.shape[0] + num_blocks = mKC.shape[0] + num_route_groups = cute.ceil_div(num_blocks, N) + q_start = q_tile_idx * M + q_len = token_count - q_start + if q_len > M: + q_len = cutlass.Int32(M) + threshold = cutlass.Float32( + mThreshold[batch_idx, q_tile_idx, head_idx] + ) + + storage = cutlass.utils.SmemAllocator().allocate(self.shared_storage_t) + if warp == 0 and lane == 0: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_Q) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_K) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_V) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_KC) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_VC) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_O) + + cg = pipeline.CooperativeGroup(pipeline.Agent.Thread) + consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_threads // 32 + ) + cta_layout_vmnk = cute.make_layout((1, 1, 1, 1)) + Q_pipeline = pipeline.PipelineTmaAsync.create( + num_stages=self.q_stage, + producer_group=cg, + consumer_group=consumer_group, + tx_count=cute.size_in_bytes( + self.Q_dtype, cute.select(Q_smem_layout, mode=[0, 1]) + ), + barrier_storage=storage.Q_barrier.data_ptr(), + cta_layout_vmnk=cta_layout_vmnk, + ) + K_pipeline = pipeline.PipelineTmaAsync.create( + num_stages=self.kv_stage, + producer_group=cg, + consumer_group=consumer_group, + tx_count=cute.size_in_bytes( + self.K_dtype, cute.select(K_smem_layout, mode=[0, 1]) + ), + barrier_storage=storage.K_barrier.data_ptr(), + cta_layout_vmnk=cta_layout_vmnk, + ) + V_pipeline = pipeline.PipelineTmaAsync.create( + num_stages=self.kv_stage, + producer_group=cg, + consumer_group=consumer_group, + tx_count=cute.size_in_bytes( + self.V_dtype, cute.select(V_smem_layout, mode=[0, 1]) + ), + barrier_storage=storage.V_barrier.data_ptr(), + cta_layout_vmnk=cta_layout_vmnk, + ) + Q_producer = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.q_stage + ) + Q_consumer = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.q_stage + ) + K_producer = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.kv_stage + ) + K_consumer = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.kv_stage + ) + V_producer = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.kv_stage + ) + V_consumer = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.kv_stage + ) + + sQ = storage.Q_smem.get_tensor( + Q_smem_layout.outer, swizzle=Q_smem_layout.inner + ) + sK = storage.K_smem.get_tensor( + K_smem_layout.outer, swizzle=K_smem_layout.inner + ) + sV = storage.V_smem.get_tensor( + V_smem_layout.outer, swizzle=V_smem_layout.inner + ) + # Q is register-resident after the prologue. Reuse its 16 KiB SMEM + # allocation for route scratch until the same allocation becomes sO + # in the epilogue. This drops the CTA below the 2-block/SM threshold + # on SM120 without changing any route reduction or synchronization. + route_f32_ptr = cute.recast_ptr( + storage.Q_smem.data_ptr(), dtype=cutlass.Float32 + ) + route_i32_ptr = cute.recast_ptr( + storage.Q_smem.data_ptr(), dtype=cutlass.Int32 + ) + route_sums = cute.make_tensor( + route_f32_ptr, cute.make_layout((4, N)) + ) + column_masks = cute.make_tensor( + route_f32_ptr + 4 * N, cute.make_layout(N) + ) + route_indices = cute.make_tensor( + route_i32_ptr + 5 * N, cute.make_layout(N) + ) + route_meta = cute.make_tensor( + route_i32_ptr + 6 * N, cute.make_layout(2) + ) + + mQ_slice = mQ[None, None, head_idx, batch_idx] + mK_slice = mK[None, None, head_idx, batch_idx] + mV_slice = mV[None, None, head_idx, batch_idx] + mO_slice = mO[None, None, head_idx, batch_idx] + mKC_slice = mKC[None, None, head_idx, batch_idx] + mVC_slice = mVC[None, None, head_idx, batch_idx] + if cutlass.const_expr(not self.debug_route_trace): + mLSE_slice = mLSE[None, head_idx, batch_idx] + + gQ = cute.local_tile( + mQ_slice, (M, D), coord=(q_tile_idx, 0) + ) + gK = cute.local_tile(mK_slice, (N, D), coord=(None, 0)) + gV = cute.local_tile(mV_slice, (DV, N), coord=(0, None)) + gKC = cute.local_tile(mKC_slice, (N, D), coord=(None, 0)) + gVC = cute.local_tile(mVC_slice, (DV, N), coord=(0, None)) + gO = cute.local_tile( + mO_slice, (M, DV), coord=(q_tile_idx, 0) + ) + + cta_coord_layout = (0, cute.make_layout(1)) + tQsQ, tQgQ = cute.nvgpu.cpasync.tma_partition( + tma_atom_Q, + *cta_coord_layout, + cute.group_modes(sQ, 0, 2), + cute.group_modes(gQ, 0, 2), + ) + tKsK, tKgK = cute.nvgpu.cpasync.tma_partition( + tma_atom_K, + *cta_coord_layout, + cute.group_modes(sK, 0, 2), + cute.group_modes(gK, 0, 2), + ) + tVsV, tVgV = cute.nvgpu.cpasync.tma_partition( + tma_atom_V, + *cta_coord_layout, + cute.group_modes(sV, 0, 2), + cute.group_modes(gV, 0, 2), + ) + tKCsK, tKCgKC = cute.nvgpu.cpasync.tma_partition( + tma_atom_KC, + *cta_coord_layout, + cute.group_modes(sK, 0, 2), + cute.group_modes(gKC, 0, 2), + ) + tVCsV, tVCgVC = cute.nvgpu.cpasync.tma_partition( + tma_atom_VC, + *cta_coord_layout, + cute.group_modes(sV, 0, 2), + cute.group_modes(gVC, 0, 2), + ) + + cS = cute.make_identity_tensor(self.tile_shape_qk[:2]) + thr_mma_qk = tiled_mma_qk.get_slice(tidx) + tSsQ = thr_mma_qk.partition_A(sQ) + tSsK = thr_mma_qk.partition_B(sK) + tSrQ = tiled_mma_qk.make_fragment_A(tSsQ[None, None, None, 0]) + tSrK = tiled_mma_qk.make_fragment_B(tSsK[None, None, None, 0]) + tSrS = cute.make_rmem_tensor( + thr_mma_qk.partition_shape_C((M, N)), self.acc_dtype + ) + tScS = thr_mma_qk.partition_C(cS) + + thr_mma_pv = tiled_mma_pv.get_slice(tidx) + tOsV = thr_mma_pv.partition_B(sV) + tOrV = tiled_mma_pv.make_fragment_B(tOsV[None, None, None, 0]) + tOrO = cute.make_rmem_tensor( + thr_mma_pv.partition_shape_C((M, DV)), self.acc_dtype + ) + + atom_copy_Q = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp( + self.Q_layout.is_m_major_a(), 4 + ), + self.Q_dtype, + ) + atom_copy_K = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp( + self.K_layout.is_n_major_b(), 4 + ), + self.K_dtype, + ) + atom_copy_V = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp( + self.V_layout.is_n_major_b(), 4 + ), + self.V_dtype, + ) + smem_copy_Q = cute.make_tiled_copy_A(atom_copy_Q, tiled_mma_qk) + smem_copy_K = cute.make_tiled_copy_B(atom_copy_K, tiled_mma_qk) + smem_copy_V = cute.make_tiled_copy_B(atom_copy_V, tiled_mma_pv) + thr_copy_Q = smem_copy_Q.get_slice(tidx) + thr_copy_K = smem_copy_K.get_slice(tidx) + thr_copy_V = smem_copy_V.get_slice(tidx) + tSsQ_copy = thr_copy_Q.partition_S(sQ) + tSrQ_copy = thr_copy_Q.retile(tSrQ) + tSsK_copy = thr_copy_K.partition_S(sK) + tOsV_copy = thr_copy_V.partition_S(sV) + + max_m_layout = cute.make_layout( + cute.size( + layout_utils.reshape_acc_to_mn(tOrO).layout, + mode=[0], + ) + ) + max_m = cute.make_rmem_tensor_like(max_m_layout, cutlass.Float32) + sum_m = cute.make_rmem_tensor_like(max_m, cutlass.Float32) + tOrO.store(cute.full_like(tOrO, 0.0, self.acc_dtype)) + max_m.store(cute.full_like(max_m, float("-inf"), cutlass.Float32)) + sum_m.store(cute.full_like(sum_m, 0.0, cutlass.Float32)) + + if warp == 0: + Q_pipeline.producer_acquire(Q_producer) + cute.copy( + tma_atom_Q, + tQgQ, + tQsQ[None, Q_producer.index], + tma_bar_ptr=Q_pipeline.producer_get_barrier(Q_producer), + ) + Q_pipeline.producer_commit(Q_producer) + Q_producer.advance() + cute.arch.sync_threads() + q_wait = Q_pipeline.consumer_try_wait(Q_consumer) + Q_pipeline.consumer_wait(Q_consumer, q_wait) + q_stage = Q_consumer.index + for k_block in cutlass.range_constexpr(cute.size(tSrQ, mode=[2])): + cute.copy( + smem_copy_Q, + tSsQ_copy[None, None, k_block, q_stage], + tSrQ_copy[None, None, k_block], + ) + Q_pipeline.consumer_release(Q_consumer) + Q_consumer.advance() + + for route_group in cutlass.range( + 0, num_route_groups, 1, unroll=1 + ): + group_start = route_group * cutlass.Int32(N) + valid_blocks = num_blocks - group_start + if valid_blocks > N: + valid_blocks = cutlass.Int32(N) + + if warp == 0: + if cutlass.const_expr(self.prefetch_next_route_k): + # P19-style terminal handoff: when the previous route + # group had an exact block, its final exact QK already + # refilled this K stage with the current group's KC. + if route_group == 0: + K_pipeline.producer_acquire(K_producer) + cute.copy( + tma_atom_KC, + tKCgKC[None, route_group], + tKCsK[None, K_producer.index], + tma_bar_ptr=K_pipeline.producer_get_barrier( + K_producer + ), + ) + K_pipeline.producer_commit(K_producer) + K_producer.advance() + else: + previous_group_exact_count = cutlass.Int32( + route_meta[0] + ) + if previous_group_exact_count == 0: + K_pipeline.producer_acquire(K_producer) + cute.copy( + tma_atom_KC, + tKCgKC[None, route_group], + tKCsK[None, K_producer.index], + tma_bar_ptr=K_pipeline.producer_get_barrier( + K_producer + ), + ) + K_pipeline.producer_commit(K_producer) + K_producer.advance() + else: + K_pipeline.producer_acquire(K_producer) + cute.copy( + tma_atom_KC, + tKCgKC[None, route_group], + tKCsK[None, K_producer.index], + tma_bar_ptr=K_pipeline.producer_get_barrier( + K_producer + ), + ) + K_pipeline.producer_commit(K_producer) + K_producer.advance() + V_pipeline.producer_acquire(V_producer) + cute.copy( + tma_atom_VC, + tVCgVC[None, route_group], + tVCsV[None, V_producer.index], + tma_bar_ptr=V_pipeline.producer_get_barrier(V_producer), + ) + V_pipeline.producer_commit(V_producer) + V_producer.advance() + + k_wait = K_pipeline.consumer_try_wait(K_consumer) + K_pipeline.consumer_wait(K_consumer, k_wait) + gemm_smem_zero_acc( + tiled_mma_qk, + tSrS, + tSrQ, + tSrK, + tSsK_copy[None, None, None, K_consumer.index], + smem_copy_K, + ) + K_pipeline.consumer_release(K_consumer) + K_consumer.advance() + + reduce_route_columns( + tSrS, + tScS, + route_sums, + warp, + lane, + q_len, + ) + cute.arch.fence_view_async_shared() + cute.arch.sync_threads() + + if warp == 0: + preceding = cutlass.Int32(0) + lane_mask_lt = cutlass.Int32(0x7FFFFFFF) >> ( + cutlass.Int32(31) - lane + ) + for word in cutlass.range_constexpr(2): + off = cutlass.Int32(word * 32) + lane + valid = off < valid_blocks + exact = False + if valid: + col_sum = ( + cutlass.Float32(route_sums[0, off]) + + cutlass.Float32(route_sums[1, off]) + + cutlass.Float32(route_sums[2, off]) + + cutlass.Float32(route_sums[3, off]) + ) + col_mean = ( + col_sum + * scale_softmax_log2e + / cutlass.Float32(q_len) + ) + kv_block = group_start + off + exact = sol_attn_route_is_exact( + q_tile_idx, + kv_block, + col_mean, + threshold, + valid, + ) + exact = exact or ( + kv_block >= sink_start_block + and kv_block < sink_end_block + ) + ballot = cutlass.Int32( + cute.arch.vote_ballot_sync(exact) + ) + column_masks[off] = ( + -cutlass.Float32.inf + if (exact or not valid) + else cutlass.Float32(0.0) + ) + rank = preceding + sol_attn_popc_b32( + ballot & lane_mask_lt + ) + if exact: + route_indices[rank] = group_start + off + preceding += sol_attn_popc_b32(ballot) + if cutlass.const_expr(self.debug_route_trace): + if lane == 0: + mLSE[ + batch_idx, + q_tile_idx, + head_idx, + route_group, + word, + ] = ballot + if lane == 0: + route_meta[0] = preceding + route_meta[1] = valid_blocks + cute.arch.fence_view_async_shared() + cute.arch.sync_threads() + + exact_count = cutlass.Int32(route_meta[0]) + has_approx = exact_count < valid_blocks + if cutlass.const_expr(self.prefetch_first_exact_k): + # Once routing identifies the first exact block, the route KC + # stage is free. Refill it before the approximate softmax/PV + # so the first exact K transfer overlaps that work. + if warp == 0 and exact_count > 0: + first_exact = cutlass.Int32(route_indices[0]) + K_pipeline.producer_acquire(K_producer) + cute.copy( + tma_atom_K, + tKgK[None, first_exact], + tKsK[None, K_producer.index], + tma_bar_ptr=K_pipeline.producer_get_barrier( + K_producer + ), + ) + K_pipeline.producer_commit(K_producer) + K_producer.advance() + v_wait = V_pipeline.consumer_try_wait(V_consumer) + V_pipeline.consumer_wait(V_consumer, v_wait) + if has_approx: + apply_route_mask(tSrS, tScS, column_masks, q_len) + row_scale = online_softmax_route( + tSrS, + tScS, + max_m, + sum_m, + scale_softmax_log2e, + group_start, + token_count, + ) + rescale_o_for_next_acc(tOrO, row_scale) + tOrP_frg = cute.make_rmem_tensor_like( + tSrS, self.K_dtype + ) + tOrP_frg.store(tSrS.load().to(self.K_dtype)) + tOrP = layout_utils.reshape_acc_to_frgA(tOrP_frg) + gemm_rs_smem( + tiled_mma_pv, + tOrO, + tOrP, + tOrV, + tOsV_copy[None, None, None, V_consumer.index], + smem_copy_V, + ) + V_pipeline.consumer_release(V_consumer) + V_consumer.advance() + + if warp == 0 and exact_count > 0: + first_exact = cutlass.Int32(route_indices[0]) + if cutlass.const_expr(not self.prefetch_first_exact_k): + K_pipeline.producer_acquire(K_producer) + cute.copy( + tma_atom_K, + tKgK[None, first_exact], + tKsK[None, K_producer.index], + tma_bar_ptr=K_pipeline.producer_get_barrier( + K_producer + ), + ) + K_pipeline.producer_commit(K_producer) + K_producer.advance() + V_pipeline.producer_acquire(V_producer) + cute.copy( + tma_atom_V, + tVgV[None, first_exact], + tVsV[None, V_producer.index], + tma_bar_ptr=V_pipeline.producer_get_barrier(V_producer), + ) + V_pipeline.producer_commit(V_producer) + V_producer.advance() + + for ordinal in cutlass.range(0, exact_count, 1, unroll=1): + exact_block = cutlass.Int32(route_indices[ordinal]) + k_wait = K_pipeline.consumer_try_wait(K_consumer) + K_pipeline.consumer_wait(K_consumer, k_wait) + gemm_smem_zero_acc( + tiled_mma_qk, + tSrS, + tSrQ, + tSrK, + tSsK_copy[None, None, None, K_consumer.index], + smem_copy_K, + ) + K_pipeline.consumer_release(K_consumer) + K_consumer.advance() + next_ordinal = ordinal + cutlass.Int32(1) + if warp == 0: + if next_ordinal < exact_count: + next_exact = cutlass.Int32( + route_indices[next_ordinal] + ) + K_pipeline.producer_acquire(K_producer) + cute.copy( + tma_atom_K, + tKgK[None, next_exact], + tKsK[None, K_producer.index], + tma_bar_ptr=K_pipeline.producer_get_barrier( + K_producer + ), + ) + K_pipeline.producer_commit(K_producer) + K_producer.advance() + else: + if cutlass.const_expr( + self.prefetch_next_route_k + ): + next_route_group = route_group + cutlass.Int32(1) + if next_route_group < num_route_groups: + # Reuse the K stage released by the final + # exact QK. The next outer prologue supplies + # VC, matching the SM90 P19 partial handoff. + K_pipeline.producer_acquire(K_producer) + cute.copy( + tma_atom_KC, + tKCgKC[None, next_route_group], + tKCsK[None, K_producer.index], + tma_bar_ptr=( + K_pipeline.producer_get_barrier( + K_producer + ) + ), + ) + K_pipeline.producer_commit(K_producer) + K_producer.advance() + block_len = token_count - exact_block * cutlass.Int32(N) + if block_len > N: + block_len = cutlass.Int32(N) + mask_exact_scores(tSrS, tScS, block_len, q_len) + row_scale = online_softmax( + tSrS, max_m, sum_m, scale_softmax_log2e + ) + rescale_o_for_next_acc(tOrO, row_scale) + tOrP_frg = cute.make_rmem_tensor_like( + tSrS, self.K_dtype + ) + tOrP_frg.store(tSrS.load().to(self.K_dtype)) + tOrP = layout_utils.reshape_acc_to_frgA(tOrP_frg) + + v_wait = V_pipeline.consumer_try_wait(V_consumer) + V_pipeline.consumer_wait(V_consumer, v_wait) + gemm_rs_smem( + tiled_mma_pv, + tOrO, + tOrP, + tOrV, + tOsV_copy[None, None, None, V_consumer.index], + smem_copy_V, + ) + V_pipeline.consumer_release(V_consumer) + V_consumer.advance() + if warp == 0 and next_ordinal < exact_count: + next_exact = cutlass.Int32(route_indices[next_ordinal]) + V_pipeline.producer_acquire(V_producer) + cute.copy( + tma_atom_V, + tVgV[None, next_exact], + tVsV[None, V_producer.index], + tma_bar_ptr=V_pipeline.producer_get_barrier( + V_producer + ), + ) + V_pipeline.producer_commit(V_producer) + V_producer.advance() + + final_ratio, lse = finalize_softmax( + max_m, sum_m, scale_softmax_log2e + ) + rescale_o_for_next_acc(tOrO, final_ratio) + if cutlass.const_expr(not self.debug_route_trace): + tScS_mn = layout_utils.reshape_acc_to_mn(tScS) + for m in cutlass.range_constexpr(cute.size(lse)): + row = tScS_mn[m, 0][0] + if tScS_mn[m, 0][1] == 0 and row < q_len: + mLSE_slice[q_start + row] = lse[m] + + tOrO_cvt = cute.make_rmem_tensor_like(tOrO, self.O_dtype) + tOrO_cvt.store(tOrO.load().to(self.O_dtype)) + sO = storage.Q_smem.get_tensor( + O_smem_layout.outer, swizzle=O_smem_layout.inner + ) + tiled_copy_O = cute.make_tiled_copy_C( + cute.make_copy_atom( + cute.nvgpu.warp.StMatrix8x8x16bOp( + self.O_layout.is_m_major_c(), 4 + ), + self.O_dtype, + ), + tiled_mma_pv, + ) + tOrO_cv = tiled_copy_O.retile(tOrO_cvt) + tOsO = tiled_copy_O.get_slice(tidx).partition_D(sO) + cute.copy(tiled_copy_O, tOrO_cv, tOsO) + cute.arch.fence_view_async_shared() + cute.arch.sync_threads() + tOsO, tOgO = cute.nvgpu.cpasync.tma_partition( + tma_atom_O, + *cta_coord_layout, + cute.group_modes(sO, 0, 2), + cute.group_modes(gO, 0, 2), + ) + if warp == 0: + cute.copy(tma_atom_O, tOsO, tOgO) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + + @cute.jit + def __call__( + self, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + o: cute.Tensor, + kc: cute.Tensor, + vc: cute.Tensor, + threshold: cute.Tensor, + lse: cute.Tensor, + softmax_scale: cutlass.Float32, + sink_start_block: cutlass.Int32, + sink_end_block: cutlass.Int32, + stream: cuda.CUstream, + ): + q_mkl, k_nkl, kc_nkl = [ + layout_utils.select(t, [1, 3, 2, 0]) + for t in (q, k, kc) + ] + v_nkl, vc_nkl = [ + layout_utils.select(t, [3, 1, 2, 0]) for t in (v, vc) + ] + o_mkl = layout_utils.select(o, [1, 3, 2, 0]) + if cutlass.const_expr(self.debug_route_trace): + lse_target = lse + else: + lse_target = layout_utils.select(lse, [1, 2, 0]) + + self.Q_dtype = q_mkl.element_type + self.K_dtype = k_nkl.element_type + self.V_dtype = v_nkl.element_type + self.O_dtype = o_mkl.element_type + self.Q_layout = utils.LayoutEnum.from_tensor(q_mkl) + self.K_layout = utils.LayoutEnum.from_tensor(k_nkl) + self.V_layout = utils.LayoutEnum.from_tensor(v_nkl) + self.O_layout = utils.LayoutEnum.from_tensor(o_mkl) + assert self.Q_dtype == cutlass.BFloat16 + assert self.K_dtype == cutlass.BFloat16 + assert self.V_dtype == cutlass.BFloat16 + + self.Q_smem_layout = sm90_utils.make_smem_layout_a( + self.Q_layout, + self.tile_shape_qk, + self.Q_dtype, + self.q_stage, + ) + self.K_smem_layout = sm90_utils.make_smem_layout_b( + self.K_layout, + self.tile_shape_qk, + self.K_dtype, + self.kv_stage, + ) + self.V_smem_layout = sm90_utils.make_smem_layout_b( + self.V_layout, + self.tile_shape_pv, + self.V_dtype, + self.kv_stage, + ) + O_smem_layout_staged = sm90_utils.make_smem_layout_epi( + self.O_dtype, + self.O_layout, + self.tile_shape_pv[:2], + 1, + ) + self.O_smem_layout = cute.select( + O_smem_layout_staged, mode=[0, 1] + ) + + @cute.struct + class SharedStorage: + Q_barrier: cute.struct.MemRange[ + cutlass.Int64, self.q_stage * 2 + ] + K_barrier: cute.struct.MemRange[ + cutlass.Int64, self.kv_stage * 2 + ] + V_barrier: cute.struct.MemRange[ + cutlass.Int64, self.kv_stage * 2 + ] + Q_smem: cute.struct.Align[ + cute.struct.MemRange[ + self.Q_dtype, cute.cosize(self.Q_smem_layout) + ], + 128, + ] + K_smem: cute.struct.Align[ + cute.struct.MemRange[ + self.K_dtype, cute.cosize(self.K_smem_layout) + ], + 128, + ] + V_smem: cute.struct.Align[ + cute.struct.MemRange[ + self.V_dtype, cute.cosize(self.V_smem_layout) + ], + 128, + ] + self.shared_storage_t = SharedStorage + + tiled_mma_qk = cute.make_tiled_mma( + cute.nvgpu.warp.MmaF16BF16Op( + self.Q_dtype, + self.acc_dtype, + (16, 8, 16), + ), + cute.make_layout((4, 1, 1)), + permutation_mnk=(64, 16, 16), + ) + tiled_mma_pv = cute.make_tiled_mma( + cute.nvgpu.warp.MmaF16BF16Op( + self.K_dtype, + self.acc_dtype, + (16, 8, 16), + ), + cute.make_layout((4, 1, 1)), + permutation_mnk=(64, 16, 16), + ) + + g2s_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() + tma_atom_Q, tma_tensor_Q = ( + cute.nvgpu.cpasync.make_tiled_tma_atom( + g2s_op, + q_mkl, + self.Q_smem_layout, + (M, D), + num_multicast=1, + ) + ) + tma_atom_K, tma_tensor_K = ( + cute.nvgpu.cpasync.make_tiled_tma_atom( + g2s_op, + k_nkl, + self.K_smem_layout, + (N, D), + num_multicast=1, + ) + ) + tma_atom_V, tma_tensor_V = ( + cute.nvgpu.cpasync.make_tiled_tma_atom( + g2s_op, + v_nkl, + self.V_smem_layout, + (DV, N), + num_multicast=1, + ) + ) + tma_atom_KC, tma_tensor_KC = ( + cute.nvgpu.cpasync.make_tiled_tma_atom( + g2s_op, + kc_nkl, + self.K_smem_layout, + (N, D), + num_multicast=1, + ) + ) + tma_atom_VC, tma_tensor_VC = ( + cute.nvgpu.cpasync.make_tiled_tma_atom( + g2s_op, + vc_nkl, + self.V_smem_layout, + (DV, N), + num_multicast=1, + ) + ) + s2g_op = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() + tma_atom_O, tma_tensor_O = ( + cute.nvgpu.cpasync.make_tiled_tma_atom( + s2g_op, + o_mkl, + self.O_smem_layout, + (M, DV), + num_multicast=1, + ) + ) + + self.kernel( + tma_tensor_Q, + tma_tensor_K, + tma_tensor_V, + tma_tensor_O, + tma_tensor_KC, + tma_tensor_VC, + threshold, + lse_target, + tma_atom_Q, + tma_atom_K, + tma_atom_V, + tma_atom_KC, + tma_atom_VC, + tma_atom_O, + tiled_mma_qk, + tiled_mma_pv, + self.Q_smem_layout, + self.K_smem_layout, + self.V_smem_layout, + self.O_smem_layout, + softmax_scale * 1.4426950408889634, + sink_start_block, + sink_end_block, + ).launch( + grid=(cute.ceil_div(q_mkl.shape[0], M), q_mkl.shape[2], q_mkl.shape[3]), + block=(self.num_threads, 1, 1), + cluster=(1, 1, 1), + smem=self.shared_storage_t.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + ) + + +@cute.jit +def gemm_smem_zero_acc( + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + tCsB: cute.Tensor, + smem_tiled_copy_B: cute.TiledCopy, +): + acc.fill(0.0) + tCrB_copy = smem_tiled_copy_B.retile(tCrB) + cute.copy( + smem_tiled_copy_B, + tCsB[None, None, 0], + tCrB_copy[None, None, 0], + ) + for k_block in cutlass.range_constexpr(cute.size(tCsB.shape[2])): + if k_block < cute.size(tCsB.shape[2]) - 1: + cute.copy( + smem_tiled_copy_B, + tCsB[None, None, k_block + 1], + tCrB_copy[None, None, k_block + 1], + ) + cute.gemm( + tiled_mma, + acc, + tCrA[None, None, k_block], + tCrB[None, None, k_block], + acc, + ) + + +@cute.jit +def gemm_rs_smem( + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + tCsB: cute.Tensor, + smem_tiled_copy_B: cute.TiledCopy, +): + tCrB_copy = smem_tiled_copy_B.retile(tCrB) + cute.copy( + smem_tiled_copy_B, + tCsB[None, None, 0], + tCrB_copy[None, None, 0], + ) + for k_block in cutlass.range_constexpr(cute.size(tCrA.shape[2])): + if k_block < cute.size(tCrA.shape[2]) - 1: + cute.copy( + smem_tiled_copy_B, + tCsB[None, None, k_block + 1], + tCrB_copy[None, None, k_block + 1], + ) + cute.gemm( + tiled_mma, + acc, + tCrA[None, None, k_block], + tCrB[None, None, k_block], + acc, + ) + + +@cute.jit +def reduce_route_columns( + scores: cute.Tensor, + coords: cute.Tensor, + route_sums: cute.Tensor, + warp: cutlass.Int32, + lane: cutlass.Int32, + q_len: cutlass.Int32, +): + """Reduce M64 score columns using the measured SM120 lane layout.""" + + scores_mn = layout_utils.reshape_acc_to_mn(scores) + coords_mn = layout_utils.reshape_acc_to_mn(coords) + row0 = coords_mn[0, 0][0] + row1 = coords_mn[1, 0][0] + valid0 = row0 < q_len + valid1 = row1 < q_len + for group in cutlass.range_constexpr(8): + n0 = group * 2 + partial0 = cutlass.Float32(0.0) + partial1 = cutlass.Float32(0.0) + if valid0: + partial0 += cutlass.Float32(scores_mn[0, n0]) + partial1 += cutlass.Float32(scores_mn[0, n0 + 1]) + if valid1: + partial0 += cutlass.Float32(scores_mn[1, n0]) + partial1 += cutlass.Float32(scores_mn[1, n0 + 1]) + for offset in (4, 8, 16): + partial0 += cute.arch.shuffle_sync_bfly(partial0, offset=offset) + partial1 += cute.arch.shuffle_sync_bfly(partial1, offset=offset) + if lane < 4: + column = cutlass.Int32(group * 8) + lane * cutlass.Int32(2) + route_sums[warp, column] = partial0 + route_sums[warp, column + 1] = partial1 + + +@cute.jit +def apply_route_mask( + scores: cute.Tensor, + coords: cute.Tensor, + column_masks: cute.Tensor, + q_len: cutlass.Int32, +): + scores_mn = layout_utils.reshape_acc_to_mn(scores) + coords_mn = layout_utils.reshape_acc_to_mn(coords) + for m in cutlass.range_constexpr(cute.size(scores_mn, mode=[0])): + valid_row = coords_mn[m, 0][0] < q_len + for n in cutlass.range_constexpr(cute.size(scores_mn, mode=[1])): + column = coords_mn[m, n][1] + scores_mn[m, n] = ( + cutlass.Float32(scores_mn[m, n]) + + cutlass.Float32(column_masks[column]) + if valid_row + else -cutlass.Float32.inf + ) + + +@cute.jit +def mask_exact_scores( + scores: cute.Tensor, + coords: cute.Tensor, + block_len: cutlass.Int32, + q_len: cutlass.Int32, +): + scores_mn = layout_utils.reshape_acc_to_mn(scores) + coords_mn = layout_utils.reshape_acc_to_mn(coords) + for m in cutlass.range_constexpr(cute.size(scores_mn, mode=[0])): + valid_row = coords_mn[m, 0][0] < q_len + for n in cutlass.range_constexpr(cute.size(scores_mn, mode=[1])): + if (not valid_row) or coords_mn[m, n][1] >= block_len: + scores_mn[m, n] = -cutlass.Float32.inf + + +@cute.jit +def online_softmax( + scores: cute.Tensor, + row_max: cute.Tensor, + row_sum: cute.Tensor, + scale_log2e: cutlass.Float32, +): + scores_mn = layout_utils.reshape_acc_to_mn(scores) + row_scale = cute.make_rmem_tensor_like(row_max, cutlass.Float32) + for m in cutlass.range_constexpr(cute.size(row_max)): + score_row = scores_mn[m, None].load() + current_max = kernel_utils.fmax_reduce( + score_row, init_val=row_max[m], arch=80 + ) + current_max = cute.arch.warp_reduction_max( + current_max, threads_in_group=4 + ) + previous_max = row_max[m] + row_max[m] = current_max + safe_max = ( + cutlass.Float32(0.0) + if current_max == -cutlass.Float32.inf + else current_max + ) + scaled_max = safe_max * scale_log2e + probabilities = cute.math.exp2( + score_row * scale_log2e - scaled_max, fastmath=True + ) + row_scale[m] = cute.math.exp2( + (previous_max - safe_max) * scale_log2e, fastmath=True + ) + row_sum[m] = kernel_utils.fadd_reduce( + probabilities, + init_val=row_sum[m] * row_scale[m], + arch=80, + ) + scores_mn[m, None].store(probabilities) + return row_scale + + +@cute.jit +def online_softmax_route( + scores: cute.Tensor, + coords: cute.Tensor, + row_max: cute.Tensor, + row_sum: cute.Tensor, + scale_log2e: cutlass.Float32, + group_start: cutlass.Int32, + token_count: cutlass.Int32, +): + scores_mn = layout_utils.reshape_acc_to_mn(scores) + coords_mn = layout_utils.reshape_acc_to_mn(coords) + row_scale = cute.make_rmem_tensor_like(row_max, cutlass.Float32) + for m in cutlass.range_constexpr(cute.size(row_max)): + score_row = scores_mn[m, None].load() + current_max = kernel_utils.fmax_reduce( + score_row, init_val=row_max[m], arch=80 + ) + current_max = cute.arch.warp_reduction_max( + current_max, threads_in_group=4 + ) + previous_max = row_max[m] + row_max[m] = current_max + safe_max = ( + cutlass.Float32(0.0) + if current_max == -cutlass.Float32.inf + else current_max + ) + probabilities = cute.math.exp2( + score_row * scale_log2e - safe_max * scale_log2e, + fastmath=True, + ) + row_scale[m] = cute.math.exp2( + (previous_max - safe_max) * scale_log2e, fastmath=True + ) + masses = cute.make_rmem_tensor_like( + scores_mn[m, None], cutlass.Float32 + ) + for n in cutlass.range_constexpr(cute.size(masses)): + block = group_start + coords_mn[m, n][1] + length = token_count - block * cutlass.Int32(N) + if length > N: + length = cutlass.Int32(N) + if length < 0: + length = cutlass.Int32(0) + masses[n] = cutlass.Float32(probabilities[n]) * cutlass.Float32( + length + ) + row_sum[m] = kernel_utils.fadd_reduce( + masses.load(), + init_val=row_sum[m] * row_scale[m], + arch=80, + ) + scores_mn[m, None].store(probabilities) + return row_scale + + +@cute.jit +def finalize_softmax( + row_max: cute.Tensor, + row_sum: cute.Tensor, + scale_log2e: cutlass.Float32, +): + row_sum.store( + kernel_utils.warp_reduce(row_sum.load(), operator.add, width=4) + ) + ratio = cute.make_rmem_tensor_like(row_sum, cutlass.Float32) + lse = cute.make_rmem_tensor_like(row_sum, cutlass.Float32) + for m in cutlass.range_constexpr(cute.size(row_sum)): + total = row_sum[m] + invalid = total == 0.0 or total != total + ratio[m] = cute.arch.rcp_approx(total if not invalid else 1.0) + lse[m] = ( + -cutlass.Float32.inf + if invalid + else ( + row_max[m] * scale_log2e + + cute.math.log2(total, fastmath=True) + ) + * 0.6931471805599453 + ) + return ratio, lse + + +@cute.jit +def rescale_o_for_next_acc( + output: cute.Tensor, + row_scale: cute.Tensor, +): + output_mn = layout_utils.reshape_acc_to_mn(output) + for m in cutlass.range_constexpr(cute.size(row_scale)): + output_mn[m, None].store( + output_mn[m, None].load() * row_scale[m] + ) + + +__all__ = ["SolAttnForwardSm120"] diff --git a/telefuser/kernel/sol_attn/sm90/__init__.py b/telefuser/kernel/sol_attn/sm90/__init__.py new file mode 100644 index 00000000..e7538ba8 --- /dev/null +++ b/telefuser/kernel/sol_attn/sm90/__init__.py @@ -0,0 +1,5 @@ +"""Hopper backend.""" + +from .kernel import make_kernel + +__all__ = ["make_kernel"] diff --git a/telefuser/kernel/sol_attn/sm90/_compat/__init__.py b/telefuser/kernel/sol_attn/sm90/_compat/__init__.py new file mode 100644 index 00000000..7e54d6ae --- /dev/null +++ b/telefuser/kernel/sol_attn/sm90/_compat/__init__.py @@ -0,0 +1 @@ +"""Local CuteDSL helper compatibility layer for the SOL_ATTN SM90 kernel.""" diff --git a/telefuser/kernel/sol_attn/sm90/_compat/activation.py b/telefuser/kernel/sol_attn/sm90/_compat/activation.py new file mode 100644 index 00000000..cd748de0 --- /dev/null +++ b/telefuser/kernel/sol_attn/sm90/_compat/activation.py @@ -0,0 +1,5 @@ +import cutlass.cute as cute + + +def sub_packed_f32x2(a, b): + return cute.arch.add_packed_f32x2(a, (-b[0], -b[1])) diff --git a/telefuser/kernel/sol_attn/sm90/_compat/copy_utils.py b/telefuser/kernel/sol_attn/sm90/_compat/copy_utils.py new file mode 100644 index 00000000..1790f149 --- /dev/null +++ b/telefuser/kernel/sol_attn/sm90/_compat/copy_utils.py @@ -0,0 +1,169 @@ +"""CuTe copy helpers used by the Hopper mainloop.""" + +from typing import Callable + +import cutlass +import cutlass.cute as cute +from cutlass import const_expr +from cutlass import pipeline +from cutlass.cute.nvgpu import cpasync +from cutlass.cutlass_dsl import dsl_user_op + + +_RAGGED_BASE = 2**30 +_RAGGED_LIMIT = 2**31 - 1 +_RAGGED_WRAP_STRIDE = 2**64 // _RAGGED_BASE + + +@dsl_user_op +def create_ragged_tensor_for_tma( + tensor: cute.Tensor, + ragged_dim: int = 0, + ptr_shift: bool = False, + *, + loc=None, + ip=None, +) -> cute.Tensor: + rank = cute.rank(tensor) + if ragged_dim < 0: + ragged_dim += rank + if ptr_shift: + shape = ( + tensor.shape[:ragged_dim] + + (_RAGGED_BASE,) + + tensor.shape[ragged_dim + 1 :] + + (_RAGGED_LIMIT,) + ) + stride = tensor.stride + (tensor.stride[ragged_dim],) + offset = ( + (None,) * ragged_dim + + (-_RAGGED_BASE,) + + (None,) * (rank - ragged_dim - 1) + ) + pointer = cute.domain_offset(offset, tensor).iterator + return cute.make_tensor( + pointer, + cute.make_layout(shape, stride=stride), + ) + + ragged_stride = tensor.stride[ragged_dim] + shape = ( + tensor.shape[:ragged_dim] + + (_RAGGED_BASE,) + + tensor.shape[ragged_dim + 1 :] + + (_RAGGED_LIMIT, _RAGGED_LIMIT) + ) + stride = ( + tensor.stride[:ragged_dim] + + (ragged_stride,) + + tensor.stride[ragged_dim + 1 :] + + (_RAGGED_WRAP_STRIDE - ragged_stride, ragged_stride) + ) + return cute.make_tensor( + tensor.iterator, + cute.make_layout(shape, stride=stride), + ) + + +def tma_get_copy_fn( + atom: cute.CopyAtom, + cta_coord: cute.Coord, + cta_layout: cute.Layout, + src_tensor: cute.Tensor, + dst_tensor: cute.Tensor, + filter_zeros: bool = False, + single_stage: bool = False, + *, + loc=None, + ip=None, + **kwargs, +) -> Callable: + source_is_smem = const_expr( + isinstance(src_tensor.iterator, cute.Pointer) + and src_tensor.memspace == cute.AddressSpace.smem + ) + smem, gmem = ( + (src_tensor, dst_tensor) + if source_is_smem + else (dst_tensor, src_tensor) + ) + smem_rank = const_expr(cute.rank(smem) - (0 if single_stage else 1)) + gmem_rank = const_expr(cute.rank(gmem) - (0 if single_stage else 1)) + smem, gmem = cpasync.tma_partition( + atom, + cta_coord, + cta_layout, + cute.group_modes(smem, 0, smem_rank), + cute.group_modes(gmem, 0, gmem_rank), + loc=loc, + ip=ip, + ) + if const_expr(filter_zeros): + smem = cute.filter_zeros(smem) + gmem = cute.filter_zeros(gmem) + source, destination = ( + (smem, gmem) if source_is_smem else (gmem, smem) + ) + + @dsl_user_op + def copy_tma( + src_idx, + dst_idx, + *, + loc=None, + ip=None, + **call_kwargs, + ): + cute.copy( + atom, + source[None, src_idx], + destination[None, dst_idx], + **call_kwargs, + **kwargs, + loc=loc, + ip=ip, + ) + + @dsl_user_op + def copy_single_stage(*, loc=None, ip=None, **call_kwargs): + cute.copy( + atom, + source, + destination, + **call_kwargs, + **kwargs, + loc=loc, + ip=ip, + ) + + return ( + copy_tma if const_expr(not single_stage) else copy_single_stage, + smem, + gmem, + ) + + +def tma_producer_copy_fn( + copy: Callable, + copy_pipeline: pipeline.PipelineAsync, +): + def copy_fn( + src_idx, + producer_state: pipeline.PipelineState, + **kwargs, + ): + copy( + src_idx=src_idx, + dst_idx=producer_state.index, + tma_bar_ptr=copy_pipeline.producer_get_barrier(producer_state), + **kwargs, + ) + + return copy_fn + + +__all__ = [ + "create_ragged_tensor_for_tma", + "tma_get_copy_fn", + "tma_producer_copy_fn", +] diff --git a/telefuser/kernel/sol_attn/sm90/_compat/cute_dsl_utils.py b/telefuser/kernel/sol_attn/sm90/_compat/cute_dsl_utils.py new file mode 100644 index 00000000..d73002f2 --- /dev/null +++ b/telefuser/kernel/sol_attn/sm90/_compat/cute_dsl_utils.py @@ -0,0 +1,191 @@ +# Copyright (c) 2025, Tri Dao. + +from typing import Tuple, get_origin +from functools import lru_cache +from dataclasses import dataclass, fields + +import os +import re + +import torch + +try: + from triton.tools.disasm import extract +except ImportError: + extract = None + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, Int64, Float16, BFloat16, Float32 +from cutlass.base_dsl.tvm_ffi_builder import spec +from cutlass.cutlass_dsl import NumericMeta + + +StaticTypes = (cutlass.Constexpr, NumericMeta, int, bool, str, float, type(None)) + + +load_cubin_module_data_og = cutlass.base_dsl.runtime.cuda.load_cubin_module_data +cute_compile_og = cute.compile + + +# Patch TVM-FFI converter to handle Constexpr type annotations as compile-time constants. +# Fields annotated with cutlass.Constexpr[T] are emitted as ConstNone (not runtime args). +# At call time, pass None for these fields; the compile-time value is baked in. +import cutlass.cute._tvm_ffi_args_spec_converter as _converter_module # noqa + +_original_convert_single_arg = _converter_module._convert_single_arg + + +def _patched_convert_single_arg(arg, arg_name, arg_type, ctx): + if arg_type is not None and get_origin(arg_type) is cutlass.Constexpr: + return spec.ConstNone(arg_name) + # If arg is a NamedTuple but arg_type doesn't have _fields (e.g. annotated as tuple), + # redirect so the converter uses the NamedTuple's own type hints. + if ( + isinstance(arg, tuple) + and hasattr(type(arg), "_fields") + and (arg_type is None or not hasattr(arg_type, "_fields")) + ): + return _original_convert_single_arg(arg, arg_name, type(arg), ctx) + return _original_convert_single_arg(arg, arg_name, arg_type, ctx) + + +_converter_module._convert_single_arg = _patched_convert_single_arg + + +torch2cute_dtype_map = { + torch.float16: Float16, + torch.bfloat16: BFloat16, + torch.float32: Float32, + torch.int32: Int32, + torch.int64: Int64, +} + + +@lru_cache +def get_device_multiprocessor_count(device_id: int = 0) -> int: + return cutlass.utils.HardwareInfo(device_id).get_device_multiprocessor_count() + + +@lru_cache +def get_max_active_clusters( + cluster_size: int, + device_capacity: Tuple[int, int] | None = None, + device_id: int = 0, +) -> int: + if device_capacity is None: + device_capacity = get_device_capacity() + if device_capacity[0] < 9: + if cluster_size != 1: + raise ValueError("SM8x kernels do not support CTA clusters; cluster_size must be 1") + return get_device_multiprocessor_count(device_id) + return cutlass.utils.HardwareInfo(device_id).get_max_active_clusters(cluster_size=cluster_size) + + +def _parse_arch_str(arch_str: str) -> Tuple[int, int]: + """Parse arch string (e.g. 'sm_90', 'sm90', '90', 'sm_100a') to (major, minor) tuple.""" + match = re.match(r"^(?:sm_?)?(\d+)(\d)([af]?)$", arch_str.strip(), re.IGNORECASE) + if not match: + raise ValueError(f"Invalid SOL_ATTN_ARCH format: {arch_str!r} (expected e.g. '90', 'sm_90')") + major, minor, _ = match.groups() + return int(major), int(minor) + + +@lru_cache +def _get_device_capacity_cached(device: torch.device = None) -> Tuple[int, int]: + """Return (major, minor) device capability. + + Override with SOL_ATTN_ARCH (e.g. 'sm_90' or '90') for CPU-only compilation + without a GPU present. + """ + arch_override = os.environ.get("SOL_ATTN_ARCH") + if arch_override is not None: + return _parse_arch_str(arch_override) + return torch.cuda.get_device_capability(device) + + +def get_device_capacity( + device: torch.device | torch.Tensor | None = None, +) -> Tuple[int, int]: + """Return (major, minor) device capability. + + Override with SOL_ATTN_ARCH (e.g. 'sm_90' or '90') for CPU-only compilation + without a GPU present. + + Accepts either a ``torch.device`` or a tensor and canonicalizes to the + underlying device before consulting the cached helper. This avoids leaking + tensors through the LRU cache key. + """ + if isinstance(device, torch.Tensor): + device = device.device + return _get_device_capacity_cached(device) + + +def _partition_fields(obj): + """Split dataclass fields into (constexpr_dict, non_constexpr_dict) by type.""" + all_fields = {field.name: getattr(obj, field.name) for field in fields(obj)} + constexpr = {n: f for n, f in all_fields.items() if isinstance(f, StaticTypes)} + non_constexpr = {n: f for n, f in all_fields.items() if not isinstance(f, StaticTypes)} + return constexpr, non_constexpr + + +def _new_from_mlir_values(self, values): + constexpr_fields, non_constexpr_fields = _partition_fields(self) + for (name, field), n_items in zip(non_constexpr_fields.items(), self._values_pos): + non_constexpr_fields[name] = cutlass.new_from_mlir_values(field, values[:n_items]) + values = values[n_items:] + return self.__class__(**non_constexpr_fields, **constexpr_fields) + + +def _namedtuple_new_from_mlir_values(self, values): + """Generic __new_from_mlir_values__ for NamedTuples. + + Applied to NamedTuple classes via the ``@mlir_namedtuple`` decorator. + + Fields that are None or Constexpr (StaticTypes) are preserved from ``self`` (the compile-time + template). Only non-static fields consume MLIR values. Multi-value fields (e.g. cute.Tensor) + consume the correct number of values via ``cutlass.new_from_mlir_values``. + + Constexpr fields (annotated ``cutlass.Constexpr[T]``) are baked into the compiled kernel via + a converter patch (see above). At call time, pass None for these fields. + """ + from cutlass.base_dsl.typing import get_mlir_types + + values = list(values) + new_fields = [] + for field_val in self: + if field_val is None or isinstance(field_val, StaticTypes): + new_fields.append(field_val) + else: + n_items = len(get_mlir_types(field_val)) + new_fields.append(cutlass.new_from_mlir_values(field_val, values[:n_items])) + values = values[n_items:] + return self.__class__(*new_fields) + + +def mlir_namedtuple(cls): + """Decorator that adds MLIR value reconstruction to a NamedTuple class. + + Usage:: + + @mlir_namedtuple + class MyArgs(NamedTuple): + tensor_arg: cute.Tensor + const_arg: cutlass.Constexpr[int] = 0 + """ + cls.__new_from_mlir_values__ = _namedtuple_new_from_mlir_values + return cls + + +@dataclass +class ParamsBase: + def __extract_mlir_values__(self): + _, non_constexpr_fields = _partition_fields(self) + values, self._values_pos = [], [] + for obj in non_constexpr_fields.values(): + obj_values = cutlass.extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + return values + + __new_from_mlir_values__ = _new_from_mlir_values diff --git a/telefuser/kernel/sol_attn/sm90/_compat/layout_utils.py b/telefuser/kernel/sol_attn/sm90/_compat/layout_utils.py new file mode 100644 index 00000000..6b0b4875 --- /dev/null +++ b/telefuser/kernel/sol_attn/sm90/_compat/layout_utils.py @@ -0,0 +1,3 @@ +"""Compatibility import for shared CuTe layout helpers.""" + +from telefuser.kernel.sol_attn.common.layout_utils import * # noqa: F403 diff --git a/telefuser/kernel/sol_attn/sm90/_compat/sm90_utils.py b/telefuser/kernel/sol_attn/sm90/_compat/sm90_utils.py new file mode 100644 index 00000000..05185783 --- /dev/null +++ b/telefuser/kernel/sol_attn/sm90/_compat/sm90_utils.py @@ -0,0 +1,173 @@ +# Copyright (c) 2025, Tri Dao. + +from typing import Literal, Type, Union, Optional + +import cutlass +import cutlass.cute as cute +import cutlass.utils.hopper_helpers as sm90_utils_og +from cutlass.cute.nvgpu import warpgroup +from cutlass.cutlass_dsl import Numeric, dsl_user_op +from cutlass import Float32, Int32, Boolean, const_expr +from cutlass.utils import LayoutEnum + + +@dsl_user_op +def make_smem_layout( + dtype: Type[Numeric], + layout: LayoutEnum, + tile: cute.Tile, + stage: Optional[int] = None, + major_mode_size: Optional[int] = None, + *, + loc=None, + ip=None, +) -> Union[cute.Layout, cute.ComposedLayout]: + shape = cute.product_each(cute.shape(tile, loc=loc, ip=ip), loc=loc, ip=ip) + if const_expr(major_mode_size is None): + major_mode_size = shape[1] if layout.is_n_major_c() else shape[0] + smem_layout_atom = warpgroup.make_smem_layout_atom( + sm90_utils_og.get_smem_layout_atom(layout, dtype, major_mode_size), + dtype, + ) + order = (1, 0, 2) if const_expr(layout.is_m_major_c()) else (0, 1, 2) + smem_layout_staged = cute.tile_to_shape( + smem_layout_atom, + cute.append(shape, stage) if const_expr(stage is not None) else shape, + order=order if const_expr(stage is not None) else order[:2], + ) + return smem_layout_staged + + +# Shared SM90/SM100 layout helper. +make_smem_layout_epi = make_smem_layout + + +def make_tiled_mma( + a_dtype: Type[Numeric], + a_major: Literal["K", "MN"], + b_major: Literal["K", "MN"], + tiler_n: int, + source: Literal["SS", "RS"] = "SS", + atom_layout_mnk: tuple = (1, 1, 1), + swap_AB: bool = False, + b_dtype: Optional[Type[Numeric]] = None, + acc_dtype: Type[Numeric] = Float32, +) -> cute.TiledMma: + """`b_dtype` defaults to `a_dtype`; pass it for mixed-precision MMAs (e.g. fp8). + `acc_dtype` defaults to Float32.""" + if b_dtype is None: + b_dtype = a_dtype + mode = {"K": cute.nvgpu.OperandMajorMode.K, "MN": cute.nvgpu.OperandMajorMode.MN} + a_mode, b_mode = mode[a_major], mode[b_major] + if swap_AB: + a_mode, b_mode = b_mode, a_mode + a_source = warpgroup.OperandSource.RMEM if source == "RS" else warpgroup.OperandSource.SMEM + return sm90_utils_og.make_trivial_tiled_mma( + a_dtype, + b_dtype, + a_mode, + b_mode, + acc_dtype, + atom_layout_mnk=atom_layout_mnk, + tiler_mn=(64, tiler_n), + a_source=a_source, + ) + + +@cute.jit +def gemm( + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + zero_init: cutlass.Constexpr[bool] = False, + wg_wait: cutlass.Constexpr[int] = 0, + # A_in_regs: cutlass.Constexpr[bool] = False, + swap_AB: cutlass.Constexpr[bool] = False, +) -> None: + if const_expr(swap_AB): + gemm(tiled_mma, acc, tCrB, tCrA, zero_init=zero_init, wg_wait=wg_wait, swap_AB=False) + else: + warpgroup.fence() + # We make a new mma_atom since we'll be modifying its attribute (accumulate). + # Otherwise the compiler complains "operand #0 does not dominate this use" + mma_atom = cute.make_mma_atom(tiled_mma.op) + mma_atom.set(warpgroup.Field.ACCUMULATE, not zero_init) + for k in cutlass.range_constexpr(cute.size(tCrA.shape[2])): + cute.gemm(mma_atom, acc, tCrA[None, None, k], tCrB[None, None, k], acc) + mma_atom.set(warpgroup.Field.ACCUMULATE, True) + warpgroup.commit_group() + if const_expr(wg_wait >= 0): + warpgroup.wait_group(wg_wait) + + +def gemm_zero_init( + tiled_mma: cute.TiledMma, + shape: cute.Shape, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + A_idx: Optional[Int32] = None, + B_idx: Optional[Int32] = None, + wg_wait: int = -1, + swap_AB: bool = False, +) -> cute.Tensor: + if const_expr(swap_AB): + return gemm_zero_init( + tiled_mma, shape[::-1], tCrB, tCrA, B_idx, A_idx, wg_wait, swap_AB=False + ) + else: + acc = cute.make_rmem_tensor(tiled_mma.partition_shape_C(shape), Float32) + rA = tCrA if const_expr(A_idx is None) else tCrA[None, None, None, A_idx] + rB = tCrB if const_expr(B_idx is None) else tCrB[None, None, None, B_idx] + gemm(tiled_mma, acc, rA, rB, zero_init=True, wg_wait=wg_wait) + return acc + + +def gemm_w_idx( + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + zero_init: Boolean, + A_idx: Optional[Int32] = None, + B_idx: Optional[Int32] = None, + wg_wait: int = -1, + swap_AB: bool = False, +) -> None: + if const_expr(swap_AB): + gemm_w_idx(tiled_mma, acc, tCrB, tCrA, zero_init, B_idx, A_idx, wg_wait, swap_AB=False) + else: + rA = tCrA if const_expr(A_idx is None) else tCrA[None, None, None, A_idx] + rB = tCrB if const_expr(B_idx is None) else tCrB[None, None, None, B_idx] + gemm(tiled_mma, acc, rA, rB, zero_init=zero_init, wg_wait=wg_wait) + + +def partition_fragment_ABC( + thr_mma: cute.ThrMma, + shape_mnk: cute.Shape, + sA: Optional[cute.Tensor], + sB: Optional[cute.Tensor], + swap_AB: bool = False, +): + is_rs = thr_mma.op.a_src == warpgroup.OperandSource.RMEM + if const_expr(not swap_AB): + acc = cute.make_rmem_tensor(thr_mma.partition_shape_C(shape_mnk[:2]), Float32) + if const_expr(not is_rs): + assert sA is not None + tCrA = thr_mma.make_fragment_A(thr_mma.partition_A(sA)) + else: + tCrA = thr_mma.make_fragment_A(thr_mma.partition_shape_A((shape_mnk[0], shape_mnk[2]))) + assert sB is not None + tCrB = thr_mma.make_fragment_B(thr_mma.partition_B(sB)) + else: + acc = cute.make_rmem_tensor( + thr_mma.partition_shape_C((shape_mnk[1], shape_mnk[0])), Float32 + ) + if const_expr(not is_rs): + assert sB is not None + tCrB = thr_mma.make_fragment_A(thr_mma.partition_A(sB)) + else: # B in rmem + tCrB = thr_mma.make_fragment_A(thr_mma.partition_shape_A((shape_mnk[1], shape_mnk[2]))) + assert sA is not None + tCrA = thr_mma.make_fragment_B(thr_mma.partition_B(sA)) + return acc, tCrA, tCrB diff --git a/telefuser/kernel/sol_attn/sm90/atoms.py b/telefuser/kernel/sol_attn/sm90/atoms.py new file mode 100644 index 00000000..e40d2a58 --- /dev/null +++ b/telefuser/kernel/sol_attn/sm90/atoms.py @@ -0,0 +1,23 @@ +"""Hopper MMA atoms used by Sol-Attn.""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32 + +from ._compat import sm90_utils + + +def make_pv_mma(tile_m: int = 64, tile_v: int = 128) -> cute.TiledMma: + return sm90_utils.make_tiled_mma( + cutlass.BFloat16, + "K", + "MN", + tile_v, + source="RS", + atom_layout_mnk=(tile_m // 64, 1, 1), + b_dtype=cutlass.BFloat16, + acc_dtype=Float32, + ) + + +__all__ = ["make_pv_mma"] diff --git a/telefuser/kernel/sol_attn/sm90/exact.py b/telefuser/kernel/sol_attn/sm90/exact.py new file mode 100644 index 00000000..46e397f1 --- /dev/null +++ b/telefuser/kernel/sol_attn/sm90/exact.py @@ -0,0 +1,243 @@ +"""Exact-block stream for the Hopper mainloop.""" + +from __future__ import annotations + +from functools import partial +from typing import Callable, Optional + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, const_expr + +from telefuser.kernel.sol_attn.common import selector + + +@cute.jit +def _consume_exact_block( + n_block: Int32, + next_n_block: Int32, + seqlen, + producer_state, + consumer_state, + load_next: Callable, + pipeline_k, + issue_load, + mma_pv: Callable, + mma_one_n_block: Callable, + mask: Callable, + score_mod: Optional[Callable], + accumulate, + mask_seqlen: cutlass.Constexpr[bool], + is_first: cutlass.Constexpr[bool], + last_n_block: Int32, + mask_last_only: cutlass.Constexpr[bool], +): + if const_expr(mask_seqlen): + next_state = mma_one_n_block( + consumer_state, + n_block=n_block, + seqlen=seqlen, + mma_pv_fn=partial(mma_pv, zero_init=not accumulate), + mask_fn=partial(mask, mask_mod=None, mask_seqlen=True), + score_mod_fn=score_mod, + is_first_n_block=is_first, + prefetch_next=True, + next_n_block=next_n_block, + kv_producer_state=producer_state, + load_K=load_next, + issue_load=issue_load, + ) + elif const_expr(mask_last_only): + next_state = mma_one_n_block( + consumer_state, + n_block=n_block, + seqlen=seqlen, + mma_pv_fn=partial(mma_pv, zero_init=not accumulate), + mask_fn=partial(mask, mask_mod=None, mask_seqlen=False), + last_block_mask_fn=partial(mask, mask_mod=None, mask_seqlen=True), + last_n_block=last_n_block, + score_mod_fn=score_mod, + is_first_n_block=is_first, + prefetch_next=True, + next_n_block=next_n_block, + kv_producer_state=producer_state, + load_K=load_next, + issue_load=issue_load, + ) + else: + next_state = mma_one_n_block( + consumer_state, + n_block=n_block, + seqlen=seqlen, + mma_pv_fn=partial(mma_pv, zero_init=not accumulate), + mask_fn=partial(mask, mask_mod=None, mask_seqlen=False), + score_mod_fn=score_mod, + is_first_n_block=is_first, + prefetch_next=True, + next_n_block=next_n_block, + kv_producer_state=producer_state, + load_K=load_next, + issue_load=issue_load, + ) + return next_state + + +@cute.jit +def consume_exact_blocks( + mask0: Int32, + mask1: Int32, + mask2: Int32, + mask3: Int32, + group_start: Int32, + seqlen, + producer_state, + consumer_state, + load_k: Callable, + load_v: Callable, + pipeline_k, + pipeline_v, + issue_load, + mma_pv: Callable, + mma_one_n_block: Callable, + mask: Callable, + score_mod: Optional[Callable], + accumulate, + scheduler_sync: Callable, + scheduler_arrive: Callable, + mask_first: cutlass.Constexpr[bool] = True, + first_is_first: cutlass.Constexpr[bool] = True, + group_words: cutlass.Constexpr[int] = 2, + last_n_block: Int32 = Int32(-1), + mask_last_only: cutlass.Constexpr[bool] = False, + first_n_block: Int32 = Int32(0), + has_exact=False, + next_route_tile: Int32 = Int32(-1), + load_next_route: Optional[Callable] = None, +): + """Consume selected blocks while overlapping K loads with softmax and PV.""" + + processed = False + if has_exact: + first = True + pending_n_block = first_n_block + + for word in cutlass.range_constexpr(group_words): + bits = selector.sol_attn_mask_word_constexpr( + mask0, mask1, mask2, mask3, word + ) + while bits != Int32(0): + lowbit = bits & (Int32(0) - bits) + next_n_block = ( + group_start + + Int32(word * 32) + + selector.sol_attn_bfind_b32(lowbit) + ) + + if issue_load: + pipeline_v.producer_acquire(producer_state) + load_v(src_idx=pending_n_block, producer_state=producer_state) + producer_state.advance() + + if first: + scheduler_sync() + consumer_state = _consume_exact_block( + pending_n_block, + next_n_block, + seqlen, + producer_state, + consumer_state, + load_k, + pipeline_k, + issue_load, + mma_pv, + mma_one_n_block, + mask, + score_mod, + accumulate, + mask_first, + first_is_first, + last_n_block, + mask_last_only, + ) + else: + consumer_state = _consume_exact_block( + pending_n_block, + next_n_block, + seqlen, + producer_state, + consumer_state, + load_k, + pipeline_k, + issue_load, + mma_pv, + mma_one_n_block, + mask, + score_mod, + accumulate, + False, + False, + last_n_block, + mask_last_only, + ) + accumulate = True + processed = True + first = False + pending_n_block = next_n_block + bits = bits & (bits - Int32(1)) + + if issue_load: + pipeline_v.producer_acquire(producer_state) + load_v(src_idx=pending_n_block, producer_state=producer_state) + producer_state.advance() + + if first: + scheduler_sync() + consumer_state = _consume_exact_block( + pending_n_block, + next_route_tile, + seqlen, + producer_state, + consumer_state, + load_next_route, + pipeline_k, + issue_load, + mma_pv, + mma_one_n_block, + mask, + score_mod, + accumulate, + mask_first, + first_is_first, + last_n_block, + mask_last_only, + ) + else: + consumer_state = _consume_exact_block( + pending_n_block, + next_route_tile, + seqlen, + producer_state, + consumer_state, + load_next_route, + pipeline_k, + issue_load, + mma_pv, + mma_one_n_block, + mask, + score_mod, + accumulate, + False, + False, + last_n_block, + mask_last_only, + ) + accumulate = True + processed = True + + if processed: + scheduler_arrive() + + return producer_state, consumer_state, accumulate, processed + + +__all__ = ["consume_exact_blocks"] diff --git a/telefuser/kernel/sol_attn/sm90/fwd.py b/telefuser/kernel/sol_attn/sm90/fwd.py new file mode 100644 index 00000000..453ab94d --- /dev/null +++ b/telefuser/kernel/sol_attn/sm90/fwd.py @@ -0,0 +1,111 @@ +"""Hopper forward operators.""" + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute + +from .mainloop import SolAttnMainloopSm90 +from .split_combine import BlockSparseAttnForwardCombine + + +class SolAttnSplitForwardSm90(SolAttnMainloopSm90): + """Run one CTA per KV split and merge the partial outputs.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._combine = BlockSparseAttnForwardCombine( + dtype=cutlass.BFloat16, + head_dim=self.tile_hdimv, + tile_m=16, + k_block_size=64, + log_max_splits=1 if self.sol_attn_num_splits == 2 else 2, + num_threads=128, + stages=4, + partial_dtype=cutlass.BFloat16, + ) + + @cute.jit + def __call__( + self, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + o: cute.Tensor, + kc: cute.Tensor, + vc: cute.Tensor, + threshold: cute.Tensor, + lse: cute.Tensor, + o_partial: cute.Tensor, + lse_partial: cute.Tensor, + softmax_scale: cutlass.Float32, + sink_range: cutlass.Int32, + stream: cuda.CUstream = None, + ): + SolAttnMainloopSm90.__call__( + self, + q, + k, + v, + o_partial, + kc, + vc, + threshold, + lse_partial, + softmax_scale, + sink_range, + stream=stream, + ) + + splits = self.sol_attn_num_splits + heads = q.shape[2] + o_partial = cute.make_tensor( + o_partial.iterator, + cute.make_layout( + ( + splits, + o_partial.shape[0], + o_partial.shape[1], + heads, + o_partial.shape[3], + ), + stride=( + heads * o_partial.stride[2], + o_partial.stride[0], + o_partial.stride[1], + o_partial.stride[2], + o_partial.stride[3], + ), + ), + ) + lse_partial = cute.make_tensor( + lse_partial.iterator, + cute.make_layout( + ( + splits, + lse_partial.shape[0], + lse_partial.shape[1], + heads, + ), + stride=( + heads * lse_partial.stride[2], + lse_partial.stride[0], + lse_partial.stride[1], + lse_partial.stride[2], + ), + ), + ) + self._combine( + o_partial, + lse_partial, + o, + lse, + None, + None, + None, + None, + None, + stream, + ) + + +__all__ = ["SolAttnSplitForwardSm90"] diff --git a/telefuser/kernel/sol_attn/sm90/kernel.py b/telefuser/kernel/sol_attn/sm90/kernel.py new file mode 100644 index 00000000..c0b3fb7b --- /dev/null +++ b/telefuser/kernel/sol_attn/sm90/kernel.py @@ -0,0 +1,47 @@ +"""Frozen Hopper kernel recipe.""" + +import cutlass + +from .fwd import SolAttnSplitForwardSm90 +from .mainloop import SolAttnMainloopSm90 + + +def make_kernel(tokens: int, kv_splits: int): + blocks = (tokens + 63) // 64 + full_groups, tail = divmod(blocks, 64) + has_full_groups = tail == 0 + has_full_blocks = tokens % 64 == 0 + kernel = SolAttnMainloopSm90 if kv_splits == 1 else SolAttnSplitForwardSm90 + + return kernel( + cutlass.BFloat16, + head_dim=128, + head_dim_v=128, + qhead_per_kvhead=1, + is_causal=False, + is_local=False, + pack_gqa=False, + tile_m=64, + tile_n=64, + num_stages=1, + num_threads=128, + sol_attn_assume_lane_group_route_reduce=( + has_full_blocks and has_full_groups + ), + sol_attn_assume_full_k_exact_blocks=has_full_blocks, + sol_attn_tail_exact_words1=0 < tail <= 8, + sol_attn_assume_full_route_groups=has_full_groups, + sol_attn_static_num_full_route_groups=( + -1 if has_full_groups else full_groups + ), + sol_attn_static_tail_valid_count=(-1 if has_full_groups else tail), + sol_attn_tail_physical_tile16=0 < tail <= 16, + sol_attn_exact_mask_seqlen_last_only=( + not has_full_blocks + ), + sol_attn_tail16_lane_group_route_reduce=tail == 16, + sol_attn_num_splits=kv_splits, + ) + + +__all__ = ["make_kernel"] diff --git a/telefuser/kernel/sol_attn/sm90/mainloop.py b/telefuser/kernel/sol_attn/sm90/mainloop.py new file mode 100644 index 00000000..47ba69f4 --- /dev/null +++ b/telefuser/kernel/sol_attn/sm90/mainloop.py @@ -0,0 +1,2312 @@ +# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. +# SM90 (Hopper) forward pass for flash attention, extracted from flash_fwd.py. + +from types import SimpleNamespace +from typing import Callable, Optional +from functools import partial + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, const_expr +from cutlass.cute.nvgpu import cpasync, warpgroup +from cutlass.utils import LayoutEnum +import cutlass.utils.hopper_helpers as sm90_utils_basic +from cutlass import pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +from cutlass.base_dsl.arch import Arch + +from ._compat import copy_utils +from ._compat import layout_utils +from ._compat import sm90_utils + +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.cute_dsl_utils import assume_tensor_aligned +from telefuser.kernel.sol_attn._vendor.flash_attn.cute import utils +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.mask import AttentionMask +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.softmax import Softmax, apply_score_mod_inner +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.seqlen_info import SeqlenInfoQK +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.block_info import BlockInfo +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.block_sparsity import BlockSparseTensors +from telefuser.kernel.sol_attn._vendor.flash_attn.cute import pipeline as pipeline_custom +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.pack_gqa import PackGQA, pack_gqa_layout, make_packgqa_tiled_tma_atom +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.named_barrier import NamedBarrierFwd +from ._compat.cute_dsl_utils import ParamsBase +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.tile_scheduler import ( + TileSchedulerArguments, + SingleTileScheduler, + SingleTileLPTScheduler, + SingleTileVarlenScheduler, +) +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.flash_fwd import FlashAttentionForwardBase +from . import atoms as sol_attn_atoms +from . import exact as exact_stream +from telefuser.kernel.sol_attn.common import selector as sol_attn_selector + + +SOL_ATTN_ROUTE_MASK_BARRIER_ID = 7 +SOL_ATTN_ROUTE_SUM_BARRIER_ID = 8 + + +class SolAttnMainloopSm90(FlashAttentionForwardBase): + def __init__( + self, + *args, + sol_attn_assume_lane_group_route_reduce: bool = False, + sol_attn_assume_full_k_exact_blocks: bool = False, + sol_attn_tail_exact_words1: bool = False, + sol_attn_assume_full_route_groups: bool = False, + sol_attn_static_num_full_route_groups: int = -1, + sol_attn_static_tail_valid_count: int = -1, + sol_attn_tail_physical_tile16: bool = False, + sol_attn_exact_mask_seqlen_last_only: bool = False, + sol_attn_tail16_lane_group_route_reduce: bool = False, + sol_attn_num_splits: int = 1, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.qk_dtype = cutlass.BFloat16 + self.pv_dtype = self.dtype + self.sol_attn_group_size = 64 + self.sol_attn_group_words = 2 + self.mma_pv_is_rs = True + self.sol_attn_mma_regs_override = 128 + self.sol_attn_warp_route_mask = True + self.sol_attn_fast_route_lens = True + self.sol_attn_early_route_mask_publish = True + self.sol_attn_lane_group_route_reduce = True + self.sol_attn_assume_lane_group_route_reduce = sol_attn_assume_lane_group_route_reduce + self.sol_attn_assume_full_k_exact_blocks = sol_attn_assume_full_k_exact_blocks + self.sol_attn_route_sum_arrive_overlap = True + self.sol_attn_route_mask_after_scale = True + self.sol_attn_assume_full_route_groups = sol_attn_assume_full_route_groups + self.sol_attn_static_num_full_route_groups = sol_attn_static_num_full_route_groups + self.sol_attn_static_tail_valid_count = sol_attn_static_tail_valid_count + self.sol_attn_tail_exact_words1 = ( + sol_attn_tail_exact_words1 and 0 < self.sol_attn_static_tail_valid_count <= 32 + ) + self.sol_attn_tail_route_mask_words1 = False + self.sol_attn_tail_physical_tile16 = ( + sol_attn_tail_physical_tile16 and 0 < self.sol_attn_static_tail_valid_count <= 16 + ) + self.sol_attn_exact_mask_seqlen_last_only = sol_attn_exact_mask_seqlen_last_only + self.sol_attn_full_route_mask_seqlen_false = True + self.sol_attn_tail16_lane_group_route_reduce = ( + sol_attn_tail16_lane_group_route_reduce and self.sol_attn_tail_physical_tile16 + ) + self.sol_attn_full_block_row_sum_prescale = False + self.sol_attn_neutral_softmax_state = True + self.sol_attn_assume_nonempty_rows = False + self.sol_attn_ballot_mask = True + self.sol_attn_approx_colmask = False + self.sol_attn_packed_route_reduction = False + self.sol_attn_num_splits = sol_attn_num_splits + self.buffer_align_bytes = 1024 + self.use_tma_KV = True + self.cluster_shape_mn = (1, 1) + if not (self.arch >= Arch.sm_90 and self.arch <= Arch.sm_90a): + raise AssertionError("The Hopper backend requires SM90") + + def _get_smem_layout_atom(self): + sQ_layout_atom = warpgroup.make_smem_layout_atom( + sm90_utils_basic.get_smem_layout_atom( + LayoutEnum.ROW_MAJOR, self.qk_dtype, self.tile_hdim + ), + self.qk_dtype, + ) + sK_layout_atom = sQ_layout_atom + sV_layout_atom = warpgroup.make_smem_layout_atom( + sm90_utils_basic.get_smem_layout_atom( + LayoutEnum.ROW_MAJOR, self.pv_dtype, self.tile_hdimv + ), + self.pv_dtype, + ) + sO_layout_atom = sV_layout_atom + if not self.mma_pv_is_rs: + sP_layout_atom = warpgroup.make_smem_layout_atom( + sm90_utils_basic.get_smem_layout_atom( + LayoutEnum.ROW_MAJOR, self.pv_dtype, self.tile_n + ), + self.pv_dtype, + ) + else: + sP_layout_atom = None + return sQ_layout_atom, sK_layout_atom, sV_layout_atom, sO_layout_atom, sP_layout_atom + + def _get_tiled_mma(self): + tiled_mma_qk = sm90_utils.make_tiled_mma( + cutlass.BFloat16, + "K", + "K", + self.tile_n, + source="SS", + atom_layout_mnk=(self.tile_m // 64, 1, 1), + b_dtype=cutlass.BFloat16, + acc_dtype=Float32, + ) + tiled_mma_pv = sol_attn_atoms.make_pv_mma( + tile_m=self.tile_m, + tile_v=self.tile_hdimv, + ) + return tiled_mma_qk, tiled_mma_pv + + @cute.jit + def sol_attn_qk_gemm_zero_init( + self, + tiled_mma: cute.TiledMma, + shape: cute.Shape, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + A_idx: Optional[Int32] = None, + B_idx: Optional[Int32] = None, + wg_wait: int = -1, + swap_AB: bool = False, + ) -> cute.Tensor: + """Run BF16 QK WGMMA directly into an FP32 accumulator.""" + + return sm90_utils.gemm_zero_init( + tiled_mma, + shape, + tCrA, + tCrB, + A_idx, + B_idx, + wg_wait, + swap_AB, + ) + + def _get_shared_storage_cls(self): + sQ_struct, sK_struct = [ + cute.struct.Align[ + cute.struct.MemRange[self.qk_dtype, cute.cosize(layout)], self.buffer_align_bytes + ] + for layout in (self.sQ_layout, self.sK_layout) + ] + sV_struct = cute.struct.Align[ + cute.struct.MemRange[self.pv_dtype, cute.cosize(self.sV_layout)], + self.buffer_align_bytes, + ] + cosize_sQV = max(cute.cosize(self.sQ_layout), cute.cosize(self.sV_layout)) + sQV_struct = cute.struct.Align[cute.struct.MemRange[self.pv_dtype, cosize_sQV], 1024] + cosize_sP = cute.cosize(self.sP_layout) if const_expr(self.sP_layout is not None) else 0 + sP_struct = cute.struct.Align[cute.struct.MemRange[self.pv_dtype, cosize_sP], 1024] + route_mask_struct = cute.struct.Align[ + cute.struct.MemRange[Int32, 4], 16 + ] + route_sums_struct = cute.struct.Align[ + cute.struct.MemRange[Float32, 4 * self.tile_n], 16 + ] + # 1 stage * 2 for Q pipeline (full + empty), self.num_stages*2 for K, self.num_stages*2 for V, + mbar_ptr_Q_struct = cute.struct.MemRange[cutlass.Int64, 1 * 2] + mbar_ptr_K_struct = cute.struct.MemRange[cutlass.Int64, self.num_stages * 2] + mbar_ptr_V_struct = cute.struct.MemRange[cutlass.Int64, self.num_stages * 2] + + @cute.struct + class SharedStorageQKV: + mbar_ptr_Q: mbar_ptr_Q_struct + mbar_ptr_K: mbar_ptr_K_struct + mbar_ptr_V: mbar_ptr_V_struct + sV: sV_struct + sQ: sQ_struct + sK: sK_struct + sP: sP_struct + route_mask: route_mask_struct + route_sums: route_sums_struct + + @cute.struct + class SharedStorageSharedQV: + mbar_ptr_Q: mbar_ptr_Q_struct + mbar_ptr_K: mbar_ptr_K_struct + mbar_ptr_V: mbar_ptr_V_struct + sQ: sQV_struct + sK: sK_struct + sP: sP_struct + route_mask: route_mask_struct + route_sums: route_sums_struct + + return SharedStorageQKV if const_expr(not self.Q_in_regs) else SharedStorageSharedQV + + @cute.jit + def sol_attn_reduce_route_sums_lane_group( + self, + acc_S_mn: cute.Tensor, + route_sums: cute.Tensor, + warp_in_mma: Int32, + lane: Int32, + ): + """Reduce route columns using the observed SM90 accumulator lane layout.""" + + for col_group in cutlass.range_constexpr(self.tile_n // 8): + base = col_group * 4 + partial0 = Float32(acc_S_mn[base]) + Float32(acc_S_mn[base + 1]) + partial1 = Float32(acc_S_mn[base + 2]) + Float32(acc_S_mn[base + 3]) + partial0 += cute.arch.shuffle_sync_down(partial0, 16) + partial1 += cute.arch.shuffle_sync_down(partial1, 16) + partial0 += cute.arch.shuffle_sync_down(partial0, 8) + partial1 += cute.arch.shuffle_sync_down(partial1, 8) + partial0 += cute.arch.shuffle_sync_down(partial0, 4) + partial1 += cute.arch.shuffle_sync_down(partial1, 4) + if lane < Int32(4): + col = Int32(8 * col_group) + lane * Int32(2) + route_sums[warp_in_mma, col] = partial0 + route_sums[warp_in_mma, col + Int32(1)] = partial1 + + @cute.jit + def sol_attn_reduce_route_sums_guarded( + self, + acc_S_mn: cute.Tensor, + route_sums: cute.Tensor, + tScS_mn: cute.Tensor, + q_start: Int32, + seqlen: SeqlenInfoQK, + warp_in_mma: Int32, + lane: Int32, + ): + """Fallback route-column reduction that ignores invalid q rows.""" + + for off in cutlass.range_constexpr(self.tile_n): + partial = Float32(0.0) + for i in cutlass.range(cute.size(acc_S_mn), unroll_full=True): + row = tScS_mn[i][0] + col = tScS_mn[i][1] + valid_row = q_start + row < seqlen.seqlen_q + if col == Int32(off) and valid_row: + partial += Float32(acc_S_mn[i]) + warp_sum = cute.arch.warp_reduction_sum(partial) + if lane == Int32(0): + route_sums[warp_in_mma, off] = warp_sum + + @cute.jit + def sol_attn_reduce_route_sums_lane_group_tail16( + self, + acc_S_mn: cute.Tensor, + route_sums: cute.Tensor, + route_col_offset: Int32, + warp_in_mma: Int32, + lane: Int32, + ): + """Reduce a physical 16-column tail route tile using the accumulator lane layout.""" + + for col_group in cutlass.range_constexpr(2): + base = col_group * 4 + partial0 = Float32(acc_S_mn[base]) + Float32(acc_S_mn[base + 1]) + partial1 = Float32(acc_S_mn[base + 2]) + Float32(acc_S_mn[base + 3]) + partial0 += cute.arch.shuffle_sync_down(partial0, 16) + partial1 += cute.arch.shuffle_sync_down(partial1, 16) + partial0 += cute.arch.shuffle_sync_down(partial0, 8) + partial1 += cute.arch.shuffle_sync_down(partial1, 8) + partial0 += cute.arch.shuffle_sync_down(partial0, 4) + partial1 += cute.arch.shuffle_sync_down(partial1, 4) + if lane < Int32(4): + col = route_col_offset + Int32(8 * col_group) + lane * Int32(2) + route_sums[warp_in_mma, col] = partial0 + route_sums[warp_in_mma, col + Int32(1)] = partial1 + + @cute.jit + def sol_attn_reduce_route_sums_static_tail( + self, + acc_S_mn: cute.Tensor, + route_sums: cute.Tensor, + tScS_mn: cute.Tensor, + q_start: Int32, + route_col_offset: Int32, + seqlen: SeqlenInfoQK, + warp_in_mma: Int32, + lane: Int32, + full_q_tile: bool, + ): + """Reduce only the compile-time-known valid columns of a static tail route group.""" + + for off in cutlass.range_constexpr(self.sol_attn_static_tail_valid_count): + route_col = route_col_offset + Int32(off) + partial = Float32(0.0) + for i in cutlass.range(cute.size(acc_S_mn), unroll_full=True): + row = tScS_mn[i][0] + col = tScS_mn[i][1] + valid_row = True + if not full_q_tile: + valid_row = q_start + row < seqlen.seqlen_q + if col == route_col and valid_row: + partial += Float32(acc_S_mn[i]) + warp_sum = cute.arch.warp_reduction_sum(partial) + if lane == Int32(0): + route_sums[warp_in_mma, route_col] = warp_sum + + @cute.jit + def sol_attn_reduce_route_sums_physical16( + self, + acc_S_mn: cute.Tensor, + route_sums: cute.Tensor, + tScS_mn: cute.Tensor, + q_start: Int32, + route_col_offset: Int32, + seqlen: SeqlenInfoQK, + warp_in_mma: Int32, + lane: Int32, + full_q_tile: bool, + ): + """Reduce a physically 16-column route tile into the 64-column route_sums buffer.""" + + for off in cutlass.range_constexpr(16): + route_col = route_col_offset + Int32(off) + partial = Float32(0.0) + for i in cutlass.range(cute.size(acc_S_mn), unroll_full=True): + row = tScS_mn[i][0] + col = tScS_mn[i][1] + valid_row = True + if not full_q_tile: + valid_row = q_start + row < seqlen.seqlen_q + if col == route_col and valid_row: + partial += Float32(acc_S_mn[i]) + warp_sum = cute.arch.warp_reduction_sum(partial) + if lane == Int32(0): + route_sums[warp_in_mma, route_col] = warp_sum + + @cute.jit + def sol_attn_build_route_mask_from_acc( + self, + acc_S: cute.Tensor, + route_sums: cute.Tensor, + tScS_mn: cute.Tensor, + m_block: Int32, + group_start_n_block: Int32, + valid_count: Int32, + route_col_offset: Int32, + seqlen: SeqlenInfoQK, + batch_idx: Int32, + head_idx: Int32, + mGlobalThresh: cute.Tensor, + softmax_scale_log2: Float32, + sink_range: Int32, + assume_full_route_group: cutlass.Constexpr[bool] = False, + physical_route_tile_n: cutlass.Constexpr[int] = 64, + route_mask_words_override: cutlass.Constexpr[int] = 0, + ): + """Build the exact mask from the distributed route QK accumulator. + + WGMMA accumulators are distributed across the 128 consumer threads. + Each consumer warp first reduces its local contribution per route + column, then one consumer thread combines the four warp partials into + the CTA-local bitmask. + """ + + tidx, _, _ = cute.arch.thread_idx() + consumer_tidx = tidx + warp_in_mma = consumer_tidx // cute.arch.WARP_SIZE + lane = cute.arch.lane_idx() + acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S) + + q_start = m_block * self.tile_m + q_len_i32 = seqlen.seqlen_q - q_start + if q_len_i32 > Int32(self.tile_m): + q_len_i32 = Int32(self.tile_m) + q_len = Float32(q_len_i32) + + full_q_tile = q_len_i32 == Int32(self.tile_m) + if const_expr(self.sol_attn_tail16_lane_group_route_reduce and physical_route_tile_n == 16): + if full_q_tile: + self.sol_attn_reduce_route_sums_lane_group_tail16( + acc_S_mn, + route_sums, + route_col_offset, + warp_in_mma, + lane, + ) + else: + self.sol_attn_reduce_route_sums_static_tail( + acc_S_mn, + route_sums, + tScS_mn, + q_start, + route_col_offset, + seqlen, + warp_in_mma, + lane, + full_q_tile, + ) + elif const_expr(physical_route_tile_n == 16): + self.sol_attn_reduce_route_sums_physical16( + acc_S_mn, + route_sums, + tScS_mn, + q_start, + route_col_offset, + seqlen, + warp_in_mma, + lane, + full_q_tile, + ) + elif const_expr(self.sol_attn_assume_lane_group_route_reduce): + self.sol_attn_reduce_route_sums_lane_group(acc_S_mn, route_sums, warp_in_mma, lane) + elif const_expr(self.sol_attn_lane_group_route_reduce): + if full_q_tile: + self.sol_attn_reduce_route_sums_lane_group(acc_S_mn, route_sums, warp_in_mma, lane) + else: + self.sol_attn_reduce_route_sums_guarded( + acc_S_mn, + route_sums, + tScS_mn, + q_start, + seqlen, + warp_in_mma, + lane, + ) + else: + for off in cutlass.range_constexpr(self.tile_n): + partial = Float32(0.0) + for i in cutlass.range(cute.size(acc_S_mn), unroll_full=True): + row = tScS_mn[i][0] + col = tScS_mn[i][1] + valid_row = q_start + row < seqlen.seqlen_q + if col == Int32(off) and valid_row: + partial += Float32(acc_S_mn[i]) + warp_sum = cute.arch.warp_reduction_sum(partial) + if lane == Int32(0): + route_sums[warp_in_mma, off] = warp_sum + + if const_expr(self.sol_attn_route_sum_arrive_overlap and self.sol_attn_warp_route_mask): + if warp_in_mma == Int32(0): + cute.arch.barrier( + barrier_id=SOL_ATTN_ROUTE_SUM_BARRIER_ID, + number_of_threads=self.num_mma_threads, + ) + else: + cute.arch.barrier_arrive( + barrier_id=SOL_ATTN_ROUTE_SUM_BARRIER_ID, + number_of_threads=self.num_mma_threads, + ) + else: + cute.arch.barrier( + barrier_id=SOL_ATTN_ROUTE_SUM_BARRIER_ID, + number_of_threads=self.num_mma_threads, + ) + + mask0 = Int32(0) + mask1 = Int32(0) + mask2 = Int32(0) + mask3 = Int32(0) + thresh = Float32(mGlobalThresh[m_block, head_idx, batch_idx]) + if const_expr(self.sol_attn_warp_route_mask): + route_mask_words = self.sol_attn_group_words + if const_expr(route_mask_words_override != 0): + route_mask_words = route_mask_words_override + if const_expr(self.sol_attn_tail_route_mask_words1 and not assume_full_route_group): + route_mask_words = 1 + build_mask = warp_in_mma == Int32(0) + if build_mask: + sink_enabled = sink_range != Int32(0) + sink_start_block = sink_range & Int32(0xFFFF) + sink_end_block = (sink_range >> Int32(16)) & Int32(0xFFFF) + if const_expr(self.sol_attn_packed_route_reduction): + off0 = lane + off1 = Int32(32) + lane + route_col0 = route_col_offset + off0 + route_col1 = route_col_offset + off1 + col_sum0 = Float32(route_sums[0, route_col0]) + Float32( + route_sums[1, route_col0] + ) + col_sum1 = Float32(route_sums[0, route_col1]) + Float32( + route_sums[1, route_col1] + ) + col_sum0 += Float32(route_sums[2, route_col0]) + col_sum1 += Float32(route_sums[2, route_col1]) + col_sum0 += Float32(route_sums[3, route_col0]) + col_sum1 += Float32(route_sums[3, route_col1]) + col_mean0 = col_sum0 * softmax_scale_log2 / q_len + col_mean1 = col_sum1 * softmax_scale_log2 / q_len + exact0 = sol_attn_selector.sol_attn_route_is_exact( + m_block, + group_start_n_block + off0, + col_mean0, + thresh, + True, + ) + exact1 = sol_attn_selector.sol_attn_route_is_exact( + m_block, + group_start_n_block + off1, + col_mean1, + thresh, + True, + ) + if sink_enabled: + exact0 = exact0 or ( + group_start_n_block + off0 >= sink_start_block + and group_start_n_block + off0 < sink_end_block + ) + exact1 = exact1 or ( + group_start_n_block + off1 >= sink_start_block + and group_start_n_block + off1 < sink_end_block + ) + word_bits0 = Int32(cute.arch.vote_ballot_sync(exact0)) + word_bits1 = Int32(cute.arch.vote_ballot_sync(exact1)) + if lane == Int32(0): + mask0 = word_bits0 + mask1 = word_bits1 + else: + for word in cutlass.range_constexpr(route_mask_words): + off = Int32(word * 32) + lane + route_col = route_col_offset + off + valid = True + if const_expr( + not ( + self.sol_attn_assume_full_route_groups + or assume_full_route_group + ) + ): + valid = off < valid_count + exact = False + if valid: + col_sum = ( + Float32(route_sums[0, route_col]) + + Float32(route_sums[1, route_col]) + + Float32(route_sums[2, route_col]) + + Float32(route_sums[3, route_col]) + ) + col_mean = col_sum * softmax_scale_log2 / q_len + exact = sol_attn_selector.sol_attn_route_is_exact( + m_block, + group_start_n_block + off, + col_mean, + thresh, + valid, + ) + if sink_enabled: + exact = exact or ( + group_start_n_block + off + >= sink_start_block + and group_start_n_block + off + < sink_end_block + ) + if const_expr(self.sol_attn_approx_colmask): + column_mask = -Float32.inf + if valid and not exact: + column_mask = Float32(0.0) + route_sums[0, route_col] = column_mask + if const_expr(self.sol_attn_ballot_mask): + word_bits = Int32(cute.arch.vote_ballot_sync(exact)) + else: + word_bits = Int32(0) + if exact: + word_bits = Int32(1) << lane + word_bits = word_bits | cute.arch.shuffle_sync_down( + word_bits, 16 + ) + word_bits = word_bits | cute.arch.shuffle_sync_down( + word_bits, 8 + ) + word_bits = word_bits | cute.arch.shuffle_sync_down( + word_bits, 4 + ) + word_bits = word_bits | cute.arch.shuffle_sync_down( + word_bits, 2 + ) + word_bits = word_bits | cute.arch.shuffle_sync_down( + word_bits, 1 + ) + if lane == Int32(0): + if const_expr(word == 0): + mask0 = word_bits + elif const_expr(word == 1): + mask1 = word_bits + elif const_expr(word == 2): + mask2 = word_bits + else: + mask3 = word_bits + else: + sink_enabled = sink_range != Int32(0) + sink_start_block = sink_range & Int32(0xFFFF) + sink_end_block = (sink_range >> Int32(16)) & Int32(0xFFFF) + for off in cutlass.range_constexpr(self.sol_attn_group_size): + route_col = route_col_offset + Int32(off) + valid = True + if const_expr( + not (self.sol_attn_assume_full_route_groups or assume_full_route_group) + ): + valid = Int32(off) < valid_count + col_sum = ( + Float32(route_sums[0, route_col]) + + Float32(route_sums[1, route_col]) + + Float32(route_sums[2, route_col]) + + Float32(route_sums[3, route_col]) + ) + if valid: + col_mean = col_sum * softmax_scale_log2 / q_len + exact = sol_attn_selector.sol_attn_route_is_exact( + m_block, + group_start_n_block + Int32(off), + col_mean, + thresh, + valid, + ) + if sink_enabled: + exact = exact or ( + group_start_n_block + Int32(off) + >= sink_start_block + and group_start_n_block + Int32(off) + < sink_end_block + ) + if exact: + mask0, mask1, mask2, mask3 = ( + sol_attn_selector.sol_attn_set_exact_bit( + mask0, mask1, mask2, mask3, Int32(off) + ) + ) + + return mask0, mask1, mask2, mask3 + + @cute.jit + def sol_attn_mask_route_approx_columns( + self, + acc_S: cute.Tensor, + route_sums: cute.Tensor, + tScS_mn: cute.Tensor, + valid_count: Int32, + route_col_offset: Int32, + mask0: Int32, + mask1: Int32, + mask2: Int32, + mask3: Int32, + assume_full_route_group: cutlass.Constexpr[bool] = False, + route_mask_words_override: cutlass.Constexpr[int] = 0, + ): + """Keep only approximate route columns in the route score tile.""" + + acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S) + if const_expr(self.sol_attn_approx_colmask): + for i in cutlass.range(cute.size(acc_S_mn), unroll_full=True): + col = tScS_mn[i][1] + acc_S_mn[i] = Float32(acc_S_mn[i]) + Float32(route_sums[0, col]) + else: + for i in cutlass.range(cute.size(acc_S_mn), unroll_full=True): + col = tScS_mn[i][1] + group_col = col - route_col_offset + valid = True + if const_expr(self.sol_attn_group_size != self.tile_n): + valid = group_col >= Int32(0) + if valid: + valid = group_col < valid_count + elif const_expr( + not (self.sol_attn_assume_full_route_groups or assume_full_route_group) + ): + valid = col < valid_count + exact = False + if valid: + route_mask_words = self.sol_attn_group_words + if const_expr(route_mask_words_override != 0): + route_mask_words = route_mask_words_override + if const_expr( + self.sol_attn_tail_route_mask_words1 + and not assume_full_route_group + ): + route_mask_words = 1 + exact = sol_attn_selector.sol_attn_test_exact_bit_limited_words( + mask0, mask1, mask2, mask3, group_col, route_mask_words + ) + if (not valid) or exact: + acc_S_mn[i] = -Float32.inf + + @cute.jit + def sol_attn_expand_route_acc_to_full_tile( + self, + acc_S: cute.Tensor, + acc_S_full_ref: cute.Tensor, + tScS_mn: cute.Tensor, + tScS_full_mn: cute.Tensor, + ) -> cute.Tensor: + """Expand a narrow physical route accumulator into a full 64-column P tile.""" + + acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S) + acc_S_full = cute.make_rmem_tensor_like(acc_S_full_ref, Float32) + acc_S_full_mn = layout_utils.reshape_acc_to_mn(acc_S_full) + for i in cutlass.range(cute.size(acc_S_full_mn), unroll_full=True): + row = tScS_full_mn[i][0] + col = tScS_full_mn[i][1] + value = -Float32.inf + if col < Int32(16): + for j in cutlass.range(cute.size(acc_S_mn), unroll_full=True): + row16 = tScS_mn[j][0] + col16 = tScS_mn[j][1] + if row == row16 and col == col16: + value = Float32(acc_S_mn[j]) + acc_S_full_mn[i] = value + return acc_S_full + + @cute.jit + def sol_attn_expand_route_prob_to_full_tile( + self, + acc_P: cute.Tensor, + acc_S_full_ref: cute.Tensor, + tScS_mn: cute.Tensor, + tScS_full_mn: cute.Tensor, + ) -> cute.Tensor: + """Expand a compact route probability tile into the full PV A fragment.""" + + acc_P_mn = layout_utils.reshape_acc_to_mn(acc_P) + acc_P_full = cute.make_rmem_tensor_like(acc_S_full_ref, Float32) + acc_P_full_mn = layout_utils.reshape_acc_to_mn(acc_P_full) + for i in cutlass.range(cute.size(acc_P_full_mn), unroll_full=True): + row = tScS_full_mn[i][0] + col = tScS_full_mn[i][1] + value = Float32(0.0) + if col < Int32(16): + for j in cutlass.range(cute.size(acc_P_mn), unroll_full=True): + row16 = tScS_mn[j][0] + col16 = tScS_mn[j][1] + if row == row16 and col == col16: + value = Float32(acc_P_mn[j]) + acc_P_full_mn[i] = value + return acc_P_full + + @cute.jit + def sol_attn_apply_route_current_lens_to_row_sum( + self, + acc_S: cute.Tensor, + tScS_mn: cute.Tensor, + group_start_n_block: Int32, + valid_count: Int32, + route_col_offset: Int32, + seqlen: SeqlenInfoQK, + softmax: Softmax, + ): + """Correct route approx denominator for VC tiles that are block sums.""" + + acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S) + last_n_block = ( + (seqlen.seqlen_k + Int32(self.tile_n - 1)) // Int32(self.tile_n) + ) - Int32(1) + tail_len = seqlen.seqlen_k - last_n_block * Int32(self.tile_n) + for r in cutlass.range(cute.size(softmax.row_sum), unroll_full=True): + extra = Float32(0.0) + for c in cutlass.range(cute.size(acc_S_mn.shape[1]), unroll_full=True): + col = tScS_mn[r, c][1] + group_col = col - route_col_offset + valid = group_col >= Int32(0) + if valid: + valid = group_col < valid_count + if valid: + kv_block_idx = group_start_n_block + group_col + current_len = Int32(self.tile_n) + if kv_block_idx == last_n_block: + current_len = tail_len + extra += Float32(acc_S_mn[r, c]) * (Float32(current_len) - Float32(1.0)) + softmax.row_sum[r] += extra + + @cute.jit + def sol_attn_apply_route_current_lens_to_row_sum_fast( + self, + acc_S: cute.Tensor, + tScS_mn: cute.Tensor, + row_sum_prev: cute.Tensor, + row_scale: cute.Tensor, + group_start_n_block: Int32, + valid_count: Int32, + route_col_offset: Int32, + seqlen: SeqlenInfoQK, + softmax: Softmax, + is_first_block: cutlass.Constexpr[bool], + ): + """Fast denominator correction for full-length route groups.""" + + last_n_block = ( + (seqlen.seqlen_k + Int32(self.tile_n - 1)) // Int32(self.tile_n) + ) - Int32(1) + tail_len = seqlen.seqlen_k - last_n_block * Int32(self.tile_n) + group_end = group_start_n_block + valid_count + full_len_group = (tail_len == Int32(self.tile_n)) or (group_end <= last_n_block) + + if full_len_group: + block_extra = Float32(self.tile_n - 1) + for r in cutlass.range(cute.size(softmax.row_sum), unroll_full=True): + prev_scaled = Float32(0.0) + if const_expr(not is_first_block): + prev_scaled = Float32(row_sum_prev[r]) * Float32(row_scale[r]) + route_row_sum = Float32(softmax.row_sum[r]) - prev_scaled + softmax.row_sum[r] += route_row_sum * block_extra + else: + # Tail blocks need per-column current_len because the last route + # column may represent fewer than tile_n values. + self.sol_attn_apply_route_current_lens_to_row_sum( + acc_S, + tScS_mn, + group_start_n_block, + valid_count, + route_col_offset, + seqlen, + softmax, + ) + + @cute.jit + def __call__( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mO: cute.Tensor, + mKC: cute.Tensor, + mVC: cute.Tensor, + mGlobalThresh: cute.Tensor, + mLSE: Optional[cute.Tensor], + softmax_scale: Float32, + sink_range: Int32, + stream: cuda.CUstream = None, + ): + """Configure and launch the Hopper Sol-Attn kernel.""" + + mCuSeqlensQ = None + mCuSeqlensK = None + mSeqUsedQ = None + mSeqUsedK = None + mPageTable = None + window_size_left = None + window_size_right = None + learnable_sink = None + blocksparse_tensors = None + piecewise_k = None + piecewise_v = None + aux_tensors = None + self.varlen_q = mCuSeqlensQ is not None or mSeqUsedQ is not None + + mQ, mK, mV, mO, mKC, mVC, mGlobalThresh = [ + assume_tensor_aligned(t) + for t in (mQ, mK, mV, mO, mKC, mVC, mGlobalThresh) + ] + if const_expr(piecewise_k is not None): + piecewise_k, piecewise_v = [ + assume_tensor_aligned(t) for t in (piecewise_k, piecewise_v) + ] + SOL_ATTN_BTHD_TRANSPOSE = [1, 3, 2, 0] + SOL_ATTN_BNH_TRANSPOSE = [1, 2, 0] + mQ, mK, mV, mO, mKC, mVC = [ + layout_utils.select(t, SOL_ATTN_BTHD_TRANSPOSE) + for t in (mQ, mK, mV, mO, mKC, mVC) + ] + mGlobalThresh = layout_utils.select( + mGlobalThresh, SOL_ATTN_BNH_TRANSPOSE + ) + if const_expr(piecewise_k is not None): + piecewise_k, piecewise_v = [ + layout_utils.select(t, SOL_ATTN_BTHD_TRANSPOSE) + for t in (piecewise_k, piecewise_v) + ] + LSE_layout_transpose = [1, 2, 0] + mLSE = ( + layout_utils.select(mLSE, LSE_layout_transpose) + if const_expr(mLSE is not None) + else None + ) + + tiled_mma_qk, tiled_mma_pv = self._get_tiled_mma() + self.num_mma_threads = tiled_mma_qk.size + self.num_threads_per_warp_group = 128 + self.num_wg_mma = self.num_mma_threads // self.num_threads_per_warp_group + assert self.num_wg_mma in [1, 2, 3] + if const_expr(self.num_wg_mma != 1): + raise NotImplementedError("SOL_ATTN SM90 path requires exactly one MMA warpgroup") + self.num_threads = self.num_threads_per_warp_group + self.num_producer_threads = 32 + self.num_Q_load_threads = self.num_threads_per_warp_group # If not TMA_Q + self.num_epilogue_threads = self.num_mma_threads + self.num_mma_regs, self.num_producer_regs = {1: (256, 56), 2: (240, 24), 3: (160, 32)}[ + self.num_wg_mma + ] + self.use_block_sparsity = cutlass.const_expr(blocksparse_tensors is not None) + self.has_piecewise_kv = cutlass.const_expr(piecewise_k is not None) + if const_expr(self.use_block_sparsity): + raise NotImplementedError("one-warpgroup SOL_ATTN path does not support block sparsity") + if const_expr(self.has_piecewise_kv): + raise NotImplementedError("one-warpgroup SOL_ATTN path does not support piecewise KV") + + self.use_scheduler_barrier = self.num_wg_mma == 2 + self.use_tma_Q = self.arch >= Arch.sm_90 and not ( + self.pack_gqa and self.tile_m % self.qhead_per_kvhead != 0 + ) + if const_expr(not self.use_tma_Q): + raise NotImplementedError("one-warpgroup SOL_ATTN path requires TMA Q/O") + # FP32 split partials require a direct register-to-global epilogue. + # A BF16 split partial matches V/O dtype and can reuse the shared-memory + # plus TMA-O epilogue. + self.use_tma_O = ( + self.sol_attn_num_splits == 1 or mO.element_type == self.dtype + ) + # Producer needs more registers when doing cp.async Q or KV loads + if const_expr(self.num_wg_mma == 2 and (not self.use_tma_Q or not self.use_tma_KV)): + self.num_mma_regs, self.num_producer_regs = 224, 40 + if const_expr(self.sol_attn_mma_regs_override is not None): + self.num_mma_regs = self.sol_attn_mma_regs_override + self.rescale_O_before_gemm = False + self._setup_attributes() + # TODO: we prob don't need most of what's in _setup_attributes + self.sQ_layout, self.sK_layout, self.sV_layout, self.sO_layout = [ + sm90_utils.make_smem_layout(mX.element_type, LayoutEnum.ROW_MAJOR, shape, stage) + for mX, shape, stage in [ + (mQ, (self.tile_m, self.tile_hdim), None), + (mK, (self.tile_n, self.tile_hdim), self.num_stages), + (mV, (self.tile_n, self.tile_hdimv), self.num_stages), + # sO always holds the BF16 PV epilogue tile. Split-KV's + # global mO is an FP32 partial workspace, so derive this + # shared-memory layout from V instead of global O. + (mV, (self.tile_m, self.tile_hdimv), None), + ] + ] + self.sP_layout = None + if const_expr(not self.mma_pv_is_rs): + self.sP_layout = sm90_utils.make_smem_layout( + mV.element_type, LayoutEnum.ROW_MAJOR, (self.tile_m, self.tile_n) + ) + SharedStorage = self._get_shared_storage_cls() + + mQ_og, mO_og = mQ, mO + if const_expr(self.pack_gqa): + nheads_kv = mK.shape[2] + mQ = pack_gqa_layout(mQ, self.qhead_per_kvhead, nheads_kv, head_idx=2) + mO = pack_gqa_layout(mO, self.qhead_per_kvhead, nheads_kv, head_idx=2) + if const_expr(mLSE is not None): + mLSE = pack_gqa_layout(mLSE, self.qhead_per_kvhead, nheads_kv, head_idx=1) + + # TMA + gmem_tiled_copy_Q = cpasync.CopyBulkTensorTileG2SOp() + gmem_tiled_copy_KV = cpasync.CopyBulkTensorTileG2SOp() # Might multicast + gmem_tiled_copy_O = cpasync.CopyBulkTensorTileS2GOp() + self.tma_copy_bytes = { + name: cute.size_in_bytes(mX.element_type, cute.select(layout, mode=[0, 1])) + for name, mX, layout in [ + ("Q", mQ, self.sQ_layout), + ("K", mK, self.sK_layout), + ("V", mV, self.sV_layout), + ] + } + make_tiled_tma_atom_fn = ( + partial(make_packgqa_tiled_tma_atom, qhead_per_kvhead=self.qhead_per_kvhead, head_idx=2) + if const_expr(self.pack_gqa) + else cpasync.make_tiled_tma_atom + ) + tma_atom_Q, tma_tensor_Q = None, None + if const_expr(self.use_tma_Q): + tma_atom_Q, tma_tensor_Q = make_tiled_tma_atom_fn( + gmem_tiled_copy_Q, + mQ_og if const_expr(self.pack_gqa) else mQ, + self.sQ_layout, + (self.tile_m, self.tile_hdim), # No mcast + ) + tma_atom_K, tma_tensor_K = None, None + tma_atom_V, tma_tensor_V = None, None + tma_atom_KC, tma_tensor_KC = None, None + tma_atom_VC, tma_tensor_VC = None, None + tma_atom_K2, tma_tensor_K2 = None, None + tma_atom_V2, tma_tensor_V2 = None, None + if const_expr(self.use_tma_KV): + tma_atom_K, tma_tensor_K = cpasync.make_tiled_tma_atom( + gmem_tiled_copy_KV, + mK, + cute.select(self.sK_layout, mode=[0, 1]), + (self.tile_n, self.tile_hdim), + 1, # No mcast for now + ) + tma_atom_V, tma_tensor_V = cpasync.make_tiled_tma_atom( + gmem_tiled_copy_KV, + mV, + cute.select(self.sV_layout, mode=[0, 1]), + (self.tile_n, self.tile_hdimv), + 1, # No mcast for now + ) + tma_atom_KC, tma_tensor_KC = cpasync.make_tiled_tma_atom( + gmem_tiled_copy_KV, + mKC, + cute.select(self.sK_layout, mode=[0, 1]), + (self.tile_n, self.tile_hdim), + 1, + ) + tma_atom_VC, tma_tensor_VC = cpasync.make_tiled_tma_atom( + gmem_tiled_copy_KV, + mVC, + cute.select(self.sV_layout, mode=[0, 1]), + (self.tile_n, self.tile_hdimv), + 1, + ) + if const_expr(self.has_piecewise_kv): + tma_atom_K2, tma_tensor_K2 = cpasync.make_tiled_tma_atom( + gmem_tiled_copy_KV, + piecewise_k, + cute.select(self.sK_layout, mode=[0, 1]), + (self.tile_n, self.tile_hdim), + 1, + ) + tma_atom_V2, tma_tensor_V2 = cpasync.make_tiled_tma_atom( + gmem_tiled_copy_KV, + piecewise_v, + cute.select(self.sV_layout, mode=[0, 1]), + (self.tile_n, self.tile_hdimv), + 1, + ) + tma_atom_O, tma_tensor_O = None, None + if const_expr(self.use_tma_O): + mO_tma = mO_og if const_expr(self.pack_gqa) else mO + if const_expr(self.varlen_q): + mO_tma = copy_utils.create_ragged_tensor_for_tma( + mO_tma, ragged_dim=0, ptr_shift=True + ) + tma_atom_O, tma_tensor_O = make_tiled_tma_atom_fn( + gmem_tiled_copy_O, + mO_tma, + self.sO_layout, + (self.tile_m, self.tile_hdimv), # No mcast + ) + if const_expr(mCuSeqlensQ is not None or mSeqUsedQ is not None): + TileScheduler = SingleTileVarlenScheduler + else: + TileScheduler = ( + SingleTileScheduler + if const_expr(not self.is_causal or self.is_local) + else SingleTileLPTScheduler + ) + tile_sched_args = TileSchedulerArguments( + cute.ceil_div(cute.size(mQ.shape[0]), self.tile_m), + cute.size(mQ.shape[2]), + cute.size(mQ.shape[3]) + if const_expr(mCuSeqlensQ is None) + else cute.size(mCuSeqlensQ.shape[0] - 1), + self.sol_attn_num_splits, + cute.size(mK.shape[0]) + if const_expr(mPageTable is None) + else mK.shape[0] * mPageTable.shape[1], + mQ.shape[1], + mV.shape[1], + total_q=cute.size(mQ.shape[0]) + if const_expr(mCuSeqlensQ is not None) + else cute.size(mQ.shape[0]) * cute.size(mQ.shape[3]), + tile_shape_mn=(self.tile_m, self.tile_n), + mCuSeqlensQ=mCuSeqlensQ, + mSeqUsedQ=mSeqUsedQ, + qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + element_size=self.dtype.width // 8, + is_persistent=False, + lpt=self.is_causal or self.is_local, + is_split_kv=self.sol_attn_num_splits > 1, + ) + tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args) + grid_dim = TileScheduler.get_grid_shape(tile_sched_params) + softmax_scale_log2, softmax_scale = utils.compute_softmax_scale_log2( + softmax_scale, self.score_mod + ) + window_size_left = Int32(window_size_left) if window_size_left is not None else None + window_size_right = Int32(window_size_right) if window_size_right is not None else None + fastdiv_mods = utils.compute_fastdiv_mods( + mQ, mK, self.qhead_per_kvhead, self.pack_gqa, aux_tensors, mPageTable + ) + + self.kernel( + tma_tensor_Q if const_expr(self.use_tma_Q) else mQ, + tma_tensor_K if const_expr(self.use_tma_KV) else mK, + tma_tensor_V if const_expr(self.use_tma_KV) else mV, + tma_tensor_KC if const_expr(self.use_tma_KV) else mKC, + tma_tensor_VC if const_expr(self.use_tma_KV) else mVC, + tma_tensor_K2 if const_expr(self.has_piecewise_kv) else None, + tma_tensor_V2 if const_expr(self.has_piecewise_kv) else None, + tma_tensor_O if const_expr(self.use_tma_O) else mO, + mGlobalThresh, + mLSE, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + mPageTable, + tma_atom_Q, + tma_atom_K, + tma_atom_V, + tma_atom_KC, + tma_atom_VC, + tma_atom_K2, + tma_atom_V2, + tma_atom_O, + softmax_scale_log2, + softmax_scale, + sink_range, + window_size_left, + window_size_right, + learnable_sink, + blocksparse_tensors, + self.sQ_layout, + self.sK_layout, + self.sV_layout, + self.sO_layout, + self.sP_layout, + self.gmem_tiled_copy_Q, + self.gmem_tiled_copy_K, + self.gmem_tiled_copy_V, + self.gmem_tiled_copy_O, + tiled_mma_qk, + tiled_mma_pv, + tile_sched_params, + TileScheduler, + SharedStorage, + fastdiv_mods, + ).launch( + grid=grid_dim, + block=[self.num_threads, 1, 1], + stream=stream, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mKC: cute.Tensor, + mVC: cute.Tensor, + mK2: Optional[cute.Tensor], + mV2: Optional[cute.Tensor], + mO: cute.Tensor, + mGlobalThresh: cute.Tensor, + mLSE: Optional[cute.Tensor], + mCuSeqlensQ: Optional[cute.Tensor], + mCuSeqlensK: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + mSeqUsedK: Optional[cute.Tensor], + mPageTable: Optional[cute.Tensor], + tma_atom_Q: Optional[cute.CopyAtom], + tma_atom_K: Optional[cute.CopyAtom], + tma_atom_V: Optional[cute.CopyAtom], + tma_atom_KC: Optional[cute.CopyAtom], + tma_atom_VC: Optional[cute.CopyAtom], + tma_atom_K2: Optional[cute.CopyAtom], + tma_atom_V2: Optional[cute.CopyAtom], + tma_atom_O: Optional[cute.CopyAtom], + softmax_scale_log2: Float32, + softmax_scale: Optional[Float32], + sink_range: Int32, + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + learnable_sink: Optional[cute.Tensor], + blocksparse_tensors: Optional[BlockSparseTensors], + sQ_layout: cute.ComposedLayout, + sK_layout: cute.ComposedLayout, + sV_layout: cute.ComposedLayout, + sO_layout: cute.ComposedLayout, + sP_layout: cute.ComposedLayout | None, + gmem_tiled_copy_Q: cute.TiledCopy, + gmem_tiled_copy_K: cute.TiledCopy, + gmem_tiled_copy_V: cute.TiledCopy, + gmem_tiled_copy_O: cute.TiledCopy, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + tile_sched_params: ParamsBase, + TileScheduler: cutlass.Constexpr[Callable], + SharedStorage: cutlass.Constexpr[Callable], + fastdiv_mods=None, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + # Prefetch tma descriptor + if warp_idx == 0: + for tma_atom in ( + tma_atom_Q, + tma_atom_K, + tma_atom_V, + tma_atom_KC, + tma_atom_VC, + tma_atom_K2, + tma_atom_V2, + tma_atom_O, + ): + if const_expr(tma_atom is not None): + cpasync.prefetch_descriptor(tma_atom) + + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # Mbarrier / pipeline init + mbar_ptr_Q = storage.mbar_ptr_Q.data_ptr() + + ThreadCooperativeGroup = partial(pipeline.CooperativeGroup, pipeline.Agent.Thread) + tma_warp = ThreadCooperativeGroup(1) + load_threads = ThreadCooperativeGroup(self.num_threads_per_warp_group) + mma_warps = ThreadCooperativeGroup(self.num_mma_threads // cute.arch.WARP_SIZE) + if const_expr(self.use_tma_Q): + pipeline_q = pipeline_custom.PipelineTmaAsync.create( + barrier_storage=mbar_ptr_Q, + num_stages=1, + producer_group=tma_warp, + consumer_group=mma_warps, + tx_count=self.tma_copy_bytes["Q"], + defer_sync=True, + ) + else: + pipeline_q = pipeline_custom.PipelineCpAsync.create( + barrier_storage=mbar_ptr_Q, + num_stages=1, + producer_group=load_threads, + consumer_group=mma_warps, + defer_sync=True, + elect_one_release=True, + syncwarp_before_release=False, + ) + + if const_expr(self.use_tma_KV): + pipeline_k = pipeline_custom.PipelineTmaAsync.create( + barrier_storage=storage.mbar_ptr_K.data_ptr(), + num_stages=self.num_stages, + producer_group=tma_warp, + consumer_group=mma_warps, + tx_count=self.tma_copy_bytes["K"], + defer_sync=True, + ) + pipeline_v = pipeline_custom.PipelineTmaAsync.create( + barrier_storage=storage.mbar_ptr_V.data_ptr(), + num_stages=self.num_stages, + producer_group=tma_warp, + consumer_group=mma_warps, + tx_count=self.tma_copy_bytes["V"], + defer_sync=True, + ) + else: + pipeline_k = pipeline_custom.PipelineCpAsync.create( + barrier_storage=storage.mbar_ptr_K.data_ptr(), + num_stages=self.num_stages, + producer_group=load_threads, + consumer_group=mma_warps, + defer_sync=True, + elect_one_release=True, + syncwarp_before_release=False, + ) + pipeline_v = pipeline_custom.PipelineCpAsync.create( + barrier_storage=storage.mbar_ptr_V.data_ptr(), + num_stages=self.num_stages, + producer_group=load_threads, + consumer_group=mma_warps, + defer_sync=True, + elect_one_release=True, + syncwarp_before_release=False, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # /////////////////////////////////////////////////////////////////////////////// + # Get shared memory buffer + # /////////////////////////////////////////////////////////////////////////////// + sQ = storage.sQ.get_tensor(sQ_layout.outer, swizzle=sQ_layout.inner) + sK = storage.sK.get_tensor(sK_layout.outer, swizzle=sK_layout.inner) + if const_expr(not self.Q_in_regs): + sV = storage.sV.get_tensor(sV_layout.outer, swizzle=sV_layout.inner) + else: + sV = storage.sQ.get_tensor( + sV_layout.outer, swizzle=sV_layout.inner, dtype=mV.element_type + ) + # Transpose view of V to tensor with layout (head_dim_v, tile_n) for tiled mma + sVt = layout_utils.transpose_view(sV) + sP = None + if const_expr(sP_layout is not None): + sP = storage.sP.get_tensor(sP_layout.outer, swizzle=sP_layout.inner) + # reuse sQ's data iterator + sO = storage.sQ.get_tensor(sO_layout.outer, swizzle=sO_layout.inner, dtype=self.dtype) + route_mask = storage.route_mask.get_tensor( + cute.make_layout((4,)) + ) + route_sums = storage.route_sums.get_tensor(cute.make_layout((4, self.tile_n))) + + block_info = BlockInfo( + self.tile_m, + self.tile_n, + self.is_causal, + self.is_local, + False, # is_split_kv + window_size_left, + window_size_right, + qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + SeqlenInfoCls = partial( + SeqlenInfoQK.create, + seqlen_q_static=mQ.shape[0] if const_expr(not self.pack_gqa) else mQ.shape[0][1], + seqlen_k_static=mK.shape[0] + if const_expr(mPageTable is None) + else mK.shape[0] * mPageTable.shape[1], + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=mCuSeqlensK, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=mSeqUsedK, + # Don't need to pass in tile_mn because we won't access offset_padded + ) + AttentionMaskCls = partial( + AttentionMask, + self.tile_m, + self.tile_n, + window_size_left=window_size_left, + window_size_right=window_size_right, + qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + TileSchedulerCls = partial(TileScheduler.create, tile_sched_params) + + # Cluster wait before starting + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + cute.arch.setmaxregister_increase(self.num_mma_regs) + self.mma_one_warpgroup_sol_attn_route_tma( + tiled_mma_qk, + tiled_mma_pv, + mQ, + mK, + mV, + mKC, + mVC, + mO, + mLSE, + sQ, + sK, + sV, + sVt, + sP, + sO, + tma_atom_Q, + tma_atom_K, + tma_atom_V, + tma_atom_KC, + tma_atom_VC, + gmem_tiled_copy_O, + tma_atom_O, + pipeline_q, + pipeline_k, + pipeline_v, + SeqlenInfoCls, + AttentionMaskCls, + TileSchedulerCls, + mGlobalThresh, + route_mask, + route_sums, + softmax_scale_log2, + softmax_scale, + sink_range, + block_info, + ) + + @cute.jit + def epilogue_one_warpgroup_tma_o( + self, + acc_O: cute.Tensor, + lse: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + sO: cute.Tensor, + seqlen: SeqlenInfoQK, + tma_atom_O: cute.CopyAtom, + tiled_mma: cute.TiledMma, + tidx: Int32, + m_block: Int32, + head_idx: Int32, + batch_idx: Int32, + ): + """One-warpgroup TMA-O epilogue with an in-CTA store owner.""" + + rO = cute.make_fragment_like(acc_O, self.dtype) + rO.store(acc_O.load().to(self.dtype)) + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.Epilogue), + number_of_threads=self.num_epilogue_threads, + ) + smem_copy_atom_O = utils.get_smem_store_atom( + self.arch.major * 10 + self.arch.minor, self.dtype + ) + smem_thr_copy_O = cute.make_tiled_copy_C(smem_copy_atom_O, tiled_mma).get_slice(tidx) + taccOrO = smem_thr_copy_O.retile(rO) + taccOsO = smem_thr_copy_O.partition_D(sO) + cute.copy(smem_copy_atom_O, taccOrO, taccOsO) + + cO = cute.make_identity_tensor((self.tile_m, self.tile_hdimv)) + if const_expr(mLSE is not None): + mLSE_cur = mLSE[None, head_idx, batch_idx] + gLSE = cute.local_tile(mLSE_cur, (self.tile_m,), (m_block,)) + gLSE_expanded_layout = cute.append( + gLSE.layout, cute.make_layout((self.tile_hdimv,), stride=(0,)) + ) + gLSE_expanded = cute.make_tensor(gLSE.iterator, gLSE_expanded_layout) + thr_mma = tiled_mma.get_slice(tidx) + taccOgLSE = layout_utils.reshape_acc_to_mn(thr_mma.partition_C(gLSE_expanded)) + taccOcO = layout_utils.reshape_acc_to_mn(thr_mma.partition_C(cO)) + t0accOcO = layout_utils.reshape_acc_to_mn(thr_mma.get_slice(0).partition_C(cO)) + if taccOcO[0][1] == 0: + for m in cutlass.range_constexpr(cute.size(taccOgLSE.shape[1])): + if ( + t0accOcO[m, 0][0] + < seqlen.seqlen_q - m_block * self.tile_m - taccOcO[0][0] + ): + taccOgLSE[m, 0] = lse[m] + + mO_cur = mO[None, None, head_idx, batch_idx] + cute.arch.fence_view_async_shared() + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.Epilogue), + number_of_threads=self.num_epilogue_threads, + ) + gO = cute.local_tile(mO_cur, (self.tile_m, self.tile_hdimv), (m_block, 0)) + store_O, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_O, 0, cute.make_layout(1), sO, gO, single_stage=True + ) + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == Int32(0): + store_O() + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + + @cute.jit + def epilogue_one_warpgroup_split_partial( + self, + acc_O: cute.Tensor, + lse: cute.Tensor, + mO: cute.Tensor, + mLSE: cute.Tensor, + seqlen: SeqlenInfoQK, + tiled_mma: cute.TiledMma, + tidx: Int32, + m_block: Int32, + partial_head_idx: Int32, + batch_idx: Int32, + ): + """Write one normalized FP32 split partial and its natural-log LSE. + + ``mO`` and ``mLSE`` use a physical split-head dimension. The caller + maps ``(split, head)`` to ``partial_head_idx``; a later combine kernel + performs the log-sum-exp weighted reduction across that dimension. + """ + + mO_cur = mO[None, None, partial_head_idx, batch_idx] + gO = cute.local_tile( + mO_cur, (self.tile_m, self.tile_hdimv), (m_block, 0) + ) + copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + Float32, + num_bits_per_copy=32, + ) + tiled_copy = cute.make_tiled_copy_C(copy_atom, tiled_mma) + rO = cute.make_rmem_tensor_like(acc_O, Float32) + rO.store(acc_O.load()) + tOrO = tiled_copy.retile(rO) + tOgO = tiled_copy.get_slice(tidx).partition_D(gO) + cute.autovec_copy(tOrO, tOgO) + + mLSE_cur = mLSE[None, partial_head_idx, batch_idx] + gLSE = cute.local_tile(mLSE_cur, (self.tile_m,), (m_block,)) + gLSE_expanded_layout = cute.append( + gLSE.layout, cute.make_layout((self.tile_hdimv,), stride=(0,)) + ) + gLSE_expanded = cute.make_tensor( + gLSE.iterator, gLSE_expanded_layout + ) + thr_mma = tiled_mma.get_slice(tidx) + taccOgLSE = layout_utils.reshape_acc_to_mn( + thr_mma.partition_C(gLSE_expanded) + ) + cO = cute.make_identity_tensor((self.tile_m, self.tile_hdimv)) + taccOcO = layout_utils.reshape_acc_to_mn(thr_mma.partition_C(cO)) + t0accOcO = layout_utils.reshape_acc_to_mn( + thr_mma.get_slice(0).partition_C(cO) + ) + if taccOcO[0][1] == 0: + for m in cutlass.range_constexpr(cute.size(taccOgLSE.shape[1])): + if ( + t0accOcO[m, 0][0] + < seqlen.seqlen_q + - m_block * self.tile_m + - taccOcO[0][0] + ): + taccOgLSE[m, 0] = lse[m] + + @cute.jit + def mma_one_warpgroup_sol_attn_route_tma( + self, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mKC: cute.Tensor, + mVC: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + sQ: cute.Tensor, + sK: cute.Tensor, + sV: cute.Tensor, + sVt: cute.Tensor, + sP: Optional[cute.Tensor], + sO: cute.Tensor, + tma_atom_Q: Optional[cute.CopyAtom], + tma_atom_K: Optional[cute.CopyAtom], + tma_atom_V: Optional[cute.CopyAtom], + tma_atom_KC: Optional[cute.CopyAtom], + tma_atom_VC: Optional[cute.CopyAtom], + gmem_tiled_copy_O: cute.TiledCopy, + tma_atom_O: Optional[cute.CopyAtom], + pipeline_q: pipeline.PipelineAsync, + pipeline_k: pipeline.PipelineAsync, + pipeline_v: pipeline.PipelineAsync, + SeqlenInfoCls: Callable, + AttentionMaskCls: Callable, + TileSchedulerCls: cutlass.Constexpr[Callable], + mGlobalThresh: cute.Tensor, + route_mask: cute.Tensor, + route_sums: cute.Tensor, + softmax_scale_log2: Float32, + softmax_scale: Float32, + sink_range: Int32, + block_info: BlockInfo, + ): + """Run the fused route, approximate, and exact attention mainloop.""" + + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if const_expr(not (self.use_tma_Q and self.use_tma_KV)): + if tidx == Int32(0) and warp_idx == Int32(0): + cute.printf("SOL_ATTN one-warpgroup path requires TMA Q/KV\n") + else: + q_producer_phase = Int32(1) + q_consumer_phase = Int32(0) + kv_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_stages + ) + kv_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_stages + ) + tile_scheduler = TileSchedulerCls() + work_tile = tile_scheduler.initial_work_tile_info() + + if work_tile.is_valid_tile: + m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx + partial_head_idx = ( + head_idx + + split_idx * mQ.shape[2] + if const_expr(self.sol_attn_num_splits > 1) + else head_idx + ) + seqlen = SeqlenInfoCls(batch_idx) + head_idx_kv = ( + head_idx // self.qhead_per_kvhead + if const_expr(not self.pack_gqa) + else head_idx + ) + mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx] + mK_cur = seqlen.offset_batch_K(mK, batch_idx, dim=3)[ + None, None, head_idx_kv + ] + mV_cur = seqlen.offset_batch_K(mV, batch_idx, dim=3)[ + None, None, head_idx_kv + ] + mKC_cur = mKC[None, None, head_idx_kv, batch_idx] + mVC_cur = mVC[None, None, head_idx_kv, batch_idx] + + gQ = cute.local_tile(mQ_cur, (self.tile_m, self.tile_hdim), (m_block, 0)) + gK = cute.local_tile(mK_cur, (self.tile_n, self.tile_hdim), (None, 0)) + gV = cute.local_tile(mV_cur, (self.tile_n, self.tile_hdimv), (None, 0)) + gKC = cute.local_tile(mKC_cur, (self.tile_n, self.tile_hdim), (None, 0)) + gVC = cute.local_tile(mVC_cur, (self.tile_n, self.tile_hdimv), (None, 0)) + + load_Q, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_Q, 0, cute.make_layout(1), gQ, sQ, single_stage=True + ) + tma_load_K_fn, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_K, 0, cute.make_layout(1), gK, sK + ) + tma_load_K_fn = copy_utils.tma_producer_copy_fn(tma_load_K_fn, pipeline_k) + tma_load_V_fn, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_V, 0, cute.make_layout(1), gV, sV + ) + tma_load_V_fn = copy_utils.tma_producer_copy_fn(tma_load_V_fn, pipeline_v) + tma_load_KC_fn, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_KC, 0, cute.make_layout(1), gKC, sK + ) + tma_load_KC_fn = copy_utils.tma_producer_copy_fn( + tma_load_KC_fn, pipeline_k + ) + tma_load_VC_fn, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_VC, 0, cute.make_layout(1), gVC, sV + ) + tma_load_VC_fn = copy_utils.tma_producer_copy_fn( + tma_load_VC_fn, pipeline_v + ) + + if warp_idx == Int32(0): + pipeline_q.producer_acquire_w_index_phase(0, q_producer_phase) + load_Q(tma_bar_ptr=pipeline_q.sync_object_full.get_barrier(0)) + + pipeline_q.consumer_wait_w_index_phase(0, q_consumer_phase) + warp_group_thread_layout = cute.make_layout( + 1, stride=self.num_threads_per_warp_group + ) + thr_mma_qk = tiled_mma_qk.get_slice(tidx) + wg_mma_qk = tiled_mma_qk.get_slice(warp_group_thread_layout(Int32(0))) + wg_mma_pv = tiled_mma_pv.get_slice(warp_group_thread_layout(Int32(0))) + _, tSrQ, tSrK = sm90_utils.partition_fragment_ABC( + wg_mma_qk, (self.tile_m, self.tile_n, self.tile_hdim), sQ, sK + ) + mma_qk_fn = partial( + self.sol_attn_qk_gemm_zero_init, + tiled_mma_qk, + (self.tile_m, self.tile_n), + tSrQ, + tSrK, + ) + acc_O, tOrP, tOrVt = sm90_utils.partition_fragment_ABC( + wg_mma_pv, (self.tile_m, self.tile_hdimv, self.tile_n), sP, sVt + ) + mma_pv_fn = partial(sm90_utils.gemm_w_idx, tiled_mma_pv, acc_O, tOrP, tOrVt) + smem_copy_atom_P = utils.get_smem_store_atom( + self.arch.major * 10 + self.arch.minor, self.dtype + ) + smem_thr_copy_P = cute.make_tiled_copy_C( + smem_copy_atom_P, tiled_mma_qk + ).get_slice(tidx) + tPsP = smem_thr_copy_P.partition_D(sP) if const_expr(sP is not None) else None + smem_copy_params = SimpleNamespace( + smem_thr_copy_P=smem_thr_copy_P, + tPsP=tPsP, + ) + acc_O.fill(0.0) + cS_route = cute.make_identity_tensor((self.tile_m, self.tile_n)) + tScS_route_mn = layout_utils.reshape_acc_to_mn( + thr_mma_qk.partition_C(cS_route) + ) + mask = AttentionMaskCls(seqlen) + mask_fn = partial( + mask.apply_mask, + batch_idx=batch_idx, + head_idx=head_idx, + m_block=m_block, + thr_mma=thr_mma_qk, + mask_causal=self.is_causal, + mask_local=self.is_local, + aux_tensors=None, + fastdiv_mods=None, + ) + score_mod_fn = None + if const_expr(self.score_mod is not None): + score_mod_fn = partial( + self.apply_score_mod, + thr_mma_qk, + batch_idx, + head_idx, + m_block, + softmax_scale=softmax_scale, + aux_tensors=None, + fastdiv_mods=None, + ) + softmax = Softmax.create( + softmax_scale_log2, + num_rows=acc_O.shape[0][0] * acc_O.shape[1], + softmax_scale=softmax_scale, + ) + if const_expr(self.sol_attn_neutral_softmax_state): + softmax.row_max.fill(-Float32.inf) + softmax.row_sum.fill(0.0) + exact_mma_one_n_block = partial( + self.mma_one_n_block, + mma_qk_fn=mma_qk_fn, + pipeline_k=pipeline_k, + pipeline_v=pipeline_v, + acc_O=acc_O, + tOrP=tOrP, + smem_copy_params=smem_copy_params, + softmax=softmax, + score_mod_fn=score_mod_fn, + score_scale_fn=None, + check_inf=not self.sol_attn_assume_nonempty_rows, + ) + n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block) + route_block_count = n_block_max - n_block_min + if const_expr(self.sol_attn_static_num_full_route_groups >= 0): + num_full_route_groups = Int32(self.sol_attn_static_num_full_route_groups) + if const_expr(self.sol_attn_static_tail_valid_count > 0): + tail_valid_count = Int32(self.sol_attn_static_tail_valid_count) + else: + tail_valid_count = Int32(0) + elif const_expr(self.sol_attn_assume_full_route_groups): + num_full_route_groups = cute.ceil_div( + route_block_count, self.sol_attn_group_size + ) + tail_valid_count = Int32(0) + else: + num_full_route_groups = route_block_count // Int32(self.sol_attn_group_size) + tail_valid_count = ( + route_block_count + - num_full_route_groups * Int32(self.sol_attn_group_size) + ) + num_route_groups = num_full_route_groups + if tail_valid_count > Int32(0): + num_route_groups += Int32(1) + if const_expr(self.sol_attn_num_splits == 1): + split_group_begin = Int32(0) + split_num_route_groups = num_route_groups + else: + groups_per_split = ( + num_route_groups + self.sol_attn_num_splits - 1 + ) // self.sol_attn_num_splits + split_group_begin = split_idx * groups_per_split + split_group_end = cutlass.min( + split_group_begin + groups_per_split, num_route_groups + ) + split_num_route_groups = cutlass.max( + split_group_end - split_group_begin, Int32(0) + ) + O_should_accumulate = self.sol_attn_neutral_softmax_state + for local_group_iter in cutlass.range( + split_num_route_groups, unroll=1 + ): + group_iter = split_group_begin + local_group_iter + group_start = n_block_min + group_iter * Int32(self.sol_attn_group_size) + route_valid_count = Int32(self.sol_attn_group_size) + if const_expr(not self.sol_attn_assume_full_route_groups): + if group_iter == num_full_route_groups and tail_valid_count > Int32(0): + route_valid_count = tail_valid_count + route_col_offset = group_start - ( + group_start // Int32(self.tile_n) + ) * Int32(self.tile_n) + route_n_block = group_start - route_col_offset + route_tile = route_n_block // Int32(self.tile_n) + has_next_route_group = ( + local_group_iter + Int32(1) < split_num_route_groups + ) + next_route_tile = Int32(-1) + if has_next_route_group: + next_group_start = group_start + Int32(self.sol_attn_group_size) + next_route_tile = next_group_start // Int32(self.tile_n) + + if warp_idx == Int32(0): + if local_group_iter == Int32(0): + pipeline_k.producer_acquire(kv_producer_state) + tma_load_KC_fn( + src_idx=route_tile, + producer_state=kv_producer_state, + ) + else: + previous_group_had_exact = ( + (route_mask[0] != Int32(0)) + or (route_mask[1] != Int32(0)) + or (route_mask[2] != Int32(0)) + or (route_mask[3] != Int32(0)) + ) + if not previous_group_had_exact: + pipeline_k.producer_acquire(kv_producer_state) + tma_load_KC_fn( + src_idx=route_tile, + producer_state=kv_producer_state, + ) + pipeline_v.producer_acquire(kv_producer_state) + tma_load_VC_fn( + src_idx=route_tile, + producer_state=kv_producer_state, + ) + kv_producer_state.advance() + + pipeline_k.consumer_wait( + kv_consumer_state, + pipeline_k.consumer_try_wait(kv_consumer_state), + ) + acc_S = mma_qk_fn(B_idx=kv_consumer_state.index, wg_wait=-1) + warpgroup.wait_group(0) + pipeline_k.consumer_release(kv_consumer_state) + mask0, mask1, mask2, mask3 = self.sol_attn_build_route_mask_from_acc( + acc_S, + route_sums, + tScS_route_mn, + m_block, + group_start, + route_valid_count, + route_col_offset, + seqlen, + batch_idx, + head_idx, + mGlobalThresh, + softmax_scale_log2, + sink_range, + False, + route_mask_words_override=2, + ) + exact_mask0 = mask0 + exact_mask1 = mask1 + exact_mask2 = mask2 + exact_mask3 = mask3 + first_exact_n_block = group_start + first_exact_exists = False + if tidx == Int32(0): + route_mask[0] = mask0 + route_mask[1] = mask1 + route_mask[2] = mask2 + route_mask[3] = mask3 + + cute.arch.barrier( + barrier_id=SOL_ATTN_ROUTE_MASK_BARRIER_ID, + number_of_threads=self.num_mma_threads, + ) + mask0 = route_mask[0] + mask1 = route_mask[1] + mask2 = route_mask[2] + mask3 = route_mask[3] + + exact_mask0 = mask0 + exact_mask1 = mask1 + exact_mask2 = mask2 + exact_mask3 = mask3 + + first_exact_exists = ( + (mask0 != Int32(0)) + or (mask1 != Int32(0)) + or (mask2 != Int32(0)) + or (mask3 != Int32(0)) + ) + if mask0 != Int32(0): + first_lowbit = mask0 & (Int32(0) - mask0) + first_exact_n_block += sol_attn_selector.sol_attn_bfind_b32( + first_lowbit + ) + exact_mask0 = mask0 & (mask0 - Int32(1)) + elif mask1 != Int32(0): + first_lowbit = mask1 & (Int32(0) - mask1) + first_exact_n_block += Int32(32) + ( + sol_attn_selector.sol_attn_bfind_b32(first_lowbit) + ) + exact_mask1 = mask1 & (mask1 - Int32(1)) + elif mask2 != Int32(0): + first_lowbit = mask2 & (Int32(0) - mask2) + first_exact_n_block += Int32(64) + ( + sol_attn_selector.sol_attn_bfind_b32(first_lowbit) + ) + exact_mask2 = mask2 & (mask2 - Int32(1)) + elif mask3 != Int32(0): + first_lowbit = mask3 & (Int32(0) - mask3) + first_exact_n_block += Int32(96) + ( + sol_attn_selector.sol_attn_bfind_b32(first_lowbit) + ) + exact_mask3 = mask3 & (mask3 - Int32(1)) + if first_exact_exists and warp_idx == Int32(0): + pipeline_k.producer_acquire(kv_producer_state) + tma_load_K_fn( + src_idx=first_exact_n_block, + producer_state=kv_producer_state, + ) + + if const_expr(self.score_mod is not None): + score_mod_fn(acc_S, n_block=route_n_block, seqlen=seqlen) + mask_fn( + acc_S, + n_block=route_n_block, + mask_mod=self.mask_mod, + mask_seqlen=not self.sol_attn_full_route_mask_seqlen_false, + ) + if const_expr(self.sol_attn_assume_full_route_groups): + route_has_approx = (mask0 != Int32(-1)) or (mask1 != Int32(-1)) + else: + valid0 = route_valid_count + if valid0 > Int32(32): + valid0 = Int32(32) + valid_bits0 = Int32(0) + if valid0 > Int32(0): + valid_bits0 = Int32(-1) + if valid0 < Int32(32): + valid_bits0 = (Int32(1) << valid0) - Int32(1) + valid1 = route_valid_count - Int32(32) + if valid1 < Int32(0): + valid1 = Int32(0) + if valid1 > Int32(32): + valid1 = Int32(32) + valid_bits1 = Int32(0) + if valid1 > Int32(0): + valid_bits1 = Int32(-1) + if valid1 < Int32(32): + valid_bits1 = (Int32(1) << valid1) - Int32(1) + route_has_approx = ( + ((mask0 & valid_bits0) != valid_bits0) + or ((mask1 & valid_bits1) != valid_bits1) + ) + self.sol_attn_mask_route_approx_columns( + acc_S, + route_sums, + tScS_route_mn, + route_valid_count, + route_col_offset, + mask0, + mask1, + mask2, + mask3, + False, + route_mask_words_override=self.sol_attn_group_words, + ) + pipeline_v.consumer_wait( + kv_consumer_state, + pipeline_v.consumer_try_wait(kv_consumer_state), + ) + if route_has_approx: + row_sum_prev = None + if const_expr( + self.sol_attn_fast_route_lens + and not self.sol_attn_full_block_row_sum_prescale + ): + row_sum_prev = cute.make_fragment_like(softmax.row_sum, Float32) + row_sum_prev.store(softmax.row_sum.load()) + if O_should_accumulate: + if const_expr(self.sol_attn_full_block_row_sum_prescale): + for r in cutlass.range( + cute.size(softmax.row_sum), unroll_full=True + ): + softmax.row_sum[r] *= Float32(1.0 / self.tile_n) + row_scale = softmax.online_softmax( + acc_S, + is_first=False, + check_inf=not self.sol_attn_assume_nonempty_rows, + ) + softmax.rescale_O(acc_O, row_scale) + if const_expr(self.sol_attn_full_block_row_sum_prescale): + for r in cutlass.range( + cute.size(softmax.row_sum), unroll_full=True + ): + softmax.row_sum[r] *= Float32(self.tile_n) + elif const_expr(self.sol_attn_fast_route_lens): + self.sol_attn_apply_route_current_lens_to_row_sum_fast( + acc_S, + tScS_route_mn, + row_sum_prev, + row_scale, + group_start, + route_valid_count, + route_col_offset, + seqlen, + softmax, + False, + ) + else: + self.sol_attn_apply_route_current_lens_to_row_sum( + acc_S, + tScS_route_mn, + group_start, + route_valid_count, + route_col_offset, + seqlen, + softmax, + ) + else: + row_scale = softmax.online_softmax( + acc_S, + is_first=True, + check_inf=not self.sol_attn_assume_nonempty_rows, + ) + if const_expr(self.sol_attn_full_block_row_sum_prescale): + for r in cutlass.range( + cute.size(softmax.row_sum), unroll_full=True + ): + softmax.row_sum[r] *= Float32(self.tile_n) + elif const_expr(self.sol_attn_fast_route_lens): + self.sol_attn_apply_route_current_lens_to_row_sum_fast( + acc_S, + tScS_route_mn, + row_sum_prev, + row_scale, + group_start, + route_valid_count, + route_col_offset, + seqlen, + softmax, + True, + ) + else: + self.sol_attn_apply_route_current_lens_to_row_sum( + acc_S, + tScS_route_mn, + group_start, + route_valid_count, + route_col_offset, + seqlen, + softmax, + ) + tOrP_acc = layout_utils.reshape_acc_to_frgA(acc_S) + tOrP_cur = ( + tOrP + if const_expr(self.mma_pv_is_rs) + else cute.make_rmem_tensor_like(tOrP_acc, self.dtype) + ) + utils.cvt_f16(tOrP_acc, tOrP_cur) + if const_expr(not self.mma_pv_is_rs): + tPrP = smem_copy_params.smem_thr_copy_P.retile(tOrP_cur) + cute.copy( + smem_copy_params.smem_thr_copy_P, + tPrP, + smem_copy_params.tPsP, + ) + cute.arch.fence_view_async_shared() + cute.arch.sync_warp() + if O_should_accumulate: + sm90_utils.gemm_w_idx( + tiled_mma_pv, + acc_O, + tOrP_cur, + tOrVt, + zero_init=False, + B_idx=kv_consumer_state.index, + wg_wait=-1, + ) + else: + sm90_utils.gemm_w_idx( + tiled_mma_pv, + acc_O, + tOrP_cur, + tOrVt, + zero_init=True, + B_idx=kv_consumer_state.index, + wg_wait=-1, + ) + warpgroup.wait_group(0) + O_should_accumulate = True + pipeline_v.consumer_release(kv_consumer_state) + kv_consumer_state.advance() + + last_n_block = Int32(-1) + if const_expr( + (not self.sol_attn_assume_full_k_exact_blocks) + or self.sol_attn_exact_mask_seqlen_last_only + ): + last_n_block = ( + (seqlen.seqlen_k + Int32(self.tile_n - 1)) // Int32(self.tile_n) + ) - Int32(1) + if O_should_accumulate: + ( + kv_producer_state, + kv_consumer_state, + O_should_accumulate, + _, + ) = exact_stream.consume_exact_blocks( + exact_mask0, + exact_mask1, + exact_mask2, + exact_mask3, + group_start, + seqlen, + kv_producer_state, + kv_consumer_state, + tma_load_K_fn, + tma_load_V_fn, + pipeline_k, + pipeline_v, + warp_idx == Int32(0), + mma_pv_fn, + exact_mma_one_n_block, + mask_fn, + score_mod_fn, + O_should_accumulate, + self.warp_scheduler_barrier_sync, + self.warp_scheduler_barrier_arrive, + not self.sol_attn_assume_full_k_exact_blocks, + False, + self.sol_attn_group_words, + last_n_block, + self.sol_attn_exact_mask_seqlen_last_only, + first_exact_n_block, + first_exact_exists, + next_route_tile, + tma_load_KC_fn, + ) + else: + ( + kv_producer_state, + kv_consumer_state, + O_should_accumulate, + _, + ) = exact_stream.consume_exact_blocks( + exact_mask0, + exact_mask1, + exact_mask2, + exact_mask3, + group_start, + seqlen, + kv_producer_state, + kv_consumer_state, + tma_load_K_fn, + tma_load_V_fn, + pipeline_k, + pipeline_v, + warp_idx == Int32(0), + mma_pv_fn, + exact_mma_one_n_block, + mask_fn, + score_mod_fn, + O_should_accumulate, + self.warp_scheduler_barrier_sync, + self.warp_scheduler_barrier_arrive, + not self.sol_attn_assume_full_k_exact_blocks, + True, + self.sol_attn_group_words, + last_n_block, + self.sol_attn_exact_mask_seqlen_last_only, + first_exact_n_block, + first_exact_exists, + next_route_tile, + tma_load_KC_fn, + ) + + pipeline_q.consumer_release_w_index(0) + final_scale = softmax.finalize(sink_val=None) + softmax.rescale_O(acc_O, final_scale) + if const_expr(self.use_tma_O): + self.epilogue_one_warpgroup_tma_o( + acc_O, + softmax.row_sum, + mO, + mLSE, + sO, + seqlen, + tma_atom_O, + tiled_mma_pv, + tidx, + m_block, + partial_head_idx, + batch_idx, + ) + else: + self.epilogue_one_warpgroup_split_partial( + acc_O, + softmax.row_sum, + mO, + mLSE, + seqlen, + tiled_mma_pv, + tidx, + m_block, + partial_head_idx, + batch_idx, + ) + + @cute.jit + def mma_one_n_block( + self, + smem_pipe_read: pipeline.PipelineState | pipeline_custom.PipelineStateSimple, + n_block: Int32, + mma_qk_fn: Callable, + mma_pv_fn: Callable, + pipeline_k: pipeline.PipelineAsync, + pipeline_v: pipeline.PipelineAsync, + acc_O: cute.Tensor, + tOrP: cute.Tensor, + smem_copy_params: SimpleNamespace, + softmax: Softmax, + seqlen: SeqlenInfoQK, + scores_scale: Optional[cute.Tensor] = None, + score_mod_fn: Optional[Callable] = None, + score_scale_fn: Optional[Callable] = None, + mask_fn: Optional[Callable] = None, + last_block_mask_fn: Optional[Callable] = None, + last_n_block: Int32 = Int32(-1), + is_first_n_block: cutlass.Constexpr = False, + check_inf: cutlass.Constexpr = True, + prefetch_next: cutlass.Constexpr = False, + next_n_block: Int32 = Int32(-1), + kv_producer_state=None, + load_K: Optional[Callable] = None, + issue_load=False, + ): + pipeline_k.consumer_wait(smem_pipe_read, pipeline_k.consumer_try_wait(smem_pipe_read)) + acc_S = mma_qk_fn(B_idx=smem_pipe_read.index, wg_wait=-1) + self.warp_scheduler_barrier_arrive() + warpgroup.wait_group(0) + pipeline_k.consumer_release(smem_pipe_read) + + # Reuse the released K stage while the current softmax and P@V run. + if const_expr(prefetch_next): + if issue_load and next_n_block >= Int32(0): + pipeline_k.producer_acquire(kv_producer_state) + load_K(src_idx=next_n_block, producer_state=kv_producer_state) + + if const_expr(score_scale_fn is not None): + score_scale_fn(acc_S, n_block=n_block) + if const_expr(score_mod_fn is not None): + score_mod_fn(acc_S, n_block=n_block, seqlen=seqlen) + if const_expr(mask_fn is not None): + mask_fn(acc_S=acc_S, n_block=n_block) + if const_expr(last_block_mask_fn is not None): + if n_block == last_n_block: + last_block_mask_fn(acc_S=acc_S, n_block=n_block) + + row_scale = softmax.online_softmax(acc_S, is_first=is_first_n_block, check_inf=check_inf) + tOrP_acc = layout_utils.reshape_acc_to_frgA(acc_S) + tOrP_cur = ( + tOrP + if const_expr(self.mma_pv_is_rs) + else cute.make_rmem_tensor_like(tOrP_acc, self.dtype) + ) + utils.cvt_f16(tOrP_acc, tOrP_cur) + if const_expr(not self.mma_pv_is_rs): + tPrP = smem_copy_params.smem_thr_copy_P.retile(tOrP_cur) + cute.copy(smem_copy_params.smem_thr_copy_P, tPrP, smem_copy_params.tPsP) + softmax.rescale_O(acc_O, row_scale) + if const_expr(not self.mma_pv_is_rs): + cute.arch.fence_view_async_shared() + cute.arch.sync_warp() + + pipeline_v.consumer_wait(smem_pipe_read, pipeline_v.consumer_try_wait(smem_pipe_read)) + self.warp_scheduler_barrier_sync() + mma_pv_fn(B_idx=smem_pipe_read.index, wg_wait=0) + pipeline_v.consumer_release(smem_pipe_read) + smem_pipe_read.advance() + return smem_pipe_read + + @cute.jit + def mma_init(self): + warp_group_idx = utils.canonical_warp_group_idx(sync=False) + if const_expr(self.use_scheduler_barrier): + if warp_group_idx == 1: + cute.arch.barrier_arrive( + barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1), + number_of_threads=2 * self.num_threads_per_warp_group, + ) + + @cute.jit + def apply_score_mod( + self, + thr_mma_qk, + batch_idx, + head_idx, + m_block, + acc_S, + n_block, + softmax_scale, + seqlen, + aux_tensors: Optional[list] = None, + fastdiv_mods=None, + ): + # Prepare index tensor + cS = cute.make_identity_tensor((self.tile_m, self.tile_n)) + cS = cute.domain_offset((m_block * self.tile_m, n_block * self.tile_n), cS) + tScS = thr_mma_qk.partition_C(cS) + + apply_score_mod_inner( + acc_S, + tScS, + self.score_mod, + batch_idx, + head_idx, + softmax_scale, + self.vec_size, + self.qk_acc_dtype, + aux_tensors, + fastdiv_mods, + seqlen_info=seqlen, + constant_q_idx=None, + qhead_per_kvhead=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + + def warp_scheduler_barrier_sync(self): + if const_expr(self.use_scheduler_barrier): + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1) + - 1 + + utils.canonical_warp_group_idx(sync=False), + number_of_threads=2 * self.num_threads_per_warp_group, + ) + + def warp_scheduler_barrier_arrive(self): + if const_expr(self.use_scheduler_barrier): + assert self.num_wg_mma in [2, 3] + cur_wg = utils.canonical_warp_group_idx(sync=False) - 1 + if const_expr(self.num_wg_mma == 2): + next_wg = 1 - cur_wg + else: + t = cur_wg + 1 + next_wg = t % self.num_wg_mma + cute.arch.barrier_arrive( + barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1) + next_wg, + number_of_threads=2 * self.num_threads_per_warp_group, + ) diff --git a/telefuser/kernel/sol_attn/sm90/split_combine.py b/telefuser/kernel/sol_attn/sm90/split_combine.py new file mode 100644 index 00000000..7813647b --- /dev/null +++ b/telefuser/kernel/sol_attn/sm90/split_combine.py @@ -0,0 +1,600 @@ +# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. +# A reimplementation of https://github.com/Dao-AILab/flash-attention/blob/main/hopper/flash_fwd_combine_kernel.h +# from Cutlass C++ to Cute-DSL. +# Vendored from cuDNN BSA v126; only helper imports are redirected to the +# equivalent private FlashAttention/CuTe source closure in this repository. +import math +from typing import Type, Optional +from functools import partial + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync +from cutlass import Float32, Int32, Boolean, const_expr + +from telefuser.kernel.sol_attn._vendor.flash_attn.cute import utils +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.cute_dsl_utils import assume_tensor_aligned +from telefuser.kernel.sol_attn._vendor.flash_attn.cute.seqlen_info import SeqlenInfo +from cutlass.cute import FastDivmodDivisor + + +class BlockSparseAttnForwardCombine: + def __init__( + self, + dtype: Type[cutlass.Numeric], + head_dim: int, + tile_m: int = 8, + k_block_size: int = 64, + log_max_splits: int = 4, + num_threads: int = 256, + stages: int = 4, + partial_dtype: Type[cutlass.Numeric] = Float32, + ): + """ + Forward combine kernel for split attention computation. + + :param dtype: output data type + :param head_dim: head dimension + :param tile_m: m block size + :param k_block_size: k block size + :param log_max_splits: log2 of maximum splits + :param num_threads: number of threads + :param varlen: whether using variable length sequences + :param stages: number of pipeline stages + """ + self.dtype = dtype + if partial_dtype not in (Float32, cutlass.BFloat16): + raise TypeError("O partial dtype must be Float32 or BFloat16") + self.partial_dtype = partial_dtype + self.head_dim = head_dim + self.tile_m = tile_m + self.k_block_size = k_block_size + self.max_splits = 1 << log_max_splits + self.num_threads = num_threads + self.is_even_k = head_dim % k_block_size == 0 + self.stages = stages + + def _setup_attributes(self): + # GMEM copy setup for O partial + universal_copy_bits = 128 + async_copy_elems = universal_copy_bits // self.partial_dtype.width + assert self.k_block_size % async_copy_elems == 0 + + k_block_gmem = 128 if self.k_block_size % 128 == 0 else (64 if self.k_block_size % 64 == 0 else 32) + gmem_threads_per_row = k_block_gmem // async_copy_elems + assert self.num_threads % gmem_threads_per_row == 0 + + # Async copy atom for O partial load + atom_async_copy_partial = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cute.nvgpu.LoadCacheMode.GLOBAL), + self.partial_dtype, + num_bits_per_copy=universal_copy_bits, + ) + tOpartial_layout = cute.make_ordered_layout( + (self.num_threads // gmem_threads_per_row, gmem_threads_per_row), + order=(1, 0), + ) + vOpartial_layout = cute.make_layout((1, async_copy_elems)) # 4 vals per load + self.gmem_tiled_copy_O_partial = cute.make_tiled_copy_tv(atom_async_copy_partial, tOpartial_layout, vOpartial_layout) + + # GMEM copy setup for final O (use universal copy for store) + atom_universal_copy = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.dtype, + num_bits_per_copy=async_copy_elems * self.dtype.width, + ) + self.gmem_tiled_copy_O = cute.make_tiled_copy_tv( + atom_universal_copy, + tOpartial_layout, + vOpartial_layout, # 4 vals per store + ) + + # LSE copy setup with async copy (alignment = 1) + lse_copy_bits = Float32.width # 1 element per copy, width is in bits + m_block_smem = ( + 128 if self.tile_m % 128 == 0 else (64 if self.tile_m % 64 == 0 else (32 if self.tile_m % 32 == 0 else (16 if self.tile_m % 16 == 0 else 8))) + ) + gmem_threads_per_row_lse = m_block_smem + assert self.num_threads % gmem_threads_per_row_lse == 0 + + # Async copy atom for LSE load + atom_async_copy_lse = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS), + Float32, + num_bits_per_copy=lse_copy_bits, + ) + tLSE_layout = cute.make_ordered_layout( + (self.num_threads // gmem_threads_per_row_lse, gmem_threads_per_row_lse), + order=(1, 0), + ) + vLSE_layout = cute.make_layout(1) + self.gmem_tiled_copy_LSE = cute.make_tiled_copy_tv(atom_async_copy_lse, tLSE_layout, vLSE_layout) + + # /////////////////////////////////////////////////////////////////////////////// + # Shared memory + # /////////////////////////////////////////////////////////////////////////////// + + # Shared memory to register copy for LSE + self.smem_threads_per_col_lse = self.num_threads // m_block_smem + assert 32 % self.smem_threads_per_col_lse == 0 # Must divide warp size + + s2r_layout_atom_lse = cute.make_ordered_layout( + (self.smem_threads_per_col_lse, self.num_threads // self.smem_threads_per_col_lse), + order=(0, 1), + ) + self.s2r_tiled_copy_LSE = cute.make_tiled_copy_tv( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Float32), + s2r_layout_atom_lse, + cute.make_layout(1), + ) + + # LSE shared memory layout with swizzling to avoid bank conflicts + # This works for kBlockMSmem = 8, 16, 32, 64, 128, no bank conflicts + if const_expr(m_block_smem == 8): + smem_lse_swizzle = cute.make_swizzle(5, 0, 5) + elif const_expr(m_block_smem == 16): + smem_lse_swizzle = cute.make_swizzle(4, 0, 4) + else: + smem_lse_swizzle = cute.make_swizzle(3, 2, 3) + smem_layout_atom_lse = cute.make_composed_layout(smem_lse_swizzle, 0, cute.make_ordered_layout((8, m_block_smem), order=(1, 0))) + self.smem_layout_lse = cute.tile_to_shape(smem_layout_atom_lse, (self.max_splits, self.tile_m), (0, 1)) + + # O partial shared memory layout (simple layout for pipeline stages) + self.smem_layout_o = cute.make_ordered_layout((self.tile_m, self.k_block_size, self.stages), order=(1, 0, 2)) + + @cute.jit + def __call__( + self, + mO_partial: cute.Tensor, + mLSE_partial: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor] = None, + cu_seqlens: Optional[cute.Tensor] = None, + seqused: Optional[cute.Tensor] = None, + num_splits_dynamic_ptr: Optional[cute.Tensor] = None, + varlen_batch_idx: Optional[cute.Tensor] = None, + semaphore_to_reset: Optional[cute.Tensor] = None, + # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI). + stream: cuda.CUstream = None, + ): + # Type checking + if const_expr(not (mO_partial.element_type == self.partial_dtype)): + raise TypeError( + "O partial tensor must match the configured partial dtype" + ) + if const_expr(not (mO.element_type == self.dtype)): + raise TypeError("O tensor must match dtype") + if const_expr(mLSE_partial.element_type not in [Float32]): + raise TypeError("LSE partial tensor must be Float32") + if const_expr(mLSE is not None and mLSE.element_type not in [Float32]): + raise TypeError("LSE tensor must be Float32") + + # Shape validation - input tensors are in user format, need to be converted to kernel format + if const_expr(len(mO_partial.shape) not in [4, 5]): + raise ValueError( + "O partial tensor must have 4 or 5 dimensions: (num_splits, batch, seqlen, nheads, headdim) or (num_splits, total_q, nheads, headdim)" + ) + if const_expr(len(mLSE_partial.shape) not in [3, 4]): + raise ValueError("LSE partial tensor must have 3 or 4 dimensions: (num_splits, batch, seqlen, nheads) or (num_splits, total_q, nheads)") + if const_expr(len(mO.shape) not in [3, 4]): + raise ValueError("O tensor must have 3 or 4 dimensions: (batch, seqlen, nheads, headdim) or (total_q, nheads, headdim)") + if const_expr(mLSE is not None and len(mLSE.shape) not in [2, 3]): + raise ValueError("LSE tensor must have 2 or 3 dimensions: (batch, seqlen, nheads) or (total_q, nheads)") + + mO_partial, mO = [assume_tensor_aligned(t) for t in (mO_partial, mO)] + # (num_splits, b, seqlen, h, d) -> (seqlen, d, num_splits, h, b) + # or (num_splits, total_q, h, d) -> (total_q, d, num_splits, h) + O_partial_layout_transpose = [2, 4, 0, 3, 1] if const_expr(cu_seqlens is None) else [1, 3, 0, 2] + # (b, seqlen, h, d) -> (seqlen, d, h, b) or (total_q, h, d) -> (total_q, d, h) + mO_partial = cute.make_tensor(mO_partial.iterator, cute.select(mO_partial.layout, mode=O_partial_layout_transpose)) + O_layout_transpose = [1, 3, 2, 0] if const_expr(cu_seqlens is None) else [0, 2, 1] + mO = cute.make_tensor(mO.iterator, cute.select(mO.layout, mode=O_layout_transpose)) + # (num_splits, b, seqlen, h) -> (seqlen, num_splits, h, b) + # or (num_splits, total_q, h) -> (total_q, num_splits, h) + LSE_partial_layout_transpose = [2, 0, 3, 1] if const_expr(cu_seqlens is None) else [1, 0, 2] + mLSE_partial = cute.make_tensor( + mLSE_partial.iterator, + cute.select(mLSE_partial.layout, mode=LSE_partial_layout_transpose), + ) + # (b, seqlen, h) -> (seqlen, h, b) or (total_q, h) -> (total_q, h) + LSE_layout_transpose = [1, 2, 0] if const_expr(cu_seqlens is None) else [0, 1] + mLSE = cute.make_tensor(mLSE.iterator, cute.select(mLSE.layout, mode=LSE_layout_transpose)) if mLSE is not None else None + + # Determine if we have variable length sequences + varlen = const_expr(cu_seqlens is not None or seqused is not None) + + self._setup_attributes() + + @cute.struct + class SharedStorage: + sLSE: cute.struct.Align[cute.struct.MemRange[Float32, cute.cosize(self.smem_layout_lse)], 128] + sMaxValidSplit: cute.struct.Align[cute.struct.MemRange[Int32, self.tile_m], 128] + sO: cute.struct.Align[cute.struct.MemRange[self.partial_dtype, cute.cosize(self.smem_layout_o)], 128] + + smem_size = SharedStorage.size_in_bytes() + + # Grid dimensions: (ceil_div(seqlen, m_block), ceil_div(head_dim, k_block), num_head * batch) + seqlen = mO_partial.shape[0] + num_head = mO_partial.shape[3] + batch_size = mO_partial.shape[4] if const_expr(cu_seqlens is None) else Int32(cu_seqlens.shape[0] - 1) + + # Create FastDivmodDivisor objects for efficient division + seqlen_divmod = FastDivmodDivisor(seqlen) + + grid_dim = ( + cute.ceil_div(seqlen * num_head, self.tile_m), + cute.ceil_div(self.head_dim, self.k_block_size), + batch_size, + ) + + self.kernel( + mO_partial, + mLSE_partial, + mO, + mLSE, + cu_seqlens, + seqused, + num_splits_dynamic_ptr, + varlen_batch_idx, + semaphore_to_reset, + SharedStorage, + self.smem_layout_lse, + self.smem_layout_o, + self.gmem_tiled_copy_O_partial, + self.gmem_tiled_copy_O, + self.gmem_tiled_copy_LSE, + self.s2r_tiled_copy_LSE, + seqlen_divmod, + varlen, + ).launch( + grid=grid_dim, + block=[self.num_threads, 1, 1], + smem=smem_size, + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mO_partial: cute.Tensor, + mLSE_partial: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + cu_seqlens: Optional[cute.Tensor], + seqused: Optional[cute.Tensor], + num_splits_dynamic_ptr: Optional[cute.Tensor], + varlen_batch_idx: Optional[cute.Tensor], + semaphore_to_reset: Optional[cute.Tensor], + SharedStorage: cutlass.Constexpr, + smem_layout_lse: cute.Layout | cute.ComposedLayout, + smem_layout_o: cute.Layout, + gmem_tiled_copy_O_partial: cute.TiledCopy, + gmem_tiled_copy_O: cute.TiledCopy, + gmem_tiled_copy_LSE: cute.TiledCopy, + s2r_tiled_copy_LSE: cute.TiledCopy, + seqlen_divmod: FastDivmodDivisor, + varlen: cutlass.Constexpr[bool], + ): + # Thread and block indices + tidx, _, _ = cute.arch.thread_idx() + m_block, k_block, maybe_virtual_batch = cute.arch.block_idx() + + # Map virtual batch index to real batch index (for persistent tile schedulers) + batch_idx = varlen_batch_idx[maybe_virtual_batch] if const_expr(varlen_batch_idx is not None) else maybe_virtual_batch + + # /////////////////////////////////////////////////////////////////////////////// + # Get shared memory buffer + # /////////////////////////////////////////////////////////////////////////////// + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + sLSE = storage.sLSE.get_tensor(smem_layout_lse) + sMaxValidSplit = storage.sMaxValidSplit.get_tensor((self.tile_m,)) + sO = storage.sO.get_tensor(smem_layout_o) + + # Handle semaphore reset after dependent grids complete. + if const_expr(semaphore_to_reset is not None): + if ( + tidx == 0 + and m_block == cute.arch.grid_dim()[0] - 1 + and k_block == cute.arch.grid_dim()[1] - 1 + and maybe_virtual_batch == cute.arch.grid_dim()[2] - 1 + ): + cute.arch.griddepcontrol_wait() + semaphore_to_reset[0] = 0 + + # Get number of splits (use maybe_virtual_batch for per-batch-slot splits) + num_splits = num_splits_dynamic_ptr[maybe_virtual_batch] if const_expr(num_splits_dynamic_ptr is not None) else mLSE_partial.shape[1] + if const_expr(cu_seqlens is None): + seqlen = mO_partial.shape[0] + offset = Int32(0) + else: + offset = cu_seqlens[batch_idx] + seqlen = seqused[batch_idx] if const_expr(seqused is not None) else cu_seqlens[batch_idx + 1] - offset + + # Extract number of heads (head index will be determined dynamically) + num_head = mO_partial.shape[3] + max_idx = seqlen * num_head + + # Early exit for single split if dynamic + if (const_expr(num_splits_dynamic_ptr is None) or num_splits > 1) and (const_expr(not varlen) or m_block * self.tile_m < max_idx): + # The BSA wrapper launches pre-schedule, partial attention, and + # combine on the same stream without CUDA dependent-grid launch + # attributes. Stream ordering already guarantees producer + # completion, so do not use griddepcontrol_wait here. + + # =============================== + # Step 1: Load LSE_partial from gmem to shared memory + # =============================== + + if const_expr(cu_seqlens is None): + mLSE_partial_cur = mLSE_partial[None, None, None, batch_idx] + else: + mLSE_partial_cur = cute.domain_offset((offset, 0, 0), mLSE_partial) + mLSE_partial_copy = cute.tiled_divide(mLSE_partial_cur, (1,)) + gmem_thr_copy_LSE = gmem_tiled_copy_LSE.get_slice(tidx) + tLSEsLSE = gmem_thr_copy_LSE.partition_D(sLSE) + # Create identity tensor for coordinate tracking + cLSE = cute.make_identity_tensor((self.max_splits, self.tile_m)) + tLSEcLSE = gmem_thr_copy_LSE.partition_S(cLSE) + + # Load LSE partial values + for m in cutlass.range(cute.size(tLSEcLSE, mode=[2]), unroll_full=True): + mi = tLSEcLSE[0, 0, m][1] # Get m coordinate + idx = m_block * self.tile_m + mi + if idx < max_idx: + # Calculate actual sequence position and head using FastDivmodDivisor + if const_expr(not varlen): + head_idx, m_idx = divmod(idx, seqlen_divmod) + else: + head_idx = idx // seqlen + m_idx = idx - head_idx * seqlen + mLSE_partial_cur_copy = mLSE_partial_copy[None, m_idx, None, head_idx] + for s in cutlass.range(cute.size(tLSEcLSE, mode=[1]), unroll_full=True): + si = tLSEcLSE[0, s, 0][0] # Get split coordinate + if si < num_splits: + cute.copy( + gmem_thr_copy_LSE, + mLSE_partial_cur_copy[None, si], + tLSEsLSE[None, s, m], + ) + else: + tLSEsLSE[None, s, m].fill(-Float32.inf) + else: + # Rows past max_idx never write O/LSE to gmem, but their sLSE slots + # still feed the max-valid-split reduction, which bounds the + # partial-O accumulation loop. Fill them with -inf so such a row + # yields max_valid_split == -1 instead of reading garbage. + for s in cutlass.range(cute.size(tLSEcLSE, mode=[1]), unroll_full=True): + tLSEsLSE[None, s, m].fill(-Float32.inf) + cute.arch.cp_async_commit_group() + + # =============================== + # Step 2: Load O_partial for pipeline stages + # =============================== + + gmem_thr_copy_O_partial = gmem_tiled_copy_O_partial.get_slice(tidx) + cO = cute.make_identity_tensor((self.tile_m, self.k_block_size)) + tOcO = gmem_thr_copy_O_partial.partition_D(cO) + tOsO_partial = gmem_thr_copy_O_partial.partition_D(sO) + if const_expr(cu_seqlens is None): + mO_partial_cur = mO_partial[None, None, None, None, batch_idx] + else: + mO_partial_cur = cute.domain_offset((offset, 0, 0, 0), mO_partial) + + # Precompute these values to avoid recomputing them in the loop + num_rows = const_expr(cute.size(tOcO, mode=[1])) + tOmidx = cute.make_rmem_tensor(num_rows, cutlass.Int32) + tOhidx = cute.make_rmem_tensor(num_rows, cutlass.Int32) + tOrOptr = cute.make_rmem_tensor(num_rows, cutlass.Int64) + for m in cutlass.range(num_rows, unroll_full=True): + mi = tOcO[0, m, 0][0] # m coordinate + idx = m_block * self.tile_m + mi + if const_expr(not varlen): + tOhidx[m], tOmidx[m] = divmod(idx, seqlen_divmod) + else: + tOhidx[m] = idx // seqlen + tOmidx[m] = idx - tOhidx[m] * seqlen + tOrOptr[m] = utils.elem_pointer(mO_partial_cur, (tOmidx[m], k_block * self.k_block_size, 0, tOhidx[m])).toint() + if idx >= max_idx: + tOhidx[m] = -1 + + tOpO = None + if const_expr(not self.is_even_k): + tOpO = cute.make_rmem_tensor(cute.size(tOcO, mode=[2]), Boolean) + for k in cutlass.range(cute.size(tOpO), unroll_full=True): + tOpO[k] = tOcO[0, 0, k][1] < mO_partial.shape[1] - k_block * self.k_block_size + + load_O_partial = partial( + self.load_O_partial, + gmem_tiled_copy_O_partial, + tOrOptr, + tOsO_partial, + tOhidx, + tOpO, + tOcO, + mO_partial_cur.layout, + ) + + # Load first few stages of O_partial + for stage in cutlass.range(self.stages - 1, unroll_full=True): + if stage < num_splits: + load_O_partial(stage, stage) + cute.arch.cp_async_commit_group() + + # =============================== + # Step 3: Load and transpose LSE from smem to registers + # =============================== + + # Wait for LSE and initial O partial stages to complete + cute.arch.cp_async_wait_group(self.stages - 1) + cute.arch.sync_threads() + + s2r_thr_copy_LSE = s2r_tiled_copy_LSE.get_slice(tidx) + ts2rsLSE = s2r_thr_copy_LSE.partition_S(sLSE) + ts2rrLSE = cute.make_rmem_tensor_like(ts2rsLSE) + cute.copy(s2r_tiled_copy_LSE, ts2rsLSE, ts2rrLSE) + + # =============================== + # Step 4: Compute final LSE along split dimension + # =============================== + + lse_sum = cute.make_rmem_tensor(cute.size(ts2rrLSE, mode=[2]), Float32) + ts2rcLSE = s2r_thr_copy_LSE.partition_D(cLSE) + # We compute the max valid split for each row to short-circuit the computation later + max_valid_split = cute.make_rmem_tensor(cute.size(ts2rrLSE, mode=[2]), Int32) + assert cute.size(ts2rrLSE, mode=[0]) == 1 + # Compute max, scales, and final LSE for each row + for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True): + # Find max LSE value across splits + threads_per_col = const_expr(self.smem_threads_per_col_lse) + lse_max = cute.arch.warp_reduction_max( + ts2rrLSE[None, None, m].load().reduce(cute.ReductionOp.MAX, init_val=-Float32.inf, reduction_profile=0), + threads_in_group=threads_per_col, + ) + # Find max valid split index + max_valid_idx = -1 + for s in cutlass.range(cute.size(ts2rrLSE, mode=[1]), unroll_full=True): + if ts2rrLSE[0, s, m] != -Float32.inf: + max_valid_idx = ts2rcLSE[0, s, 0][0] # Get split coordinate + max_valid_split[m] = cute.arch.warp_reduction_max(max_valid_idx, threads_in_group=threads_per_col) + # Compute exp scales and sum + lse_max_cur = 0.0 if lse_max == -Float32.inf else lse_max # In case all local LSEs are -inf + LOG2_E = math.log2(math.e) + lse_sum_cur = 0.0 + for s in cutlass.range(cute.size(ts2rrLSE, mode=[1]), unroll_full=True): + scale = cute.math.exp2(ts2rrLSE[0, s, m] * LOG2_E - (lse_max_cur * LOG2_E), fastmath=True) + lse_sum_cur += scale + ts2rrLSE[0, s, m] = scale # Store scale for later use + lse_sum_cur = cute.arch.warp_reduction_sum(lse_sum_cur, threads_in_group=threads_per_col) + lse_sum[m] = cute.math.log(lse_sum_cur, fastmath=True) + lse_max + # Normalize scales + inv_sum = 0.0 if (lse_sum_cur == 0.0 or lse_sum_cur != lse_sum_cur) else 1.0 / lse_sum_cur + ts2rrLSE[None, None, m].store(ts2rrLSE[None, None, m].load() * inv_sum) + # Store the scales exp(lse - lse_logsum) back to smem + cute.copy(s2r_tiled_copy_LSE, ts2rrLSE, ts2rsLSE) + + # Store max valid split to smem + for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True): + if ts2rcLSE[0, 0, m][0] == 0: # Only thread responsible for s=0 writes + mi = ts2rcLSE[0, 0, m][1] + if mi < self.tile_m: + sMaxValidSplit[mi] = max_valid_split[m] + + # =============================== + # Step 5: Store final LSE to gmem + # =============================== + + if const_expr(mLSE is not None): + if const_expr(cu_seqlens is None): + mLSE_cur = mLSE[None, None, batch_idx] + else: + mLSE_cur = cute.domain_offset((offset, 0), mLSE) + if k_block == 0: # Only first k_block writes LSE when mLSE is provided + for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True): + if ts2rcLSE[0, 0, m][0] == 0: # Only thread responsible for s=0 writes + mi = ts2rcLSE[0, 0, m][1] + idx = m_block * self.tile_m + mi + if idx < max_idx: + if const_expr(not varlen): + head_idx, m_idx = divmod(idx, seqlen_divmod) + else: + head_idx = idx // seqlen + m_idx = idx - head_idx * seqlen + mLSE_cur[m_idx, head_idx] = lse_sum[m] + + # =============================== + # Step 6: Read O_partial and accumulate final O + # =============================== + + cute.arch.sync_threads() + + # Get max valid split for this thread + thr_max_valid_split = sMaxValidSplit[tOcO[0, 0, 0][0]] + for m in cutlass.range(1, cute.size(tOcO, mode=[1]), unroll_full=True): + thr_max_valid_split = max(thr_max_valid_split, sMaxValidSplit[tOcO[0, m, 0][0]]) + + tOrO_partial = cute.make_rmem_tensor_like(tOsO_partial[None, None, None, 0]) + tOrO = cute.make_rmem_tensor_like(tOrO_partial, Float32) + tOrO.fill(0.0) + + stage_load = self.stages - 1 + stage_compute = 0 + + # Main accumulation loop + for s in cutlass.range(thr_max_valid_split + 1, unroll=4): + # Get scales for this split + scale = cute.make_rmem_tensor(num_rows, Float32) + for m in cutlass.range(num_rows, unroll_full=True): + scale[m] = sLSE[s, tOcO[0, m, 0][0]] # Get scale from smem + + # Load next stage if needed + split_to_load = s + self.stages - 1 + if split_to_load <= thr_max_valid_split: + load_O_partial(split_to_load, stage_load) + cute.arch.cp_async_commit_group() + stage_load = 0 if stage_load == self.stages - 1 else stage_load + 1 + + # Wait for the current stage to be ready + cute.arch.cp_async_wait_group(self.stages - 1) + # We don't need __syncthreads() because each thread is just reading its own data from smem + # Copy from smem to registers + cute.autovec_copy(tOsO_partial[None, None, None, stage_compute], tOrO_partial) + stage_compute = 0 if stage_compute == self.stages - 1 else stage_compute + 1 + + # Accumulate scaled partial results + for m in cutlass.range(num_rows, unroll_full=True): + if tOhidx[m] >= 0 and scale[m] > 0.0: + tOrO[None, m, None].store(tOrO[None, m, None].load() + scale[m] * tOrO_partial[None, m, None].load().to(Float32)) + + # =============================== + # Step 7: Write final O to gmem + # =============================== + + rO = cute.make_rmem_tensor_like(tOrO, self.dtype) + rO.store(tOrO.load().to(self.dtype)) + if const_expr(cu_seqlens is None): + mO_cur = mO[None, None, None, batch_idx] + else: + mO_cur = cute.domain_offset((offset, 0, 0), mO) + mO_cur = utils.domain_offset_aligned((0, k_block * self.k_block_size, 0), mO_cur) + elems_per_store = const_expr(cute.size(gmem_tiled_copy_O.layout_tv_tiled[1])) + gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx) + # Write final results + for m in cutlass.range(num_rows, unroll_full=True): + if tOhidx[m] >= 0: + mO_cur_copy = cute.tiled_divide(mO_cur[tOmidx[m], None, tOhidx[m]], (elems_per_store,)) + for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True): + k_idx = tOcO[0, 0, k][1] // elems_per_store + if const_expr(self.is_even_k) or tOpO[k]: + cute.copy(gmem_thr_copy_O, rO[None, m, k], mO_cur_copy[None, k_idx]) + + @cute.jit + def load_O_partial( + self, + gmem_tiled_copy_O_partial: cute.TiledCopy, + tOrOptr: cute.Tensor, + tOsO_partial: cute.Tensor, + tOhidx: cute.Tensor, + tOpO: Optional[cute.Tensor], + tOcO: cute.Tensor, + mO_cur_partial_layout: cute.Layout, + split: Int32, + stage: Int32, + ) -> None: + elems_per_load = const_expr(cute.size(gmem_tiled_copy_O_partial.layout_tv_tiled[1])) + tOsO_partial_cur = tOsO_partial[None, None, None, stage] + for m in cutlass.range(cute.size(tOcO, [1]), unroll_full=True): + if tOhidx[m] >= 0: + o_gmem_ptr = cute.make_ptr(tOsO_partial.element_type, tOrOptr[m], cute.AddressSpace.gmem, assumed_align=16) + mO_partial_cur = cute.make_tensor(o_gmem_ptr, cute.slice_(mO_cur_partial_layout, (0, None, None, 0))) + mO_partial_cur_copy = cute.tiled_divide(mO_partial_cur, (elems_per_load,)) + for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True): + k_idx = tOcO[0, 0, k][1] // elems_per_load + if const_expr(tOpO is None) or tOpO[k]: + cute.copy( + gmem_tiled_copy_O_partial, + mO_partial_cur_copy[None, k_idx, split], + tOsO_partial_cur[None, m, k], + ) diff --git a/telefuser/kernel/sol_attn/triton_ref/__init__.py b/telefuser/kernel/sol_attn/triton_ref/__init__.py new file mode 100644 index 00000000..f8234aab --- /dev/null +++ b/telefuser/kernel/sol_attn/triton_ref/__init__.py @@ -0,0 +1,5 @@ +"""Triton reference implementation.""" + +from .fwd import sol_attn + +__all__ = ["sol_attn"] diff --git a/telefuser/kernel/sol_attn/triton_ref/fwd.py b/telefuser/kernel/sol_attn/triton_ref/fwd.py new file mode 100644 index 00000000..67edcc57 --- /dev/null +++ b/telefuser/kernel/sol_attn/triton_ref/fwd.py @@ -0,0 +1,475 @@ +"""Architecture-aware Triton Sol-Attn forward implementation.""" + +import torch +import triton +import triton.language as tl + +try: + from triton.tools.tensor_descriptor import TensorDescriptor +except ImportError: # Pointer path remains usable without descriptors. + TensorDescriptor = None + +from telefuser.kernel.sol_attn.interface import ( + _sink_block_range, + _validate_inputs, +) + +from .preprocess import prepare as prepare_ptr + + +BLOCK = 64 +GROUP = 32 + + +def _use_tma(device) -> bool: + """Use descriptor-backed Triton I/O when the GPU supports TMA.""" + + capability = torch.cuda.get_device_capability(device) + return capability[0] >= 9 and TensorDescriptor is not None + + +@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 _forward_tma( + q_desc, + k_desc, + v_desc, + kc_desc, + vc_desc, + threshold, + o_desc, + scale, + T, + sink_start_block, + sink_end_block, + HAS_SINK: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + NT: tl.constexpr, + BV: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + GROUP_SIZE: tl.constexpr, +): + v_tile, q_block, batch_head = ( + tl.program_id(0), + tl.program_id(1), + tl.program_id(2), + ) + batch, head = batch_head // H, batch_head % H + group_offsets = tl.max_contiguous( + tl.arange(0, GROUP_SIZE), + GROUP_SIZE, + ) + token_offsets = tl.max_contiguous( + tl.arange(0, BLOCK_SIZE), + BLOCK_SIZE, + ) + q_start = q_block * BLOCK_SIZE + q = q_desc.load([batch, q_start, head, 0]).reshape([BLOCK_SIZE, D]) + q_len = tl.minimum(BLOCK_SIZE, T - q_start).to(tl.float32) + + output = tl.zeros([BLOCK_SIZE, BV], dtype=tl.float32) + row_sum = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) + row_max = tl.full((BLOCK_SIZE,), -float("inf"), tl.float32) + scale_log2 = scale * 1.4426950408889634 + tail_length = T - (NT - 1) * BLOCK_SIZE + route_threshold = tl.load( + threshold + (batch * NT + q_block) * H + head + ) + + for group_start in range(0, NT, GROUP_SIZE): + block_indices = group_start + group_offsets + valid = block_indices < NT + kc = kc_desc.load( + [batch, group_start, head, 0] + ).reshape([GROUP_SIZE, D]) + vc = vc_desc.load( + [batch, group_start, head, v_tile * BV] + ).reshape([GROUP_SIZE, BV]) + scores = tl.dot(q, kc.T).to(tl.float32) * scale_log2 + exact = ( + (tl.sum(scores, axis=0) / q_len > route_threshold) + | (tl.abs(q_block - block_indices) <= 1) + ) + if HAS_SINK: + exact = exact | ( + (block_indices >= sink_start_block) + & (block_indices < sink_end_block) + ) + exact = exact & valid + + approximate = valid & ~exact + approximate_scores = tl.where( + approximate[None, :], scores, -float("inf") + ) + new_max = tl.maximum(row_max, tl.max(approximate_scores, axis=1)) + alpha = tl.math.exp2( + tl.where(row_max == new_max, 0.0, row_max - new_max) + ) + approximate_probability = tl.where( + approximate[None, :], + tl.math.exp2(approximate_scores - new_max[:, None]), + 0.0, + ) + output = output * alpha[:, None] + tl.dot( + approximate_probability.to(vc.dtype), + vc, + ) + lengths = tl.where( + block_indices == NT - 1, tail_length, BLOCK_SIZE + ).to(tl.float32) + row_sum = row_sum * alpha + tl.sum( + approximate_probability * lengths[None, :], + axis=1, + ) + row_max = new_max + + exact_offsets = tl.where(exact, group_offsets, GROUP_SIZE) + for _ in range(tl.sum(exact.to(tl.int32))): + offset = tl.min(exact_offsets) + block = group_start + offset + exact_offsets = tl.where( + group_offsets == offset, + GROUP_SIZE, + exact_offsets, + ) + kv_start = block * BLOCK_SIZE + k = k_desc.load( + [batch, kv_start, head, 0] + ).reshape([BLOCK_SIZE, D]) + exact_scores = tl.dot(q, k.T).to(tl.float32) * scale_log2 + exact_scores += tl.where( + (kv_start + token_offsets)[None, :] < T, + 0.0, + -float("inf"), + ) + new_max = tl.maximum(row_max, tl.max(exact_scores, axis=1)) + alpha = tl.math.exp2(row_max - new_max) + exact_probability = tl.math.exp2( + exact_scores - new_max[:, None] + ) + row_sum = row_sum * alpha + tl.sum( + exact_probability, + axis=1, + ) + v = v_desc.load( + [batch, kv_start, head, v_tile * BV] + ).reshape([BLOCK_SIZE, BV]) + output = output * alpha[:, None] + tl.dot( + exact_probability.to(v.dtype), + v, + ) + row_max = new_max + + o_desc.store( + [batch, q_start, head, v_tile * BV], + (output / row_sum[:, None]).to(tl.bfloat16)[None, :, None, :], + ) + + +@triton.autotune( + configs=[ + triton.Config({"BV": 128}, num_warps=4, num_stages=1), + triton.Config({"BV": 128}, num_warps=8, num_stages=1), + triton.Config({"BV": 128}, num_warps=4, num_stages=2), + triton.Config({"BV": 64}, num_warps=4, num_stages=1), + ], + key=["T"], +) +@triton.jit +def _forward_ptr( + q_ptr, + k_ptr, + v_ptr, + kc_ptr, + vc_ptr, + threshold_ptr, + o_ptr, + scale, + T, + NPAD, + sink_start_block, + sink_end_block, + HAS_SINK: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + NT: tl.constexpr, + BV: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + GROUP_SIZE: tl.constexpr, +): + v_tile, q_block, batch_head = ( + tl.program_id(0), + tl.program_id(1), + tl.program_id(2), + ) + batch, head = batch_head // H, batch_head % H + group_offsets = tl.max_contiguous( + tl.arange(0, GROUP_SIZE), + GROUP_SIZE, + ) + token_offsets = tl.max_contiguous( + tl.arange(0, BLOCK_SIZE), + BLOCK_SIZE, + ) + dims = tl.arange(0, D) + value_dims = v_tile * BV + tl.arange(0, BV) + q_tokens = q_block * BLOCK_SIZE + token_offsets + q_valid = q_tokens < T + q_offsets = ( + ((batch * T + q_tokens[:, None]).to(tl.int64) * H + head) * D + + dims[None, :] + ) + q = tl.load(q_ptr + q_offsets, mask=q_valid[:, None], other=0.0) + q_len = tl.minimum(BLOCK_SIZE, T - q_block * BLOCK_SIZE).to(tl.float32) + + output = tl.zeros([BLOCK_SIZE, BV], dtype=tl.float32) + row_sum = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) + row_max = tl.full((BLOCK_SIZE,), -float("inf"), tl.float32) + scale_log2 = scale * 1.4426950408889634 + route_threshold = tl.load( + threshold_ptr + (batch * NT + q_block) * H + head + ) + + for group_start in range(0, NT, GROUP_SIZE): + block_indices = group_start + group_offsets + valid = block_indices < NT + kc_offsets = ( + ((batch * NPAD + block_indices[:, None]) * H + head) * D + + dims[None, :] + ) + vc_offsets = ( + ((batch * NPAD + block_indices[:, None]) * H + head) * D + + value_dims[None, :] + ) + kc = tl.load(kc_ptr + kc_offsets) + vc = tl.load(vc_ptr + vc_offsets) + scores = tl.dot(q, kc.T).to(tl.float32) * scale_log2 + exact = ( + (tl.sum(scores, axis=0) / q_len > route_threshold) + | (tl.abs(q_block - block_indices) <= 1) + ) + if HAS_SINK: + exact = exact | ( + (block_indices >= sink_start_block) + & (block_indices < sink_end_block) + ) + exact = exact & valid + + approximate = valid & ~exact + has_approximate = tl.sum(approximate.to(tl.int32), axis=0) > 0 + approximate_scores = tl.where( + approximate[None, :], + scores, + -float("inf"), + ) + safe_scores = tl.where(has_approximate, approximate_scores, 0.0) + candidate_max = tl.maximum(row_max, tl.max(safe_scores, axis=1)) + new_max = tl.where(has_approximate, candidate_max, row_max) + alpha = tl.math.exp2( + tl.where(has_approximate, row_max - new_max, 0.0) + ) + probability = tl.math.exp2( + safe_scores + - tl.where(has_approximate, new_max, 0.0)[:, None] + ) + probability = tl.where( + has_approximate & approximate[None, :], + probability, + 0.0, + ) + output = output * alpha[:, None] + tl.dot( + probability.to(vc.dtype), + vc, + ) + lengths = tl.minimum( + BLOCK_SIZE, + tl.maximum(0, T - block_indices * BLOCK_SIZE), + ).to(tl.float32) + row_sum = row_sum * alpha + tl.sum( + probability * lengths[None, :], + axis=1, + ) + row_max = new_max + + exact_offsets = tl.where(exact, group_offsets, GROUP_SIZE) + num_exact = tl.sum(exact.to(tl.int32), axis=0) + for _ in range(num_exact): + offset = tl.min(exact_offsets) + block = group_start + offset + exact_offsets = tl.where( + group_offsets == offset, + GROUP_SIZE, + exact_offsets, + ) + kv_tokens = block * BLOCK_SIZE + token_offsets + kv_valid = kv_tokens < T + k_offsets = ( + ((batch * T + kv_tokens[:, None]).to(tl.int64) * H + head) + * D + + dims[None, :] + ) + k = tl.load( + k_ptr + k_offsets, + mask=kv_valid[:, None], + other=0.0, + ) + exact_scores = tl.dot(q, k.T).to(tl.float32) * scale_log2 + exact_scores += tl.where( + kv_valid[None, :], + 0.0, + -float("inf"), + ) + new_max = tl.maximum(row_max, tl.max(exact_scores, axis=1)) + alpha = tl.math.exp2(row_max - new_max) + exact_probability = tl.math.exp2( + exact_scores - new_max[:, None] + ) + row_sum = row_sum * alpha + tl.sum( + exact_probability, + axis=1, + ) + v_offsets = ( + ((batch * T + kv_tokens[:, None]).to(tl.int64) * H + head) + * D + + value_dims[None, :] + ) + v = tl.load( + v_ptr + v_offsets, + mask=kv_valid[:, None], + other=0.0, + ) + output = output * alpha[:, None] + tl.dot( + exact_probability.to(v.dtype), + v, + ) + row_max = new_max + + output_offsets = ( + ((batch * T + q_tokens[:, None]).to(tl.int64) * H + head) * D + + value_dims[None, :] + ) + tl.store( + o_ptr + output_offsets, + (output / row_sum[:, None]).to(tl.bfloat16), + mask=q_valid[:, None], + ) + + +def sol_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + scale: float | None = None, + tau: float = 1.0, + thresh_type: str = "diag", + sink_tokens: int = 0, + sink_start: int | None = None, +) -> torch.Tensor: + """Run Triton Sol-Attn on contiguous BF16 BTHD inputs.""" + + arch = _validate_inputs( + q, + k, + v, + thresh_type, + sink_tokens, + sink_start, + ) + if arch[0] < 8: + raise RuntimeError( + "Triton Sol-Attn requires an NVIDIA GPU with compute " + f"capability >= 8.0; got SM{arch[0]}{arch[1]}" + ) + scale = q.shape[-1] ** -0.5 if scale is None else float(scale) + tau = float(tau) + batch, tokens, heads, head_dim = q.shape + blocks = triton.cdiv(tokens, BLOCK) + sink_start_block, sink_end_block = _sink_block_range( + tokens, + sink_start, + sink_tokens, + ) + use_tma = _use_tma(q.device) + + if use_tma: + # Keep the original descriptor-backed preprocessing on TMA devices. + # The pointer preprocessing below exists only for older architectures. + from ..preprocess import prepare as prepare_tma + + kc, vc, threshold = prepare_tma( + q, + k, + v, + scale=scale, + tau=tau, + thresh_type=thresh_type, + ) + output = torch.empty_like(v) + block_shape = [1, BLOCK, 1, head_dim] + summary_shape = [1, GROUP, 1, head_dim] + _forward_tma[(1, blocks, batch * heads)]( + TensorDescriptor.from_tensor(q, block_shape), + TensorDescriptor.from_tensor(k, block_shape), + TensorDescriptor.from_tensor(v, block_shape), + TensorDescriptor.from_tensor(kc, summary_shape), + TensorDescriptor.from_tensor(vc, summary_shape), + threshold, + TensorDescriptor.from_tensor(output, block_shape), + scale, + tokens, + sink_start_block, + sink_end_block, + sink_tokens > 0, + heads, + head_dim, + blocks, + head_dim, + BLOCK, + GROUP, + ) + return output + + kc, vc, threshold = prepare_ptr( + q, + k, + v, + scale=scale, + tau=tau, + thresh_type=thresh_type, + tokens=tokens, + ) + output = torch.empty_like(v) + grid = lambda meta: (head_dim // meta["BV"], blocks, batch * heads) + _forward_ptr[grid]( + q, + k, + v, + kc, + vc, + threshold, + output, + scale, + tokens, + kc.shape[1], + sink_start_block, + sink_end_block, + HAS_SINK=sink_tokens > 0, + H=heads, + D=head_dim, + NT=blocks, + BLOCK_SIZE=BLOCK, + GROUP_SIZE=GROUP, + ) + return output + + +__all__ = ["sol_attn"] diff --git a/telefuser/kernel/sol_attn/triton_ref/preprocess.py b/telefuser/kernel/sol_attn/triton_ref/preprocess.py new file mode 100644 index 00000000..6341736c --- /dev/null +++ b/telefuser/kernel/sol_attn/triton_ref/preprocess.py @@ -0,0 +1,389 @@ +"""Pointer preprocessing for Triton Sol-Attn when TMA is unavailable.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +BLOCK_SIZE = 64 +HEAD_DIM = 128 +THRESHOLD_GROUP_SIZE = 64 +SUMMARY_PAD = 64 + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=warps, num_stages=stages) + for warps in (4, 8) + for stages in (1, 2) + ], + key=["T"], +) +@triton.jit +def _reduce_kv_kernel( + k, + v, + kc, + vc, + T, + TP, + NPAD, + H: tl.constexpr, + D: tl.constexpr, + BLOCK: tl.constexpr, +): + block, batch_head = tl.program_id(0), tl.program_id(1) + batch, head = batch_head // H, batch_head % H + tokens = block * BLOCK + tl.arange(0, BLOCK) + dims = tl.arange(0, D) + valid = tokens < T + offsets = ( + ((batch * TP + tokens[:, None]).to(tl.int64) * H + head) * D + + dims[None, :] + ) + k_values = tl.load(k + offsets, mask=valid[:, None], other=0.0) + v_values = tl.load(v + offsets, mask=valid[:, None], other=0.0) + block_len = tl.minimum(BLOCK, T - block * BLOCK).to(tl.float32) + summary_offsets = ( + ((batch * NPAD + block) * H + head) * D + dims + ) + tl.store(kc + summary_offsets, tl.sum(k_values, axis=0) / block_len) + tl.store(vc + summary_offsets, tl.sum(v_values, axis=0)) + + +@triton.jit +def _reduce_kc_stats_kernel( + kc, + kc_mean, + kc_var_diag, + NPAD, + H: tl.constexpr, + N: tl.constexpr, + D: tl.constexpr, + GROUP: tl.constexpr, +): + batch_head = tl.program_id(0) + batch, head = batch_head // H, batch_head % H + blocks = tl.max_contiguous(tl.arange(0, GROUP), GROUP) + dims = tl.arange(0, D) + total = tl.zeros((D,), dtype=tl.float32) + total_sq = tl.zeros((D,), dtype=tl.float32) + count = tl.full((), 0.0, dtype=tl.float32) + for start in range(0, N, GROUP): + block_indices = start + blocks + valid = block_indices < N + offsets = ( + ((batch * NPAD + block_indices[:, None]) * H + head) * D + + dims[None, :] + ) + values = tl.load( + kc + offsets, + mask=valid[:, None], + other=0.0, + ).to(tl.float32) + 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) + tl.store(kc_mean + batch_head * D + dims, mean) + tl.store(kc_var_diag + batch_head * D + dims, variance) + + +@triton.jit +def _diag_threshold_kernel( + q, + kc_mean, + kc_var_diag, + threshold, + scale, + T, + TP, + H: tl.constexpr, + N: tl.constexpr, + D: tl.constexpr, + BLOCK: 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 + tokens = q_block * BLOCK + tl.arange(0, BLOCK) + dims = tl.arange(0, D) + valid = tokens < T + offsets = ( + ((batch * TP + tokens[:, None]).to(tl.int64) * H + head) * D + + dims[None, :] + ) + q_values = tl.load(q + offsets, mask=valid[:, None], other=0.0) + q_len = tl.minimum(BLOCK, T - q_block * BLOCK).to(tl.float32) + q_centroid = tl.sum(q_values.to(tl.float32), axis=0) / q_len + mean_kc = tl.load(kc_mean + batch_head * D + dims) + var_kc = tl.load(kc_var_diag + batch_head * D + dims) + log2_scale = 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( + threshold + (batch * N + q_block) * H + head, + mean + TAU * std, + ) + + +@triton.jit +def _pool_query_kernel( + q, + q_bar, + T, + TP, + H: tl.constexpr, + N: tl.constexpr, + D: tl.constexpr, + BLOCK: tl.constexpr, +): + q_block, batch_head = tl.program_id(0), tl.program_id(1) + batch, head = batch_head // H, batch_head % H + tokens = q_block * BLOCK + tl.arange(0, BLOCK) + dims = tl.arange(0, D) + valid = tokens < T + offsets = ( + ((batch * TP + tokens[:, None]).to(tl.int64) * H + head) * D + + dims[None, :] + ) + values = tl.load(q + offsets, mask=valid[:, None], other=0.0) + q_len = tl.minimum(BLOCK, T - q_block * BLOCK).to(tl.float32) + centroid = tl.sum(values.to(tl.float32), axis=0) / q_len + tl.store(q_bar + (batch_head * N + q_block) * D + dims, centroid) + + +@triton.jit +def _exact_fused_threshold_kernel( + q_bar, + kc_mean, + kc_second_moment, + threshold, + scale, + H: tl.constexpr, + N: tl.constexpr, + D: tl.constexpr, + BLOCK_M: 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) + dims = tl.arange(0, D) + valid_rows = rows < N + q_centroid = tl.load( + q_bar + (batch_head * N + rows[:, None]) * D + dims[None, :], + mask=valid_rows[:, None], + other=0.0, + ) + mean_kc = tl.load(kc_mean + batch_head * D + dims) + second_moment = tl.load( + kc_second_moment + + batch_head * D * D + + dims[:, None] * D + + dims[None, :] + ) + 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 = scale * 1.4426950408889634 + mean = raw_mean * log2_scale + variance = tl.maximum( + raw_second_moment - raw_mean * raw_mean, + 0.0, + ) * (log2_scale * log2_scale) + result = mean + TAU * tl.sqrt(variance + 1.0e-6) + batch, head = batch_head // H, batch_head % H + tl.store( + threshold + (batch * N + rows) * H + head, + result, + mask=valid_rows, + ) + + +def _reduce_kv( + k: torch.Tensor, + v: torch.Tensor, + *, + tokens: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + batch, padded_tokens, heads, head_dim = k.shape + tokens = padded_tokens if tokens is None else int(tokens) + blocks = triton.cdiv(tokens, BLOCK_SIZE) + padded_blocks = triton.cdiv(blocks, SUMMARY_PAD) * SUMMARY_PAD + kc = torch.zeros( + (batch, padded_blocks, heads, head_dim), + device=k.device, + dtype=torch.bfloat16, + ) + vc = torch.zeros_like(kc) + _reduce_kv_kernel[(blocks, batch * heads)]( + k, + v, + kc, + vc, + tokens, + padded_tokens, + padded_blocks, + heads, + head_dim, + BLOCK_SIZE, + ) + return kc, vc + + +def _compute_diag_threshold( + q: torch.Tensor, + kc: torch.Tensor, + *, + tau: float, + scale: float, + tokens: int | None = None, +) -> torch.Tensor: + batch, padded_tokens, heads, head_dim = q.shape + tokens = padded_tokens if tokens is None else int(tokens) + blocks = triton.cdiv(tokens, BLOCK_SIZE) + batch_heads = batch * heads + kc_mean = torch.empty( + (batch_heads, head_dim), + device=q.device, + dtype=torch.float32, + ) + kc_var_diag = torch.empty_like(kc_mean) + threshold = torch.empty( + (batch, blocks, heads), + device=q.device, + dtype=torch.float32, + ) + _reduce_kc_stats_kernel[(batch_heads,)]( + kc, + kc_mean, + kc_var_diag, + kc.shape[1], + heads, + blocks, + head_dim, + THRESHOLD_GROUP_SIZE, + num_warps=4, + num_stages=2, + ) + _diag_threshold_kernel[(blocks, batch_heads)]( + q, + kc_mean, + kc_var_diag, + threshold, + scale, + tokens, + padded_tokens, + heads, + blocks, + head_dim, + BLOCK_SIZE, + tau, + num_warps=4, + num_stages=2, + ) + return threshold + + +def _compute_exact_threshold( + q: torch.Tensor, + kc: torch.Tensor, + *, + tau: float, + scale: float, + tokens: int | None = None, +) -> torch.Tensor: + batch, padded_tokens, heads, head_dim = q.shape + tokens = padded_tokens if tokens is None else int(tokens) + blocks = triton.cdiv(tokens, BLOCK_SIZE) + batch_heads = batch * heads + kc_bh = kc[:, :blocks].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, + ) + threshold = torch.empty( + (batch, blocks, heads), + device=q.device, + dtype=torch.float32, + ) + _pool_query_kernel[(blocks, batch_heads)]( + q, + q_bar, + tokens, + padded_tokens, + heads, + blocks, + head_dim, + BLOCK_SIZE, + 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, + threshold, + scale, + heads, + blocks, + head_dim, + block_m, + tau, + num_warps=4, + num_stages=1, + ) + return threshold + + +def prepare( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + tau: float, + scale: float, + thresh_type: str = "diag", + tokens: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + kc, vc = _reduce_kv(k, v, tokens=tokens) + if thresh_type == "exact": + threshold = _compute_exact_threshold( + q, + kc, + tau=tau, + scale=scale, + tokens=tokens, + ) + else: + threshold = _compute_diag_threshold( + q, + kc, + tau=tau, + scale=scale, + tokens=tokens, + ) + return kc, vc, threshold + + +__all__ = ["prepare"] diff --git a/telefuser/models/wan_video_dit.py b/telefuser/models/wan_video_dit.py index c87b085e..9048055b 100755 --- a/telefuser/models/wan_video_dit.py +++ b/telefuser/models/wan_video_dit.py @@ -10,7 +10,7 @@ from torch.distributed.device_mesh import DeviceMesh from telefuser.core.base_model import BaseModel -from telefuser.core.config import AttentionConfig, AttnImplType, OffloadConfig +from telefuser.core.config import AttentionConfig, AttnImplType, OffloadConfig, SparseAttentionConfig from telefuser.distributed.device_mesh import ( get_pp_group, get_pp_rank, @@ -37,7 +37,6 @@ from telefuser.offload.async_offload import AsyncOffloadManager from telefuser.ops.attention import MaskMap, SparseAttentionState from telefuser.ops.attention import attention as attn_func -from telefuser.ops.attention import long_context_attention as long_attn_func from telefuser.ops.normalization import LayerNorm, RMSNorm, fused_scale_shift, modulate from telefuser.ops.rotary import apply_rotary_emb from telefuser.utils.logging import logger @@ -116,6 +115,11 @@ def __init__(self, dim: int, num_heads: int, eps: float = 1e-6): self.norm_k = RMSNorm(dim, eps=eps) self.usp_flag = False + def _resolve_attention_config(self, sparse_state: SparseAttentionState | None) -> AttentionConfig: + if sparse_state is None and self.attention_config.is_sparse(): + return AttentionConfig.dense_attention(AttnImplType.FLASH_ATTN_2) + return self.attention_config + def async_usp_forward( self, x: torch.Tensor, @@ -145,32 +149,15 @@ def async_usp_forward( q = rearrange(q, "b n s d -> (b s) n d", s=seqlen, n=self.num_heads) k = rearrange(k, "b n s d -> (b s) n d", s=seqlen, n=self.num_heads) v = rearrange(v, "b n s d -> (b s) n d", s=seqlen, n=self.num_heads) - - from telefuser.core.config import AttentionConfig - - attention_config = AttentionConfig( - attn_impl=AttnImplType.RADIAL_ATTN, - sparse_config=sparse_state.config, - ) - x = attn_func( - q, - k, - v, - attention_config=attention_config, - sparse_state=sparse_state, - input_layout="BSND", - output_layout="BSND", - ) - else: - if self.attention_config.is_sparse(): - from telefuser.core.config import AttentionConfig - - dense_config = AttentionConfig.dense_attention(AttnImplType.FLASH_ATTN_2) - x = attn_func(q, k, v, attention_config=dense_config, input_layout="BSND", output_layout="BSND") - else: - x = attn_func( - q, k, v, attention_config=self.attention_config, input_layout="BSND", output_layout="BSND" - ) + x = attn_func( + q, + k, + v, + attention_config=self._resolve_attention_config(sparse_state), + sparse_state=sparse_state, + input_layout="BSND", + output_layout="BSND", + ) out_wait = ulysses_gather_heads(x, group, num_heads=self.num_heads) out = out_wait() out = rearrange(out, "b s n d -> b s (n d)", n=self.num_heads) @@ -205,75 +192,20 @@ def default_forward( q = rearrange(q, "b s (n d) -> b s n d", n=self.num_heads) k = rearrange(k, "b s (n d) -> b s n d", n=self.num_heads) v = rearrange(v, "b s (n d) -> b s n d", n=self.num_heads) - if self.usp_flag: - if sparse_state is not None and sparse_state.config.sparse_impl == "radial": - from telefuser.core.config import AttentionConfig - - attention_config = AttentionConfig( - attn_impl=AttnImplType.RADIAL_ATTN, sparse_config=sparse_state.config - ) - x = long_attn_func( - q, - k, - v, - attention_config=attention_config, - sparse_state=sparse_state, - input_layout="BSND", - output_layout="BSND", - device_mesh=device_mesh, - ) - else: - if self.attention_config.is_sparse(): - from telefuser.core.config import AttentionConfig - - dense_config = AttentionConfig.dense_attention(AttnImplType.FLASH_ATTN_2) - x = long_attn_func( - q, - k, - v, - input_layout="BSND", - output_layout="BSND", - device_mesh=device_mesh, - attention_config=dense_config, - ) - else: - x = long_attn_func( - q, - k, - v, - input_layout="BSND", - output_layout="BSND", - device_mesh=device_mesh, - attention_config=self.attention_config, - ) - elif sparse_state is not None and sparse_state.config.sparse_impl == "radial": + if sparse_state is not None and sparse_state.config.sparse_impl == "radial": seqlen = q.shape[2] q = rearrange(q, "b s n d -> (b s) n d", s=seqlen, n=self.num_heads) k = rearrange(k, "b s n d -> (b s) n d", s=seqlen, n=self.num_heads) v = rearrange(v, "b s n d -> (b s) n d", s=seqlen, n=self.num_heads) - - from telefuser.core.config import AttentionConfig - - attention_config = AttentionConfig(attn_impl=AttnImplType.RADIAL_ATTN, sparse_config=sparse_state.config) - x = attn_func( - q, - k, - v, - attention_config=attention_config, - sparse_state=sparse_state, - input_layout="BSND", - output_layout="BSND", - ) - else: - if self.attention_config.is_sparse(): - from telefuser.core.config import AttentionConfig - - dense_config = AttentionConfig.dense_attention(AttnImplType.FLASH_ATTN_2) - x = attn_func(q, k, v, input_layout="BSND", output_layout="BSND", attention_config=dense_config) - else: - x = attn_func( - q, k, v, input_layout="BSND", output_layout="BSND", attention_config=self.attention_config - ) + x = attn_func( + q, + k, + v, + attention_config=self._resolve_attention_config(sparse_state), + sparse_state=sparse_state, + input_layout="BSND", + output_layout="BSND", + ) x = rearrange(x, "b s n d -> b s (n d)", n=self.num_heads) return self.o(x) @@ -569,11 +501,14 @@ def forward_blocks( freqs_sin: torch.Tensor, sparse_state: SparseAttentionState | None = None, ) -> torch.Tensor: + x, t_mod, freqs_cos, freqs_sin = self._apply_sol_token_order( + x, t_mod, freqs_cos, freqs_sin, sparse_state, reorder_tokens=True + ) for block_id, block in enumerate(self.blocks): if sparse_state is not None: sparse_state.update(layer_idx=block_id) x = block(x, context, t_mod, freqs_cos, freqs_sin, sparse_state=sparse_state, device_mesh=self.device_mesh) - return x + return self._restore_sol_token_order(x, sparse_state) def forward_blocks_pp( self, @@ -588,10 +523,48 @@ def forward_blocks_pp( Only processes the blocks assigned to this stage. """ + x, t_mod, freqs_cos, freqs_sin = self._apply_sol_token_order( + x, + t_mod, + freqs_cos, + freqs_sin, + sparse_state, + reorder_tokens=self.pp_start_idx == 0, + ) for block_id, block in enumerate(self.blocks[self.pp_start_idx : self.pp_end_idx]): if sparse_state is not None: sparse_state.update(layer_idx=self.pp_start_idx + block_id) x = block(x, context, t_mod, freqs_cos, freqs_sin, sparse_state=sparse_state, device_mesh=self.device_mesh) + return self._restore_sol_token_order(x, sparse_state, enabled=self.pp_end_idx == len(self.blocks)) + + def _apply_sol_token_order( + self, + x: torch.Tensor, + t_mod: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + sparse_state: SparseAttentionState | None, + *, + reorder_tokens: bool, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if sparse_state is None or sparse_state.config.sparse_impl != "sol": + return x, t_mod, freqs_cos, freqs_sin + perm = self.sol_morton_perm.to(x.device) + if reorder_tokens: + x = x.index_select(1, perm) + if t_mod.shape[1] == perm.numel(): + t_mod = t_mod.index_select(1, perm) + return x, t_mod, freqs_cos.index_select(0, perm), freqs_sin.index_select(0, perm) + + def _restore_sol_token_order( + self, + x: torch.Tensor, + sparse_state: SparseAttentionState | None, + *, + enabled: bool = True, + ) -> torch.Tensor: + if enabled and sparse_state is not None and sparse_state.config.sparse_impl == "sol": + return x.index_select(1, self.sol_morton_inverse.to(x.device)) return x def enable_usp(self): @@ -1010,6 +983,55 @@ def enable_radial_attention( self.sparse_attention_state = SparseAttentionState(config=sparse_config, mask_map=mask_map, model_type="wan") logger.info(f"Radial attention initialized: video_token_num={video_token_num}") + def enable_sol_attention( + self, + height: int, + width: int, + num_frames: int, + sparse_config: SparseAttentionConfig, + ) -> None: + """Enable Sol-Attn with the official Wan Morton3D token ordering.""" + if sparse_config.sparse_impl != "sol": + raise ValueError("Sol-Attn requires a Sol sparse attention config") + latent_frames = (num_frames - 1) // 4 + 1 + grid = ( + latent_frames // self.patch_size[0], + height // (8 * self.patch_size[1]), + width // (8 * self.patch_size[2]), + ) + perm, inverse = self._morton3d_permutation(grid) + self.register_buffer("sol_morton_perm", perm, persistent=False) + self.register_buffer("sol_morton_inverse", inverse, persistent=False) + logger.info( + f"Enabling Sol-Attn: grid={grid}, dense_layers={sparse_config.dense_layers}, " + f"dense_timesteps={sparse_config.dense_timesteps}, tau={sparse_config.sol_tau}" + ) + self.sparse_attention_state = SparseAttentionState(config=sparse_config, mask_map=None, model_type="wan") + + @staticmethod + def _morton3d_permutation(grid: tuple[int, int, int]) -> tuple[torch.Tensor, torch.Tensor]: + """Return the canonical x/y/z-interleaved Morton permutation and inverse.""" + frames, height, width = grid + total = frames * height * width + linear = torch.arange(total, dtype=torch.long) + frame_area = height * width + z = linear // frame_area + remainder = linear - z * frame_area + y = remainder // width + x = remainder - y * width + + def part1by2(value: torch.Tensor) -> torch.Tensor: + value = value & 0x1FFFFF + value = (value | (value << 32)) & 0x1F00000000FFFF + value = (value | (value << 16)) & 0x1F0000FF0000FF + value = (value | (value << 8)) & 0x100F00F00F00F00F + value = (value | (value << 4)) & 0x10C30C30C30C30C3 + return (value | (value << 2)) & 0x1249249249249249 + + code = part1by2(x) | (part1by2(y) << 1) | (part1by2(z) << 2) + perm = linear[torch.argsort(code)] + return perm, torch.argsort(perm) + def create_sparse_state(self, numeral_timestep: int = 0, layer_idx: int = 0) -> SparseAttentionState | None: """Create/update sparse attention state for current step.""" if not hasattr(self, "sparse_attention_state"): diff --git a/telefuser/ops/attention/attention_impl.py b/telefuser/ops/attention/attention_impl.py index 9375fdde..eac66305 100755 --- a/telefuser/ops/attention/attention_impl.py +++ b/telefuser/ops/attention/attention_impl.py @@ -2,7 +2,7 @@ Supports multiple attention implementations: - Dense: TORCH_SDPA, TORCH_CUDNN, FLASH_ATTN_2/3/4, SAGE_ATTN variants, SPARGE_ATTN -- Sparse: RADIAL_ATTN, LOCAL_SPARSE_ATTN +- Sparse: RADIAL_ATTN, LOCAL_SPARSE_ATTN, SOL_ATTN Note: Attention functions are decorated with @torch.compiler.disable because: 1. SageAttention requires static tensor shapes for quantization scales @@ -33,6 +33,7 @@ FLASH_ATTN_4_AVAILABLE, SAGE_ATTN_AVAILABLE, SDPA_AVAILABLE, + SOL_ATTN_AVAILABLE, flash_attn2, flash_attn3, flash_attn4, @@ -40,6 +41,7 @@ get_lse_fallback_impl, sageattention, sdpa_attn_cudnn, + sol_attn, sparge_attn, supports_return_lse, ) @@ -135,7 +137,7 @@ class SparseAttentionState: def __init__( self, config: SparseAttentionConfig, - mask_map: MaskMap, + mask_map: MaskMap | None, model_type: str = "wan", ) -> None: self.config = config @@ -160,6 +162,15 @@ def get_sparsity_type(self) -> str: return "dense" if self.should_use_dense() else self.config.sparse_impl +def _resolve_sol_kv_splits(q: Tensor, kv_splits: int | str) -> int: + """Match the official Sol-Engine automatic split policy.""" + if kv_splits != "auto": + return int(kv_splits) + if torch.cuda.get_device_capability(q.device) == (9, 0) and q.shape[1] >= 65536: + return 4 + return 1 + + @torch.compiler.disable def attention( q: Tensor, @@ -216,14 +227,19 @@ def attention( raise ValueError("packed sequence attention requires TORCH_SDPA or FLASH_ATTN_4") # Handle sparse attention - if attn_impl in (AttnImplType.RADIAL_ATTN, AttnImplType.LOCAL_SPARSE_ATTN): + if attn_impl in (AttnImplType.RADIAL_ATTN, AttnImplType.LOCAL_SPARSE_ATTN, AttnImplType.SOL_ATTN): if sparse_state is None: msg = "Sparse attention requires sparse_state, falling back to FLASH_ATTN_2" if msg not in _warned_attn_fallback: _warned_attn_fallback.add(msg) logger.warning(msg) attn_impl = AttnImplType.FLASH_ATTN_2 + elif attn_impl == AttnImplType.SOL_ATTN: + if sparse_state.should_use_dense(): + attn_impl = AttnImplType.FLASH_ATTN_2 else: + if sparse_state.mask_map is None: + raise RuntimeError("Radial attention requires a mask map") return radial_attention( query=q, key=k, @@ -324,7 +340,17 @@ def attention( output = _packed_sdpa(q, k, v, sequence_lengths, scale=scale, is_causal=is_causal) # Sage Attention variants - elif sequence_lengths is None and SAGE_ATTN_AVAILABLE and sageattention is not None: + elif ( + sequence_lengths is None + and attn_impl + in { + AttnImplType.SAGE_ATTN_2_8_8, + AttnImplType.SAGE_ATTN_2_8_16, + AttnImplType.SAGE_ATTN_2_8_8_SM90, + } + and SAGE_ATTN_AVAILABLE + and sageattention is not None + ): # SageAttention tensor_layout: "NHD" for BSND, "HND" for BNSD sage_tensor_layout = "NHD" if current_layout == "BSND" else "HND" if attn_impl == AttnImplType.SAGE_ATTN_2_8_8: @@ -371,6 +397,39 @@ def attention( elif sequence_lengths is None and attn_impl == AttnImplType.SPARGE_ATTN: output = sparge_attn(q, k, v, attn_mask=attn_mask, scale=scale) + # Sol-Attn + elif attn_impl == AttnImplType.SOL_ATTN and SOL_ATTN_AVAILABLE and sol_attn is not None: + eligible = ( + attn_mask is None + and not is_causal + and not return_lse + and attention_config.dropout == 0.0 + and q.shape == k.shape == v.shape + and q.ndim == 4 + and q.shape[-1] == 128 + and q.dtype == torch.bfloat16 + and q.is_cuda + ) + if eligible: + sparse_config = attention_config.sparse_config + if sparse_config is None: + raise RuntimeError("Sol-Attn requires sparse attention configuration") + try: + output = sol_attn( + q.contiguous(), + k.contiguous(), + v.contiguous(), + scale=scale, + tau=sparse_config.sol_tau, + thresh_type=sparse_config.sol_threshold_type, + kv_splits=_resolve_sol_kv_splits(q, sparse_config.sol_kv_splits), + ) + except (RuntimeError, TypeError, ValueError) as error: + msg = "Sol-Attn execution failed, falling back to TORCH_SDPA" + if msg not in _warned_attn_fallback: + _warned_attn_fallback.add(msg) + logger.warning("%s: %s", msg, error) + # Fallback to SDPA if output is None: msg = f"Attention {attn_impl} not available, falling back to TORCH_SDPA" diff --git a/telefuser/ops/attention/backends.py b/telefuser/ops/attention/backends.py index 37f09f4e..1b76942b 100644 --- a/telefuser/ops/attention/backends.py +++ b/telefuser/ops/attention/backends.py @@ -23,6 +23,7 @@ SAGE_ATTN_AVAILABLE = False SPARGE_ATTN_AVAILABLE = False FLASHINFER_AVAILABLE = False +SOL_ATTN_AVAILABLE = False # Backend function references (populated on successful import) flash_attn2: Callable | None = None @@ -32,6 +33,7 @@ sageattention: object | None = None spas_sage2_attn_meansim_cuda: Callable | None = None flashinfer: object | None = None +sol_attn: Callable | None = None def _try_import_flash_attn() -> None: @@ -139,12 +141,33 @@ def _try_import_flashinfer() -> None: pass +def _try_import_sol_attn() -> None: + """Import TeleFuser's built-in Sol-Attn kernel.""" + global SOL_ATTN_AVAILABLE, sol_attn + + SOL_ATTN_AVAILABLE = False + sol_attn = None + module_name = "telefuser.kernel.sol_attn" + try: + if importlib.util.find_spec(module_name) is None: + return + candidate = getattr(importlib.import_module(module_name), "sol_attn", None) + except (ModuleNotFoundError, ImportError, RuntimeError) as error: + logger.debug("Built-in Sol-Attn backend unavailable: %s", error) + return + if callable(candidate): + sol_attn = candidate + SOL_ATTN_AVAILABLE = True + logger.debug("Built-in Sol-Attn kernel available") + + # Initialize all backends _try_import_flash_attn() _try_import_sdpa() _try_import_sage_attn() _try_import_sparge_attn() _try_import_flashinfer() +_try_import_sol_attn() def supports_return_lse(attn_impl: str) -> bool: @@ -210,12 +233,14 @@ def sparge_attn( "SAGE_ATTN_AVAILABLE", "SPARGE_ATTN_AVAILABLE", "FLASHINFER_AVAILABLE", + "SOL_ATTN_AVAILABLE", "flash_attn2", "flash_attn3", "flash_attn4", "flash_attn4_varlen", "sageattention", "flashinfer", + "sol_attn", "supports_return_lse", "get_lse_fallback_impl", "sdpa_attn_cudnn", diff --git a/telefuser/pipelines/wan_video/single_dit_denoising.py b/telefuser/pipelines/wan_video/single_dit_denoising.py index 829fcb95..f69915f1 100644 --- a/telefuser/pipelines/wan_video/single_dit_denoising.py +++ b/telefuser/pipelines/wan_video/single_dit_denoising.py @@ -193,10 +193,11 @@ def process( if ref_latent is not None: input_latent = torch.cat([input_latent, ref_latent], dim=1) - # Create sparse_state if radial attention is enabled + # Create sparse state when the configured attention backend needs it. sparse_state = None if has_sparse_attention: - numeral_timestep = num_inference_steps - progress_id - 1 + sparse_impl = self.dit.sparse_attention_state.config.sparse_impl + numeral_timestep = progress_id if sparse_impl == "sol" else num_inference_steps - progress_id - 1 sparse_state = self.dit.create_sparse_state( numeral_timestep=numeral_timestep, layer_idx=0, # Updated per layer in forward_blocks diff --git a/telefuser/pipelines/wan_video/wan21_video.py b/telefuser/pipelines/wan_video/wan21_video.py index 12d11ffa..a2f3955a 100644 --- a/telefuser/pipelines/wan_video/wan21_video.py +++ b/telefuser/pipelines/wan_video/wan21_video.py @@ -165,20 +165,31 @@ def __call__( attention_config = self.config.dit_config.attention_config if attention_config.is_sparse(): sparse_config = attention_config.sparse_config - self.denoise_stage.dit.enable_radial_attention( - height=height, - width=width, - num_frames=num_frames, - dense_layers=sparse_config.dense_layers, - dense_timesteps=sparse_config.dense_timesteps, - decay_factor=sparse_config.decay_factor, - use_sage_attention=sparse_config.use_sage_attention, - ) + if sparse_config is None: + raise ValueError("Sparse attention requires sparse configuration") + if sparse_config.sparse_impl == "radial": + self.denoise_stage.dit.enable_radial_attention( + height=height, + width=width, + num_frames=num_frames, + dense_layers=sparse_config.dense_layers, + dense_timesteps=sparse_config.dense_timesteps, + decay_factor=sparse_config.decay_factor, + use_sage_attention=sparse_config.use_sage_attention, + ) + elif sparse_config.sparse_impl == "sol": + self.denoise_stage.dit.enable_sol_attention( + height=height, + width=width, + num_frames=num_frames, + sparse_config=sparse_config, + ) + else: + raise ValueError(f"Unsupported Wan sparse attention: {sparse_config.sparse_impl}") logger.info( f"Sparse attention enabled ({sparse_config.sparse_impl}): " f"dense_layers={sparse_config.dense_layers}, " - f"dense_timesteps={sparse_config.dense_timesteps}, " - f"decay_factor={sparse_config.decay_factor}" + f"dense_timesteps={sparse_config.dense_timesteps}" ) # denoise diff --git a/tests/unit/kernel/test_sol_attn.py b/tests/unit/kernel/test_sol_attn.py new file mode 100644 index 00000000..4d730e16 --- /dev/null +++ b/tests/unit/kernel/test_sol_attn.py @@ -0,0 +1,62 @@ +from pathlib import Path + +import pytest +import torch + +try: + import tomllib +except ModuleNotFoundError: # Python 3.10 + import tomli as tomllib + +from telefuser.kernel.sol_attn.interface import _backend_for_arch, _validate_inputs, sol_attn + + +def test_sol_attn_is_packaged_with_telefuser() -> None: + project_root = Path(__file__).parents[3] + with (project_root / "pyproject.toml").open("rb") as config_file: + package_data = tomllib.load(config_file)["tool"]["setuptools"]["package-data"] + + assert package_data["telefuser.kernel.sol_attn"] == [ + "THIRD_PARTY_NOTICES.md", + "sm100/LICENSE.flash-attention", + ] + assert (project_root / "telefuser" / "kernel" / "sol_attn" / "__init__.py").is_file() + assert not (project_root / "tf-kernel" / "tf_kernel" / "_sol_attn").exists() + + +def test_backend_selection_prefers_cute_and_falls_back_to_triton() -> None: + assert _backend_for_arch((9, 0), cute_available=True) == "cute_sm90" + assert _backend_for_arch((10, 0), cute_available=True) == "cute_sm100" + assert _backend_for_arch((12, 0), cute_available=True) == "cute_sm120" + assert _backend_for_arch((9, 0), cute_available=False) == "triton" + assert _backend_for_arch((8, 0), cute_available=True) == "triton" + + with pytest.raises(RuntimeError, match="compute capability >= 8.0"): + _backend_for_arch((7, 5), cute_available=False) + + +def test_input_validation_rejects_cpu_tensors() -> None: + q = torch.randn(1, 64, 2, 128, dtype=torch.bfloat16) + + with pytest.raises(ValueError, match="same CUDA device"): + _validate_inputs(q, q, q, "diag") + + +@pytest.mark.gpu +def test_sol_attn_dense_limit_matches_sdpa_on_sm90() -> None: + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0): + pytest.skip("Sol-Attn SM90 correctness test requires H100") + + torch.manual_seed(0) + q = torch.randn(1, 256, 2, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + + output = sol_attn(q, k, v, tau=-1000.0, thresh_type="diag", kv_splits=1) + expected = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + ).transpose(1, 2) + + torch.testing.assert_close(output, expected, atol=0.05, rtol=0.02) diff --git a/tests/unit/models/test_wan_video_sol_attention.py b/tests/unit/models/test_wan_video_sol_attention.py new file mode 100644 index 00000000..d0c8afd6 --- /dev/null +++ b/tests/unit/models/test_wan_video_sol_attention.py @@ -0,0 +1,132 @@ +from unittest.mock import patch + +import pytest +import torch + +from telefuser.core.config import AttentionConfig, AttnImplType, SparseAttentionConfig +from telefuser.models.wan_video_dit import SelfAttention, WanModel, precompute_freqs_cis_3d +from telefuser.ops.attention import SparseAttentionState, attention_impl + + +def test_wan_model_enables_sol_attention_state() -> None: + model = WanModel.__new__(WanModel) + torch.nn.Module.__init__(model) + model.patch_size = (1, 2, 2) + + sparse_config = SparseAttentionConfig( + sparse_impl="sol", + dense_layers=2, + dense_timesteps=3, + sol_tau=0.75, + sol_threshold_type="exact", + sol_kv_splits=2, + ) + model.enable_sol_attention(height=64, width=64, num_frames=5, sparse_config=sparse_config) + + state = model.create_sparse_state(numeral_timestep=4, layer_idx=5) + assert state is not None + assert state.mask_map is None + assert state.numeral_timestep == 4 + assert state.layer_idx == 5 + assert state.config is sparse_config + perm = model.sol_morton_perm + inverse = model.sol_morton_inverse + torch.testing.assert_close(perm.index_select(0, inverse), torch.arange(perm.numel())) + + +def test_wan_sol_attention_uses_official_dense_guards() -> None: + model = WanModel.__new__(WanModel) + torch.nn.Module.__init__(model) + model.patch_size = (1, 2, 2) + config = AttentionConfig.sol_attention() + assert config.sparse_config is not None + model.enable_sol_attention(height=480, width=832, num_frames=81, sparse_config=config.sparse_config) + + state = model.create_sparse_state(numeral_timestep=0, layer_idx=1) + assert state is not None and state.should_use_dense() + state.update(numeral_timestep=10, layer_idx=0) + assert state.should_use_dense() + state.update(numeral_timestep=10, layer_idx=1) + assert not state.should_use_dense() + assert model.sol_morton_perm.numel() == 21 * 30 * 52 + + +def test_wan_sol_token_order_round_trip() -> None: + model = WanModel.__new__(WanModel) + torch.nn.Module.__init__(model) + model.patch_size = (1, 2, 2) + config = AttentionConfig.sol_attention() + assert config.sparse_config is not None + model.enable_sol_attention(height=64, width=64, num_frames=5, sparse_config=config.sparse_config) + state = model.create_sparse_state() + + tokens = model.sol_morton_perm.numel() + x = torch.arange(tokens).reshape(1, tokens, 1) + t_mod = x.clone() + freqs_cos = torch.arange(tokens).reshape(tokens, 1) + freqs_sin = -freqs_cos + + ordered = model._apply_sol_token_order(x, t_mod, freqs_cos, freqs_sin, state, reorder_tokens=True) + perm = model.sol_morton_perm + torch.testing.assert_close(ordered[0], x.index_select(1, perm)) + torch.testing.assert_close(ordered[1], t_mod.index_select(1, perm)) + torch.testing.assert_close(ordered[2], freqs_cos.index_select(0, perm)) + torch.testing.assert_close(ordered[3], freqs_sin.index_select(0, perm)) + torch.testing.assert_close(model._restore_sol_token_order(ordered[0], state), x) + + +def test_wan_self_attention_dispatches_sol_through_public_ops() -> None: + module = SelfAttention(dim=128, num_heads=1) + sparse_config = SparseAttentionConfig(sparse_impl="sol", dense_timesteps=0) + module.attention_config = AttentionConfig(attn_impl=AttnImplType.SOL_ATTN, sparse_config=sparse_config) + state = SparseAttentionState(sparse_config, mask_map=None) + captured = {} + + def fake_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs) -> torch.Tensor: + captured.update(kwargs) + return q + + x = torch.randn(1, 4, 128) + with ( + patch("telefuser.models.wan_video_dit.rope_apply", side_effect=lambda tensor, *_args: tensor), + patch("telefuser.models.wan_video_dit.attn_func", side_effect=fake_attention), + ): + output = module.default_forward(x, torch.empty(0), torch.empty(0), sparse_state=state) + + assert output.shape == x.shape + assert captured["attention_config"].attn_impl is AttnImplType.SOL_ATTN + assert captured["attention_config"].sparse_config is sparse_config + assert captured["sparse_state"] is state + + +@pytest.mark.gpu +def test_wan_self_attention_executes_sol_on_h100(monkeypatch: pytest.MonkeyPatch) -> None: + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0): + pytest.skip("Wan Sol-Attn execution test requires H100") + + assert attention_impl.SOL_ATTN_AVAILABLE + assert attention_impl.sol_attn is not None + kernel_calls = 0 + sol_attn = attention_impl.sol_attn + + def tracked_sol_attn(*args, **kwargs): + nonlocal kernel_calls + kernel_calls += 1 + return sol_attn(*args, **kwargs) + + monkeypatch.setattr(attention_impl, "sol_attn", tracked_sol_attn) + + module = SelfAttention(dim=128, num_heads=1).eval().cuda().to(torch.bfloat16) + x = torch.randn(1, 256, 128, device="cuda", dtype=torch.bfloat16) + freqs = precompute_freqs_cis_3d(128) + freqs_cos = torch.cat([freq.real for freq in freqs], dim=-1)[:256].cuda() + freqs_sin = torch.cat([freq.imag for freq in freqs], dim=-1)[:256].cuda() + sparse_config = SparseAttentionConfig(sparse_impl="sol", dense_timesteps=0, sol_tau=-1000.0) + module.attention_config = AttentionConfig(attn_impl=AttnImplType.SOL_ATTN, sparse_config=sparse_config) + state = SparseAttentionState(sparse_config, mask_map=None) + + output = module(x, freqs_cos, freqs_sin, sparse_state=state) + + assert output.shape == x.shape + assert torch.isfinite(output).all() + assert kernel_calls == 1 diff --git a/tests/unit/ops/test_sol_attention.py b/tests/unit/ops/test_sol_attention.py new file mode 100644 index 00000000..ab663fd1 --- /dev/null +++ b/tests/unit/ops/test_sol_attention.py @@ -0,0 +1,130 @@ +from types import ModuleType +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from telefuser.core.config import AttentionConfig, AttnImplType, SparseAttentionConfig +from telefuser.ops.attention import attention_impl, backends +from telefuser.ops.attention.attention_impl import SparseAttentionState + + +def test_sol_attention_config_defaults_and_validation() -> None: + config = AttentionConfig.sol_attention() + + assert config.attn_impl is AttnImplType.SOL_ATTN + assert config.is_sparse() + assert config.sparse_config == SparseAttentionConfig( + sparse_impl="sol", + dense_timesteps=10, + dense_layers=1, + sol_tau=1.0, + sol_threshold_type="diag", + sol_kv_splits="auto", + ) + + with pytest.raises(ValueError, match="threshold type"): + AttentionConfig.sol_attention(threshold_type="unknown") + with pytest.raises(ValueError, match="KV splits"): + AttentionConfig.sol_attention(kv_splits=3) + + +def test_sol_attention_loads_from_telefuser_kernel() -> None: + imported_modules: list[str] = [] + kernel_module = ModuleType("telefuser.kernel.sol_attn") + kernel_module.sol_attn = MagicMock() + previous_available = backends.SOL_ATTN_AVAILABLE + previous_backend = backends.sol_attn + + def import_module(name: str) -> ModuleType: + imported_modules.append(name) + return kernel_module + + try: + with ( + patch("telefuser.ops.attention.backends.importlib.util.find_spec", return_value=object()), + patch("telefuser.ops.attention.backends.importlib.import_module", side_effect=import_module), + ): + backends._try_import_sol_attn() + + assert imported_modules == ["telefuser.kernel.sol_attn"] + assert backends.SOL_ATTN_AVAILABLE is True + assert backends.sol_attn is kernel_module.sol_attn + finally: + backends.SOL_ATTN_AVAILABLE = previous_available + backends.sol_attn = previous_backend + + +def test_sol_attention_ineligible_input_falls_back_to_sdpa() -> None: + q = torch.randn(1, 8, 2, 16) + kernel = MagicMock() + config = AttentionConfig.sol_attention() + assert config.sparse_config is not None + state = SparseAttentionState(config.sparse_config, mask_map=None) + + with ( + patch.object(attention_impl, "SOL_ATTN_AVAILABLE", True), + patch.object(attention_impl, "sol_attn", kernel), + ): + output = attention_impl.attention(q, q, q, attention_config=config, sparse_state=state) + + expected = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), + q.transpose(1, 2), + q.transpose(1, 2), + ).transpose(1, 2) + torch.testing.assert_close(output, expected) + kernel.assert_not_called() + + +def test_sol_attention_dense_guard_does_not_call_kernel() -> None: + q = torch.randn(1, 8, 2, 16) + kernel = MagicMock() + config = AttentionConfig.sol_attention(dense_timesteps=1) + assert config.sparse_config is not None + state = SparseAttentionState(config.sparse_config, mask_map=None) + + with ( + patch.object(attention_impl, "SOL_ATTN_AVAILABLE", True), + patch.object(attention_impl, "sol_attn", kernel), + ): + output = attention_impl.attention(q, q, q, attention_config=config, sparse_state=state) + + assert output.shape == q.shape + kernel.assert_not_called() + + +@pytest.mark.gpu +def test_sol_attention_public_ops_matches_sdpa_on_h100(monkeypatch: pytest.MonkeyPatch) -> None: + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0): + pytest.skip("Sol-Attn public ops test requires H100") + + assert attention_impl.SOL_ATTN_AVAILABLE + assert attention_impl.sol_attn is not None + kernel_calls = 0 + sol_attn = attention_impl.sol_attn + + def tracked_sol_attn(*args, **kwargs): + nonlocal kernel_calls + kernel_calls += 1 + return sol_attn(*args, **kwargs) + + monkeypatch.setattr(attention_impl, "sol_attn", tracked_sol_attn) + + torch.manual_seed(0) + q = torch.randn(1, 256, 2, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + config = AttentionConfig.sol_attention(dense_timesteps=0, dense_layers=0, tau=-1000.0) + assert config.sparse_config is not None + state = SparseAttentionState(config.sparse_config, mask_map=None) + + output = attention_impl.attention(q, k, v, attention_config=config, sparse_state=state) + assert kernel_calls == 1 + expected = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + ).transpose(1, 2) + + torch.testing.assert_close(output, expected, atol=0.05, rtol=0.02)