Add Sol-Attn Kernel Builder source

Pinned to NVlabs/Sana commit 8a26fb0ec9e353125ead798cb2e312d5ce48cded; packaging changes are limited to version-isolated relative imports.

.gitignore

@@ -0,0 +1,4 @@
+.hf-venv/
+build/
+__pycache__/
+*.pyc

CARD.md

@@ -0,0 +1,57 @@
+---
+library_name: kernels
+{% if license %}license: {{ license }}
+{% endif %}tags:
+- kernels
+- cuda
+- attention
+- triton
+- cute-dsl
+---
+
+# Sol-Attn
+
+Sol-Attn accelerates image and video generation with on-the-fly attention
+sparsification. The public API dispatches to CuTe DSL kernels on SM90, SM100,
+and SM120, and to Triton on SM80 and SM89 or when CuTe DSL is unavailable.
+
+## Usage
+
+```python
+from kernels import get_kernel
+
+kernel = get_kernel("{{ repo_id }}", version={{ version }})
+
+out = kernel.sol_attn(
+    q,  # Contiguous BF16 CUDA tensor [batch, tokens, heads, 128].
+    k,  # Same shape, dtype, layout, and device as q.
+    v,  # Same shape, dtype, layout, and device as q.
+    tau=1.0,
+    thresh_type="exact",
+)
+```
+
+The released implementation is noncausal and forward-only. Q/K/V must have
+the same BTHD shape. An optional exact KV sink is available through
+`sink_start` and `sink_tokens`.
+
+## Backends
+
+| Architecture | Example GPU | Backend |
+|---|---|---|
+| SM90 | H100 | CuTe DSL |
+| SM100 | GB200 | CuTe DSL |
+| SM120 | RTX 5090 | CuTe DSL |
+| SM80 / SM89 | A100 / RTX 4090 | Triton |
+
+## Paper
+
+[Accelerating Video Generation Inference via On-the-Fly Attention
+Sparsification](https://arxiv.org/abs/2607.24027)
+
+## Source
+
+The implementation is maintained in
+[`NVlabs/Sana`](https://github.com/NVlabs/Sana/tree/sol-engine/techniques/sparse_backends/sol_attn).
+This release is pinned to commit
+[`8a26fb0`](https://github.com/NVlabs/Sana/commit/8a26fb0ec9e353125ead798cb2e312d5ce48cded).

README.md

@@ -0,0 +1,40 @@
+---
+license: apache-2.0
+tags:
+- kernels
+- cuda
+- attention
+- triton
+- cute-dsl
+---
+
+# Sol-Attn Kernel Builder Source
+
+This repository contains the Hugging Face Kernel Builder packaging for
+[Sol-Attn](https://github.com/NVlabs/Sana/tree/sol-engine/techniques/sparse_backends/sol_attn).
+The kernel implementation is pinned to NVIDIA's `NVlabs/Sana` commit
+[`8a26fb0`](https://github.com/NVlabs/Sana/commit/8a26fb0ec9e353125ead798cb2e312d5ce48cded).
+
+The packaged Python sources preserve that implementation. The only source
+relocation required by Kernel Hub is converting internal `sol_attn.*` imports
+to package-relative imports, so the kernel remains loadable under the
+version-isolated module name assigned by `kernels.get_kernel(...)`.
+
+Published builds are loaded from
+[`Efficient-Large-Model/Sol-Attn`](https://huggingface.co/Efficient-Large-Model/Sol-Attn):
+
+```python
+from kernels import get_kernel
+
+kernel = get_kernel("Efficient-Large-Model/Sol-Attn", version=1)
+
+out = kernel.sol_attn(
+    q,  # Contiguous BF16 CUDA tensor [batch, tokens, heads, 128].
+    k,  # Same shape, dtype, layout, and device as q.
+    v,  # Same shape, dtype, layout, and device as q.
+    tau=1.0,
+    thresh_type="exact",
+)
+```
+
+See [SOURCE.md](SOURCE.md) for provenance and the verification command.

SOURCE.md

@@ -0,0 +1,23 @@
+# Source provenance
+
+- Upstream repository: <https://github.com/NVlabs/Sana>
+- Upstream branch: `sol-engine`
+- Upstream commit: `8a26fb0ec9e353125ead798cb2e312d5ce48cded`
+- Upstream path: `techniques/sparse_backends/sol_attn`
+- Kernel Hub target: `Efficient-Large-Model/Sol-Attn`
+
+`torch-ext/sol_attn` contains the complete upstream package at the pinned
+commit. Its only packaging-level changes are semantically equivalent
+absolute-to-relative internal imports. Kernel Hub loads each kernel version
+under an isolated module name, so hard-coded `sol_attn.*` package imports
+would otherwise escape that namespace.
+
+Given a checkout of the pinned Sana commit, verify the file inventory, all
+non-import source text, and import semantics with:
+
+```bash
+python tools/verify_upstream.py /path/to/Sana
+```
+
+The publishing process does not add files, workflows, or commits to
+`NVlabs/Sana`.

benchmarks/benchmark.py

@@ -0,0 +1,58 @@
+import importlib
+
+import torch
+from kernels.benchmark import Benchmark
+
+
+class SolAttnBenchmark(Benchmark):
+    seed: int = 42
+
+    def _setup(self, tokens, heads):
+        self.q = torch.randn(
+            1,
+            tokens,
+            heads,
+            128,
+            device=self.device,
+            dtype=torch.bfloat16,
+        )
+        self.k = torch.randn_like(self.q)
+        self.v = torch.randn_like(self.q)
+        module = importlib.import_module(f"{self.kernel.__name__}.triton_ref")
+        self.triton_sol_attn = module.sol_attn
+
+    def _run(self):
+        self.out = self.kernel.sol_attn(
+            self.q,
+            self.k,
+            self.v,
+            tau=1.0,
+            thresh_type="exact",
+        )
+
+    def _reference(self):
+        return self.triton_sol_attn(
+            self.q,
+            self.k,
+            self.v,
+            tau=1.0,
+            thresh_type="exact",
+        )
+
+    def setup_base(self):
+        self._setup(tokens=4096, heads=8)
+
+    def benchmark_base(self):
+        self._run()
+
+    def verify_base(self):
+        return self._reference()
+
+    def setup_video(self):
+        self._setup(tokens=16384, heads=16)
+
+    def benchmark_video(self):
+        self._run()
+
+    def verify_video(self):
+        return self._reference()

build.toml

@@ -0,0 +1,21 @@
+[general]
+name = "sol-attn"
+version = 1
+edition = 5
+license = "Apache-2.0"
+backends = ["cuda"]
+python-depends = ["tvm-ffi"]
+upstream = "https://github.com/NVlabs/Sana.git"
+source = "https://huggingface.co/Efficient-Large-Model/Sol-Attn-Kernel-Source"
+
+[general.cuda]
+minver = "12.8"
+python-depends = ["nvidia-cutlass-dsl"]
+
+[general.hub]
+repo-id = "Efficient-Large-Model/Sol-Attn"
+
+[torch-noarch]
+cuda-capabilities = ["8.0", "8.9", "9.0", "10.0", "12.0"]
+
+[kernel]

example.py

@@ -0,0 +1,18 @@
+import torch
+from kernels import get_kernel
+
+
+kernel = get_kernel("Efficient-Large-Model/Sol-Attn", version=1)
+
+q = torch.randn(1, 4096, 16, 128, device="cuda", dtype=torch.bfloat16)
+k = torch.randn_like(q)
+v = torch.randn_like(q)
+
+out = kernel.sol_attn(
+    q,
+    k,
+    v,
+    tau=1.0,
+    thresh_type="exact",
+)
+print(out.shape)

flake.lock

@@ -0,0 +1,117 @@
+{
+  "nodes": {
+    "flake-compat": {
+      "locked": {
+        "lastModified": 1767039857,
+        "narHash": "sha256-vNpUSpF5Nuw8xvDLj2KCwwksIbjua2LZCqhV1LNRDns=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "5edf11c44bc78a0d334f6334cdaf7d60d732daab",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "kernel-builder": {
+      "inputs": {
+        "flake-compat": "flake-compat",
+        "flake-utils": "flake-utils",
+        "nixpkgs": "nixpkgs",
+        "rust-overlay": "rust-overlay"
+      },
+      "locked": {
+        "lastModified": 1785782967,
+        "narHash": "sha256-5Esv3PsgzKfP6I3Te1n4Ge4kQG69WQETPJoxGB9WKpk=",
+        "owner": "huggingface",
+        "repo": "kernels",
+        "rev": "a6564d1f481adcbd3273099c0f4432e7b833e846",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "repo": "kernels",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1783284758,
+        "narHash": "sha256-tiQ8/qi8I45OOaBBYlVbXoAVkeQzvvTQOv5I45rMw5o=",
+        "owner": "NixOS",
+        "repo": "nixpkgs",
+        "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d",
+        "type": "github"
+      },
+      "original": {
+        "owner": "NixOS",
+        "repo": "nixpkgs",
+        "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "kernel-builder": "kernel-builder"
+      }
+    },
+    "rust-overlay": {
+      "inputs": {
+        "nixpkgs": [
+          "kernel-builder",
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1783320166,
+        "narHash": "sha256-l7C/OsjcnWDOk2K3ssj+SBduwL67LashjBqis9+t468=",
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "rev": "20ee15370c9256669d66968b89ee20a4b0a4e673",
+        "type": "github"
+      },
+      "original": {
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "type": "github"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}

flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "Flake for Sol-Attn kernels";
+
+  inputs = {
+    kernel-builder.url = "github:huggingface/kernels";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

tests/test_sol_attn.py

@@ -0,0 +1,89 @@
+import importlib
+
+import pytest
+import torch
+import torch.nn.functional as F
+from kernels import get_kernel
+
+
+kernel = get_kernel("Efficient-Large-Model/Sol-Attn", version=1)
+
+
+def _inputs(tokens=256, heads=4):
+    torch.manual_seed(42)
+    q = torch.randn(
+        1,
+        tokens,
+        heads,
+        128,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    return q, torch.randn_like(q), torch.randn_like(q)
+
+
+@pytest.mark.kernels_ci
+def test_backend_dispatch_contract():
+    interface = importlib.import_module(f"{kernel.__name__}.interface")
+
+    assert interface._backend_for_arch((8, 0), cute_available=True) == "triton"
+    assert interface._backend_for_arch((8, 9), cute_available=True) == "triton"
+    assert interface._backend_for_arch((9, 0), cute_available=True) == "cute_sm90"
+    assert interface._backend_for_arch((10, 0), cute_available=True) == "cute_sm100"
+    assert interface._backend_for_arch((12, 0), cute_available=True) == "cute_sm120"
+    assert interface._backend_for_arch((9, 0), cute_available=False) == "triton"
+    assert interface._backend_for_arch((10, 0), cute_available=False) == "triton"
+    assert interface._backend_for_arch((12, 0), cute_available=False) == "triton"
+
+
+@pytest.mark.kernels_ci
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+def test_full_sink_matches_sdpa():
+    capability = torch.cuda.get_device_capability()
+    if capability[0] < 8:
+        pytest.skip("Sol-Attn requires compute capability 8.0 or newer")
+
+    q, k, v = _inputs()
+    expected = F.scaled_dot_product_attention(
+        q.transpose(1, 2),
+        k.transpose(1, 2),
+        v.transpose(1, 2),
+    ).transpose(1, 2)
+    actual = kernel.sol_attn(
+        q,
+        k,
+        v,
+        tau=1.0,
+        thresh_type="exact",
+        sink_start=0,
+        sink_tokens=q.shape[1],
+    )
+
+    torch.testing.assert_close(actual, expected, atol=2e-2, rtol=3e-2)
+
+
+@pytest.mark.kernels_ci
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+def test_selected_backend_matches_triton_reference():
+    capability = torch.cuda.get_device_capability()
+    if capability[0] < 8:
+        pytest.skip("Sol-Attn requires compute capability 8.0 or newer")
+
+    triton_ref = importlib.import_module(f"{kernel.__name__}.triton_ref")
+    q, k, v = _inputs()
+    expected = triton_ref.sol_attn(
+        q,
+        k,
+        v,
+        tau=1.0,
+        thresh_type="exact",
+    )
+    actual = kernel.sol_attn(
+        q,
+        k,
+        v,
+        tau=1.0,
+        thresh_type="exact",
+    )
+
+    torch.testing.assert_close(actual, expected, atol=2e-2, rtol=3e-2)

tools/verify_upstream.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+"""Verify that the packaged sources only relocate internal imports."""
+
+from __future__ import annotations
+
+import ast
+import importlib.util
+import subprocess
+import sys
+from pathlib import Path
+
+
+PINNED_COMMIT = "8a26fb0ec9e353125ead798cb2e312d5ce48cded"
+UPSTREAM_SUBDIR = Path("techniques/sparse_backends/sol_attn")
+IGNORED_PARTS = {"__pycache__"}
+
+
+def source_files(root: Path) -> dict[Path, Path]:
+    return {
+        path.relative_to(root): path
+        for path in root.rglob("*")
+        if path.is_file() and not (set(path.parts) & IGNORED_PARTS)
+    }
+
+
+def module_package(relative_path: Path) -> str:
+    parent_parts = relative_path.with_suffix("").parts[:-1]
+    return ".".join(("sol_attn", *parent_parts))
+
+
+def normalized_imports(tree: ast.AST, package: str):
+    imports = []
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Import):
+            imports.extend(
+                (alias.name, alias.asname or alias.name.split(".")[0])
+                for alias in node.names
+            )
+        elif isinstance(node, ast.ImportFrom):
+            module = node.module or ""
+            if node.level:
+                relative = "." * node.level + module
+                module = importlib.util.resolve_name(relative, package)
+            imports.extend(
+                (
+                    f"{module}.{alias.name}" if module else alias.name,
+                    alias.asname or alias.name,
+                )
+                for alias in node.names
+            )
+    return sorted(imports)
+
+
+def without_imports(text: str) -> str:
+    tree = ast.parse(text)
+    lines = text.splitlines(keepends=True)
+    for node in ast.walk(tree):
+        if isinstance(node, (ast.Import, ast.ImportFrom)):
+            for index in range(node.lineno - 1, node.end_lineno):
+                lines[index] = "\n" if lines[index].endswith("\n") else ""
+    return "".join(lines)
+
+
+def main() -> int:
+    if len(sys.argv) != 2:
+        print("usage: verify_upstream.py /path/to/Sana", file=sys.stderr)
+        return 2
+
+    checkout = Path(sys.argv[1]).resolve()
+    commit = subprocess.check_output(
+        ["git", "-C", str(checkout), "rev-parse", "HEAD"],
+        text=True,
+    ).strip()
+    if commit != PINNED_COMMIT:
+        raise SystemExit(f"expected Sana {PINNED_COMMIT}, found {commit}")
+
+    upstream_root = checkout / UPSTREAM_SUBDIR
+    packaged_root = Path(__file__).resolve().parents[1] / "torch-ext/sol_attn"
+    upstream = source_files(upstream_root)
+    packaged = source_files(packaged_root)
+    if upstream.keys() != packaged.keys():
+        missing = sorted(upstream.keys() - packaged.keys())
+        extra = sorted(packaged.keys() - upstream.keys())
+        raise SystemExit(f"source inventory mismatch; missing={missing}, extra={extra}")
+
+    relocated = []
+    for relative_path in sorted(upstream):
+        upstream_bytes = upstream[relative_path].read_bytes()
+        packaged_bytes = packaged[relative_path].read_bytes()
+        if upstream_bytes == packaged_bytes:
+            continue
+        if relative_path.suffix != ".py":
+            raise SystemExit(f"non-Python source differs: {relative_path}")
+
+        upstream_text = upstream_bytes.decode()
+        packaged_text = packaged_bytes.decode()
+        if without_imports(upstream_text) != without_imports(packaged_text):
+            raise SystemExit(f"non-import source differs: {relative_path}")
+
+        package = module_package(relative_path)
+        upstream_imports = normalized_imports(ast.parse(upstream_text), package)
+        packaged_imports = normalized_imports(ast.parse(packaged_text), package)
+        if upstream_imports != packaged_imports:
+            raise SystemExit(f"import semantics differ: {relative_path}")
+        relocated.append(relative_path)
+
+    print(
+        "Verified pinned NVlabs/Sana source; only semantically equivalent "
+        f"import relocation is present in {len(relocated)} files."
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

torch-ext/sol_attn/THIRD_PARTY_NOTICES.md

@@ -0,0 +1,15 @@
+# Third-party notices
+
+The files under `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
+`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.

torch-ext/sol_attn/__init__.py

@@ -0,0 +1,5 @@
+"""Sol-Attn."""
+
+from .interface import sol_attn
+
+__all__ = ["sol_attn"]

torch-ext/sol_attn/_vendor/__init__.py

@@ -0,0 +1 @@
+"""Private source dependencies bundled with Sol-Attn."""

torch-ext/sol_attn/_vendor/flash_attn/__init__.py

@@ -0,0 +1 @@
+"""Local FlashAttention Cute shim for the SOL_ATTN SM90 release."""

torch-ext/sol_attn/_vendor/flash_attn/cute/__init__.py

@@ -0,0 +1 @@
+"""Vendored FlashAttention Cute Python helpers used by SOL_ATTN SM90."""

torch-ext/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()

torch-ext/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 ...._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))

torch-ext/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 ...._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

torch-ext/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())

torch-ext/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

torch-ext/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 ....sm90._compat import copy_utils
+from ....sm90._compat import layout_utils
+
+from ...._vendor.flash_attn.cute import ampere_helpers as sm80_utils
+from ...._vendor.flash_attn.cute.cute_dsl_utils import assume_tensor_aligned
+from ...._vendor.flash_attn.cute import utils
+from ...._vendor.flash_attn.cute.mask import AttentionMask
+from ...._vendor.flash_attn.cute.softmax import Softmax
+from ...._vendor.flash_attn.cute.seqlen_info import SeqlenInfoQK
+from ...._vendor.flash_attn.cute.block_info import BlockInfo
+from ...._vendor.flash_attn.cute.pack_gqa import PackGQA
+from ...._vendor.flash_attn.cute.named_barrier import NamedBarrierFwd
+from ...._vendor.flash_attn.cute.block_sparsity import BlockSparseTensors
+from ...._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}")

torch-ext/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 ....sm90._compat import layout_utils
+from . import utils
+from ...._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,
+                    )

torch-ext/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()

torch-ext/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 ....sm90._compat import layout_utils
+from . import 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,
+                    )

torch-ext/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)

torch-ext/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 ....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,
+        )

torch-ext/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 ....sm90._compat import layout_utils
+from . import utils
+from ....sm90._compat.cute_dsl_utils import ParamsBase
+from ...._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]

torch-ext/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 ....sm90._compat.cute_dsl_utils import ParamsBase
+
+from . import utils
+from ...._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)

torch-ext/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 ....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]

torch-ext/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"]

torch-ext/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",
+]

torch-ext/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"]

torch-ext/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",
+]

