From 04cc2cc4d5c879c06405de3325ee9c029ad6a7a2 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Wed, 5 Aug 2026 14:12:38 +0000 Subject: [PATCH] feat(attention): integrate Sol-Attn for Wan2.1 Vendor the official Sol-Attn kernels into TeleFuser and expose them through the compile-aware public attention ops with dense fallbacks. Adapt Wan2.1 to the official Morton3D token order, dense warmup schedule, layer guard, threshold policy, and automatic SM90 KV splitting. Keep the runtime integration compact by sharing sparse dispatch and token-order helpers. Add kernel, ops, model, and pipeline coverage, including H100 numerical checks and Morton round-trip validation. Document configuration, packaging, compatibility, and Wan2.1 usage. Verification: 29 related tests passed; ruff check and format checks passed; git diff --check passed. --- docs/en/attention.md | 37 +- docs/zh/attention.md | 36 +- examples/wan_video/README.md | 11 + pyproject.toml | 10 +- telefuser/core/config.py | 48 +- .../kernel/sol_attn/THIRD_PARTY_NOTICES.md | 23 + telefuser/kernel/sol_attn/__init__.py | 5 + telefuser/kernel/sol_attn/_vendor/__init__.py | 1 + .../sol_attn/_vendor/flash_attn/__init__.py | 1 + .../_vendor/flash_attn/cute/__init__.py | 1 + .../_vendor/flash_attn/cute/ampere_helpers.py | 103 + .../_vendor/flash_attn/cute/block_info.py | 139 + .../_vendor/flash_attn/cute/block_sparsity.py | 463 ++++ .../_vendor/flash_attn/cute/cute_dsl_utils.py | 129 + .../_vendor/flash_attn/cute/fast_math.py | 21 + .../_vendor/flash_attn/cute/flash_fwd.py | 1218 +++++++++ .../sol_attn/_vendor/flash_attn/cute/mask.py | 712 +++++ .../_vendor/flash_attn/cute/named_barrier.py | 47 + .../_vendor/flash_attn/cute/pack_gqa.py | 263 ++ .../_vendor/flash_attn/cute/pipeline.py | 402 +++ .../_vendor/flash_attn/cute/seqlen_info.py | 290 +++ .../_vendor/flash_attn/cute/softmax.py | 639 +++++ .../_vendor/flash_attn/cute/tile_scheduler.py | 1087 ++++++++ .../sol_attn/_vendor/flash_attn/cute/utils.py | 800 ++++++ telefuser/kernel/sol_attn/common/__init__.py | 5 + .../kernel/sol_attn/common/layout_utils.py | 130 + telefuser/kernel/sol_attn/common/runtime.py | 14 + telefuser/kernel/sol_attn/common/selector.py | 169 ++ telefuser/kernel/sol_attn/interface.py | 399 +++ telefuser/kernel/sol_attn/preprocess.py | 463 ++++ .../sol_attn/sm100/LICENSE.flash-attention | 29 + telefuser/kernel/sol_attn/sm100/__init__.py | 5 + telefuser/kernel/sol_attn/sm100/kernel.py | 5 + telefuser/kernel/sol_attn/sm100/mainloop.py | 1762 +++++++++++++ telefuser/kernel/sol_attn/sm100/math.py | 29 + telefuser/kernel/sol_attn/sm100/softmax.py | 156 ++ telefuser/kernel/sol_attn/sm100/tmem.py | 138 + telefuser/kernel/sol_attn/sm120/__init__.py | 5 + telefuser/kernel/sol_attn/sm120/kernel.py | 19 + telefuser/kernel/sol_attn/sm120/mainloop.py | 1172 +++++++++ telefuser/kernel/sol_attn/sm90/__init__.py | 5 + .../kernel/sol_attn/sm90/_compat/__init__.py | 1 + .../sol_attn/sm90/_compat/activation.py | 5 + .../sol_attn/sm90/_compat/copy_utils.py | 169 ++ .../sol_attn/sm90/_compat/cute_dsl_utils.py | 191 ++ .../sol_attn/sm90/_compat/layout_utils.py | 3 + .../sol_attn/sm90/_compat/sm90_utils.py | 173 ++ telefuser/kernel/sol_attn/sm90/atoms.py | 23 + telefuser/kernel/sol_attn/sm90/exact.py | 243 ++ telefuser/kernel/sol_attn/sm90/fwd.py | 111 + telefuser/kernel/sol_attn/sm90/kernel.py | 47 + telefuser/kernel/sol_attn/sm90/mainloop.py | 2312 +++++++++++++++++ .../kernel/sol_attn/sm90/split_combine.py | 600 +++++ .../kernel/sol_attn/triton_ref/__init__.py | 5 + telefuser/kernel/sol_attn/triton_ref/fwd.py | 475 ++++ .../kernel/sol_attn/triton_ref/preprocess.py | 389 +++ telefuser/models/wan_video_dit.py | 210 +- telefuser/ops/attention/attention_impl.py | 67 +- telefuser/ops/attention/backends.py | 25 + .../wan_video/single_dit_denoising.py | 5 +- telefuser/pipelines/wan_video/wan21_video.py | 33 +- tests/unit/kernel/test_sol_attn.py | 62 + .../models/test_wan_video_sol_attention.py | 132 + tests/unit/ops/test_sol_attention.py | 130 + 64 files changed, 16272 insertions(+), 130 deletions(-) create mode 100644 telefuser/kernel/sol_attn/THIRD_PARTY_NOTICES.md create mode 100644 telefuser/kernel/sol_attn/__init__.py create mode 100644 telefuser/kernel/sol_attn/_vendor/__init__.py create mode 100644 telefuser/kernel/sol_attn/_vendor/flash_attn/__init__.py create mode 100644 telefuser/kernel/sol_attn/_vendor/flash_attn/cute/__init__.py create mode 100644 telefuser/kernel/sol_attn/_vendor/flash_attn/cute/ampere_helpers.py create mode 100644 telefuser/kernel/sol_attn/_vendor/flash_attn/cute/block_info.py create mode 100644 telefuser/kernel/sol_attn/_vendor/flash_attn/cute/block_sparsity.py create mode 100644 telefuser/kernel/sol_attn/_vendor/flash_attn/cute/cute_dsl_utils.py create mode 100644 telefuser/kernel/sol_attn/_vendor/flash_attn/cute/fast_math.py create mode 100644 telefuser/kernel/sol_attn/_vendor/flash_attn/cute/flash_fwd.py create mode 100644 telefuser/kernel/sol_attn/_vendor/flash_attn/cute/mask.py create mode 100644 telefuser/kernel/sol_attn/_vendor/flash_attn/cute/named_barrier.py create mode 100644 telefuser/kernel/sol_attn/_vendor/flash_attn/cute/pack_gqa.py create mode 100644 telefuser/kernel/sol_attn/_vendor/flash_attn/cute/pipeline.py create mode 100644 telefuser/kernel/sol_attn/_vendor/flash_attn/cute/seqlen_info.py create mode 100644 telefuser/kernel/sol_attn/_vendor/flash_attn/cute/softmax.py create mode 100644 telefuser/kernel/sol_attn/_vendor/flash_attn/cute/tile_scheduler.py create mode 100644 telefuser/kernel/sol_attn/_vendor/flash_attn/cute/utils.py create mode 100644 telefuser/kernel/sol_attn/common/__init__.py create mode 100644 telefuser/kernel/sol_attn/common/layout_utils.py create mode 100644 telefuser/kernel/sol_attn/common/runtime.py create mode 100644 telefuser/kernel/sol_attn/common/selector.py create mode 100644 telefuser/kernel/sol_attn/interface.py create mode 100644 telefuser/kernel/sol_attn/preprocess.py create mode 100644 telefuser/kernel/sol_attn/sm100/LICENSE.flash-attention create mode 100644 telefuser/kernel/sol_attn/sm100/__init__.py create mode 100644 telefuser/kernel/sol_attn/sm100/kernel.py create mode 100644 telefuser/kernel/sol_attn/sm100/mainloop.py create mode 100644 telefuser/kernel/sol_attn/sm100/math.py create mode 100644 telefuser/kernel/sol_attn/sm100/softmax.py create mode 100644 telefuser/kernel/sol_attn/sm100/tmem.py create mode 100644 telefuser/kernel/sol_attn/sm120/__init__.py create mode 100644 telefuser/kernel/sol_attn/sm120/kernel.py create mode 100644 telefuser/kernel/sol_attn/sm120/mainloop.py create mode 100644 telefuser/kernel/sol_attn/sm90/__init__.py create mode 100644 telefuser/kernel/sol_attn/sm90/_compat/__init__.py create mode 100644 telefuser/kernel/sol_attn/sm90/_compat/activation.py create mode 100644 telefuser/kernel/sol_attn/sm90/_compat/copy_utils.py create mode 100644 telefuser/kernel/sol_attn/sm90/_compat/cute_dsl_utils.py create mode 100644 telefuser/kernel/sol_attn/sm90/_compat/layout_utils.py create mode 100644 telefuser/kernel/sol_attn/sm90/_compat/sm90_utils.py create mode 100644 telefuser/kernel/sol_attn/sm90/atoms.py create mode 100644 telefuser/kernel/sol_attn/sm90/exact.py create mode 100644 telefuser/kernel/sol_attn/sm90/fwd.py create mode 100644 telefuser/kernel/sol_attn/sm90/kernel.py create mode 100644 telefuser/kernel/sol_attn/sm90/mainloop.py create mode 100644 telefuser/kernel/sol_attn/sm90/split_combine.py create mode 100644 telefuser/kernel/sol_attn/triton_ref/__init__.py create mode 100644 telefuser/kernel/sol_attn/triton_ref/fwd.py create mode 100644 telefuser/kernel/sol_attn/triton_ref/preprocess.py create mode 100644 tests/unit/kernel/test_sol_attn.py create mode 100644 tests/unit/models/test_wan_video_sol_attention.py create mode 100644 tests/unit/ops/test_sol_attention.py 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)