torch-ext/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"]

torch-ext/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"]

torch-ext/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.

torch-ext/sol_attn/sm100/__init__.py

@@ -0,0 +1,5 @@
+"""Blackwell backend."""
+
+from .kernel import forward
+
+__all__ = ["forward"]

torch-ext/sol_attn/sm100/kernel.py

@@ -0,0 +1,5 @@
+"""Blackwell kernel entry."""
+
+from .mainloop import forward
+
+__all__ = ["forward"]

torch-ext/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
+from .._vendor.flash_attn.cute import pipeline as fa_pipeline
+from .._vendor.flash_attn.cute import 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 .._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 ..common import layout_utils
+from ..common.selector import (
+    sol_attn_popc_b32,
+    sol_attn_route_is_exact,
+)
+from .tmem import (
+    _add_physical_tmem_base,
+    _zero_based_tmem_tensor,
+    load_m64_o_fp32_256b,
+    tcgen05_wait_st,
+)
+
+
+M = 64
+N_MEMBER = 64
+N_PACK_HALF = 128
+D = 128
+DV = 128
+THREADS = 192
+PAIR_STAGES = 1
+TMEM_COLS = 256
+PAIR_SCORE_OFFSET = 0
+PAIR_P_OFFSET = 64
+O_OFFSET = 128
+PACK_QK_INST = (M, N_PACK_HALF, 16)
+PACK_QK_TILE = (M, N_PACK_HALF, D)
+PACK_PV_INST = (M, DV, 16)
+PACK_PV_TILE = (M, DV, N_PACK_HALF)
+PACK_QK_QUARTER_INST = (M, N_MEMBER, 16)
+PACK_PV_QUARTER_INST = (M, 64, 16)
+PACK_QK_GATHER_TILE = (M, N_MEMBER, 64)
+PACK_PV_GATHER_TILE = (M, 64, 64)
+LOG2E = math.log2(math.e)
+LN2 = math.log(2.0)
+SEMANTIC_ROW_OFFSET = 16
+LOGICAL_GROUP_SIZE = 256
+ROUTE_TILE_SIZE = 128
+ROUTE_HALVES_PER_GROUP = LOGICAL_GROUP_SIZE // ROUTE_TILE_SIZE
+ROUTE_MASK_WORDS = 4
+# masks[0:4], current-half exact count, append base, cumulative exact count,
+# logical-terminal-half flag
+PACKET_WORDS = 8
+ROUTE_INDEX_CAPACITY = LOGICAL_GROUP_SIZE
+PAIR_P_CHUNKS = 4
+PAIR_P_CHUNK_PACKED_COLUMNS = (N_PACK_HALF // 2) // PAIR_P_CHUNKS
+PAIR_P_PACKED_REGISTERS_PER_THREAD_PER_CHUNK = 8
+O_PACKED_STORE_VALUES_PER_WORD = 2
+O_PACKED_STORE_ALIGNMENT_BYTES = 4
+O_PACKED_STORE_WRITER_THREADS = 4 * 32
+O_ROWS_PER_OWNER_THREAD = 2
+O_PACKED_WORDS_PER_ROW_PER_THREAD = 16
+O_PACKED_COLUMN_STRIDE = 8
+
+@dsl_user_op
+def _cvt_bf16x2_f32(
+    hi: Float32,
+    lo: Float32,
+    *,
+    loc=None,
+    ip=None,
+) -> Int32:
+    """Round two FP32 values and pack them as ``{lo, hi}`` BF16 bits."""
+
+    return Int32(
+        llvm.inline_asm(
+            T.i32(),
+            [
+                Float32(hi).ir_value(loc=loc, ip=ip),
+                Float32(lo).ir_value(loc=loc, ip=ip),
+            ],
+            "cvt.rn.bf16x2.f32 $0, $1, $2;",
+            "=r,f,f",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
+@dsl_user_op
+def _store_global_u32_inline(
+    ptr: cute.Pointer,
+    value: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    """Store one aligned same-row BF16 pair as a single 32-bit word."""
+
+    llvm.inline_asm(
+        None,
+        [
+            ptr.toint().ir_value(),
+            Int32(value).ir_value(loc=loc, ip=ip),
+        ],
+        "st.global.u32 [$0], $1;",
+        "l,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+@dsl_user_op
+def _prmt_b32(
+    a: Int32,
+    b: Int32,
+    sel: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> Int32:
+    """Select four bytes from packed words ``a`` and ``b``."""
+
+    return Int32(
+        llvm.inline_asm(
+            T.i32(),
+            [
+                Int32(a).ir_value(loc=loc, ip=ip),
+                Int32(b).ir_value(loc=loc, ip=ip),
+                Int32(sel).ir_value(loc=loc, ip=ip),
+            ],
+            "prmt.b32 $0, $1, $2, $3;",
+            "=r,r,r,r",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
+@cute.jit
+def _store_pair_probability_chunked_tmemp(
+    o_template: cute.Tensor,
+    probabilities: cute.Tensor,
+    tmem_base: Int32,
+    p_offset: Int32,
+    owner_tidx: Int32,
+):
+    """Store M64xN128 BF16 P as four live-range-bounded x8 chunks.
+
+    Probabilities remain FP32 until each x8 fragment is converted, and every
+    chunk waits for its St16x64b store before the fragment goes out of scope.
+    """
+
+    assert o_template.element_type == Float32
+    assert cute.size(o_template) == M * DV
+    p_chunk_layout = cute.composition(
+        o_template.layout,
+        cute.make_layout((M, PAIR_P_CHUNK_PACKED_COLUMNS)),
+    )
+    relative_chunk = _zero_based_tmem_tensor(Float32, p_chunk_layout)
+    store_atom = cute.make_copy_atom(
+        tcgen05.copy.St16x64bOp(tcgen05.copy.Repetition(8)),
+        Float32,
+    )
+    tiled_store = tcgen05.make_tmem_copy(store_atom, relative_chunk)
+    thread_store = tiled_store.get_slice(owner_tidx)
+    destination_relative = thread_store.partition_D(relative_chunk)
+    destination = _add_physical_tmem_base(
+        destination_relative, tmem_base + p_offset
+    )
+    p_store_coordinates = thread_store.partition_S(
+        cute.make_identity_tensor((M, PAIR_P_CHUNK_PACKED_COLUMNS))
+    )
+    lane = owner_tidx % Int32(32)
+
+    for chunk_idx in cutlass.range_constexpr(PAIR_P_CHUNKS):
+        p_store_registers = cute.make_rmem_tensor(
+            p_store_coordinates.shape, Float32
+        )
+        assert (
+            cute.size(p_store_registers)
+            == PAIR_P_PACKED_REGISTERS_PER_THREAD_PER_CHUNK
+        )
+        assert (
+            cute.size(probabilities)
+            == 2 * cute.size(p_store_registers) * PAIR_P_CHUNKS
+        )
+        p_store_words = cute.make_tensor(
+            cute.recast_ptr(p_store_registers.iterator, dtype=Int32),
+            p_store_registers.layout,
+        )
+        probability_base = chunk_idx * (2 * cute.size(p_store_registers))
+        for i in cutlass.range(
+            cute.size(p_store_registers), unroll_full=True
+        ):
+            low = probability_base + i * 2
+            high = low + 1
+            own = _cvt_bf16x2_f32(
+                Float32(probabilities[high]),
+                Float32(probabilities[low]),
+            )
+            peer = cute.arch.shuffle_sync_bfly(own, offset=2)
+            if (lane & Int32(2)) == Int32(0):
+                p_store_words[i] = _prmt_b32(
+                    own, peer, Int32(0x5410)
+                )
+            else:
+                p_store_words[i] = _prmt_b32(
+                    own, peer, Int32(0x3276)
+                )
+
+        destination_chunk = cute.make_tensor(
+            destination.iterator
+            + chunk_idx * PAIR_P_CHUNK_PACKED_COLUMNS,
+            destination.layout,
+        )
+        cute.copy(tiled_store, p_store_registers, destination_chunk)
+        tcgen05_wait_st()
+
+    cute.arch.fence_view_async_tmem_store()
+
+
+@cute.jit
+def _load_pack_k_half(
+    tma_atom_pack_k: cute.CopyAtom,
+    tPackKgK: cute.Tensor,
+    tPackKsK: cute.Tensor,
+    block0: Int32,
+    block1: Int32,
+    quarter0: Int32,
+    barrier,
+):
+    """Gather one canonical N128 K tile as K0/N0,K0/N1,K1/N0,K1/N1."""
+
+    cute.copy(
+        tma_atom_pack_k,
+        tPackKgK[(None, block0, Int32(0))],
+        tPackKsK[(None, quarter0)],
+        tma_bar_ptr=barrier,
+    )
+    cute.copy(
+        tma_atom_pack_k,
+        tPackKgK[(None, block1, Int32(0))],
+        tPackKsK[(None, quarter0 + Int32(1))],
+        tma_bar_ptr=barrier,
+    )
+    cute.copy(
+        tma_atom_pack_k,
+        tPackKgK[(None, block0, Int32(1))],
+        tPackKsK[(None, quarter0 + Int32(2))],
+        tma_bar_ptr=barrier,
+    )
+    cute.copy(
+        tma_atom_pack_k,
+        tPackKgK[(None, block1, Int32(1))],
+        tPackKsK[(None, quarter0 + Int32(3))],
+        tma_bar_ptr=barrier,
+    )
+
+
+@cute.jit
+def _load_pack_v_half(
+    tma_atom_pack_v: cute.CopyAtom,
+    tPackVgV: cute.Tensor,
+    tPackVsV: cute.Tensor,
+    block0: Int32,
+    block1: Int32,
+    quarter0: Int32,
+    barrier,
+):
+    """Gather one canonical N128 V tile as D0/N0,D0/N1,D1/N0,D1/N1."""
+
+    cute.copy(
+        tma_atom_pack_v,
+        tPackVgV[(None, Int32(0), block0)],
+        tPackVsV[(None, quarter0)],
+        tma_bar_ptr=barrier,
+    )
+    cute.copy(
+        tma_atom_pack_v,
+        tPackVgV[(None, Int32(0), block1)],
+        tPackVsV[(None, quarter0 + Int32(1))],
+        tma_bar_ptr=barrier,
+    )
+    cute.copy(
+        tma_atom_pack_v,
+        tPackVgV[(None, Int32(1), block0)],
+        tPackVsV[(None, quarter0 + Int32(2))],
+        tma_bar_ptr=barrier,
+    )
+    cute.copy(
+        tma_atom_pack_v,
+        tPackVgV[(None, Int32(1), block1)],
+        tPackVsV[(None, quarter0 + Int32(3))],
+        tma_bar_ptr=barrier,
+    )
+
+
+@cute.struct
+class SharedStorage:
+    q_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2]
+    pack_k_mbar_ptr: cute.struct.MemRange[
+        cutlass.Int64, PAIR_STAGES * 2
+    ]
+    pack_v_mbar_ptr: cute.struct.MemRange[
+        cutlass.Int64, PAIR_STAGES * 2
+    ]
+    pair_score_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2]
+    pair_o_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2]
+    final_stats: cute.struct.Align[
+        cute.struct.MemRange[Float32, M * 2], 128
+    ]
+    route_partial: cute.struct.Align[
+        cute.struct.MemRange[Float32, 4 * ROUTE_TILE_SIZE], 16
+    ]
+    column_masks: cute.struct.Align[
+        cute.struct.MemRange[Float32, ROUTE_TILE_SIZE], 16
+    ]
+    route_packet: cute.struct.Align[
+        cute.struct.MemRange[Int32, PACKET_WORDS], 16
+    ]
+    tmem_holding_buf: Int32
+    # Owner-warp 0 lane 0 appends both N128 route masks.  The full-CTA
+    # pre-exact join publishes the completed list to warp 5; no HBM indices.
+    route_indices: cute.struct.Align[
+        cute.struct.MemRange[Int32, ROUTE_INDEX_CAPACITY], 16
+    ]
+
+
+@cute.kernel
+def _sol_attn_sm100_bf16_kernel(
+    tiled_pack_qk: cute.TiledMma,
+    tiled_pack_pv: cute.TiledMma,
+    tma_atom_q: cute.CopyAtom,
+    mQ_mkl: cute.Tensor,
+    tma_atom_pack_k: cute.CopyAtom,
+    mPackK_nkl: cute.Tensor,
+    tma_atom_pack_v: cute.CopyAtom,
+    mPackV_nkl: cute.Tensor,
+    tma_atom_kc: cute.CopyAtom,
+    mKC_nkl: cute.Tensor,
+    tma_atom_vc: cute.CopyAtom,
+    mVC_nkl: cute.Tensor,
+    mThreshold_bnh: cute.Tensor,
+    mO_bthd: cute.Tensor,
+    mLSE_bth: cute.Tensor,
+    token_count: Int32,
+    route_valid_total: Int32,
+    num_route_tiles: Int32,
+    softmax_scale: Float32,
+    sink_start_block: Int32,
+    sink_end_block: Int32,
+    q_layout: cute.ComposedLayout,
+    pack_k_layout: cute.ComposedLayout,
+    pack_k_gather_layout: cute.ComposedLayout,
+    pack_p_layout: cute.ComposedLayout,
+    pack_v_layout: cute.ComposedLayout,
+    pack_v_gather_layout: cute.ComposedLayout,
+    route_k_layout: cute.ComposedLayout,
+    route_v_layout: cute.ComposedLayout,
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+    q_block_idx_raw, head_idx_raw, batch_idx_raw = cute.arch.block_idx()
+    q_block_idx = Int32(q_block_idx_raw)
+    head_idx = Int32(head_idx_raw)
+    batch_idx = Int32(batch_idx_raw)
+    softmax_scale_log2 = softmax_scale * Float32(LOG2E)
+
+    smem = utils.SmemAllocator()
+    storage = smem.allocate(SharedStorage)
+    sFinalStats = storage.final_stats.get_tensor(
+        cute.make_layout((M, 2))
+    )
+    route_partial = storage.route_partial.get_tensor(
+        cute.make_layout((4, ROUTE_TILE_SIZE))
+    )
+    column_masks = storage.column_masks.get_tensor(
+        cute.make_layout((ROUTE_TILE_SIZE,))
+    )
+    route_packet = storage.route_packet.get_tensor(
+        cute.make_layout((PACKET_WORDS,))
+    )
+    route_indices = storage.route_indices.get_tensor(
+        cute.make_layout((ROUTE_INDEX_CAPACITY,))
+    )
+    sQ = smem.allocate_tensor(
+        element_type=BFloat16,
+        layout=q_layout.outer,
+        byte_alignment=128,
+        swizzle=q_layout.inner,
+    )
+    sPackK = smem.allocate_tensor(
+        element_type=BFloat16,
+        layout=pack_k_layout.outer,
+        byte_alignment=128,
+        swizzle=pack_k_layout.inner,
+    )
+    sPackV = smem.allocate_tensor(
+        element_type=BFloat16,
+        layout=pack_v_layout.outer,
+        byte_alignment=128,
+        swizzle=pack_v_layout.inner,
+    )
+    # One independent physical N128 K stage and one N128 V stage.  Every
+    # runtime route/exact transaction stays in this completion domain.
+    sPackKGather = cute.make_tensor(
+        cute.recast_ptr(
+            sPackK.iterator, pack_k_gather_layout.inner, BFloat16
+        ),
+        pack_k_gather_layout.outer,
+    )
+    sPackVGather = cute.make_tensor(
+        cute.recast_ptr(
+            sPackV.iterator, pack_v_gather_layout.inner, BFloat16
+        ),
+        pack_v_gather_layout.outer,
+    )
+    # KC/VC and exact K/V have disjoint lifetimes within each runtime group.
+    # They reuse the same independent N128 K and V allocations without a
+    # cross-operand alias barrier.
+    sKC = cute.make_tensor(
+        cute.recast_ptr(sPackK.iterator, route_k_layout.inner, BFloat16),
+        route_k_layout.outer,
+    )
+    sVC = cute.make_tensor(
+        cute.recast_ptr(sPackV.iterator, route_v_layout.inner, BFloat16),
+        route_v_layout.outer,
+    )
+
+    tmem_barrier = pipeline.NamedBarrier(barrier_id=1, num_threads=THREADS)
+    score_loaded_barrier = pipeline.NamedBarrier(
+        barrier_id=2, num_threads=4 * 32
+    )
+    final_stats_ready_barrier = pipeline.NamedBarrier(
+        barrier_id=3, num_threads=4 * 32
+    )
+    pack_score_loaded_barrier = pipeline.NamedBarrier(
+        barrier_id=4, num_threads=4 * 32
+    )
+    route_packet_ready_barrier = pipeline.NamedBarrier(
+        barrier_id=5, num_threads=5 * 32
+    )
+    exact_pair_p_ready_barrier = pipeline.NamedBarrier(
+        barrier_id=6, num_threads=5 * 32
+    )
+    tmem = utils.TmemAllocator(
+        storage.tmem_holding_buf.ptr,
+        barrier_for_retrieve=tmem_barrier,
+    )
+    tmem.allocate(TMEM_COLS)
+
+    one_thread = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1)
+    pack_owner_threads = pipeline.CooperativeGroup(
+        pipeline.Agent.Thread, 4 * 32
+    )
+    q_bytes = cute.size_in_bytes(
+        BFloat16, cute.select(q_layout, mode=[0, 1, 2])
+    )
+    route_k_bytes = cute.size_in_bytes(
+        BFloat16, cute.select(route_k_layout, mode=[0, 1, 2])
+    )
+    route_v_bytes = cute.size_in_bytes(
+        BFloat16, cute.select(route_v_layout, mode=[0, 1, 2])
+    )
+    pack_k_bytes = cute.size_in_bytes(
+        BFloat16, cute.select(pack_k_layout, mode=[0, 1, 2])
+    )
+    pack_v_bytes = cute.size_in_bytes(
+        BFloat16, cute.select(pack_v_layout, mode=[0, 1, 2])
+    )
+    assert route_k_bytes == pack_k_bytes
+    assert route_v_bytes == pack_v_bytes
+    q_pipe = fa_pipeline.PipelineTmaUmma.create(
+        num_stages=1,
+        producer_group=one_thread,
+        consumer_group=one_thread,
+        tx_count=q_bytes,
+        barrier_storage=storage.q_mbar_ptr.data_ptr(),
+    )
+    pack_k_pipe = fa_pipeline.PipelineTmaUmma.create(
+        num_stages=PAIR_STAGES,
+        producer_group=one_thread,
+        consumer_group=one_thread,
+        tx_count=pack_k_bytes,
+        barrier_storage=storage.pack_k_mbar_ptr.data_ptr(),
+    )
+    pack_v_pipe = fa_pipeline.PipelineTmaUmma.create(
+        num_stages=PAIR_STAGES,
+        producer_group=one_thread,
+        consumer_group=one_thread,
+        tx_count=pack_v_bytes,
+        barrier_storage=storage.pack_v_mbar_ptr.data_ptr(),
+    )
+    pair_score_pipe = fa_pipeline.PipelineUmmaAsync.create(
+        num_stages=1,
+        producer_group=one_thread,
+        consumer_group=pack_owner_threads,
+        barrier_storage=storage.pair_score_mbar_ptr.data_ptr(),
+    )
+    pair_o_pipe = fa_pipeline.PipelineUmmaAsync.create(
+        num_stages=1,
+        producer_group=one_thread,
+        consumer_group=pack_owner_threads,
+        barrier_storage=storage.pair_o_mbar_ptr.data_ptr(),
+    )
+
+    mQ_cur = mQ_mkl[None, None, head_idx, batch_idx]
+    mPackK_cur = mPackK_nkl[None, None, head_idx, batch_idx]
+    mPackV_cur = mPackV_nkl[None, None, head_idx, batch_idx]
+    mKC_cur = mKC_nkl[None, None, head_idx, batch_idx]
+    mVC_cur = mVC_nkl[None, None, head_idx, batch_idx]
+    gQ = cute.local_tile(mQ_cur, (M, D), (None, 0))
+    gPackK = cute.local_tile(
+        mPackK_cur, (N_MEMBER, 64), (None, None)
+    )
+    gPackV = cute.local_tile(
+        mPackV_cur, (64, N_MEMBER), (None, None)
+    )
+    gKC = cute.local_tile(mKC_cur, (N_PACK_HALF, D), (None, 0))
+    gVC = cute.local_tile(mVC_cur, (DV, N_PACK_HALF), (0, None))
+    thr_pack_qk = tiled_pack_qk.get_slice(0)
+    thr_pack_pv = tiled_pack_pv.get_slice(0)
+    tCgQ = thr_pack_qk.partition_A(gQ)
+    tCgKC = thr_pack_qk.partition_B(gKC)
+    tCgVC = thr_pack_pv.partition_B(gVC)
+    tCrKC = tiled_pack_qk.make_fragment_B(sKC)
+    tCrVC = tiled_pack_pv.make_fragment_B(sVC)
+    tCrPackQ = tiled_pack_qk.make_fragment_A(sQ)
+    tCrPackK = tiled_pack_qk.make_fragment_B(sPackK)
+    tCrPackV = tiled_pack_pv.make_fragment_B(sPackV)
+
+    tQsQ, tQgQ = cpasync.tma_partition(
+        tma_atom_q,
+        0,
+        cute.make_layout(1),
+        cute.group_modes(sQ, 0, 3),
+        cute.group_modes(tCgQ, 0, 3),
+    )
+    tPackKsK, tPackKgK = cpasync.tma_partition(
+        tma_atom_pack_k,
+        0,
+        cute.make_layout(1),
+        cute.group_modes(sPackKGather, 0, 3),
+        cute.group_modes(gPackK, 0, 2),
+    )
+    tPackVsV, tPackVgV = cpasync.tma_partition(
+        tma_atom_pack_v,
+        0,
+        cute.make_layout(1),
+        cute.group_modes(sPackVGather, 0, 3),
+        cute.group_modes(gPackV, 0, 2),
+    )
+    tKCsKC, tKCgKC = cpasync.tma_partition(
+        tma_atom_kc,
+        0,
+        cute.make_layout(1),
+        cute.group_modes(sKC, 0, 3),
+        cute.group_modes(tCgKC, 0, 3),
+    )
+    tVCsVC, tVCgVC = cpasync.tma_partition(
+        tma_atom_vc,
+        0,
+        cute.make_layout(1),
+        cute.group_modes(sVC, 0, 3),
+        cute.group_modes(tCgVC, 0, 3),
+    )
+
+    pack_score_shape = tiled_pack_qk.partition_shape_C(
+        PACK_QK_TILE[:2]
+    )
+    pack_score_template = tiled_pack_qk.make_fragment_C(pack_score_shape)
+    pack_o_shape = tiled_pack_pv.partition_shape_C(PACK_PV_TILE[:2])
+    pack_o_template = tiled_pack_pv.make_fragment_C(pack_o_shape)
+
+    tmem.wait_for_alloc()
+    tmem_ptr = tmem.retrieve_ptr(Float32)
+    # The 256-column allocation leaves the second half of SM TMEM available to
+    # another CTA.  The live allocation remains owned after permit release.
+    tmem.relinquish_alloc_permit()
+    tmem_base = tmem_ptr.toint()
+    pair_tScore = cute.make_tensor(
+        cute.make_ptr(
+            Float32,
+            tmem_base + Int32(PAIR_SCORE_OFFSET),
+            cute.AddressSpace.tmem,
+            assumed_align=16,
+        ),
+        pack_score_template.layout,
+    )
+    pair_tO = cute.make_tensor(
+        cute.make_ptr(
+            Float32,
+            tmem_base + Int32(O_OFFSET),
+            cute.AddressSpace.tmem,
+            assumed_align=16,
+        ),
+        pack_o_template.layout,
+    )
+    # make_fragment_A drops the physical TMEM allocation base and addresses
+    # packed BF16 columns in half-column units. Restore both facts so
+    # 2*tmem_base + 2*PAIR_P_OFFSET names columns 64..127.
+    pair_tP_storage = cute.make_tensor(
+        pair_tScore.iterator, pack_p_layout.outer
+    )
+    pair_tP_base = tiled_pack_pv.make_fragment_A(pair_tP_storage)[
+        None, None, None, 0
+    ]
+    pair_tP = cute.make_tensor(
+        pair_tP_base.iterator
+        + tmem_base
+        + tmem_base
+        + Int32(PAIR_P_OFFSET * 2),
+        pair_tP_base.layout,
+    )
+    q_producer = fa_pipeline.make_pipeline_state(
+        pipeline.PipelineUserType.Producer, 1
+    )
+    q_consumer = fa_pipeline.make_pipeline_state(
+        pipeline.PipelineUserType.Consumer, 1
+    )
+    pack_k_producer = pipeline.make_pipeline_state(
+        pipeline.PipelineUserType.Producer, PAIR_STAGES
+    )
+    pack_k_consumer = pipeline.make_pipeline_state(
+        pipeline.PipelineUserType.Consumer, PAIR_STAGES
+    )
+    pack_v_producer = pipeline.make_pipeline_state(
+        pipeline.PipelineUserType.Producer, PAIR_STAGES
+    )
+    pack_v_consumer = pipeline.make_pipeline_state(
+        pipeline.PipelineUserType.Consumer, PAIR_STAGES
+    )
+    pair_score_producer = fa_pipeline.make_pipeline_state(
+        pipeline.PipelineUserType.Producer, 1
+    )
+    pair_score_consumer = fa_pipeline.make_pipeline_state(
+        pipeline.PipelineUserType.Consumer, 1
+    )
+    pair_o_producer = fa_pipeline.make_pipeline_state(
+        pipeline.PipelineUserType.Producer, 1
+    )
+    pair_o_consumer = fa_pipeline.make_pipeline_state(
+        pipeline.PipelineUserType.Consumer, 1
+    )
+    route_start_base = Int32(0)
+    q_len = token_count - q_block_idx * Int32(M)
+    if q_len > Int32(M):
+        q_len = Int32(M)
+    threshold = Float32(
+        mThreshold_bnh[batch_idx, q_block_idx, head_idx]
+    )
+
+    if warp_idx == Int32(5):
+        cpasync.prefetch_descriptor(tma_atom_q)
+        cpasync.prefetch_descriptor(tma_atom_pack_k)
+        cpasync.prefetch_descriptor(tma_atom_pack_v)
+        cpasync.prefetch_descriptor(tma_atom_kc)
+        cpasync.prefetch_descriptor(tma_atom_vc)
+
+        q_pipe.producer_acquire(q_producer)
+        q_barrier = q_pipe.producer_get_barrier(q_producer)
+        cute.copy(
+            tma_atom_q,
+            tQgQ[(None, q_block_idx)],
+            tQsQ[(None, q_producer.index)],
+            tma_bar_ptr=q_barrier,
+        )
+        q_producer.advance()
+
+    is_owner = warp_idx >= Int32(1) and warp_idx <= Int32(4)
+    is_score_consumer = warp_idx <= Int32(4)
+    owner_tidx = tidx - Int32(32)
+
+    # One register-resident online state and one TMEM-O initialization bit span
+    # every route/exact transaction in every runtime group.
+    running_max = -Float32.inf
+    running_sum = Float32(0.0)
+    owner_o_initialized = Int32(0)
+    mma_o_initialized = Int32(0)
+
+    if warp_idx == Int32(0):
+        q_pipe.consumer_wait(q_consumer)
+
+    # The outer loop owns one logical G256 exact-index lifetime.  The inner
+    # loop consumes each physical score/PV half immediately; it appends only
+    # integer indices, never a second score or probability fragment.
+    num_logical_groups = (
+        num_route_tiles + Int32(ROUTE_HALVES_PER_GROUP - 1)
+    ) // Int32(ROUTE_HALVES_PER_GROUP)
+    # BEGIN_G256_CURSOR_UNIFORM_INDUCTION
+    # arch_make_warp_uniform is a lowering hint, not a value broadcast. Both
+    # values are CTA-invariant integer scalars before the hint.
+    logical_group_idx = cute.arch.make_warp_uniform(Int32(0))
+    remaining_group_tiles = cute.arch.make_warp_uniform(num_route_tiles)
+    while logical_group_idx < num_logical_groups:
+        is_final_logical_group = (
+            logical_group_idx + Int32(1) == num_logical_groups
+        )
+        group_route_tile_base = logical_group_idx * Int32(
+            ROUTE_HALVES_PER_GROUP
+        )
+        physical_halves_this_group = remaining_group_tiles
+        if physical_halves_this_group > Int32(ROUTE_HALVES_PER_GROUP):
+            physical_halves_this_group = Int32(ROUTE_HALVES_PER_GROUP)
+
+        for half_idx in cutlass.range(
+            physical_halves_this_group, unroll=1
+        ):
+            route_tile_idx = cute.arch.make_warp_uniform(
+                group_route_tile_base + half_idx
+            )
+            is_final_route_tile = (
+                route_tile_idx + Int32(1) == num_route_tiles
+            )
+            is_logical_terminal_half = (
+                half_idx + Int32(1) == physical_halves_this_group
+            )
+            route_start = cute.arch.make_warp_uniform(
+                route_start_base
+                + route_tile_idx * Int32(ROUTE_TILE_SIZE)
+            )
+            remaining_route_count = cute.arch.make_warp_uniform(
+                route_valid_total
+                - route_tile_idx * Int32(ROUTE_TILE_SIZE)
+            )
+            valid_route_count = remaining_route_count
+            if valid_route_count > Int32(ROUTE_TILE_SIZE):
+                valid_route_count = Int32(ROUTE_TILE_SIZE)
+            if valid_route_count < Int32(0):
+                valid_route_count = Int32(0)
+
+            # One native N128 route transaction shares the independent K/V stages
+            # with the exact-pair engine.  Route and exact are separated by
+            # a full-CTA phase boundary, so no K<->V alias handoff is required.
+            if warp_idx == Int32(5):
+                pack_k_pipe.producer_acquire(pack_k_producer)
+                route_k_barrier = pack_k_pipe.producer_get_barrier(
+                    pack_k_producer
+                )
+                cute.copy(
+                    tma_atom_kc,
+                    tKCgKC[(None, route_tile_idx)],
+                    tKCsKC[(None, pack_k_producer.index)],
+                    tma_bar_ptr=route_k_barrier,
+                )
+                pack_k_producer.advance()
+
+                pack_v_pipe.producer_acquire(pack_v_producer)
+                route_v_barrier = pack_v_pipe.producer_get_barrier(
+                    pack_v_producer
+                )
+                cute.copy(
+                    tma_atom_vc,
+                    tVCgVC[(None, route_tile_idx)],
+                    tVCsVC[(None, pack_v_producer.index)],
+                    tma_bar_ptr=route_v_barrier,
+                )
+                pack_v_producer.advance()
+
+            if warp_idx == Int32(0):
+                pack_k_pipe.consumer_wait(pack_k_consumer)
+                pair_score_pipe.producer_acquire(pair_score_producer)
+                mma_utils.gemm(
+                    tiled_pack_qk,
+                    pair_tScore,
+                    tCrPackQ[None, None, None, q_consumer.index],
+                    tCrKC[None, None, None, pack_k_consumer.index],
+                    zero_init=True,
+                )
+                pair_score_pipe.producer_commit(pair_score_producer)
+                pair_score_producer.advance()
+                pack_k_pipe.consumer_release(pack_k_consumer)
+                pack_k_consumer.advance()
+
+            # BEGIN_RUNTIME_GROUP_BODY
+
+            # Route generation: four physical owner warps reduce the native N128
+            # score tile into one four-word mask.  HBM receives only the diagnostic
+            # copy; the compacted exact stream remains resident in SMEM.
+            if is_owner:
+                pair_score_pipe.consumer_wait(pair_score_consumer)
+                score_raw, score_coords = _load_pair_score(
+                    pack_score_template,
+                    thr_pack_qk,
+                    tmem_base,
+                    Int32(PAIR_SCORE_OFFSET),
+                    owner_tidx,
+                )
+                pack_score_loaded_barrier.arrive_and_wait()
+                pair_score_pipe.consumer_release(pair_score_consumer)
+                pair_score_consumer.advance()
+                owner_warp = owner_tidx // Int32(32)
+                lane = owner_tidx % Int32(32)
+                semantic_row = (
+                    score_coords[0][0] + Int32(SEMANTIC_ROW_OFFSET)
+                ) & Int32(M - 1)
+                row_valid = semantic_row < q_len
+                lane_col_parity = (lane // Int32(2)) % Int32(2)
+                # Column-pair reduction: parity-0 lanes carry column 2*pair and
+                # parity-1 lanes carry column 2*pair+1.  The XOR-1/16/8/4
+                # butterfly tree never crosses lane column-parity classes
+                # ((l^k)//2 keeps (l//2)%2 for k in {1,16,8,4}), so one tree
+                # reduces both columns at once; every surviving addition chain
+                # sees the same zero-padded operand streams, and the removed
+                # chains only ever accumulated 0.0. Writer lanes 0 and 2 equal
+                # 2*(col%2).
+                for pair_idx in cutlass.range_constexpr(
+                    0, ROUTE_TILE_SIZE // 2, 2
+                ):
+                    my_col0 = Int32(2 * pair_idx) + lane_col_parity
+                    partial0 = Float32(0.0)
+                    if row_valid and my_col0 < valid_route_count:
+                        partial0 = Float32(score_raw[pair_idx])
+                    my_col1 = Int32(2 * (pair_idx + 1)) + lane_col_parity
+                    partial1 = Float32(0.0)
+                    if row_valid and my_col1 < valid_route_count:
+                        partial1 = Float32(score_raw[pair_idx + 1])
+
+                    raw_partial0 = partial0
+                    raw_partial1 = partial1
+                    scaled0, scaled1 = cute.arch.mul_packed_f32x2(
+                        (raw_partial0, raw_partial1),
+                        (softmax_scale_log2, softmax_scale_log2),
+                    )
+                    peer_scaled0 = cute.arch.shuffle_sync_bfly(
+                        scaled0, offset=1
+                    )
+                    peer_scaled1 = cute.arch.shuffle_sync_bfly(
+                        scaled1, offset=1
+                    )
+                    partial0, partial1 = cute.arch.fma_packed_f32x2(
+                        (raw_partial0, raw_partial1),
+                        (softmax_scale_log2, softmax_scale_log2),
+                        (peer_scaled0, peer_scaled1),
+                    )
+                    peer0 = cute.arch.shuffle_sync_bfly(
+                        partial0, offset=16
+                    )
+                    peer1 = cute.arch.shuffle_sync_bfly(
+                        partial1, offset=16
+                    )
+                    partial0, partial1 = cute.arch.add_packed_f32x2(
+                        (partial0, partial1), (peer0, peer1)
+                    )
+                    peer0 = cute.arch.shuffle_sync_bfly(
+                        partial0, offset=8
+                    )
+                    peer1 = cute.arch.shuffle_sync_bfly(
+                        partial1, offset=8
+                    )
+                    partial0, partial1 = cute.arch.add_packed_f32x2(
+                        (partial0, partial1), (peer0, peer1)
+                    )
+                    peer0 = cute.arch.shuffle_sync_bfly(
+                        partial0, offset=4
+                    )
+                    peer1 = cute.arch.shuffle_sync_bfly(
+                        partial1, offset=4
+                    )
+                    partial0, partial1 = cute.arch.add_packed_f32x2(
+                        (partial0, partial1), (peer0, peer1)
+                    )
+                    if lane == Int32(0):
+                        route_partial[owner_warp, 2 * pair_idx] = partial0
+                        route_partial[owner_warp, 2 * (pair_idx + 1)] = (
+                            partial1
+                        )
+                    if lane == Int32(2):
+                        route_partial[owner_warp, 2 * pair_idx + 1] = partial0
+                        route_partial[
+                            owner_warp, 2 * (pair_idx + 1) + 1
+                        ] = partial1
+
+                cute.arch.fence_view_async_shared()
+                score_loaded_barrier.arrive_and_wait()
+                if owner_warp == Int32(0):
+                    mask0 = Int32(0)
+                    mask1 = Int32(0)
+                    mask2 = Int32(0)
+                    mask3 = Int32(0)
+
+                    # Half 0 starts a fresh G256 stream and half 1 appends to
+                    # lane 0's cumulative packet word.  The preceding packet
+                    # barrier makes the base warp-uniform before the vote.
+                    append_base = Int32(0)
+                    if half_idx != Int32(0):
+                        append_base = Int32(route_packet[6])
+
+                    # A positive signed shift avoids materializing 1<<31:
+                    # lane 0 gets zero and lane 31 gets 0x7fffffff.
+                    lane_mask_lt = Int32(0x7FFFFFFF) >> (
+                        Int32(31) - lane
+                    )
+                    preceding_word_count = Int32(0)
+                    for word in cutlass.range_constexpr(ROUTE_MASK_WORDS):
+                        off = Int32(word * 32) + lane
+                        valid = off < valid_route_count
+                        exact_pred = False
+                        if valid:
+                            pair_02 = Float32(route_partial[0, off]) + Float32(
+                                route_partial[2, off]
+                            )
+                            pair_13 = Float32(route_partial[1, off]) + Float32(
+                                route_partial[3, off]
+                            )
+                            col_mean = (pair_02 + pair_13) / Float32(q_len)
+                            exact_pred = sol_attn_route_is_exact(
+                                q_block_idx,
+                                route_start + off,
+                                col_mean,
+                                threshold,
+                                valid,
+                            )
+                            # Sink is a KV-only contract. Text queries remain
+                            # a caller-side dense operation in MMDiT models.
+                            exact_pred = (
+                                exact_pred
+                                or (
+                                    route_start + off >= sink_start_block
+                                    and route_start + off < sink_end_block
+                                )
+                            )
+                        word_mask = Int32(
+                            cute.arch.vote_ballot_sync(exact_pred)
+                        )
+                        # Site 2: preserve the route decision and its four
+                        # ordered ballots, but materialize the resulting
+                        # approximate-column mask exactly once.  Dedicated
+                        # SMEM holds the two N64 mask halves so the reduction
+                        # scratch remains non-aliasing for ptxas scheduling.
+                        # The existing shared fence and owner barrier below
+                        # publish them to every score owner.
+                        if valid and not exact_pred:
+                            column_masks[off] = Float32(0.0)
+                        else:
+                            column_masks[off] = -Float32.inf
+                        lane_rank = (
+                            append_base
+                            + preceding_word_count
+                            + sol_attn_popc_b32(word_mask & lane_mask_lt)
+                        )
+                        if exact_pred:
+                            route_indices[lane_rank] = route_start + off
+                        if cutlass.const_expr(word == 0):
+                            mask0 = word_mask
+                        elif cutlass.const_expr(word == 1):
+                            mask1 = word_mask
+                        elif cutlass.const_expr(word == 2):
+                            mask2 = word_mask
+                        else:
+                            mask3 = word_mask
+                        preceding_word_count = (
+                            preceding_word_count
+                            + sol_attn_popc_b32(word_mask)
+                        )
+
+                    # Every selected lane has a unique rank; lane 0 publishes
+                    # the packet after reconvergence.
+                    exact_count = preceding_word_count
+                    if lane == Int32(0):
+                        route_rank = append_base + exact_count
+
+                        route_packet[0] = mask0
+                        route_packet[1] = mask1
+                        route_packet[2] = mask2
+                        route_packet[3] = mask3
+                        route_packet[4] = exact_count
+                        route_packet[5] = append_base
+                        route_packet[6] = route_rank
+                        terminal_half_word = Int32(0)
+                        if is_logical_terminal_half:
+                            terminal_half_word = Int32(1)
+                        route_packet[7] = terminal_half_word
+                        cute.arch.fence_view_async_shared()
+
+                # The selector packet is now immutable.  Reuse the already resident
+                # route scores for the non-exact transaction; no offset list or second
+                # route-score load is introduced.
+                score_loaded_barrier.arrive_and_wait()
+                route_exact_count = Int32(route_packet[4])
+                has_route_approx = route_exact_count < valid_route_count
+                if has_route_approx:
+                    row_mask = -Float32.inf
+                    if row_valid:
+                        row_mask = Float32(0.0)
+                    # Route generation has consumed every raw score.  Apply
+                    # the shared mask in place so raw and masked N128
+                    # fragments never overlap in registers; the same object
+                    # remains available for the later route-mass scratch.
+                    route_scores = score_raw
+                    assert cute.size(score_raw) % 2 == 0
+                    for i in cutlass.range_constexpr(
+                        0, cute.size(score_raw), 2
+                    ):
+                        group_col0 = score_coords[i][1]
+                        group_col1 = score_coords[i + 1][1]
+                        mask0 = Float32(column_masks[group_col0])
+                        mask1 = Float32(column_masks[group_col1])
+                        mask0, mask1 = cute.arch.add_packed_f32x2(
+                            (mask0, mask1), (row_mask, row_mask)
+                        )
+                        mask0, mask1 = cute.arch.add_packed_f32x2(
+                            (
+                                Float32(score_raw[i]),
+                                Float32(score_raw[i + 1]),
+                            ),
+                            (mask0, mask1),
+                        )
+                        route_scores[i] = mask0
+                        route_scores[i + 1] = mask1
+
+                    local_max = fa_utils.fmax_reduce(
+                        route_scores.load(), arch=100
+                    )
+                    local_max = Float32(local_max) * softmax_scale
+                    peer_max = cute.arch.shuffle_sync_bfly(local_max, offset=2)
+                    pair_max = local_max
+                    if peer_max > pair_max:
+                        pair_max = peer_max
+
+                    old_max = running_max
+                    old_sum = running_sum
+                    new_max = old_max
+                    if old_max == -Float32.inf or pair_max > old_max:
+                        new_max = pair_max
+                    row_alpha = Float32(0.0)
+                    if old_max != -Float32.inf:
+                        row_alpha = cute.math.exp2(
+                            (old_max - new_max) * Float32(LOG2E),
+                            fastmath=True,
+                        )
+
+                    route_probabilities = cute.make_rmem_tensor(
+                        route_scores.shape, Float32
+                    )
+                    if new_max == -Float32.inf:
+                        for i in cutlass.range(
+                            cute.size(route_scores), unroll_full=True
+                        ):
+                            route_probabilities[i] = Float32(0.0)
+                    else:
+                        for i in cutlass.range(
+                            cute.size(route_scores), unroll_full=True
+                        ):
+                            route_probabilities[i] = cute.math.exp2(
+                                Float32(route_scores[i]) * softmax_scale_log2
+                                - new_max * Float32(LOG2E),
+                                fastmath=True,
+                            )
+                    # ``route_scores`` is dead after the exponentials above.  Use
+                    # it as mass scratch so the compiler does not need a second
+                    # full N128-shaped fragment while probabilities remain live
+                    # for the chunked TMEM-P store below.  Keeping the same shape,
+                    # index order, and fadd_reduce preserves floating-point
+                    # reduction order and every phase edge.
+                    assert cute.size(route_probabilities) % 2 == 0
+                    for i in cutlass.range_constexpr(
+                        0, cute.size(route_probabilities), 2
+                    ):
+                        block_idx0 = route_start + score_coords[i][1]
+                        raw_length0 = (
+                            token_count - block_idx0 * Int32(N_MEMBER)
+                        )
+                        block_length0 = max(
+                            Int32(0), min(raw_length0, Int32(N_MEMBER))
+                        )
+                        block_idx1 = route_start + score_coords[i + 1][1]
+                        raw_length1 = (
+                            token_count - block_idx1 * Int32(N_MEMBER)
+                        )
+                        block_length1 = max(
+                            Int32(0), min(raw_length1, Int32(N_MEMBER))
+                        )
+                        mass0, mass1 = cute.arch.mul_packed_f32x2(
+                            (
+                                Float32(route_probabilities[i]),
+                                Float32(route_probabilities[i + 1]),
+                            ),
+                            (
+                                Float32(block_length0),
+                                Float32(block_length1),
+                            ),
+                        )
+                        route_scores[i] = mass0
+                        route_scores[i + 1] = mass1
+                    current_sum = fa_utils.fadd_reduce(
+                        route_scores.load(), arch=100
+                    )
+                    current_sum += cute.arch.shuffle_sync_bfly(
+                        current_sum, offset=2
+                    )
+                    # KC is a block mean and VC a valid-token sum.  Route mass uses
+                    # the true block length while PV still consumes p*VC once.
+                    running_sum = old_sum * row_alpha + current_sum
+                    running_max = new_max
+                    if owner_o_initialized != Int32(0):
+                        _rescale_pair_o(
+                            pack_o_template,
+                            thr_pack_pv,
+                            tmem_base,
+                            Int32(O_OFFSET),
+                            owner_tidx,
+                            row_alpha,
+                        )
+                    _store_pair_probability_chunked_tmemp(
+                        pack_o_template,
+                        route_probabilities,
+                        tmem_base,
+                        Int32(PAIR_P_OFFSET),
+                        owner_tidx,
+                    )
+                    owner_o_initialized = Int32(1)
+            # Publish the mask/P decision to warp 0.  The route PV is deliberately
+            # drained before exact work so all-exact, all-approx, odd, and
+            # partial-tail paths share one phase boundary.
+            if is_score_consumer:
+                route_packet_ready_barrier.arrive_and_wait()
+            if warp_idx == Int32(0):
+                route_exact_count = Int32(route_packet[4])
+                route_has_approx = route_exact_count < valid_route_count
+                pack_v_pipe.consumer_wait(pack_v_consumer)
+                if route_has_approx:
+                    mma_utils.gemm(
+                        tiled_pack_pv,
+                        pair_tO,
+                        pair_tP,
+                        tCrVC[None, None, None, pack_v_consumer.index],
+                        zero_init=mma_o_initialized == Int32(0),
+                    )
+                    # Half 0 is followed by half-1 route QK.  The terminal
+                    # route half is followed by exact QK0 whenever the fused
+                    # G256 index stream is nonempty.  Those score completions
+                    # prove this PV complete; only a final route-only CTA needs
+                    # an explicit O completion here.
+                    if (
+                        is_final_route_tile
+                        and Int32(route_packet[6]) == Int32(0)
+                    ):
+                        pair_o_pipe.producer_commit(pair_o_producer)
+                    mma_o_initialized = Int32(1)
+                pack_v_pipe.consumer_release(pack_v_consumer)
+                pack_v_consumer.advance()
+            if is_owner:
+                cumulative_exact_count = Int32(route_packet[6])
+                if (
+                    is_final_route_tile
+                    and cumulative_exact_count == Int32(0)
+                ):
+                    pair_o_pipe.consumer_wait(pair_o_consumer)
+
+            # route_packet may be reused by the next physical half without a
+            # CTA join.  Warp 0 reads this half's packet before it can issue
+            # next-half QK; owner-warp 0 cannot overwrite the packet until
+            # that QK's pair-score completion has released all owners.
+
+        # Both route halves have published their packet/index data and drained
+        # approximate PV.  This is the only CTA-wide pre-exact join in the
+        # logical G256 group; it publishes the combined list to warp 5.
+        cute.arch.barrier()
+        # The cumulative count covers half 0 followed by half 1.  Pairing this
+        # one ordered stream removes cross-half odd padding without retaining
+        # either physical score fragment.
+        exact_block_count = Int32(route_packet[6])
+        exact_pair_count = (exact_block_count + Int32(1)) // Int32(2)
+        pair_count = exact_pair_count
+        has_pair_exact = exact_block_count > Int32(0)
+
+        # BEGIN_GENERAL_N128_PAIR
+        # Every executable exact count, including a logical-group terminal
+        # exact1, stays in the N128 domain.
+
+        # Warp 5 streams one physical N128 K stage and one physical N128 V
+        # stage.  A missing odd peer duplicates block0 only for the physical
+        # transaction; owners mask all upper-64 scores before softmax.
+        if warp_idx == Int32(5) and has_pair_exact:
+            for pair_idx in cutlass.range(pair_count, unroll=1):
+                ordinal0 = pair_idx * Int32(2)
+                block0 = Int32(route_indices[ordinal0])
+                block1 = block0
+                if ordinal0 + Int32(1) < exact_block_count:
+                    block1 = Int32(route_indices[ordinal0 + Int32(1)])
+
+                pack_k_pipe.producer_acquire(pack_k_producer)
+                pair_k_barrier = pack_k_pipe.producer_get_barrier(
+                    pack_k_producer
+                )
+                _load_pack_k_half(
+                    tma_atom_pack_k,
+                    tPackKgK,
+                    tPackKsK,
+                    block0,
+                    block1,
+                    pack_k_producer.index * Int32(4),
+                    pair_k_barrier,
+                )
+                pack_k_producer.advance()
+
+                pack_v_pipe.producer_acquire(pack_v_producer)
+                pair_v_barrier = pack_v_pipe.producer_get_barrier(
+                    pack_v_producer
+                )
+                _load_pack_v_half(
+                    tma_atom_pack_v,
+                    tPackVgV,
+                    tPackVsV,
+                    block0,
+                    block1,
+                    pack_v_producer.index * Int32(4),
+                    pair_v_barrier,
+                )
+                pack_v_producer.advance()
+
+        if warp_idx == Int32(0) and has_pair_exact:
+            # QK0 prologue.  K and score cursors advance exactly once per QK;
+            # neither V nor O state is touched until the steady-state PV path.
+            pack_k_pipe.consumer_wait(pack_k_consumer)
+            pair_score_pipe.producer_acquire(pair_score_producer)
+            mma_utils.gemm(
+                tiled_pack_qk,
+                pair_tScore,
+                tCrPackQ[None, None, None, q_consumer.index],
+                tCrPackK[None, None, None, pack_k_consumer.index],
+                zero_init=True,
+            )
+            pair_score_pipe.producer_commit(pair_score_producer)
+            pair_score_producer.advance()
+            # PipelineTmaUmma release is tcgen05-completion-backed.
+            pack_k_pipe.consumer_release(pack_k_consumer)
+            pack_k_consumer.advance()
+
+            for pair_idx in cutlass.range(pair_count, unroll=1):
+                # P aliases the drained upper half of S.  PV must therefore be
+                # issued before QK(i+1) overwrites S.  Both instructions are
+                # emitted back-to-back by warp 0, retaining the full-G128
+                # tcgen05 dependency order without its K/V alias barriers.
+                pack_v_pipe.consumer_wait(pack_v_consumer)
+                # All four owners have completed their synchronous chunked
+                # TMEM stores and the helper's TMEM store fence before this
+                # five-warp rendezvous releases the single MMA warp.
+                exact_pair_p_ready_barrier.arrive_and_wait()
+                mma_utils.gemm(
+                    tiled_pack_pv,
+                    pair_tO,
+                    pair_tP,
+                    tCrPackV[None, None, None, pack_v_consumer.index],
+                    zero_init=mma_o_initialized == Int32(0),
+                )
+                # QK(i+1) completion dominates PV(i) completion for every
+                # nonterminal transaction on this tcgen05 issuer.  Commit one
+                # explicit O-full generation only for the CTA's final PV.
+                if (
+                    is_final_logical_group
+                    and pair_idx + Int32(1) == pair_count
+                ):
+                    pair_o_pipe.producer_commit(pair_o_producer)
+                mma_o_initialized = Int32(1)
+                pack_v_pipe.consumer_release(pack_v_consumer)
+                pack_v_consumer.advance()
+
+                if pair_idx + Int32(1) < pair_count:
+                    pack_k_pipe.consumer_wait(pack_k_consumer)
+                    pair_score_pipe.producer_acquire(pair_score_producer)
+                    mma_utils.gemm(
+                        tiled_pack_qk,
+                        pair_tScore,
+                        tCrPackQ[None, None, None, q_consumer.index],
+                        tCrPackK[
+                            None, None, None, pack_k_consumer.index
+                        ],
+                        zero_init=True,
+                    )
+                    pair_score_pipe.producer_commit(pair_score_producer)
+                    pair_score_producer.advance()
+                    pack_k_pipe.consumer_release(pack_k_consumer)
+                    pack_k_consumer.advance()
+
+        if is_owner and has_pair_exact:
+            exact_owner_warp = owner_tidx // Int32(32)
+            exact_lane = owner_tidx % Int32(32)
+            for pair_idx in cutlass.range(pair_count, unroll=1):
+                ordinal0 = pair_idx * Int32(2)
+                block0 = Int32(route_indices[ordinal0])
+                has_peer = ordinal0 + Int32(1) < exact_block_count
+                block1 = block0
+                if has_peer:
+                    block1 = Int32(route_indices[ordinal0 + Int32(1)])
+                valid0 = token_count - block0 * Int32(N_MEMBER)
+                valid1 = Int32(0)
+                if has_peer:
+                    valid1 = token_count - block1 * Int32(N_MEMBER)
+                # Keep packed-select integer min/max lowering and exact-pair
+                # bookkeeping unchanged.
+                valid0 = max(Int32(0), min(valid0, Int32(N_MEMBER)))
+                valid1 = max(Int32(0), min(valid1, Int32(N_MEMBER)))
+
+                # Site 1: owner warp 0 builds two 64-column gates once for
+                # this exact N128 pair.  The existing score-load barrier below
+                # both protects the S/P alias and publishes these stores; no
+                # barrier or shared allocation is added.
+                if exact_owner_warp == Int32(0):
+                    for cohort in cutlass.range_constexpr(4):
+                        column = Int32(cohort * 32) + exact_lane
+                        if cutlass.const_expr(cohort < 2):
+                            if column >= valid0:
+                                column_masks[column] = -Float32.inf
+                            else:
+                                column_masks[column] = Float32(0.0)
+                        else:
+                            if column - Int32(N_MEMBER) >= valid1:
+                                column_masks[column] = -Float32.inf
+                            else:
+                                column_masks[column] = Float32(0.0)
+                    cute.arch.fence_view_async_shared()
+
+                pair_score_pipe.consumer_wait(pair_score_consumer)
+                # Keep the exact ae9 score-load helper and fragment scope.
+                pair_scores, pair_coords = _load_pair_score(
+                    pack_score_template,
+                    thr_pack_qk,
+                    tmem_base,
+                    Int32(PAIR_SCORE_OFFSET),
+                    owner_tidx,
+                )
+                # Every owner retires the complete score load before the
+                # packed P store aliases columns 64..127 of S.
+                pack_score_loaded_barrier.arrive_and_wait()
+                pair_score_pipe.consumer_release(pair_score_consumer)
+                pair_score_consumer.advance()
+
+                semantic_row = (
+                    pair_coords[0][0] + Int32(SEMANTIC_ROW_OFFSET)
+                ) & Int32(M - 1)
+                row_valid = semantic_row < q_len
+                row_mask = -Float32.inf
+                if row_valid:
+                    row_mask = Float32(0.0)
+                assert cute.size(pair_scores) % 2 == 0
+                for i in cutlass.range_constexpr(
+                    0, cute.size(pair_scores), 2
+                ):
+                    column0 = pair_coords[i][1]
+                    column1 = pair_coords[i + 1][1]
+                    mask0 = Float32(column_masks[column0])
+                    mask1 = Float32(column_masks[column1])
+                    mask0, mask1 = cute.arch.add_packed_f32x2(
+                        (mask0, mask1), (row_mask, row_mask)
+                    )
+                    mask0, mask1 = cute.arch.add_packed_f32x2(
+                        (
+                            Float32(pair_scores[i]),
+                            Float32(pair_scores[i + 1]),
+                        ),
+                        (mask0, mask1),
+                    )
+                    pair_scores[i] = mask0
+                    pair_scores[i + 1] = mask1
+
+                probabilities, next_max, next_sum, row_alpha = (
+                    _online_update_pair(
+                        pair_scores,
+                        running_max,
+                        running_sum,
+                        softmax_scale,
+                    )
+                )
+                # For i>0, pair-score completion comes from QK(i), issued
+                # after PV(i-1) on the same tcgen05 issuer.  The score wait and
+                # load above therefore retire PV(i-1) before this O rescale.
+                # Pair0 similarly follows either route QK or route PV->QK0.
+                if owner_o_initialized != Int32(0):
+                    _rescale_pair_o(
+                        pack_o_template,
+                        thr_pack_pv,
+                        tmem_base,
+                        Int32(O_OFFSET),
+                        owner_tidx,
+                        row_alpha,
+                    )
+                # The one TMEM P image is free once PV(i-1) completes.  Keep
+                # probabilities FP32 until the live-range-bounded chunked R2T.
+                _store_pair_probability_chunked_tmemp(
+                    pack_o_template,
+                    probabilities,
+                    tmem_base,
+                    Int32(PAIR_P_OFFSET),
+                    owner_tidx,
+                )
+                # The preceding helper performs tcgen05.wait::st for every
+                # chunk and a TMEM-store fence.  Publish P to warp 0 with one
+                # uniform generation shared by warps 0-4; warp 5 is excluded.
+                exact_pair_p_ready_barrier.arrive_and_wait()
+                running_max = next_max
+                running_sum = next_sum
+                owner_o_initialized = Int32(1)
+
+            # There is no successor QK after the CTA's final exact PV.  Keep
+            # exactly one completion-backed wait before the epilogue; all
+            # earlier groups flow into a successor route QK completion.
+            if is_final_logical_group and pair_count > Int32(0):
+                pair_o_pipe.consumer_wait(pair_o_consumer)
+
+        # route_indices reuse HB proof for the next logical group:
+        # (1) warp 5 reads both indices before producing each pair's K/V, and
+        #     final-pair score completion therefore dominates its last read;
+        # (2) all owner index reads precede the final exact-P NamedBarrier;
+        # (3) owner-warp0/lane0 is the sole next-group writer and reaches it
+        #     only after that same exact loop.  For exact_count==0 there are no
+        #     readers.  Therefore no group-tail CTA barrier is required.
+
+        # Cross-group progress is carried by the existing K/V buffer-free
+        # phases and pair-score ready phase.  There is no CTA-wide group-tail
+        # join: the next producer acquire cannot overwrite a live K/V stage,
+        # and the next owner score load cannot precede QK completion.
+        # END_GENERAL_N128_PAIR
+
+        logical_group_idx = cute.arch.make_warp_uniform(
+            logical_group_idx + Int32(1)
+        )
+        remaining_group_tiles = cute.arch.make_warp_uniform(
+            remaining_group_tiles - Int32(ROUTE_HALVES_PER_GROUP)
+        )
+        # END_RUNTIME_GROUP_BODY
+    # END_G256_CURSOR_UNIFORM_INDUCTION
+
+    if warp_idx == Int32(0):
+        q_pipe.consumer_release(q_consumer)
+        q_consumer.advance()
+
+    if is_owner:
+        lane = owner_tidx % Int32(32)
+        owner_warp = owner_tidx // Int32(32)
+        owner_row = (
+            owner_warp * Int32(16)
+            + lane // Int32(4)
+            + (lane % Int32(2)) * Int32(8)
+            + Int32(SEMANTIC_ROW_OFFSET)
+        ) & Int32(M - 1)
+        # Register state remains owner-local for the entire exact stream.  It
+        # is published only once here because the final Ld16x256b epilogue
+        # remaps rows differently from the Ld16x64b xor-2 score ownership.
+        if (lane & Int32(2)) == Int32(0):
+            sFinalStats[owner_row, 0] = running_sum
+            sFinalStats[owner_row, 1] = running_max
+        cute.arch.fence_view_async_shared()
+        final_stats_ready_barrier.arrive_and_wait()
+
+        o_regs, o_coords = load_m64_o_fp32_256b(
+            pack_o_template,
+            thr_pack_pv,
+            tmem_base,
+            owner_tidx,
+        )
+        assert cute.size(o_regs) == 64
+        assert cute.size(o_coords) == 64
+
+        # B7's device inversion proves that 4*w/4*w+1 belong to one
+        # semantic row and 4*w+2/4*w+3 to its row-plus-eight peer.  Hoist
+        # validity, final-sum LDS, reciprocal, and row base once per stratum.
+        semantic_row0 = (
+            owner_warp * Int32(16)
+            + lane // Int32(4)
+            + Int32(SEMANTIC_ROW_OFFSET)
+        ) & Int32(M - 1)
+        semantic_row1 = (semantic_row0 + Int32(8)) & Int32(M - 1)
+        even_col_base = (lane % Int32(4)) * Int32(2)
+
+        if semantic_row0 < q_len:
+            inv_sum0 = cute.arch.rcp_approx(
+                Float32(sFinalStats[semantic_row0, 0])
+            )
+            query_idx0 = q_block_idx * Int32(M) + semantic_row0
+            destination_row0 = cute.domain_offset(
+                (batch_idx, query_idx0, head_idx, Int32(0)), mO_bthd
+            )
+            for word_i in cutlass.range(
+                O_PACKED_WORDS_PER_ROW_PER_THREAD, unroll_full=True
+            ):
+                even_i = word_i * 4
+                odd_i = even_i + 1
+                even_value = Float32(o_regs[even_i]) * inv_sum0
+                odd_value = Float32(o_regs[odd_i]) * inv_sum0
+                packed_word = _cvt_bf16x2_f32(
+                    Float32(odd_value), Float32(even_value)
+                )
+                even_col = (
+                    even_col_base + word_i * O_PACKED_COLUMN_STRIDE
+                )
+                _store_global_u32_inline(
+                    destination_row0.iterator + even_col, packed_word
+                )
+
+        if semantic_row1 < q_len:
+            inv_sum1 = cute.arch.rcp_approx(
+                Float32(sFinalStats[semantic_row1, 0])
+            )
+            query_idx1 = q_block_idx * Int32(M) + semantic_row1
+            destination_row1 = cute.domain_offset(
+                (batch_idx, query_idx1, head_idx, Int32(0)), mO_bthd
+            )
+            for word_i in cutlass.range(
+                O_PACKED_WORDS_PER_ROW_PER_THREAD, unroll_full=True
+            ):
+                even_i = word_i * 4 + 2
+                odd_i = even_i + 1
+                even_value = Float32(o_regs[even_i]) * inv_sum1
+                odd_value = Float32(o_regs[odd_i]) * inv_sum1
+                packed_word = _cvt_bf16x2_f32(
+                    Float32(odd_value), Float32(even_value)
+                )
+                even_col = (
+                    even_col_base + word_i * O_PACKED_COLUMN_STRIDE
+                )
+                _store_global_u32_inline(
+                    destination_row1.iterator + even_col, packed_word
+                )
+
+        if (lane & Int32(2)) == Int32(0) and owner_row < q_len:
+            query_idx = q_block_idx * Int32(M) + owner_row
+            mLSE_bth[batch_idx, query_idx, head_idx] = (
+                running_max
+                + cute.math.log2(running_sum, fastmath=True) * Float32(LN2)
+            )
+
+    cute.arch.barrier()
+    tmem.free(tmem_ptr)
+
+
+@cute.jit
+def _sol_attn_sm100_bf16_host(
+    q: cute.Tensor,
+    k: cute.Tensor,
+    v: cute.Tensor,
+    o: cute.Tensor,
+    kc: cute.Tensor,
+    vc: cute.Tensor,
+    threshold: cute.Tensor,
+    lse: cute.Tensor,
+    softmax_scale: Float32,
+    sink_start_block: Int32,
+    sink_end_block: Int32,
+    stream: cuda.CUstream = None,
+):
+    q, k, v, o, kc, vc = tuple(
+        assume_tensor_aligned(t) for t in (q, k, v, o, kc, vc)
+    )
+    q_mkl, k_nkl, kc_nkl = [
+        layout_utils.select(t, [1, 3, 2, 0]) for t in (q, k, kc)
+    ]
+    v_nkl, vc_nkl = [
+        layout_utils.select(t, [3, 1, 2, 0]) for t in (v, vc)
+    ]
+    token_count = cute.size(q_mkl.shape[0])
+    num_blocks = cute.size(kc_nkl.shape[0])
+    num_heads = cute.size(q_mkl.shape[2])
+    num_batches = cute.size(q_mkl.shape[3])
+    num_route_tiles = cute.ceil_div(num_blocks, ROUTE_TILE_SIZE)
+    pack_qk_op = tcgen05.MmaF16BF16Op(
+        BFloat16,
+        Float32,
+        PACK_QK_INST,
+        tcgen05.CtaGroup.ONE,
+        tcgen05.OperandSource.SMEM,
+        cute.nvgpu.OperandMajorMode.K,
+        cute.nvgpu.OperandMajorMode.K,
+    )
+    tiled_pack_qk = cute.make_tiled_mma(pack_qk_op)
+    pack_pv_op = tcgen05.MmaF16BF16Op(
+        BFloat16,
+        Float32,
+        PACK_PV_INST,
+        tcgen05.CtaGroup.ONE,
+        tcgen05.OperandSource.TMEM,
+        cute.nvgpu.OperandMajorMode.K,
+        cute.nvgpu.OperandMajorMode.MN,
+    )
+    tiled_pack_pv = cute.make_tiled_mma(pack_pv_op)
+    pack_qk_quarter_op = tcgen05.MmaF16BF16Op(
+        BFloat16,
+        Float32,
+        PACK_QK_QUARTER_INST,
+        tcgen05.CtaGroup.ONE,
+        tcgen05.OperandSource.SMEM,
+        cute.nvgpu.OperandMajorMode.K,
+        cute.nvgpu.OperandMajorMode.K,
+    )
+    tiled_pack_qk_gather = cute.make_tiled_mma(pack_qk_quarter_op)
+    pack_pv_quarter_op = tcgen05.MmaF16BF16Op(
+        BFloat16,
+        Float32,
+        PACK_PV_QUARTER_INST,
+        tcgen05.CtaGroup.ONE,
+        tcgen05.OperandSource.TMEM,
+        cute.nvgpu.OperandMajorMode.K,
+        cute.nvgpu.OperandMajorMode.MN,
+    )
+    tiled_pack_pv_gather = cute.make_tiled_mma(pack_pv_quarter_op)
+    q_layout = sm100_utils.make_smem_layout_a(
+        tiled_pack_qk, PACK_QK_TILE, BFloat16, 1
+    )
+    pack_k_layout = sm100_utils.make_smem_layout_b(
+        tiled_pack_qk, PACK_QK_TILE, BFloat16, PAIR_STAGES
+    )
+    pack_v_layout = sm100_utils.make_smem_layout_b(
+        tiled_pack_pv, PACK_PV_TILE, BFloat16, PAIR_STAGES
+    )
+    pack_k_gather_layout = sm100_utils.make_smem_layout_b(
+        tiled_pack_qk_gather,
+        PACK_QK_GATHER_TILE,
+        BFloat16,
+        PAIR_STAGES * 4,
+    )
+    pack_v_gather_layout = sm100_utils.make_smem_layout_b(
+        tiled_pack_pv_gather,
+        PACK_PV_GATHER_TILE,
+        BFloat16,
+        PAIR_STAGES * 4,
+    )
+    pack_p_layout = sm100_utils.make_smem_layout_a(
+        tiled_pack_pv, PACK_PV_TILE, BFloat16, 1
+    )
+    route_k_layout = sm100_utils.make_smem_layout_b(
+        tiled_pack_qk, PACK_QK_TILE, BFloat16, PAIR_STAGES
+    )
+    route_v_layout = sm100_utils.make_smem_layout_b(
+        tiled_pack_pv, PACK_PV_TILE, BFloat16, PAIR_STAGES
+    )
+    copy_op = cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE)
+    q_tma_atom, q_tma_tensor = cute.nvgpu.make_tiled_tma_atom_A(
+        copy_op,
+        q_mkl,
+        cute.select(q_layout, mode=[0, 1, 2]),
+        PACK_QK_TILE,
+        tiled_pack_qk,
+    )
+    pack_k_tma_layout = cute.make_composed_layout(
+        pack_k_gather_layout.inner,
+        0,
+        cute.make_layout((64, 64), stride=(64, 1)),
+    )
+    pack_k_tma_atom, pack_k_tma_tensor = cpasync.make_tiled_tma_atom(
+        copy_op,
+        k_nkl,
+        pack_k_tma_layout,
+        (64, 64),
+    )
+    pack_v_tma_layout = cute.make_composed_layout(
+        pack_v_gather_layout.inner,
+        0,
+        cute.make_layout((64, 64), stride=(1, 64)),
+    )
+    pack_v_tma_atom, pack_v_tma_tensor = cpasync.make_tiled_tma_atom(
+        copy_op,
+        v_nkl,
+        pack_v_tma_layout,
+        (64, 64),
+    )
+    kc_tma_atom, kc_tma_tensor = cute.nvgpu.make_tiled_tma_atom_B(
+        copy_op,
+        kc_nkl,
+        cute.select(route_k_layout, mode=[0, 1, 2]),
+        PACK_QK_TILE,
+        tiled_pack_qk,
+    )
+    vc_tma_atom, vc_tma_tensor = cute.nvgpu.make_tiled_tma_atom_B(
+        copy_op,
+        vc_nkl,
+        cute.select(route_v_layout, mode=[0, 1, 2]),
+        PACK_PV_TILE,
+        tiled_pack_pv,
+    )
+    _sol_attn_sm100_bf16_kernel(
+        tiled_pack_qk,
+        tiled_pack_pv,
+        q_tma_atom,
+        q_tma_tensor,
+        pack_k_tma_atom,
+        pack_k_tma_tensor,
+        pack_v_tma_atom,
+        pack_v_tma_tensor,
+        kc_tma_atom,
+        kc_tma_tensor,
+        vc_tma_atom,
+        vc_tma_tensor,
+        threshold,
+        o,
+        lse,
+        Int32(token_count),
+        Int32(num_blocks),
+        Int32(num_route_tiles),
+        softmax_scale,
+        sink_start_block,
+        sink_end_block,
+        q_layout,
+        pack_k_layout,
+        pack_k_gather_layout,
+        pack_p_layout,
+        pack_v_layout,
+        pack_v_gather_layout,
+        route_k_layout,
+        route_v_layout,
+    ).launch(
+        grid=(num_blocks, num_heads, num_batches),
+        block=(THREADS, 1, 1),
+        stream=stream,
+        min_blocks_per_mp=2,
+    )
+
+
+@cute.jit
+def forward(
+    q: cute.Tensor,
+    k: cute.Tensor,
+    v: cute.Tensor,
+    o: cute.Tensor,
+    kc: cute.Tensor,
+    vc: cute.Tensor,
+    threshold: cute.Tensor,
+    lse: cute.Tensor,
+    softmax_scale: Float32,
+    sink_start_block: Int32,
+    sink_end_block: Int32,
+    stream: cuda.CUstream = None,
+):
+    return _sol_attn_sm100_bf16_host(
+        q,
+        k,
+        v,
+        o,
+        kc,
+        vc,
+        threshold,
+        lse,
+        softmax_scale,
+        sink_start_block,
+        sink_end_block,
+        stream,
+    )
+
+
+__all__ = ["forward"]

torch-ext/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"]

torch-ext/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 .._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",
+]

torch-ext/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",
+]

torch-ext/sol_attn/sm120/__init__.py

@@ -0,0 +1,5 @@
+"""GeForce Blackwell (SM120) backend."""
+
+from .kernel import make_kernel
+
+__all__ = ["make_kernel"]

torch-ext/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"]

torch-ext/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 .._vendor.flash_attn.cute import utils as kernel_utils
+from ..common import layout_utils
+from ..common.selector import (
+    sol_attn_popc_b32,
+    sol_attn_route_is_exact,
+)
+
+
+M = 64
+N = 64
+D = 128
+DV = 128
+THREADS = 128
+STAGES = 1
+
+
+class SolAttnForwardSm120:
+    """M64/N64 warp-MMA Sol-Attn kernel for BF16 D128 inputs."""
+
+    def __init__(
+        self,
+        *,
+        debug_route_trace: bool = False,
+        prefetch_first_exact_k: bool = True,
+        prefetch_next_route_k: bool = True,
+    ):
+        self.dtype = cutlass.BFloat16
+        self.acc_dtype = cutlass.Float32
+        self.tile_shape_qk = (M, N, D)
+        self.tile_shape_pv = (M, DV, N)
+        self.num_threads = THREADS
+        self.q_stage = 1
+        self.kv_stage = STAGES
+        self.debug_route_trace = debug_route_trace
+        self.prefetch_first_exact_k = prefetch_first_exact_k
+        self.prefetch_next_route_k = prefetch_next_route_k
+
+    @cute.kernel
+    def kernel(
+        self,
+        mQ: cute.Tensor,
+        mK: cute.Tensor,
+        mV: cute.Tensor,
+        mO: cute.Tensor,
+        mKC: cute.Tensor,
+        mVC: cute.Tensor,
+        mThreshold: cute.Tensor,
+        mLSE: cute.Tensor,
+        tma_atom_Q: cute.CopyAtom,
+        tma_atom_K: cute.CopyAtom,
+        tma_atom_V: cute.CopyAtom,
+        tma_atom_KC: cute.CopyAtom,
+        tma_atom_VC: cute.CopyAtom,
+        tma_atom_O: cute.CopyAtom,
+        tiled_mma_qk: cute.TiledMma,
+        tiled_mma_pv: cute.TiledMma,
+        Q_smem_layout: cute.ComposedLayout,
+        K_smem_layout: cute.ComposedLayout,
+        V_smem_layout: cute.ComposedLayout,
+        O_smem_layout: cute.ComposedLayout,
+        scale_softmax_log2e: cutlass.Float32,
+        sink_start_block: cutlass.Int32,
+        sink_end_block: cutlass.Int32,
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        lane = cute.arch.lane_idx()
+        warp = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+        q_tile_idx, head_idx, batch_idx = cute.arch.block_idx()
+        q_tile_idx = cute.arch.make_warp_uniform(q_tile_idx)
+        head_idx = cute.arch.make_warp_uniform(head_idx)
+        batch_idx = cute.arch.make_warp_uniform(batch_idx)
+
+        token_count = mK.shape[0]
+        num_blocks = mKC.shape[0]
+        num_route_groups = cute.ceil_div(num_blocks, N)
+        q_start = q_tile_idx * M
+        q_len = token_count - q_start
+        if q_len > M:
+            q_len = cutlass.Int32(M)
+        threshold = cutlass.Float32(
+            mThreshold[batch_idx, q_tile_idx, head_idx]
+        )
+
+        storage = cutlass.utils.SmemAllocator().allocate(self.shared_storage_t)
+        if warp == 0 and lane == 0:
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_Q)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_K)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_V)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_KC)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_VC)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_O)
+
+        cg = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+        consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, self.num_threads // 32
+        )
+        cta_layout_vmnk = cute.make_layout((1, 1, 1, 1))
+        Q_pipeline = pipeline.PipelineTmaAsync.create(
+            num_stages=self.q_stage,
+            producer_group=cg,
+            consumer_group=consumer_group,
+            tx_count=cute.size_in_bytes(
+                self.Q_dtype, cute.select(Q_smem_layout, mode=[0, 1])
+            ),
+            barrier_storage=storage.Q_barrier.data_ptr(),
+            cta_layout_vmnk=cta_layout_vmnk,
+        )
+        K_pipeline = pipeline.PipelineTmaAsync.create(
+            num_stages=self.kv_stage,
+            producer_group=cg,
+            consumer_group=consumer_group,
+            tx_count=cute.size_in_bytes(
+                self.K_dtype, cute.select(K_smem_layout, mode=[0, 1])
+            ),
+            barrier_storage=storage.K_barrier.data_ptr(),
+            cta_layout_vmnk=cta_layout_vmnk,
+        )
+        V_pipeline = pipeline.PipelineTmaAsync.create(
+            num_stages=self.kv_stage,
+            producer_group=cg,
+            consumer_group=consumer_group,
+            tx_count=cute.size_in_bytes(
+                self.V_dtype, cute.select(V_smem_layout, mode=[0, 1])
+            ),
+            barrier_storage=storage.V_barrier.data_ptr(),
+            cta_layout_vmnk=cta_layout_vmnk,
+        )
+        Q_producer = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.q_stage
+        )
+        Q_consumer = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.q_stage
+        )
+        K_producer = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.kv_stage
+        )
+        K_consumer = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.kv_stage
+        )
+        V_producer = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, self.kv_stage
+        )
+        V_consumer = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.kv_stage
+        )
+
+        sQ = storage.Q_smem.get_tensor(
+            Q_smem_layout.outer, swizzle=Q_smem_layout.inner
+        )
+        sK = storage.K_smem.get_tensor(
+            K_smem_layout.outer, swizzle=K_smem_layout.inner
+        )
+        sV = storage.V_smem.get_tensor(
+            V_smem_layout.outer, swizzle=V_smem_layout.inner
+        )
+        # Q is register-resident after the prologue.  Reuse its 16 KiB SMEM
+        # allocation for route scratch until the same allocation becomes sO
+        # in the epilogue.  This drops the CTA below the 2-block/SM threshold
+        # on SM120 without changing any route reduction or synchronization.
+        route_f32_ptr = cute.recast_ptr(
+            storage.Q_smem.data_ptr(), dtype=cutlass.Float32
+        )
+        route_i32_ptr = cute.recast_ptr(
+            storage.Q_smem.data_ptr(), dtype=cutlass.Int32
+        )
+        route_sums = cute.make_tensor(
+            route_f32_ptr, cute.make_layout((4, N))
+        )
+        column_masks = cute.make_tensor(
+            route_f32_ptr + 4 * N, cute.make_layout(N)
+        )
+        route_indices = cute.make_tensor(
+            route_i32_ptr + 5 * N, cute.make_layout(N)
+        )
+        route_meta = cute.make_tensor(
+            route_i32_ptr + 6 * N, cute.make_layout(2)
+        )
+
+        mQ_slice = mQ[None, None, head_idx, batch_idx]
+        mK_slice = mK[None, None, head_idx, batch_idx]
+        mV_slice = mV[None, None, head_idx, batch_idx]
+        mO_slice = mO[None, None, head_idx, batch_idx]
+        mKC_slice = mKC[None, None, head_idx, batch_idx]
+        mVC_slice = mVC[None, None, head_idx, batch_idx]
+        if cutlass.const_expr(not self.debug_route_trace):
+            mLSE_slice = mLSE[None, head_idx, batch_idx]
+
+        gQ = cute.local_tile(
+            mQ_slice, (M, D), coord=(q_tile_idx, 0)
+        )
+        gK = cute.local_tile(mK_slice, (N, D), coord=(None, 0))
+        gV = cute.local_tile(mV_slice, (DV, N), coord=(0, None))
+        gKC = cute.local_tile(mKC_slice, (N, D), coord=(None, 0))
+        gVC = cute.local_tile(mVC_slice, (DV, N), coord=(0, None))
+        gO = cute.local_tile(
+            mO_slice, (M, DV), coord=(q_tile_idx, 0)
+        )
+
+        cta_coord_layout = (0, cute.make_layout(1))
+        tQsQ, tQgQ = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_Q,
+            *cta_coord_layout,
+            cute.group_modes(sQ, 0, 2),
+            cute.group_modes(gQ, 0, 2),
+        )
+        tKsK, tKgK = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_K,
+            *cta_coord_layout,
+            cute.group_modes(sK, 0, 2),
+            cute.group_modes(gK, 0, 2),
+        )
+        tVsV, tVgV = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_V,
+            *cta_coord_layout,
+            cute.group_modes(sV, 0, 2),
+            cute.group_modes(gV, 0, 2),
+        )
+        tKCsK, tKCgKC = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_KC,
+            *cta_coord_layout,
+            cute.group_modes(sK, 0, 2),
+            cute.group_modes(gKC, 0, 2),
+        )
+        tVCsV, tVCgVC = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_VC,
+            *cta_coord_layout,
+            cute.group_modes(sV, 0, 2),
+            cute.group_modes(gVC, 0, 2),
+        )
+
+        cS = cute.make_identity_tensor(self.tile_shape_qk[:2])
+        thr_mma_qk = tiled_mma_qk.get_slice(tidx)
+        tSsQ = thr_mma_qk.partition_A(sQ)
+        tSsK = thr_mma_qk.partition_B(sK)
+        tSrQ = tiled_mma_qk.make_fragment_A(tSsQ[None, None, None, 0])
+        tSrK = tiled_mma_qk.make_fragment_B(tSsK[None, None, None, 0])
+        tSrS = cute.make_rmem_tensor(
+            thr_mma_qk.partition_shape_C((M, N)), self.acc_dtype
+        )
+        tScS = thr_mma_qk.partition_C(cS)
+
+        thr_mma_pv = tiled_mma_pv.get_slice(tidx)
+        tOsV = thr_mma_pv.partition_B(sV)
+        tOrV = tiled_mma_pv.make_fragment_B(tOsV[None, None, None, 0])
+        tOrO = cute.make_rmem_tensor(
+            thr_mma_pv.partition_shape_C((M, DV)), self.acc_dtype
+        )
+
+        atom_copy_Q = cute.make_copy_atom(
+            cute.nvgpu.warp.LdMatrix8x8x16bOp(
+                self.Q_layout.is_m_major_a(), 4
+            ),
+            self.Q_dtype,
+        )
+        atom_copy_K = cute.make_copy_atom(
+            cute.nvgpu.warp.LdMatrix8x8x16bOp(
+                self.K_layout.is_n_major_b(), 4
+            ),
+            self.K_dtype,
+        )
+        atom_copy_V = cute.make_copy_atom(
+            cute.nvgpu.warp.LdMatrix8x8x16bOp(
+                self.V_layout.is_n_major_b(), 4
+            ),
+            self.V_dtype,
+        )
+        smem_copy_Q = cute.make_tiled_copy_A(atom_copy_Q, tiled_mma_qk)
+        smem_copy_K = cute.make_tiled_copy_B(atom_copy_K, tiled_mma_qk)
+        smem_copy_V = cute.make_tiled_copy_B(atom_copy_V, tiled_mma_pv)
+        thr_copy_Q = smem_copy_Q.get_slice(tidx)
+        thr_copy_K = smem_copy_K.get_slice(tidx)
+        thr_copy_V = smem_copy_V.get_slice(tidx)
+        tSsQ_copy = thr_copy_Q.partition_S(sQ)
+        tSrQ_copy = thr_copy_Q.retile(tSrQ)
+        tSsK_copy = thr_copy_K.partition_S(sK)
+        tOsV_copy = thr_copy_V.partition_S(sV)
+
+        max_m_layout = cute.make_layout(
+            cute.size(
+                layout_utils.reshape_acc_to_mn(tOrO).layout,
+                mode=[0],
+            )
+        )
+        max_m = cute.make_rmem_tensor_like(max_m_layout, cutlass.Float32)
+        sum_m = cute.make_rmem_tensor_like(max_m, cutlass.Float32)
+        tOrO.store(cute.full_like(tOrO, 0.0, self.acc_dtype))
+        max_m.store(cute.full_like(max_m, float("-inf"), cutlass.Float32))
+        sum_m.store(cute.full_like(sum_m, 0.0, cutlass.Float32))
+
+        if warp == 0:
+            Q_pipeline.producer_acquire(Q_producer)
+            cute.copy(
+                tma_atom_Q,
+                tQgQ,
+                tQsQ[None, Q_producer.index],
+                tma_bar_ptr=Q_pipeline.producer_get_barrier(Q_producer),
+            )
+            Q_pipeline.producer_commit(Q_producer)
+            Q_producer.advance()
+        cute.arch.sync_threads()
+        q_wait = Q_pipeline.consumer_try_wait(Q_consumer)
+        Q_pipeline.consumer_wait(Q_consumer, q_wait)
+        q_stage = Q_consumer.index
+        for k_block in cutlass.range_constexpr(cute.size(tSrQ, mode=[2])):
+            cute.copy(
+                smem_copy_Q,
+                tSsQ_copy[None, None, k_block, q_stage],
+                tSrQ_copy[None, None, k_block],
+            )
+        Q_pipeline.consumer_release(Q_consumer)
+        Q_consumer.advance()
+
+        for route_group in cutlass.range(
+            0, num_route_groups, 1, unroll=1
+        ):
+            group_start = route_group * cutlass.Int32(N)
+            valid_blocks = num_blocks - group_start
+            if valid_blocks > N:
+                valid_blocks = cutlass.Int32(N)
+
+            if warp == 0:
+                if cutlass.const_expr(self.prefetch_next_route_k):
+                    # P19-style terminal handoff: when the previous route
+                    # group had an exact block, its final exact QK already
+                    # refilled this K stage with the current group's KC.
+                    if route_group == 0:
+                        K_pipeline.producer_acquire(K_producer)
+                        cute.copy(
+                            tma_atom_KC,
+                            tKCgKC[None, route_group],
+                            tKCsK[None, K_producer.index],
+                            tma_bar_ptr=K_pipeline.producer_get_barrier(
+                                K_producer
+                            ),
+                        )
+                        K_pipeline.producer_commit(K_producer)
+                        K_producer.advance()
+                    else:
+                        previous_group_exact_count = cutlass.Int32(
+                            route_meta[0]
+                        )
+                        if previous_group_exact_count == 0:
+                            K_pipeline.producer_acquire(K_producer)
+                            cute.copy(
+                                tma_atom_KC,
+                                tKCgKC[None, route_group],
+                                tKCsK[None, K_producer.index],
+                                tma_bar_ptr=K_pipeline.producer_get_barrier(
+                                    K_producer
+                                ),
+                            )
+                            K_pipeline.producer_commit(K_producer)
+                            K_producer.advance()
+                else:
+                    K_pipeline.producer_acquire(K_producer)
+                    cute.copy(
+                        tma_atom_KC,
+                        tKCgKC[None, route_group],
+                        tKCsK[None, K_producer.index],
+                        tma_bar_ptr=K_pipeline.producer_get_barrier(
+                            K_producer
+                        ),
+                    )
+                    K_pipeline.producer_commit(K_producer)
+                    K_producer.advance()
+                V_pipeline.producer_acquire(V_producer)
+                cute.copy(
+                    tma_atom_VC,
+                    tVCgVC[None, route_group],
+                    tVCsV[None, V_producer.index],
+                    tma_bar_ptr=V_pipeline.producer_get_barrier(V_producer),
+                )
+                V_pipeline.producer_commit(V_producer)
+                V_producer.advance()
+
+            k_wait = K_pipeline.consumer_try_wait(K_consumer)
+            K_pipeline.consumer_wait(K_consumer, k_wait)
+            gemm_smem_zero_acc(
+                tiled_mma_qk,
+                tSrS,
+                tSrQ,
+                tSrK,
+                tSsK_copy[None, None, None, K_consumer.index],
+                smem_copy_K,
+            )
+            K_pipeline.consumer_release(K_consumer)
+            K_consumer.advance()
+
+            reduce_route_columns(
+                tSrS,
+                tScS,
+                route_sums,
+                warp,
+                lane,
+                q_len,
+            )
+            cute.arch.fence_view_async_shared()
+            cute.arch.sync_threads()
+
+            if warp == 0:
+                preceding = cutlass.Int32(0)
+                lane_mask_lt = cutlass.Int32(0x7FFFFFFF) >> (
+                    cutlass.Int32(31) - lane
+                )
+                for word in cutlass.range_constexpr(2):
+                    off = cutlass.Int32(word * 32) + lane
+                    valid = off < valid_blocks
+                    exact = False
+                    if valid:
+                        col_sum = (
+                            cutlass.Float32(route_sums[0, off])
+                            + cutlass.Float32(route_sums[1, off])
+                            + cutlass.Float32(route_sums[2, off])
+                            + cutlass.Float32(route_sums[3, off])
+                        )
+                        col_mean = (
+                            col_sum
+                            * scale_softmax_log2e
+                            / cutlass.Float32(q_len)
+                        )
+                        kv_block = group_start + off
+                        exact = sol_attn_route_is_exact(
+                            q_tile_idx,
+                            kv_block,
+                            col_mean,
+                            threshold,
+                            valid,
+                        )
+                        exact = exact or (
+                            kv_block >= sink_start_block
+                            and kv_block < sink_end_block
+                        )
+                    ballot = cutlass.Int32(
+                        cute.arch.vote_ballot_sync(exact)
+                    )
+                    column_masks[off] = (
+                        -cutlass.Float32.inf
+                        if (exact or not valid)
+                        else cutlass.Float32(0.0)
+                    )
+                    rank = preceding + sol_attn_popc_b32(
+                        ballot & lane_mask_lt
+                    )
+                    if exact:
+                        route_indices[rank] = group_start + off
+                    preceding += sol_attn_popc_b32(ballot)
+                    if cutlass.const_expr(self.debug_route_trace):
+                        if lane == 0:
+                            mLSE[
+                                batch_idx,
+                                q_tile_idx,
+                                head_idx,
+                                route_group,
+                                word,
+                            ] = ballot
+                if lane == 0:
+                    route_meta[0] = preceding
+                    route_meta[1] = valid_blocks
+                    cute.arch.fence_view_async_shared()
+            cute.arch.sync_threads()
+
+            exact_count = cutlass.Int32(route_meta[0])
+            has_approx = exact_count < valid_blocks
+            if cutlass.const_expr(self.prefetch_first_exact_k):
+                # Once routing identifies the first exact block, the route KC
+                # stage is free.  Refill it before the approximate softmax/PV
+                # so the first exact K transfer overlaps that work.
+                if warp == 0 and exact_count > 0:
+                    first_exact = cutlass.Int32(route_indices[0])
+                    K_pipeline.producer_acquire(K_producer)
+                    cute.copy(
+                        tma_atom_K,
+                        tKgK[None, first_exact],
+                        tKsK[None, K_producer.index],
+                        tma_bar_ptr=K_pipeline.producer_get_barrier(
+                            K_producer
+                        ),
+                    )
+                    K_pipeline.producer_commit(K_producer)
+                    K_producer.advance()
+            v_wait = V_pipeline.consumer_try_wait(V_consumer)
+            V_pipeline.consumer_wait(V_consumer, v_wait)
+            if has_approx:
+                apply_route_mask(tSrS, tScS, column_masks, q_len)
+                row_scale = online_softmax_route(
+                    tSrS,
+                    tScS,
+                    max_m,
+                    sum_m,
+                    scale_softmax_log2e,
+                    group_start,
+                    token_count,
+                )
+                rescale_o_for_next_acc(tOrO, row_scale)
+                tOrP_frg = cute.make_rmem_tensor_like(
+                    tSrS, self.K_dtype
+                )
+                tOrP_frg.store(tSrS.load().to(self.K_dtype))
+                tOrP = layout_utils.reshape_acc_to_frgA(tOrP_frg)
+                gemm_rs_smem(
+                    tiled_mma_pv,
+                    tOrO,
+                    tOrP,
+                    tOrV,
+                    tOsV_copy[None, None, None, V_consumer.index],
+                    smem_copy_V,
+                )
+            V_pipeline.consumer_release(V_consumer)
+            V_consumer.advance()
+
+            if warp == 0 and exact_count > 0:
+                first_exact = cutlass.Int32(route_indices[0])
+                if cutlass.const_expr(not self.prefetch_first_exact_k):
+                    K_pipeline.producer_acquire(K_producer)
+                    cute.copy(
+                        tma_atom_K,
+                        tKgK[None, first_exact],
+                        tKsK[None, K_producer.index],
+                        tma_bar_ptr=K_pipeline.producer_get_barrier(
+                            K_producer
+                        ),
+                    )
+                    K_pipeline.producer_commit(K_producer)
+                    K_producer.advance()
+                V_pipeline.producer_acquire(V_producer)
+                cute.copy(
+                    tma_atom_V,
+                    tVgV[None, first_exact],
+                    tVsV[None, V_producer.index],
+                    tma_bar_ptr=V_pipeline.producer_get_barrier(V_producer),
+                )
+                V_pipeline.producer_commit(V_producer)
+                V_producer.advance()
+
+            for ordinal in cutlass.range(0, exact_count, 1, unroll=1):
+                exact_block = cutlass.Int32(route_indices[ordinal])
+                k_wait = K_pipeline.consumer_try_wait(K_consumer)
+                K_pipeline.consumer_wait(K_consumer, k_wait)
+                gemm_smem_zero_acc(
+                    tiled_mma_qk,
+                    tSrS,
+                    tSrQ,
+                    tSrK,
+                    tSsK_copy[None, None, None, K_consumer.index],
+                    smem_copy_K,
+                )
+                K_pipeline.consumer_release(K_consumer)
+                K_consumer.advance()
+                next_ordinal = ordinal + cutlass.Int32(1)
+                if warp == 0:
+                    if next_ordinal < exact_count:
+                        next_exact = cutlass.Int32(
+                            route_indices[next_ordinal]
+                        )
+                        K_pipeline.producer_acquire(K_producer)
+                        cute.copy(
+                            tma_atom_K,
+                            tKgK[None, next_exact],
+                            tKsK[None, K_producer.index],
+                            tma_bar_ptr=K_pipeline.producer_get_barrier(
+                                K_producer
+                            ),
+                        )
+                        K_pipeline.producer_commit(K_producer)
+                        K_producer.advance()
+                    else:
+                        if cutlass.const_expr(
+                            self.prefetch_next_route_k
+                        ):
+                            next_route_group = route_group + cutlass.Int32(1)
+                            if next_route_group < num_route_groups:
+                                # Reuse the K stage released by the final
+                                # exact QK.  The next outer prologue supplies
+                                # VC, matching the SM90 P19 partial handoff.
+                                K_pipeline.producer_acquire(K_producer)
+                                cute.copy(
+                                    tma_atom_KC,
+                                    tKCgKC[None, next_route_group],
+                                    tKCsK[None, K_producer.index],
+                                    tma_bar_ptr=(
+                                        K_pipeline.producer_get_barrier(
+                                            K_producer
+                                        )
+                                    ),
+                                )
+                                K_pipeline.producer_commit(K_producer)
+                                K_producer.advance()
+                block_len = token_count - exact_block * cutlass.Int32(N)
+                if block_len > N:
+                    block_len = cutlass.Int32(N)
+                mask_exact_scores(tSrS, tScS, block_len, q_len)
+                row_scale = online_softmax(
+                    tSrS, max_m, sum_m, scale_softmax_log2e
+                )
+                rescale_o_for_next_acc(tOrO, row_scale)
+                tOrP_frg = cute.make_rmem_tensor_like(
+                    tSrS, self.K_dtype
+                )
+                tOrP_frg.store(tSrS.load().to(self.K_dtype))
+                tOrP = layout_utils.reshape_acc_to_frgA(tOrP_frg)
+
+                v_wait = V_pipeline.consumer_try_wait(V_consumer)
+                V_pipeline.consumer_wait(V_consumer, v_wait)
+                gemm_rs_smem(
+                    tiled_mma_pv,
+                    tOrO,
+                    tOrP,
+                    tOrV,
+                    tOsV_copy[None, None, None, V_consumer.index],
+                    smem_copy_V,
+                )
+                V_pipeline.consumer_release(V_consumer)
+                V_consumer.advance()
+                if warp == 0 and next_ordinal < exact_count:
+                    next_exact = cutlass.Int32(route_indices[next_ordinal])
+                    V_pipeline.producer_acquire(V_producer)
+                    cute.copy(
+                        tma_atom_V,
+                        tVgV[None, next_exact],
+                        tVsV[None, V_producer.index],
+                        tma_bar_ptr=V_pipeline.producer_get_barrier(
+                            V_producer
+                        ),
+                    )
+                    V_pipeline.producer_commit(V_producer)
+                    V_producer.advance()
+
+        final_ratio, lse = finalize_softmax(
+            max_m, sum_m, scale_softmax_log2e
+        )
+        rescale_o_for_next_acc(tOrO, final_ratio)
+        if cutlass.const_expr(not self.debug_route_trace):
+            tScS_mn = layout_utils.reshape_acc_to_mn(tScS)
+            for m in cutlass.range_constexpr(cute.size(lse)):
+                row = tScS_mn[m, 0][0]
+                if tScS_mn[m, 0][1] == 0 and row < q_len:
+                    mLSE_slice[q_start + row] = lse[m]
+
+        tOrO_cvt = cute.make_rmem_tensor_like(tOrO, self.O_dtype)
+        tOrO_cvt.store(tOrO.load().to(self.O_dtype))
+        sO = storage.Q_smem.get_tensor(
+            O_smem_layout.outer, swizzle=O_smem_layout.inner
+        )
+        tiled_copy_O = cute.make_tiled_copy_C(
+            cute.make_copy_atom(
+                cute.nvgpu.warp.StMatrix8x8x16bOp(
+                    self.O_layout.is_m_major_c(), 4
+                ),
+                self.O_dtype,
+            ),
+            tiled_mma_pv,
+        )
+        tOrO_cv = tiled_copy_O.retile(tOrO_cvt)
+        tOsO = tiled_copy_O.get_slice(tidx).partition_D(sO)
+        cute.copy(tiled_copy_O, tOrO_cv, tOsO)
+        cute.arch.fence_view_async_shared()
+        cute.arch.sync_threads()
+        tOsO, tOgO = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_O,
+            *cta_coord_layout,
+            cute.group_modes(sO, 0, 2),
+            cute.group_modes(gO, 0, 2),
+        )
+        if warp == 0:
+            cute.copy(tma_atom_O, tOsO, tOgO)
+            cute.arch.cp_async_bulk_commit_group()
+            cute.arch.cp_async_bulk_wait_group(0, read=True)
+
+    @cute.jit
+    def __call__(
+        self,
+        q: cute.Tensor,
+        k: cute.Tensor,
+        v: cute.Tensor,
+        o: cute.Tensor,
+        kc: cute.Tensor,
+        vc: cute.Tensor,
+        threshold: cute.Tensor,
+        lse: cute.Tensor,
+        softmax_scale: cutlass.Float32,
+        sink_start_block: cutlass.Int32,
+        sink_end_block: cutlass.Int32,
+        stream: cuda.CUstream,
+    ):
+        q_mkl, k_nkl, kc_nkl = [
+            layout_utils.select(t, [1, 3, 2, 0])
+            for t in (q, k, kc)
+        ]
+        v_nkl, vc_nkl = [
+            layout_utils.select(t, [3, 1, 2, 0]) for t in (v, vc)
+        ]
+        o_mkl = layout_utils.select(o, [1, 3, 2, 0])
+        if cutlass.const_expr(self.debug_route_trace):
+            lse_target = lse
+        else:
+            lse_target = layout_utils.select(lse, [1, 2, 0])
+
+        self.Q_dtype = q_mkl.element_type
+        self.K_dtype = k_nkl.element_type
+        self.V_dtype = v_nkl.element_type
+        self.O_dtype = o_mkl.element_type
+        self.Q_layout = utils.LayoutEnum.from_tensor(q_mkl)
+        self.K_layout = utils.LayoutEnum.from_tensor(k_nkl)
+        self.V_layout = utils.LayoutEnum.from_tensor(v_nkl)
+        self.O_layout = utils.LayoutEnum.from_tensor(o_mkl)
+        assert self.Q_dtype == cutlass.BFloat16
+        assert self.K_dtype == cutlass.BFloat16
+        assert self.V_dtype == cutlass.BFloat16
+
+        self.Q_smem_layout = sm90_utils.make_smem_layout_a(
+            self.Q_layout,
+            self.tile_shape_qk,
+            self.Q_dtype,
+            self.q_stage,
+        )
+        self.K_smem_layout = sm90_utils.make_smem_layout_b(
+            self.K_layout,
+            self.tile_shape_qk,
+            self.K_dtype,
+            self.kv_stage,
+        )
+        self.V_smem_layout = sm90_utils.make_smem_layout_b(
+            self.V_layout,
+            self.tile_shape_pv,
+            self.V_dtype,
+            self.kv_stage,
+        )
+        O_smem_layout_staged = sm90_utils.make_smem_layout_epi(
+            self.O_dtype,
+            self.O_layout,
+            self.tile_shape_pv[:2],
+            1,
+        )
+        self.O_smem_layout = cute.select(
+            O_smem_layout_staged, mode=[0, 1]
+        )
+
+        @cute.struct
+        class SharedStorage:
+            Q_barrier: cute.struct.MemRange[
+                cutlass.Int64, self.q_stage * 2
+            ]
+            K_barrier: cute.struct.MemRange[
+                cutlass.Int64, self.kv_stage * 2
+            ]
+            V_barrier: cute.struct.MemRange[
+                cutlass.Int64, self.kv_stage * 2
+            ]
+            Q_smem: cute.struct.Align[
+                cute.struct.MemRange[
+                    self.Q_dtype, cute.cosize(self.Q_smem_layout)
+                ],
+                128,
+            ]
+            K_smem: cute.struct.Align[
+                cute.struct.MemRange[
+                    self.K_dtype, cute.cosize(self.K_smem_layout)
+                ],
+                128,
+            ]
+            V_smem: cute.struct.Align[
+                cute.struct.MemRange[
+                    self.V_dtype, cute.cosize(self.V_smem_layout)
+                ],
+                128,
+            ]
+        self.shared_storage_t = SharedStorage
+
+        tiled_mma_qk = cute.make_tiled_mma(
+            cute.nvgpu.warp.MmaF16BF16Op(
+                self.Q_dtype,
+                self.acc_dtype,
+                (16, 8, 16),
+            ),
+            cute.make_layout((4, 1, 1)),
+            permutation_mnk=(64, 16, 16),
+        )
+        tiled_mma_pv = cute.make_tiled_mma(
+            cute.nvgpu.warp.MmaF16BF16Op(
+                self.K_dtype,
+                self.acc_dtype,
+                (16, 8, 16),
+            ),
+            cute.make_layout((4, 1, 1)),
+            permutation_mnk=(64, 16, 16),
+        )
+
+        g2s_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp()
+        tma_atom_Q, tma_tensor_Q = (
+            cute.nvgpu.cpasync.make_tiled_tma_atom(
+                g2s_op,
+                q_mkl,
+                self.Q_smem_layout,
+                (M, D),
+                num_multicast=1,
+            )
+        )
+        tma_atom_K, tma_tensor_K = (
+            cute.nvgpu.cpasync.make_tiled_tma_atom(
+                g2s_op,
+                k_nkl,
+                self.K_smem_layout,
+                (N, D),
+                num_multicast=1,
+            )
+        )
+        tma_atom_V, tma_tensor_V = (
+            cute.nvgpu.cpasync.make_tiled_tma_atom(
+                g2s_op,
+                v_nkl,
+                self.V_smem_layout,
+                (DV, N),
+                num_multicast=1,
+            )
+        )
+        tma_atom_KC, tma_tensor_KC = (
+            cute.nvgpu.cpasync.make_tiled_tma_atom(
+                g2s_op,
+                kc_nkl,
+                self.K_smem_layout,
+                (N, D),
+                num_multicast=1,
+            )
+        )
+        tma_atom_VC, tma_tensor_VC = (
+            cute.nvgpu.cpasync.make_tiled_tma_atom(
+                g2s_op,
+                vc_nkl,
+                self.V_smem_layout,
+                (DV, N),
+                num_multicast=1,
+            )
+        )
+        s2g_op = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp()
+        tma_atom_O, tma_tensor_O = (
+            cute.nvgpu.cpasync.make_tiled_tma_atom(
+                s2g_op,
+                o_mkl,
+                self.O_smem_layout,
+                (M, DV),
+                num_multicast=1,
+            )
+        )
+
+        self.kernel(
+            tma_tensor_Q,
+            tma_tensor_K,
+            tma_tensor_V,
+            tma_tensor_O,
+            tma_tensor_KC,
+            tma_tensor_VC,
+            threshold,
+            lse_target,
+            tma_atom_Q,
+            tma_atom_K,
+            tma_atom_V,
+            tma_atom_KC,
+            tma_atom_VC,
+            tma_atom_O,
+            tiled_mma_qk,
+            tiled_mma_pv,
+            self.Q_smem_layout,
+            self.K_smem_layout,
+            self.V_smem_layout,
+            self.O_smem_layout,
+            softmax_scale * 1.4426950408889634,
+            sink_start_block,
+            sink_end_block,
+        ).launch(
+            grid=(cute.ceil_div(q_mkl.shape[0], M), q_mkl.shape[2], q_mkl.shape[3]),
+            block=(self.num_threads, 1, 1),
+            cluster=(1, 1, 1),
+            smem=self.shared_storage_t.size_in_bytes(),
+            stream=stream,
+            min_blocks_per_mp=1,
+        )
+
+
+@cute.jit
+def gemm_smem_zero_acc(
+    tiled_mma: cute.TiledMma,
+    acc: cute.Tensor,
+    tCrA: cute.Tensor,
+    tCrB: cute.Tensor,
+    tCsB: cute.Tensor,
+    smem_tiled_copy_B: cute.TiledCopy,
+):
+    acc.fill(0.0)
+    tCrB_copy = smem_tiled_copy_B.retile(tCrB)
+    cute.copy(
+        smem_tiled_copy_B,
+        tCsB[None, None, 0],
+        tCrB_copy[None, None, 0],
+    )
+    for k_block in cutlass.range_constexpr(cute.size(tCsB.shape[2])):
+        if k_block < cute.size(tCsB.shape[2]) - 1:
+            cute.copy(
+                smem_tiled_copy_B,
+                tCsB[None, None, k_block + 1],
+                tCrB_copy[None, None, k_block + 1],
+            )
+        cute.gemm(
+            tiled_mma,
+            acc,
+            tCrA[None, None, k_block],
+            tCrB[None, None, k_block],
+            acc,
+        )
+
+
+@cute.jit
+def gemm_rs_smem(
+    tiled_mma: cute.TiledMma,
+    acc: cute.Tensor,
+    tCrA: cute.Tensor,
+    tCrB: cute.Tensor,
+    tCsB: cute.Tensor,
+    smem_tiled_copy_B: cute.TiledCopy,
+):
+    tCrB_copy = smem_tiled_copy_B.retile(tCrB)
+    cute.copy(
+        smem_tiled_copy_B,
+        tCsB[None, None, 0],
+        tCrB_copy[None, None, 0],
+    )
+    for k_block in cutlass.range_constexpr(cute.size(tCrA.shape[2])):
+        if k_block < cute.size(tCrA.shape[2]) - 1:
+            cute.copy(
+                smem_tiled_copy_B,
+                tCsB[None, None, k_block + 1],
+                tCrB_copy[None, None, k_block + 1],
+            )
+        cute.gemm(
+            tiled_mma,
+            acc,
+            tCrA[None, None, k_block],
+            tCrB[None, None, k_block],
+            acc,
+        )
+
+
+@cute.jit
+def reduce_route_columns(
+    scores: cute.Tensor,
+    coords: cute.Tensor,
+    route_sums: cute.Tensor,
+    warp: cutlass.Int32,
+    lane: cutlass.Int32,
+    q_len: cutlass.Int32,
+):
+    """Reduce M64 score columns using the measured SM120 lane layout."""
+
+    scores_mn = layout_utils.reshape_acc_to_mn(scores)
+    coords_mn = layout_utils.reshape_acc_to_mn(coords)
+    row0 = coords_mn[0, 0][0]
+    row1 = coords_mn[1, 0][0]
+    valid0 = row0 < q_len
+    valid1 = row1 < q_len
+    for group in cutlass.range_constexpr(8):
+        n0 = group * 2
+        partial0 = cutlass.Float32(0.0)
+        partial1 = cutlass.Float32(0.0)
+        if valid0:
+            partial0 += cutlass.Float32(scores_mn[0, n0])
+            partial1 += cutlass.Float32(scores_mn[0, n0 + 1])
+        if valid1:
+            partial0 += cutlass.Float32(scores_mn[1, n0])
+            partial1 += cutlass.Float32(scores_mn[1, n0 + 1])
+        for offset in (4, 8, 16):
+            partial0 += cute.arch.shuffle_sync_bfly(partial0, offset=offset)
+            partial1 += cute.arch.shuffle_sync_bfly(partial1, offset=offset)
+        if lane < 4:
+            column = cutlass.Int32(group * 8) + lane * cutlass.Int32(2)
+            route_sums[warp, column] = partial0
+            route_sums[warp, column + 1] = partial1
+
+
+@cute.jit
+def apply_route_mask(
+    scores: cute.Tensor,
+    coords: cute.Tensor,
+    column_masks: cute.Tensor,
+    q_len: cutlass.Int32,
+):
+    scores_mn = layout_utils.reshape_acc_to_mn(scores)
+    coords_mn = layout_utils.reshape_acc_to_mn(coords)
+    for m in cutlass.range_constexpr(cute.size(scores_mn, mode=[0])):
+        valid_row = coords_mn[m, 0][0] < q_len
+        for n in cutlass.range_constexpr(cute.size(scores_mn, mode=[1])):
+            column = coords_mn[m, n][1]
+            scores_mn[m, n] = (
+                cutlass.Float32(scores_mn[m, n])
+                + cutlass.Float32(column_masks[column])
+                if valid_row
+                else -cutlass.Float32.inf
+            )
+
+
+@cute.jit
+def mask_exact_scores(
+    scores: cute.Tensor,
+    coords: cute.Tensor,
+    block_len: cutlass.Int32,
+    q_len: cutlass.Int32,
+):
+    scores_mn = layout_utils.reshape_acc_to_mn(scores)
+    coords_mn = layout_utils.reshape_acc_to_mn(coords)
+    for m in cutlass.range_constexpr(cute.size(scores_mn, mode=[0])):
+        valid_row = coords_mn[m, 0][0] < q_len
+        for n in cutlass.range_constexpr(cute.size(scores_mn, mode=[1])):
+            if (not valid_row) or coords_mn[m, n][1] >= block_len:
+                scores_mn[m, n] = -cutlass.Float32.inf
+
+
+@cute.jit
+def online_softmax(
+    scores: cute.Tensor,
+    row_max: cute.Tensor,
+    row_sum: cute.Tensor,
+    scale_log2e: cutlass.Float32,
+):
+    scores_mn = layout_utils.reshape_acc_to_mn(scores)
+    row_scale = cute.make_rmem_tensor_like(row_max, cutlass.Float32)
+    for m in cutlass.range_constexpr(cute.size(row_max)):
+        score_row = scores_mn[m, None].load()
+        current_max = kernel_utils.fmax_reduce(
+            score_row, init_val=row_max[m], arch=80
+        )
+        current_max = cute.arch.warp_reduction_max(
+            current_max, threads_in_group=4
+        )
+        previous_max = row_max[m]
+        row_max[m] = current_max
+        safe_max = (
+            cutlass.Float32(0.0)
+            if current_max == -cutlass.Float32.inf
+            else current_max
+        )
+        scaled_max = safe_max * scale_log2e
+        probabilities = cute.math.exp2(
+            score_row * scale_log2e - scaled_max, fastmath=True
+        )
+        row_scale[m] = cute.math.exp2(
+            (previous_max - safe_max) * scale_log2e, fastmath=True
+        )
+        row_sum[m] = kernel_utils.fadd_reduce(
+            probabilities,
+            init_val=row_sum[m] * row_scale[m],
+            arch=80,
+        )
+        scores_mn[m, None].store(probabilities)
+    return row_scale
+
+
+@cute.jit
+def online_softmax_route(
+    scores: cute.Tensor,
+    coords: cute.Tensor,
+    row_max: cute.Tensor,
+    row_sum: cute.Tensor,
+    scale_log2e: cutlass.Float32,
+    group_start: cutlass.Int32,
+    token_count: cutlass.Int32,
+):
+    scores_mn = layout_utils.reshape_acc_to_mn(scores)
+    coords_mn = layout_utils.reshape_acc_to_mn(coords)
+    row_scale = cute.make_rmem_tensor_like(row_max, cutlass.Float32)
+    for m in cutlass.range_constexpr(cute.size(row_max)):
+        score_row = scores_mn[m, None].load()
+        current_max = kernel_utils.fmax_reduce(
+            score_row, init_val=row_max[m], arch=80
+        )
+        current_max = cute.arch.warp_reduction_max(
+            current_max, threads_in_group=4
+        )
+        previous_max = row_max[m]
+        row_max[m] = current_max
+        safe_max = (
+            cutlass.Float32(0.0)
+            if current_max == -cutlass.Float32.inf
+            else current_max
+        )
+        probabilities = cute.math.exp2(
+            score_row * scale_log2e - safe_max * scale_log2e,
+            fastmath=True,
+        )
+        row_scale[m] = cute.math.exp2(
+            (previous_max - safe_max) * scale_log2e, fastmath=True
+        )
+        masses = cute.make_rmem_tensor_like(
+            scores_mn[m, None], cutlass.Float32
+        )
+        for n in cutlass.range_constexpr(cute.size(masses)):
+            block = group_start + coords_mn[m, n][1]
+            length = token_count - block * cutlass.Int32(N)
+            if length > N:
+                length = cutlass.Int32(N)
+            if length < 0:
+                length = cutlass.Int32(0)
+            masses[n] = cutlass.Float32(probabilities[n]) * cutlass.Float32(
+                length
+            )
+        row_sum[m] = kernel_utils.fadd_reduce(
+            masses.load(),
+            init_val=row_sum[m] * row_scale[m],
+            arch=80,
+        )
+        scores_mn[m, None].store(probabilities)
+    return row_scale
+
+
+@cute.jit
+def finalize_softmax(
+    row_max: cute.Tensor,
+    row_sum: cute.Tensor,
+    scale_log2e: cutlass.Float32,
+):
+    row_sum.store(
+        kernel_utils.warp_reduce(row_sum.load(), operator.add, width=4)
+    )
+    ratio = cute.make_rmem_tensor_like(row_sum, cutlass.Float32)
+    lse = cute.make_rmem_tensor_like(row_sum, cutlass.Float32)
+    for m in cutlass.range_constexpr(cute.size(row_sum)):
+        total = row_sum[m]
+        invalid = total == 0.0 or total != total
+        ratio[m] = cute.arch.rcp_approx(total if not invalid else 1.0)
+        lse[m] = (
+            -cutlass.Float32.inf
+            if invalid
+            else (
+                row_max[m] * scale_log2e
+                + cute.math.log2(total, fastmath=True)
+            )
+            * 0.6931471805599453
+        )
+    return ratio, lse
+
+
+@cute.jit
+def rescale_o_for_next_acc(
+    output: cute.Tensor,
+    row_scale: cute.Tensor,
+):
+    output_mn = layout_utils.reshape_acc_to_mn(output)
+    for m in cutlass.range_constexpr(cute.size(row_scale)):
+        output_mn[m, None].store(
+            output_mn[m, None].load() * row_scale[m]
+        )
+
+
+__all__ = ["SolAttnForwardSm120"]

torch-ext/sol_attn/sm90/__init__.py

@@ -0,0 +1,5 @@
+"""Hopper backend."""
+
+from .kernel import make_kernel
+
+__all__ = ["make_kernel"]

torch-ext/sol_attn/sm90/_compat/__init__.py

@@ -0,0 +1 @@
+"""Local CuteDSL helper compatibility layer for the SOL_ATTN SM90 kernel."""

torch-ext/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]))

torch-ext/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",
+]