Build uploaded using `kernels`.

build/torch-universal/__init__.py

@@ -0,0 +1,21 @@
+"""Flash Attention CUTE (CUDA Template Engine) implementation."""
+
+__version__ = "0.1.0"
+
+import cutlass.cute as cute
+
+from .interface import (
+    flash_attn_func,
+    flash_attn_varlen_func,
+)
+
+from .cute_dsl_utils import cute_compile_patched
+
+# Patch cute.compile to optionally dump SASS
+cute.compile = cute_compile_patched
+
+
+__all__ = [
+    "flash_attn_func",
+    "flash_attn_varlen_func",
+]

build/torch-universal/_ops.py

@@ -0,0 +1,8 @@
+import torch
+ops = torch.ops._flash_attn_cute_cae4d59_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flash_attn_cute_cae4d59_dirty::{op_name}"

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

build/torch-universal/barrier.py

@@ -0,0 +1,71 @@
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32
+from cutlass.cutlass_dsl import T, dsl_user_op
+from cutlass._mlir.dialects import llvm
+
+
+@dsl_user_op
+def ld_acquire(lock_ptr: cute.Pointer, *, loc=None, ip=None) -> cutlass.Int32:
+    lock_ptr_i64 = lock_ptr.toint(loc=loc, ip=ip).ir_value()
+    state = llvm.inline_asm(
+        T.i32(),
+        [lock_ptr_i64],
+        "ld.global.acquire.gpu.b32 $0, [$1];",
+        "=r,l",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+    return cutlass.Int32(state)
+
+
+@dsl_user_op
+def red_relaxed(
+    lock_ptr: cute.Pointer, val: cutlass.Constexpr[Int32], *, loc=None, ip=None
+) -> None:
+    lock_ptr_i64 = lock_ptr.toint(loc=loc, ip=ip).ir_value()
+    llvm.inline_asm(
+        None,
+        [lock_ptr_i64, Int32(val).ir_value(loc=loc, ip=ip)],
+        "red.relaxed.gpu.global.add.s32 [$0], $1;",
+        "l,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+@dsl_user_op
+def red_release(
+    lock_ptr: cute.Pointer, val: cutlass.Constexpr[Int32], *, loc=None, ip=None
+) -> None:
+    lock_ptr_i64 = lock_ptr.toint(loc=loc, ip=ip).ir_value()
+    llvm.inline_asm(
+        None,
+        [lock_ptr_i64, Int32(val).ir_value(loc=loc, ip=ip)],
+        "red.release.gpu.global.add.s32 [$0], $1;",
+        "l,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+@cute.jit
+def wait_eq(lock_ptr: cute.Pointer, thread_idx: int | Int32, flag_offset: int, val: Int32) -> None:
+    flag_ptr = lock_ptr + flag_offset
+    if thread_idx == 0:
+        read_val = Int32(0)
+        while read_val != val:
+            read_val = ld_acquire(flag_ptr)
+
+
+@cute.jit
+def arrive_inc(
+    lock_ptr: cute.Pointer, thread_idx: int | Int32, flag_offset: int, val: cutlass.Constexpr[Int32]
+) -> None:
+    flag_ptr = lock_ptr + flag_offset
+    if thread_idx == 0:
+        red_release(flag_ptr, val)
+        # red_relaxed(flag_ptr, val)

build/torch-universal/benchmark.py

@@ -0,0 +1,268 @@
+# Copyright (c) 2023, Tri Dao.
+"""Useful functions for writing test code."""
+
+import torch
+import torch.utils.benchmark as benchmark
+
+
+def benchmark_forward(
+    fn, *inputs, repeats=10, desc="", verbose=True, amp=False, amp_dtype=torch.float16, **kwinputs
+):
+    """Use Pytorch Benchmark on the forward pass of an arbitrary function."""
+    if verbose:
+        print(desc, "- Forward pass")
+
+    def amp_wrapper(*inputs, **kwinputs):
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            fn(*inputs, **kwinputs)
+
+    t = benchmark.Timer(
+        stmt="fn_amp(*inputs, **kwinputs)",
+        globals={"fn_amp": amp_wrapper, "inputs": inputs, "kwinputs": kwinputs},
+        num_threads=torch.get_num_threads(),
+    )
+    m = t.timeit(repeats)
+    if verbose:
+        print(m)
+    return t, m
+
+
+def benchmark_backward(
+    fn,
+    *inputs,
+    grad=None,
+    repeats=10,
+    desc="",
+    verbose=True,
+    amp=False,
+    amp_dtype=torch.float16,
+    **kwinputs,
+):
+    """Use Pytorch Benchmark on the backward pass of an arbitrary function."""
+    if verbose:
+        print(desc, "- Backward pass")
+    with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+        y = fn(*inputs, **kwinputs)
+        if type(y) is tuple:
+            y = y[0]
+    if grad is None:
+        grad = torch.randn_like(y)
+    else:
+        if grad.shape != y.shape:
+            raise RuntimeError("Grad shape does not match output shape")
+
+    def f(*inputs, y, grad):
+        # Set .grad to None to avoid extra operation of gradient accumulation
+        for x in inputs:
+            if isinstance(x, torch.Tensor):
+                x.grad = None
+        y.backward(grad, retain_graph=True)
+
+    t = benchmark.Timer(
+        stmt="f(*inputs, y=y, grad=grad)",
+        globals={"f": f, "inputs": inputs, "y": y, "grad": grad},
+        num_threads=torch.get_num_threads(),
+    )
+    m = t.timeit(repeats)
+    if verbose:
+        print(m)
+    return t, m
+
+
+def benchmark_combined(
+    fn,
+    *inputs,
+    grad=None,
+    repeats=10,
+    desc="",
+    verbose=True,
+    amp=False,
+    amp_dtype=torch.float16,
+    **kwinputs,
+):
+    """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
+    if verbose:
+        print(desc, "- Forward + Backward pass")
+    with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+        y = fn(*inputs, **kwinputs)
+        if type(y) is tuple:
+            y = y[0]
+    if grad is None:
+        grad = torch.randn_like(y)
+    else:
+        if grad.shape != y.shape:
+            raise RuntimeError("Grad shape does not match output shape")
+
+    def f(grad, *inputs, **kwinputs):
+        for x in inputs:
+            if isinstance(x, torch.Tensor):
+                x.grad = None
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            y = fn(*inputs, **kwinputs)
+            if type(y) is tuple:
+                y = y[0]
+        y.backward(grad, retain_graph=True)
+
+    t = benchmark.Timer(
+        stmt="f(grad, *inputs, **kwinputs)",
+        globals={"f": f, "fn": fn, "inputs": inputs, "grad": grad, "kwinputs": kwinputs},
+        num_threads=torch.get_num_threads(),
+    )
+    m = t.timeit(repeats)
+    if verbose:
+        print(m)
+    return t, m
+
+
+def benchmark_fwd_bwd(
+    fn,
+    *inputs,
+    grad=None,
+    repeats=10,
+    desc="",
+    verbose=True,
+    amp=False,
+    amp_dtype=torch.float16,
+    **kwinputs,
+):
+    """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
+    return (
+        benchmark_forward(
+            fn,
+            *inputs,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+        benchmark_backward(
+            fn,
+            *inputs,
+            grad=grad,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+    )
+
+
+def benchmark_all(
+    fn,
+    *inputs,
+    grad=None,
+    repeats=10,
+    desc="",
+    verbose=True,
+    amp=False,
+    amp_dtype=torch.float16,
+    **kwinputs,
+):
+    """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
+    return (
+        benchmark_forward(
+            fn,
+            *inputs,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+        benchmark_backward(
+            fn,
+            *inputs,
+            grad=grad,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+        benchmark_combined(
+            fn,
+            *inputs,
+            grad=grad,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+    )
+
+
+def pytorch_profiler(
+    fn,
+    *inputs,
+    trace_filename=None,
+    backward=False,
+    amp=False,
+    amp_dtype=torch.float16,
+    cpu=False,
+    verbose=True,
+    **kwinputs,
+):
+    """Wrap benchmark functions in Pytorch profiler to see CUDA information."""
+    if backward:
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            out = fn(*inputs, **kwinputs)
+            if type(out) is tuple:
+                out = out[0]
+            g = torch.randn_like(out)
+    for _ in range(30):  # Warm up
+        if backward:
+            for x in inputs:
+                if isinstance(x, torch.Tensor):
+                    x.grad = None
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            out = fn(*inputs, **kwinputs)
+            if type(out) is tuple:
+                out = out[0]
+        # Backward should be done outside autocast
+        if backward:
+            out.backward(g, retain_graph=True)
+    activities = ([torch.profiler.ProfilerActivity.CPU] if cpu else []) + [
+        torch.profiler.ProfilerActivity.CUDA
+    ]
+    with torch.profiler.profile(
+        activities=activities,
+        record_shapes=True,
+        # profile_memory=True,
+        with_stack=True,
+    ) as prof:
+        if backward:
+            for x in inputs:
+                if isinstance(x, torch.Tensor):
+                    x.grad = None
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            out = fn(*inputs, **kwinputs)
+            if type(out) is tuple:
+                out = out[0]
+        if backward:
+            out.backward(g, retain_graph=True)
+    if verbose:
+        # print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=50))
+        print(prof.key_averages().table(row_limit=50))
+    if trace_filename is not None:
+        prof.export_chrome_trace(trace_filename)
+
+
+def benchmark_memory(fn, *inputs, desc="", verbose=True, **kwinputs):
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
+    torch.cuda.synchronize()
+    fn(*inputs, **kwinputs)
+    torch.cuda.synchronize()
+    mem = torch.cuda.max_memory_allocated() / ((2**20) * 1000)
+    if verbose:
+        print(f"{desc} max memory: {mem}GB")
+    torch.cuda.empty_cache()
+    return mem

build/torch-universal/benchmark_mask_mod.py

@@ -0,0 +1,686 @@
+"""
+FlashAttention benchmarking script with Flex Attention-style
+mask mod support and varlen sequences.
+"""
+
+from dataclasses import dataclass
+import math
+from typing import Any, Dict, Optional, Tuple
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+import numpy as np
+import torch
+
+from flash_fwd import FlashAttentionForwardSm90
+from mask_definitions import (
+    get_mask_pair,
+    random_doc_id_tensor,
+)
+from .block_sparsity import (
+    compute_block_sparsity,
+    BlockSparseTensorsTorch,
+    to_cute_block_sparse_tensors,
+)
+
+
+@dataclass
+class BenchmarkConfig:
+    """Benchmark configuration"""
+
+    # Model parameters
+    headdim: int
+    headdim_v: int
+    nheads: int
+    nheads_kv: int
+    dtype: torch.dtype
+
+    # Sequence parameters
+    batch_size: int = 2
+    seqlen_q: int = 8192
+    seqlen_k: int = 8192
+
+    # Varlen parameters
+    use_varlen: bool = False
+    min_seqlen_q: Optional[int] = None  # If None, use seqlen_q // 2
+    max_seqlen_q: Optional[int] = None  # If None, use seqlen_q
+    min_seqlen_k: Optional[int] = None  # If None, use seqlen_k // 2
+    max_seqlen_k: Optional[int] = None  # If None, use seqlen_k
+
+    # Mask parameters
+    use_mask_mod: bool = True
+    mask_mod_name: str = "causal"
+    has_aux_tensors: bool = mask_mod_name == "document"
+
+    # Sliding window parameter (used when mask_mod_name == "sliding_window")
+    window_size: int = 128
+
+    # Attention parameters
+    causal: bool = False
+    is_local: bool = False
+    window_left: Optional[int] = 128  # For base Flash Attention local
+    window_right: Optional[int] = 0  # For base Flash Attention local
+    softcap: Optional[float] = None
+    use_learnable_sink: bool = False
+
+    # Kernel configuration
+    tile_m: int = 128
+    tile_n: int = 128
+    num_stages: int = 2
+    num_threads: int = 384
+    intra_wg_overlap: bool = True
+    mma_pv_is_rs: bool = True
+
+    # Benchmark parameters
+    warmup_iters: int = 5
+    benchmark_iters: int = 20
+    verbose: bool = False
+    seed: int = 42
+
+
+class FlashAttentionBenchmark:
+    def __init__(self, config: BenchmarkConfig):
+        self.config = config
+
+        torch.manual_seed(config.seed)
+        np.random.seed(config.seed)
+
+        # Verify SM90 compute capability
+        compute_capability = torch.cuda.get_device_capability()
+        assert compute_capability >= (9, 0), (
+            f"Requires SM90+, got SM{compute_capability[0]}{compute_capability[1]}"
+        )
+        # causal overrides use_mask_mod
+        if config.causal:
+            config.use_mask_mod = False
+
+        if config.use_mask_mod:
+            self.mask_mod_cute, self.mask_mod_flex = get_mask_pair(
+                config.mask_mod_name,
+                seqlen_q=config.seqlen_q,
+                seqlen_k=config.seqlen_k,
+                window_size=config.window_size,
+            )
+        else:
+            self.mask_mod_cute = None
+            self.mask_mod_flex = None
+
+        self._validate_config()
+
+    def _validate_config(self):
+        config = self.config
+
+        assert config.headdim <= 256, "headdim must be <= 256"
+        assert config.headdim_v <= 256, "headdim_v must be <= 256"
+        assert config.nheads % config.nheads_kv == 0, "nheads must be divisible by nheads_kv"
+
+        alignment = 16 // config.dtype.itemsize
+        assert config.headdim % alignment == 0, f"headdim must be divisible by {alignment}"
+        assert config.headdim_v % alignment == 0, f"headdim_v must be divisible by {alignment}"
+
+        # Validate is_local configuration
+        if config.is_local:
+            assert config.window_left is not None or config.window_right is not None, (
+                "When is_local=True, at least one of window_left or window_right must be set"
+            )
+            assert not config.use_mask_mod, (
+                "Cannot use both is_local and use_mask_mod simultaneously"
+            )
+            assert not config.causal, "Cannot use both is_local and causal simultaneously"
+
+        # Validate mask_mod configuration
+        if config.use_mask_mod and config.mask_mod_name == "sliding_window":
+            assert config.window_size > 0, (
+                "window_size must be positive when using sliding_window mask"
+            )
+
+    def _generate_varlen_seqlens(self, min_len: int, max_len: int) -> Tuple[torch.Tensor, int]:
+        """Generate random sequence lengths and compute cumulative lengths."""
+        seqlens = torch.randint(
+            min_len, max_len + 1, (self.config.batch_size,), dtype=torch.int32, device="cuda"
+        )
+        cu_seqlens = torch.cat(
+            [
+                torch.zeros(1, dtype=torch.int32, device="cuda"),
+                torch.cumsum(seqlens, dtype=torch.int32, dim=0),
+            ]
+        )
+
+        total_tokens = cu_seqlens[-1].item()
+        return cu_seqlens, total_tokens
+
+    def _create_tensors(self) -> Dict[str, torch.Tensor]:
+        config = self.config
+        device = "cuda"
+
+        if config.use_varlen:
+            # Set defaults for varlen range
+            min_q = config.min_seqlen_q if config.min_seqlen_q is not None else config.seqlen_q // 2
+            max_q = config.max_seqlen_q if config.max_seqlen_q is not None else config.seqlen_q
+            min_k = config.min_seqlen_k if config.min_seqlen_k is not None else config.seqlen_k // 2
+            max_k = config.max_seqlen_k if config.max_seqlen_k is not None else config.seqlen_k
+
+            # Generate cu_seqlens
+            cu_seqlens_q, total_q = self._generate_varlen_seqlens(min_q, max_q)
+            cu_seqlens_k, total_k = self._generate_varlen_seqlens(min_k, max_k)
+
+            # Varlen shape: (total_tokens, nheads, headdim)
+            q = torch.randn(
+                total_q, config.nheads, config.headdim, dtype=config.dtype, device=device
+            )
+            k = torch.randn(
+                total_k, config.nheads_kv, config.headdim, dtype=config.dtype, device=device
+            )
+            v = torch.randn(
+                total_k, config.nheads_kv, config.headdim_v, dtype=config.dtype, device=device
+            )
+            out = torch.empty(
+                total_q, config.nheads, config.headdim_v, dtype=config.dtype, device=device
+            )
+            lse = torch.empty(config.nheads, total_q, dtype=torch.float32, device=device)
+
+            tensors = {
+                "q": q.contiguous(),
+                "k": k.contiguous(),
+                "v": v.contiguous(),
+                "out": out.contiguous(),
+                "lse": lse.contiguous(),
+                "cu_seqlens_q": cu_seqlens_q.contiguous(),
+                "cu_seqlens_k": cu_seqlens_k.contiguous(),
+            }
+
+            if config.verbose:
+                print(f"Varlen: total_q={total_q}, total_k={total_k}")
+                print(f"Q seqlens: {cu_seqlens_q[1:] - cu_seqlens_q[:-1]}")
+                print(f"K seqlens: {cu_seqlens_k[1:] - cu_seqlens_k[:-1]}")
+        else:
+            # Standard shape: (batch, seqlen, nheads, headdim)
+            q = torch.randn(
+                config.batch_size,
+                config.seqlen_q,
+                config.nheads,
+                config.headdim,
+                dtype=config.dtype,
+                device=device,
+            )
+            k = torch.randn(
+                config.batch_size,
+                config.seqlen_k,
+                config.nheads_kv,
+                config.headdim,
+                dtype=config.dtype,
+                device=device,
+            )
+            v = torch.randn(
+                config.batch_size,
+                config.seqlen_k,
+                config.nheads_kv,
+                config.headdim_v,
+                dtype=config.dtype,
+                device=device,
+            )
+            out = torch.empty(
+                config.batch_size,
+                config.seqlen_q,
+                config.nheads,
+                config.headdim_v,
+                dtype=config.dtype,
+                device=device,
+            )
+            lse = torch.empty(
+                config.batch_size,
+                config.nheads,
+                config.seqlen_q,
+                dtype=torch.float32,
+                device=device,
+            )
+
+            tensors = {
+                "q": q.contiguous(),
+                "k": k.contiguous(),
+                "v": v.contiguous(),
+                "out": out.contiguous(),
+                "lse": lse.contiguous(),
+            }
+
+        if config.use_learnable_sink:
+            learnable_sink = torch.rand(config.nheads, dtype=torch.bfloat16, device=device)
+
+            tensors["learnable_sink"] = learnable_sink.contiguous()
+
+        # Compute block sparsity when using mask_mod
+        if config.use_mask_mod:
+            if config.mask_mod_name == "document":
+                doc_id = random_doc_id_tensor(
+                    config.batch_size, config.nheads, config.seqlen_q, device=device
+                )
+                tensors["aux_tensors"] = [doc_id.contiguous()]
+            full_cnt, full_idx, mask_cnt, mask_idx = compute_block_sparsity(
+                config=self.config,
+                mask_mod_flex=self.mask_mod_flex,
+                device=device,
+                cu_seqlens_q=tensors.get("cu_seqlens_q"),
+                cu_seqlens_k=tensors.get("cu_seqlens_k"),
+                aux_tensors=tensors.get("aux_tensors"),
+            )
+
+            if all(t is not None for t in [full_cnt, full_idx, mask_cnt, mask_idx]):
+                tensors["block_sparse_tensors"] = BlockSparseTensorsTorch(
+                    mask_block_cnt=mask_cnt.contiguous(),
+                    mask_block_idx=mask_idx.contiguous(),
+                    full_block_cnt=full_cnt.contiguous(),
+                    full_block_idx=full_idx.contiguous(),
+                )
+
+                if config.verbose:
+                    total_full = full_cnt.sum().item()
+                    total_partial = mask_cnt.sum().item()
+
+                    if config.use_varlen:
+                        # Compute max possible blocks across all sequences
+                        max_blocks = 0
+                        for i in range(config.batch_size):
+                            seq_len_q = (
+                                tensors["cu_seqlens_q"][i + 1] - tensors["cu_seqlens_q"][i]
+                            ).item()
+                            seq_len_k = (
+                                tensors["cu_seqlens_k"][i + 1] - tensors["cu_seqlens_k"][i]
+                            ).item()
+                            n_blocks_q = (seq_len_q + config.tile_m - 1) // config.tile_m
+                            n_blocks_k = (seq_len_k + config.tile_n - 1) // config.tile_n
+                            max_blocks += n_blocks_q * n_blocks_k * config.nheads
+                    else:
+                        n_blocks_k = (config.seqlen_k + config.tile_n - 1) // config.tile_n
+                        n_blocks_q = (config.seqlen_q + config.tile_m - 1) // config.tile_m
+                        max_blocks = n_blocks_k * n_blocks_q * config.nheads * config.batch_size
+
+                    skipped = max_blocks - total_full - total_partial
+                    print(
+                        f"Block stats: Full={total_full}, Partial={total_partial}, "
+                        f"Skipped={skipped}/{max_blocks}"
+                    )
+
+        return tensors
+
+    def _compile_kernel(self, tensors: Dict[str, torch.Tensor]) -> Tuple[Any, tuple]:
+        config = self.config
+
+        dtype_map = {
+            torch.float16: cutlass.Float16,
+            torch.bfloat16: cutlass.BFloat16,
+            torch.float32: cutlass.Float32,
+        }
+        cute_dtype = dtype_map[config.dtype]
+
+        qhead_per_kvhead = config.nheads // config.nheads_kv
+        kernel = FlashAttentionForwardSm90(
+            cute_dtype,
+            config.headdim,
+            config.headdim_v,
+            qhead_per_kvhead,
+            is_causal=config.causal,
+            is_local=config.is_local,
+            pack_gqa=False,
+            tile_m=config.tile_m,
+            tile_n=config.tile_n,
+            num_stages=config.num_stages,
+            num_threads=config.num_threads,
+            intra_wg_overlap=config.intra_wg_overlap,
+            mma_pv_is_rs=config.mma_pv_is_rs,
+            mask_mod=self.mask_mod_cute,
+            Q_in_regs=False,
+            has_aux_tensors=config.has_aux_tensors,
+        )
+
+        softmax_scale = 1.0 / math.sqrt(config.headdim)
+        current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+
+        # Convert tensors to cute
+        q_cute = from_dlpack(tensors["q"].detach(), assumed_align=16).mark_layout_dynamic(
+            leading_dim=tensors["q"].ndim - 1
+        )
+        k_cute = from_dlpack(tensors["k"].detach(), assumed_align=16).mark_layout_dynamic(
+            leading_dim=tensors["k"].ndim - 1
+        )
+        v_cute = from_dlpack(tensors["v"].detach(), assumed_align=16).mark_layout_dynamic(
+            leading_dim=tensors["v"].ndim - 1
+        )
+        out_cute = from_dlpack(tensors["out"].detach(), assumed_align=16).mark_layout_dynamic(
+            leading_dim=tensors["out"].ndim - 1
+        )
+        lse_cute = from_dlpack(tensors["lse"].detach(), assumed_align=4).mark_layout_dynamic(
+            leading_dim=tensors["lse"].ndim - 1
+        )
+
+        # Varlen tensors
+        cu_seqlens_q_cute = (
+            from_dlpack(tensors["cu_seqlens_q"].detach(), assumed_align=4).mark_layout_dynamic(
+                leading_dim=0
+            )
+            if "cu_seqlens_q" in tensors
+            else None
+        )
+        cu_seqlens_k_cute = (
+            from_dlpack(tensors["cu_seqlens_k"].detach(), assumed_align=4).mark_layout_dynamic(
+                leading_dim=0
+            )
+            if "cu_seqlens_k" in tensors
+            else None
+        )
+        learnable_sink_cute = (
+            from_dlpack(tensors["learnable_sink"].detach(), assumed_align=4).mark_layout_dynamic(
+                leading_dim=0
+            )
+            if "learnable_sink" in tensors
+            else None
+        )
+
+        blocksparse_tensors_cute = (
+            to_cute_block_sparse_tensors(tensors["block_sparse_tensors"])
+            if "block_sparse_tensors" in tensors
+            else None
+        )
+
+        if "aux_tensors" in tensors:
+            aux_tensors_cute = []
+            for i in range(len(tensors["aux_tensors"])):
+                buf = from_dlpack(tensors["aux_tensors"][i].detach(), assumed_align=4)
+                aux_tensors_cute.append(buf.mark_layout_dynamic(leading_dim=2))
+
+        else:
+            aux_tensors_cute = None
+
+        # Window parameters for is_local
+        window_left_cute = (
+            cutlass.Int32(config.window_left) if config.window_left is not None else None
+        )
+        window_right_cute = (
+            cutlass.Int32(config.window_right) if config.window_right is not None else None
+        )
+
+        compiled = cute.compile(
+            kernel,
+            q_cute,
+            k_cute,
+            v_cute,
+            out_cute,
+            lse_cute,
+            softmax_scale,
+            current_stream,
+            cu_seqlens_q_cute,
+            cu_seqlens_k_cute,
+            None,  # seqused_q
+            None,  # seqused_k
+            None,  # page_table
+            window_left_cute,
+            window_right_cute,
+            learnable_sink_cute,
+            blocksparse_tensors_cute,
+            aux_tensors_cute,
+            # None,
+        )
+
+        args = (
+            q_cute,
+            k_cute,
+            v_cute,
+            out_cute,
+            lse_cute,
+            softmax_scale,
+            current_stream,
+            cu_seqlens_q_cute,
+            cu_seqlens_k_cute,
+            None,
+            None,
+            None,
+            window_left_cute,
+            window_right_cute,
+            learnable_sink_cute,
+            blocksparse_tensors_cute,
+            aux_tensors_cute,
+            # None,
+        )
+
+        return compiled, args
+
+    def _calculate_flops(self, tensors: Dict[str, torch.Tensor]) -> float:
+        config = self.config
+
+        # Estimate sparsity for known mask patterns
+        if config.is_local:
+            # Local attention with window_left and window_right
+            window_left = config.window_left if config.window_left is not None else 0
+            window_right = config.window_right if config.window_right is not None else 0
+            total_window = window_left + window_right + 1  # +1 for current position
+            sparsity_ratio = min(1.0, total_window / config.seqlen_k)
+        elif config.use_mask_mod:
+            if config.mask_mod_name in ["identity", "identity_partial"]:
+                sparsity_ratio = 1.0
+            elif config.mask_mod_name in ["causal", "block_causal"]:
+                sparsity_ratio = 0.5
+            elif config.mask_mod_name == "sliding_window":
+                # Use configured window size
+                sparsity_ratio = min(1.0, config.window_size / config.seqlen_k)
+            elif config.mask_mod_name == "block_diagonal":
+                block_size = 64
+                num_blocks = (config.seqlen_k + block_size - 1) // block_size
+                sparsity_ratio = 1.0 / num_blocks if num_blocks > 1 else 1.0
+            elif config.mask_mod_name == "document":
+                vals = tensors["aux_tensors"][0]
+                val_mask = torch.ones_like(vals, dtype=torch.bool)
+                val_mask[..., 1:] = vals[..., 1:] != vals[..., :-1]
+                total = torch.where(val_mask, vals.square(), 0).sum()
+                sparsity_ratio = total / (config.seqlen_q * config.seqlen_k)
+            else:
+                sparsity_ratio = 1.0
+        elif config.causal:
+            sparsity_ratio = 0.5
+        else:
+            sparsity_ratio = 1.0
+
+        if config.use_varlen:
+            # Compute FLOPs per sequence and sum
+            total_flops = 0
+            cu_q = tensors["cu_seqlens_q"]
+            cu_k = tensors["cu_seqlens_k"]
+            for i in range(config.batch_size):
+                seq_len_q = (cu_q[i + 1] - cu_q[i]).item()
+                seq_len_k = (cu_k[i + 1] - cu_k[i]).item()
+
+                # Adjust sparsity for local attention in varlen case
+                if config.is_local:
+                    window_left = config.window_left if config.window_left is not None else 0
+                    window_right = config.window_right if config.window_right is not None else 0
+                    total_window = window_left + window_right + 1
+                    seq_sparsity = min(1.0, total_window / seq_len_k)
+                elif config.use_mask_mod and config.mask_mod_name == "sliding_window":
+                    seq_sparsity = min(1.0, config.window_size / seq_len_k)
+                else:
+                    seq_sparsity = sparsity_ratio
+
+                num_cells = int(seq_len_q * seq_len_k * seq_sparsity)
+
+                if config.headdim == config.headdim_v:
+                    flops_this_seq = 4 * config.nheads * num_cells * config.headdim
+                else:
+                    flops_this_seq = (
+                        2 * config.nheads * num_cells * config.headdim
+                        + 2 * config.nheads * num_cells * config.headdim_v
+                    )
+                total_flops += flops_this_seq
+            return total_flops
+        else:
+            num_cells = int(config.seqlen_q * config.seqlen_k * sparsity_ratio)
+            if config.headdim == config.headdim_v:
+                flops_per_batch = 4 * config.nheads * num_cells * config.headdim
+            else:
+                flops_per_batch = (
+                    2 * config.nheads * num_cells * config.headdim
+                    + 2 * config.nheads * num_cells * config.headdim_v
+                )
+            return flops_per_batch * config.batch_size
+
+    def benchmark(self) -> Dict[str, Any]:
+        config = self.config
+
+        tensors = self._create_tensors()
+        compiled_kernel, args = self._compile_kernel(tensors)
+
+        # Warmup
+        for _ in range(config.warmup_iters):
+            compiled_kernel(*args)
+        torch.cuda.synchronize()
+
+        # Benchmark
+        times = []
+        for _ in range(config.benchmark_iters):
+            start = torch.cuda.Event(enable_timing=True)
+            end = torch.cuda.Event(enable_timing=True)
+
+            start.record()
+            compiled_kernel(*args)
+            end.record()
+            torch.cuda.synchronize()
+
+            times.append(start.elapsed_time(end))
+
+        times_tensor = torch.tensor(times)
+        mean_time = times_tensor.mean().item()
+        std_time = times_tensor.std().item() if len(times) > 1 else 0.0
+
+        total_flops = self._calculate_flops(tensors)
+        tflops = total_flops / (mean_time * 1e-3) / 1e12
+
+        # Bandwidth calculation
+        bytes_per_element = config.dtype.itemsize
+        if config.use_varlen:
+            total_q = tensors["q"].shape[0]
+            total_k = tensors["k"].shape[0]
+            memory_accessed = (
+                total_q * config.nheads * config.headdim * bytes_per_element
+                + total_k * config.nheads_kv * config.headdim * bytes_per_element
+                + total_k * config.nheads_kv * config.headdim_v * bytes_per_element
+                + total_q * config.nheads * config.headdim_v * bytes_per_element
+            )
+        else:
+            memory_accessed = (
+                config.batch_size
+                * config.seqlen_q
+                * config.nheads
+                * config.headdim
+                * bytes_per_element
+                + config.batch_size
+                * config.seqlen_k
+                * config.nheads_kv
+                * config.headdim
+                * bytes_per_element
+                + config.batch_size
+                * config.seqlen_k
+                * config.nheads_kv
+                * config.headdim_v
+                * bytes_per_element
+                + config.batch_size
+                * config.seqlen_q
+                * config.nheads
+                * config.headdim_v
+                * bytes_per_element
+            )
+        bandwidth_gbps = memory_accessed / (mean_time * 1e-3) / 1e9
+
+        results = {
+            "mean_time_ms": mean_time,
+            "std_time_ms": std_time,
+            "tflops": tflops,
+            "bandwidth_gbps": bandwidth_gbps,
+        }
+
+        if config.verbose:
+            self._print_results(results)
+
+        return results
+
+    def _print_results(self, results: Dict[str, Any]):
+        config = self.config
+
+        # Basic configuration
+        if config.use_varlen:
+            print(
+                f"Shape: B={config.batch_size} (varlen), HD={config.headdim}, "
+                f"NH={config.nheads}, NKV={config.nheads_kv}"
+            )
+        else:
+            print(
+                f"Shape: B={config.batch_size}, Q={config.seqlen_q}, K={config.seqlen_k}, "
+                f"HD={config.headdim}, NH={config.nheads}, NKV={config.nheads_kv}"
+            )
+
+        # Attention pattern
+        attn_info = []
+        if config.causal:
+            attn_info.append("causal")
+        if config.is_local:
+            window_info = f"local(L={config.window_left},R={config.window_right})"
+            attn_info.append(window_info)
+        if config.use_mask_mod:
+            if config.mask_mod_name == "sliding_window":
+                attn_info.append(f"mask_mod={config.mask_mod_name}(w={config.window_size})")
+            else:
+                attn_info.append(f"mask_mod={config.mask_mod_name}")
+        if config.use_varlen:
+            attn_info.append("varlen")
+        if attn_info:
+            print(f"Attention: {', '.join(attn_info)}")
+
+        # Performance metrics
+        print(f"Time: {results['mean_time_ms']:.3f} ± {results['std_time_ms']:.3f} ms")
+        print(f"Throughput: {results['tflops']:.2f} TFLOPS")
+        print(f"Bandwidth: {results['bandwidth_gbps']:.1f} GB/s")
+
+
+if __name__ == "__main__":
+    B = 2
+    config = BenchmarkConfig(
+        headdim=128,
+        headdim_v=128,
+        nheads=16,
+        nheads_kv=16,
+        dtype=torch.bfloat16,
+        batch_size=B,
+        # batch_size=1,
+        seqlen_q=16384 // B,
+        # seqlen_q=128,
+        seqlen_k=16384 // B,
+        # seqlen_k=192,
+        use_varlen=False,
+        use_mask_mod=True,
+        mask_mod_name="causal",
+        window_size=128,  # Configurable window size for mask_mod
+        use_learnable_sink=False,
+        causal=False,
+        is_local=False,
+        verbose=True,
+    )
+
+    # Example 2: Base Flash Attention Local
+    # config = BenchmarkConfig(
+    #     headdim=64,
+    #     headdim_v=64,
+    #     nheads=64,
+    #     nheads_kv=8,
+    #     dtype=torch.bfloat16,
+    #     batch_size=2,
+    #     seqlen_q=8192,
+    #     seqlen_k=8192,
+    #     use_varlen=False,
+    #     use_mask_mod=False,
+    #     causal=False,
+    #     is_local=True,
+    #     window_left=128,   # Left window size for base local attention
+    #     window_right=0,    # Right window size for base local attention
+    #     verbose=True,
+    # )
+
+    benchmark = FlashAttentionBenchmark(config)
+    results = benchmark.benchmark()

build/torch-universal/blackwell_helpers.py

@@ -0,0 +1,753 @@
+# Copyright (c) 2025, Tri Dao.
+from typing import Optional, Tuple
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32, Boolean, const_expr
+from cutlass.cute.nvgpu import tcgen05
+from cutlass._mlir.dialects import llvm
+
+from . import mma_sm100_desc as sm100_desc
+from .utils import parse_swizzle_from_pointer
+
+
+@cute.jit
+def gemm_w_idx(
+    tiled_mma: cute.TiledMma,
+    acc: cute.Tensor,
+    tCrA: cute.Tensor,
+    tCrB: cute.Tensor,
+    A_idx: Optional[Int32] = None,
+    B_idx: Optional[Int32] = None,
+    zero_init: bool | Boolean = False,
+    swap_AB: bool = False,
+) -> None:
+    if const_expr(swap_AB):
+        return gemm_w_idx(
+            tiled_mma, acc, tCrB, tCrA, B_idx, A_idx, zero_init=zero_init, swap_AB=False
+        )
+    else:
+        rA = tCrA if const_expr(A_idx is None) else tCrA[None, None, None, A_idx]
+        rB = tCrB if const_expr(B_idx is None) else tCrB[None, None, None, B_idx]
+        mma_atom = cute.make_mma_atom(tiled_mma.op)
+        for k in cutlass.range_constexpr(cute.size(tCrA.shape[2])):
+            mma_atom.set(tcgen05.Field.ACCUMULATE, not zero_init or k != 0)
+            cute.gemm(mma_atom, acc, rA[None, None, k], rB[None, None, k], acc)
+
+
+@cute.jit
+def gemm_ptx_w_idx(
+    tiled_mma: cute.TiledMma,
+    acc: cute.Tensor,
+    tCrA: cute.Tensor,
+    tCrB: cute.Tensor,
+    sA: Optional[cute.Tensor],
+    sB: cute.Tensor,
+    A_idx: Optional[Int32] = None,
+    B_idx: Optional[Int32] = None,
+    zero_init: bool | Boolean = False,
+    **kwargs,
+) -> None:
+    rA = tCrA if const_expr(A_idx is None) else tCrA[None, None, None, A_idx]
+    rB = tCrB if const_expr(B_idx is None) else tCrB[None, None, None, B_idx]
+    sA_cur = None
+    if const_expr(sA is not None):
+        sA_cur = sA if const_expr(A_idx is None) else sA[None, None, None, A_idx]
+    sB_cur = sB if const_expr(B_idx is None) else sB[None, None, None, B_idx]
+    mma_atom = cute.make_mma_atom(tiled_mma.op)
+    acc_tmem_addr = acc.iterator.toint()
+    gemm_ptx_partial(
+        mma_atom.op, acc_tmem_addr, rA, rB, sA_cur, sB_cur, zero_init=zero_init, **kwargs
+    )
+
+
+@cute.jit
+def gemm(
+    tiled_mma: cute.TiledMma,
+    acc: cute.Tensor,
+    tCrA: cute.Tensor,
+    tCrB: cute.Tensor,
+    zero_init: bool | Boolean = False,
+) -> cute.TiledMma:
+    for k in cutlass.range_constexpr(cute.size(tCrA.shape[2])):
+        tiled_mma.set(tcgen05.Field.ACCUMULATE, not zero_init or k != 0)
+        cute.gemm(tiled_mma, acc, tCrA[None, None, k], tCrB[None, None, k], acc)
+    return tiled_mma
+
+
+def i64_to_i32x2(i: int) -> Tuple[int, int]:
+    """Convert a 64-bit integer to a tuple of two 32-bit integers."""
+    return i & 0xFFFF_FFFF, (i >> 32) & 0xFFFF_FFFF
+
+
+@cute.jit
+def gemm_ptx(
+    op: cute.nvgpu.tcgen05.mma.MmaOp,
+    acc: cute.Tensor,
+    tCrA: cute.Tensor,
+    tCrB: cute.Tensor,
+    sA: Optional[cute.Tensor],
+    sB: cute.Tensor,
+    zero_init: bool | Boolean = False,
+) -> None:
+    is_ts = op.a_src == cute.nvgpu.tcgen05.OperandSource.TMEM
+    if const_expr(not is_ts):
+        assert sA is not None, "sA must be provided when a_src is not TMEM"
+    sA_layout = sA.layout if sA is not None else None
+    sB_layout = sB.layout
+    idesc: int = const_expr(sm100_desc.mma_op_to_idesc(op))
+    if const_expr(not is_ts):
+        sA_swizzle = parse_swizzle_from_pointer(sA.iterator)
+        smem_desc_base_a: int = const_expr(
+            sm100_desc.make_smem_desc_base(
+                cute.recast_layout(128, op.a_dtype.width, sA_layout[0]),
+                sA_swizzle,
+                sm100_desc.Major.K
+                if const_expr(op.a_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K)
+                else sm100_desc.Major.MN,
+            )
+        )
+        smem_desc_base_a_lo, smem_desc_a_hi = i64_to_i32x2(smem_desc_base_a)
+        smem_desc_base_a_lo = const_expr(smem_desc_base_a_lo)
+        smem_desc_a_hi = const_expr(smem_desc_a_hi)
+    else:
+        smem_desc_base_a = None
+        smem_desc_base_a_lo, smem_desc_a_hi = None, None
+    sB_swizzle = parse_swizzle_from_pointer(sB.iterator)
+    smem_desc_base_b: int = const_expr(
+        sm100_desc.make_smem_desc_base(
+            cute.recast_layout(128, op.b_dtype.width, sB_layout[0]),
+            sB_swizzle,
+            sm100_desc.Major.K
+            if const_expr(op.b_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K)
+            else sm100_desc.Major.MN,
+        )
+    )
+    smem_desc_base_b_lo, smem_desc_b_hi = i64_to_i32x2(smem_desc_base_b)
+    smem_desc_base_b_lo = const_expr(smem_desc_base_b_lo)
+    smem_desc_b_hi = const_expr(smem_desc_b_hi)
+
+    if const_expr(not is_ts):
+        smem_desc_start_a_lo = Int32(smem_desc_base_a_lo) | sm100_desc.make_smem_desc_start_addr(
+            sA[None, None, 0].iterator
+        )
+    else:
+        smem_desc_start_a_lo = None
+    smem_desc_start_b_lo = Int32(smem_desc_base_b_lo) | sm100_desc.make_smem_desc_start_addr(
+        sB[None, None, 0].iterator
+    )
+    for k in cutlass.range_constexpr(cute.size(tCrA.shape[2])):
+        if const_expr(not is_ts):
+            smem_desc_a_lo = smem_desc_start_a_lo + (
+                (cute.crd2idx((0, 0, k), sA_layout) * sA.element_type.width // 8) >> 4
+            )
+        smem_desc_b_lo = smem_desc_start_b_lo + (
+            (cute.crd2idx((0, 0, k), sB_layout) * sB.element_type.width // 8) >> 4
+        )
+        # with cute.arch.elect_one():
+        #     cute.printf("smem_desc_a_lo = {}, smem_desc_b_lo = {}", smem_desc_a_lo, smem_desc_b_lo)
+        #     cute.printf("smem_desc_a_lo_correct = {}, smem_desc_b_lo_correct = {}", smem_desc_a_lo_correct, smem_desc_b_lo_correct)
+        with cute.arch.elect_one():
+            if const_expr(not is_ts):
+                llvm.inline_asm(
+                    None,
+                    [
+                        acc.iterator.toint().ir_value(),
+                        smem_desc_a_lo.ir_value(),
+                        smem_desc_b_lo.ir_value(),
+                        Int32(not zero_init or k != 0).ir_value(),
+                    ],
+                    "{\n\t"
+                    ".reg .pred p;\n\t"
+                    ".reg .b64 smem_desc_a, smem_desc_b;\n\t"
+                    ".reg .b32 idesc;\n\t"
+                    f"mov.b32 idesc, {hex(idesc)};\n\t"
+                    f"mov.b64 smem_desc_a, {{$1, {hex(smem_desc_a_hi)}}};\n\t"
+                    f"mov.b64 smem_desc_b, {{$2, {hex(smem_desc_b_hi)}}};\n\t"
+                    "setp.ne.b32 p, $3, 0;\n\t"
+                    f"tcgen05.mma.cta_group::1.kind::f16 [$0], smem_desc_a, smem_desc_b, idesc, p;\n\t"
+                    "}\n",
+                    "r,r,r,r",
+                    has_side_effects=True,
+                    is_align_stack=False,
+                    asm_dialect=llvm.AsmDialect.AD_ATT,
+                )
+            else:
+                llvm.inline_asm(
+                    None,
+                    [
+                        acc.iterator.toint().ir_value(),
+                        tCrA[None, None, k].iterator.toint().ir_value(),
+                        smem_desc_b_lo.ir_value(),
+                        Int32(not zero_init or k != 0).ir_value(),
+                    ],
+                    "{\n\t"
+                    ".reg .pred p;\n\t"
+                    ".reg .b64 smem_desc_b;\n\t"
+                    f"mov.b64 smem_desc_b, {{$2, {hex(smem_desc_b_hi)}}};\n\t"
+                    "setp.ne.b32 p, $3, 0;\n\t"
+                    f"tcgen05.mma.cta_group::1.kind::f16 [$0], [$1], smem_desc_b, {hex(idesc)}, p;\n\t"
+                    "}\n",
+                    "r,r,r,r",
+                    has_side_effects=True,
+                    is_align_stack=False,
+                    asm_dialect=llvm.AsmDialect.AD_ATT,
+                )
+
+
+@cute.jit
+def gemm_ptx_loop(
+    op: cute.nvgpu.tcgen05.mma.MmaOp,
+    acc: cute.Tensor,
+    tCrA: cute.Tensor,
+    tCrB: cute.Tensor,
+    sA: Optional[cute.Tensor],
+    sB: cute.Tensor,
+    zero_init: bool | Boolean = False,
+) -> None:
+    is_ts = op.a_src == cute.nvgpu.tcgen05.OperandSource.TMEM
+    if const_expr(not is_ts):
+        assert sA is not None, "sA must be provided when a_src is not TMEM"
+    sA_layout = sA.layout if sA is not None else tCrA.layout
+    sB_layout = sB.layout
+    idesc: int = const_expr(sm100_desc.mma_op_to_idesc(op))
+    if const_expr(not is_ts):
+        sA_swizzle = parse_swizzle_from_pointer(sA.iterator)
+        smem_desc_base_a: int = const_expr(
+            sm100_desc.make_smem_desc_base(
+                cute.recast_layout(128, op.a_dtype.width, sA_layout[0]),
+                sA_swizzle,
+                sm100_desc.Major.K
+                if const_expr(op.a_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K)
+                else sm100_desc.Major.MN,
+            )
+        )
+        smem_desc_base_a_lo, smem_desc_a_hi = i64_to_i32x2(smem_desc_base_a)
+        smem_desc_base_a_lo = const_expr(smem_desc_base_a_lo)
+        smem_desc_a_hi = const_expr(smem_desc_a_hi)
+    else:
+        smem_desc_base_a = None
+        smem_desc_base_a_lo, smem_desc_a_hi = None, None
+    sB_swizzle = parse_swizzle_from_pointer(sB.iterator)
+    smem_desc_base_b: int = const_expr(
+        sm100_desc.make_smem_desc_base(
+            cute.recast_layout(128, op.b_dtype.width, sB_layout[0]),
+            sB_swizzle,
+            sm100_desc.Major.K
+            if const_expr(op.b_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K)
+            else sm100_desc.Major.MN,
+        )
+    )
+    smem_desc_base_b_lo, smem_desc_b_hi = i64_to_i32x2(smem_desc_base_b)
+    smem_desc_base_b_lo = const_expr(smem_desc_base_b_lo)
+    smem_desc_b_hi = const_expr(smem_desc_b_hi)
+
+    if const_expr(not is_ts):
+        offset_a = [
+            (cute.crd2idx((0, 0, k), sA_layout) * sA.element_type.width // 8) >> 4
+            for k in cutlass.range_constexpr(cute.size(tCrA.shape[2]))
+        ]
+    else:
+        offset_a = [
+            cute.crd2idx((0, 0, k), sA_layout) * op.a_dtype.width // 32
+            for k in cutlass.range_constexpr(cute.size(tCrA.shape[2]))
+        ]
+    offset_a_diff = [
+        offset_a[k] - offset_a[k - 1] for k in cutlass.range_constexpr(1, cute.size(tCrA.shape[2]))
+    ]
+    offset_b = [
+        (cute.crd2idx((0, 0, k), sB_layout) * sB.element_type.width // 8) >> 4
+        for k in cutlass.range_constexpr(cute.size(tCrB.shape[2]))
+    ]
+    offset_b_diff = [
+        offset_b[k] - offset_b[k - 1] for k in cutlass.range_constexpr(1, cute.size(tCrB.shape[2]))
+    ]
+
+    if const_expr(not is_ts):
+        smem_desc_start_a_lo = Int32(
+            smem_desc_base_a_lo | sm100_desc.make_smem_desc_start_addr(sA[None, None, 0].iterator)
+        )
+    else:
+        smem_desc_start_a_lo = None
+    smem_desc_start_b_lo = Int32(
+        smem_desc_base_b_lo | sm100_desc.make_smem_desc_start_addr(sB[None, None, 0].iterator)
+    )
+    pred_str = "p" if isinstance(zero_init, Boolean) else "0" if zero_init else "1"
+    if const_expr(not is_ts):
+        llvm.inline_asm(
+            None,
+            [
+                acc.iterator.toint().ir_value(),
+                Int32(cute.arch.make_warp_uniform(smem_desc_start_a_lo)).ir_value(),
+                Int32(cute.arch.make_warp_uniform(smem_desc_start_b_lo)).ir_value(),
+                Int32(not zero_init).ir_value(),
+            ],
+            "{\n\t"
+            ".reg .pred leader_thread;\n\t"
+            ".reg .pred p;\n\t"
+            ".reg .b32 idesc;\n\t"
+            ".reg .b32 smem_desc_a_lo, smem_desc_b_lo;\n\t"
+            ".reg .b32 smem_desc_a_hi, smem_desc_b_hi;\n\t"
+            ".reg .b64 smem_desc_a, smem_desc_b;\n\t"
+            "elect.sync _|leader_thread, -1;\n\t"
+            f"mov.b32 idesc, {hex(idesc)};\n\t"
+            "mov.b32 smem_desc_a_lo, $1;\n\t"
+            "mov.b32 smem_desc_b_lo, $2;\n\t"
+            f"mov.b32 smem_desc_a_hi, {hex(smem_desc_a_hi)};\n\t"
+            f"mov.b32 smem_desc_b_hi, {hex(smem_desc_b_hi)};\n\t"
+            f"mov.b64 smem_desc_a, {{smem_desc_a_lo, smem_desc_a_hi}};\n\t"
+            f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t"
+            "setp.ne.b32 p, $3, 0;\n\t"
+            f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [$0], smem_desc_a, smem_desc_b, idesc, {pred_str};\n\t"
+            + "".join(
+                (
+                    f"add.u32 smem_desc_a_lo, smem_desc_a_lo, {hex(offset_a_diff[k - 1])};\n\t"
+                    f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t"
+                    f"mov.b64 smem_desc_a, {{smem_desc_a_lo, smem_desc_a_hi}};\n\t"
+                    f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t"
+                    f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [$0], smem_desc_a, smem_desc_b, idesc, 1;\n\t"
+                )
+                for k in cutlass.range_constexpr(1, cute.size(tCrA.shape[2]))
+            )
+            + "}\n",
+            "r,r,r,r",
+            has_side_effects=True,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    else:
+        llvm.inline_asm(
+            None,
+            [
+                acc.iterator.toint().ir_value(),
+                Int32(tCrA[None, None, 0].iterator.toint()).ir_value(),
+                Int32(smem_desc_start_b_lo).ir_value(),
+                Int32(not zero_init).ir_value(),
+            ],
+            "{\n\t"
+            ".reg .pred leader_thread;\n\t"
+            ".reg .pred p;\n\t"
+            ".reg .b32 idesc;\n\t"
+            ".reg .b32 tmem_a;\n\t"
+            ".reg .b32 smem_desc_b_lo;\n\t"
+            ".reg .b32 smem_desc_b_hi;\n\t"
+            ".reg .b64 smem_desc_b;\n\t"
+            "elect.sync _|leader_thread, -1;\n\t"
+            f"mov.b32 idesc, {hex(idesc)};\n\t"
+            "mov.b32 tmem_a, $1;\n\t"
+            "mov.b32 smem_desc_b_lo, $2;\n\t"
+            f"mov.b32 smem_desc_b_hi, {hex(smem_desc_b_hi)};\n\t"
+            f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t"
+            "setp.ne.b32 p, $3, 0;\n\t"
+            f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [$0], [tmem_a], smem_desc_b, idesc, {pred_str};\n\t"
+            + "".join(
+                (
+                    # f"add.u32 tmem_a, tmem_a, {hex(offset_a_diff[k - 1])};\n\t"
+                    f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t"
+                    f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t"
+                    # f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [$0], [tmem_a], smem_desc_b, idesc, 1;\n\t"
+                    f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [$0], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t"
+                )
+                for k in cutlass.range_constexpr(1, cute.size(tCrA.shape[2]))
+            )
+            + "}\n",
+            "r,r,r,r",
+            has_side_effects=True,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+
+
+@cute.jit
+def gemm_ptx_partial(
+    op: cute.nvgpu.tcgen05.mma.MmaOp,
+    acc_tmem_addr: Int32,
+    tCrA: cute.Tensor,
+    tCrB: cute.Tensor,
+    sA: Optional[cute.Tensor],
+    sB: cute.Tensor,
+    mbar_ptr: Optional[cutlass.Pointer] = None,
+    mbar_phase: Optional[Int32] = None,
+    zero_init: bool | Boolean = False,
+    # sA_offset: Int32 = 0,
+    # acc_offset: Int32 = 0,
+    tA_addr: Optional[Int32] = None,
+) -> None:
+    # acc_tmem_addr += acc_offset
+    is_ts = op.a_src == cute.nvgpu.tcgen05.OperandSource.TMEM
+    if const_expr(not is_ts):
+        assert sA is not None, "sA must be provided when a_src is not TMEM"
+    sA_layout = sA.layout if sA is not None else tCrA.layout
+    sB_layout = sB.layout
+    idesc: int = const_expr(sm100_desc.mma_op_to_idesc(op))
+    if const_expr(not is_ts):
+        sA_swizzle = parse_swizzle_from_pointer(sA.iterator)
+        smem_desc_base_a: int = const_expr(
+            sm100_desc.make_smem_desc_base(
+                cute.recast_layout(128, op.a_dtype.width, sA_layout[0]),
+                sA_swizzle,
+                sm100_desc.Major.K
+                if const_expr(op.a_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K)
+                else sm100_desc.Major.MN,
+            )
+        )
+        smem_desc_base_a_lo, smem_desc_a_hi = i64_to_i32x2(smem_desc_base_a)
+        smem_desc_base_a_lo = const_expr(smem_desc_base_a_lo)
+        smem_desc_a_hi = const_expr(smem_desc_a_hi)
+    else:
+        smem_desc_base_a = None
+        smem_desc_base_a_lo, smem_desc_a_hi = None, None
+    sB_swizzle = parse_swizzle_from_pointer(sB.iterator)
+    smem_desc_base_b: int = const_expr(
+        sm100_desc.make_smem_desc_base(
+            cute.recast_layout(128, op.b_dtype.width, sB_layout[0]),
+            sB_swizzle,
+            sm100_desc.Major.K
+            if const_expr(op.b_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K)
+            else sm100_desc.Major.MN,
+        )
+    )
+    smem_desc_base_b_lo, smem_desc_b_hi = i64_to_i32x2(smem_desc_base_b)
+    smem_desc_base_b_lo = const_expr(smem_desc_base_b_lo)
+    smem_desc_b_hi = const_expr(smem_desc_b_hi)
+
+    tCrA_layout = (
+        tCrA.layout
+        if const_expr(not is_ts)
+        else cute.recast_layout(32, tCrA.element_type.width, tCrA.layout)
+    )
+    offset_a = [cute.crd2idx((0, 0, k), tCrA_layout) for k in range(cute.size(tCrA.shape[2]))]
+    offset_a_diff = [offset_a[k] - offset_a[k - 1] for k in range(1, cute.size(tCrA.shape[2]))]
+    offset_b = [cute.crd2idx((0, 0, k), tCrB.layout) for k in range(cute.size(tCrB.shape[2]))]
+    offset_b_diff = [offset_b[k] - offset_b[k - 1] for k in range(1, cute.size(tCrB.shape[2]))]
+
+    if const_expr(not is_ts):
+        smem_desc_start_a_lo = Int32(
+            smem_desc_base_a_lo | sm100_desc.make_smem_desc_start_addr(sA[None, None, 0].iterator)
+        )
+        # ) + sA_offset
+    else:
+        smem_desc_start_a_lo = None
+    smem_desc_start_b_lo = Int32(
+        smem_desc_base_b_lo | sm100_desc.make_smem_desc_start_addr(sB[None, None, 0].iterator)
+    )
+    pred_str = "p" if isinstance(zero_init, Boolean) else "0" if zero_init else "1"
+    if const_expr(not is_ts):
+        assert mbar_ptr is None, "mbar_ptr must be None when a_src is not TMEM"
+        llvm.inline_asm(
+            None,
+            [
+                # acc.iterator.toint().ir_value(),
+                Int32(cute.arch.make_warp_uniform(smem_desc_start_a_lo)).ir_value(),
+                Int32(cute.arch.make_warp_uniform(smem_desc_start_b_lo)).ir_value(),
+                Int32(not zero_init).ir_value(),
+                Int32(cute.arch.make_warp_uniform(acc_tmem_addr)).ir_value(),
+            ],
+            "{\n\t"
+            ".reg .pred leader_thread;\n\t"
+            ".reg .pred p;\n\t"
+            ".reg .b32 idesc;\n\t"
+            ".reg .b32 tmem_acc;\n\t"
+            ".reg .b32 smem_desc_a_lo_start, smem_desc_b_lo_start;\n\t"
+            ".reg .b32 smem_desc_a_lo, smem_desc_b_lo;\n\t"
+            ".reg .b32 smem_desc_a_hi, smem_desc_b_hi;\n\t"
+            ".reg .b64 smem_desc_a, smem_desc_b;\n\t"
+            "elect.sync _|leader_thread, -1;\n\t"
+            f"mov.b32 idesc, {hex(idesc)};\n\t"
+            # f"mov.b32 tmem_acc, {hex(acc_tmem_addr)};\n\t"
+            f"mov.b32 tmem_acc, $3;\n\t"
+            "mov.b32 smem_desc_a_lo_start, $0;\n\t"
+            "mov.b32 smem_desc_b_lo_start, $1;\n\t"
+            f"mov.b32 smem_desc_a_hi, {hex(smem_desc_a_hi)};\n\t"
+            f"mov.b32 smem_desc_b_hi, {hex(smem_desc_b_hi)};\n\t"
+            f"mov.b64 smem_desc_a, {{smem_desc_a_lo_start, smem_desc_a_hi}};\n\t"
+            f"mov.b64 smem_desc_b, {{smem_desc_b_lo_start, smem_desc_b_hi}};\n\t"
+            "setp.ne.b32 p, $2, 0;\n\t"
+            f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [tmem_acc], smem_desc_a, smem_desc_b, idesc, {pred_str};\n\t"
+            + "".join(
+                (
+                    # f"add.u32 smem_desc_a_lo, smem_desc_a_lo, {hex(offset_a_diff[k - 1])};\n\t"
+                    # f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t"
+                    f"add.u32 smem_desc_a_lo, smem_desc_a_lo_start, {hex(offset_a[k])};\n\t"
+                    f"add.u32 smem_desc_b_lo, smem_desc_b_lo_start, {hex(offset_b[k])};\n\t"
+                    f"mov.b64 smem_desc_a, {{smem_desc_a_lo, smem_desc_a_hi}};\n\t"
+                    f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t"
+                    f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [tmem_acc], smem_desc_a, smem_desc_b, idesc, 1;\n\t"
+                )
+                for k in range(1, cute.size(tCrA.shape[2]))
+            )
+            + "}\n",
+            # "r,r,r",
+            "r,r,r,r",
+            has_side_effects=True,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    else:
+        # For TS gemm, somehow tCrA.iterator.toint() returns 0 no matter what, so we need to
+        # explicitly pass in the tA_addr for correctness.
+        tA_addr = tCrA[None, None, 0].iterator.toint() if tA_addr is None else tA_addr
+        input_args = [
+            # Int32(cute.arch.make_warp_uniform(tCrA[None, None, 0].iterator.toint())).ir_value(),
+            Int32(cute.arch.make_warp_uniform(tA_addr)).ir_value(),
+            Int32(cute.arch.make_warp_uniform(smem_desc_start_b_lo)).ir_value(),
+            Int32(not zero_init).ir_value(),
+            Int32(cute.arch.make_warp_uniform(acc_tmem_addr)).ir_value(),
+        ]
+        if const_expr(mbar_ptr is not None):
+            assert mbar_phase is not None, "mbar_phase must be provided when mbar_ptr is not None"
+            input_args.append(mbar_ptr.toint().ir_value())
+            input_args.append(Int32(mbar_phase).ir_value())
+            mbar_wait_str = (
+                ".reg .pred P1; \n\t"
+                "LAB_WAIT: \n\t"
+                "mbarrier.try_wait.parity.shared::cta.b64 P1, [$4], $5, 10000000; \n\t"
+                "@P1 bra DONE; \n\t"
+                "bra     LAB_WAIT; \n\t"
+                "DONE: \n\t"
+            )
+        else:
+            mbar_wait_str = ""
+        llvm.inline_asm(
+            None,
+            # [
+            #     # acc.iterator.toint().ir_value(),
+            #     Int32(tCrA[None, None, 0].iterator.toint()).ir_value(),
+            #     Int32(smem_desc_start_b_lo).ir_value(),
+            #     Int32(not zero_init).ir_value(),
+            # ],
+            input_args,
+            "{\n\t"
+            ".reg .pred leader_thread;\n\t"
+            ".reg .pred p;\n\t"
+            ".reg .b32 idesc;\n\t"
+            ".reg .b32 tmem_acc;\n\t"
+            ".reg .b32 tmem_a;\n\t"
+            ".reg .b32 smem_desc_b_lo_start;\n\t"
+            ".reg .b32 smem_desc_b_lo;\n\t"
+            ".reg .b32 smem_desc_b_hi;\n\t"
+            ".reg .b64 smem_desc_b;\n\t"
+            "elect.sync _|leader_thread, -1;\n\t"
+            f"mov.b32 idesc, {hex(idesc)};\n\t"
+            # f"mov.b32 tmem_acc, {hex(acc_tmem_addr)};\n\t"
+            f"mov.b32 tmem_acc, $3;\n\t"
+            f"mov.b32 tmem_a, $0;\n\t"
+            f"mov.b32 smem_desc_b_lo_start, $1;\n\t"
+            f"mov.b32 smem_desc_b_hi, {hex(smem_desc_b_hi)};\n\t"
+            f"mov.b64 smem_desc_b, {{smem_desc_b_lo_start, smem_desc_b_hi}};\n\t"
+            "setp.ne.b32 p, $2, 0;\n\t"
+            f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [tmem_acc], [tmem_a], smem_desc_b, idesc, {pred_str};\n\t"
+            + "".join(
+                (
+                    # f"add.u32 tmem_a, tmem_a, {hex(offset_a_diff[k - 1])};\n\t"
+                    # f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t"
+                    f"add.u32 smem_desc_b_lo, smem_desc_b_lo_start, {hex(offset_b[k])};\n\t"
+                    f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t"
+                    # f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [tmem_acc], [tmem_a], smem_desc_b, idesc, 1;\n\t"
+                    f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t"
+                )
+                for k in range(
+                    1,
+                    cute.size(tCrA.shape[2])
+                    if const_expr(mbar_ptr is None)
+                    else cute.size(tCrA.shape[2]) // 4 * 3,
+                )
+            )
+            + mbar_wait_str
+            + (
+                "".join(
+                    (
+                        f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t"
+                        f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t"
+                        f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t"
+                    )
+                    for k in range(cute.size(tCrA.shape[2]) // 4 * 3, cute.size(tCrA.shape[2]))
+                )
+                if const_expr(mbar_ptr is not None)
+                else ""
+            )
+            + "}\n",
+            "r,r,r,r" if const_expr(mbar_ptr is None) else "r,r,r,r,r,r",
+            has_side_effects=True,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+
+
+@cute.jit
+def gemm_ptx_partial1(
+    op: cute.nvgpu.tcgen05.mma.MmaOp,
+    acc_tmem_addr: cutlass.Constexpr[int],
+    tCrA: cute.Tensor,
+    tCrB: cute.Tensor,
+    sA_base_addr_for_desc: Int32,
+    sA_addr_offset_for_desc: cutlass.Constexpr[int],
+    sA_stage: Int32,
+    sB_base_addr_for_desc: Int32,
+    sB_addr_offset_for_desc: cutlass.Constexpr[int],
+    sB_stage: Int32,
+    sA_layout: Optional[cute.Layout],
+    sB_layout: Optional[cute.Layout],
+    sA_swizzle: Optional[cute.Swizzle],
+    sB_swizzle: cute.Swizzle,
+    zero_init: bool | Boolean = False,
+) -> None:
+    is_ts = op.a_src == cute.nvgpu.tcgen05.OperandSource.TMEM
+    if const_expr(not is_ts):
+        assert sA_layout is not None, "sA_layout must be provided when a_src is not TMEM"
+        assert sA_swizzle is not None, "sA_swizzle must be provided when a_src is not TMEM"
+    idesc: int = const_expr(sm100_desc.mma_op_to_idesc(op))
+    if const_expr(not is_ts):
+        smem_desc_base_a: int = const_expr(
+            sm100_desc.make_smem_desc_base(
+                cute.recast_layout(128, op.a_dtype.width, sA_layout[0]),
+                sA_swizzle,
+                sm100_desc.Major.K
+                if const_expr(op.a_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K)
+                else sm100_desc.Major.MN,
+            )
+        )
+        smem_desc_base_a_lo, smem_desc_a_hi = i64_to_i32x2(smem_desc_base_a)
+        smem_desc_base_a_lo = const_expr(smem_desc_base_a_lo)
+        smem_desc_a_hi = const_expr(smem_desc_a_hi)
+    else:
+        smem_desc_base_a = None
+        smem_desc_base_a_lo, smem_desc_a_hi = None, None
+    smem_desc_base_b: int = const_expr(
+        sm100_desc.make_smem_desc_base(
+            cute.recast_layout(128, op.b_dtype.width, sB_layout[0]),
+            sB_swizzle,
+            sm100_desc.Major.K
+            if const_expr(op.b_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K)
+            else sm100_desc.Major.MN,
+        )
+    )
+    smem_desc_base_b_lo, smem_desc_b_hi = i64_to_i32x2(smem_desc_base_b)
+    smem_desc_base_b_lo = const_expr(smem_desc_base_b_lo)
+    smem_desc_b_hi = const_expr(smem_desc_b_hi)
+    mask = [Int32(0)] * 4
+
+    if const_expr(not is_ts):
+        offset_a = [
+            (cute.crd2idx((0, 0, k), sA_layout) * op.a_dtype.width // 8) >> 4
+            for k in range(cute.size(tCrA.shape[2]))
+        ]
+    else:
+        offset_a = [
+            cute.crd2idx((0, 0, k), sA_layout) * op.a_dtype.width // 32
+            for k in range(cute.size(tCrA.shape[2]))
+        ]
+    offset_a_diff = [offset_a[k] - offset_a[k - 1] for k in range(1, cute.size(tCrA.shape[2]))]
+    offset_b = [
+        (cute.crd2idx((0, 0, k), sB_layout) * op.b_dtype.width // 8) >> 4
+        for k in range(cute.size(tCrB.shape[2]))
+    ]
+    offset_b_diff = [offset_b[k] - offset_b[k - 1] for k in range(1, cute.size(tCrB.shape[2]))]
+
+    if const_expr(not is_ts):
+        # smem_desc_start_a_lo = Int32(smem_desc_base_a_lo | sm100_desc.make_smem_desc_start_addr(sA[None, None, 0].iterator))
+        smem_desc_start_a_lo = const_expr(smem_desc_base_a_lo)
+    else:
+        smem_desc_start_a_lo = None
+    # smem_desc_start_b_lo = Int32(smem_desc_base_b_lo | sm100_desc.make_smem_desc_start_addr(sB[None, None, 0].iterator))
+    smem_desc_start_b_lo = const_expr(smem_desc_base_b_lo)
+    pred_str = "p" if isinstance(zero_init, Boolean) else "0" if zero_init else "1"
+    if const_expr(not is_ts):
+        llvm.inline_asm(
+            None,
+            [
+                # acc.iterator.toint().ir_value(),
+                # Int32(cute.arch.make_warp_uniform(smem_desc_start_a_lo)).ir_value(),
+                Int32(sA_base_addr_for_desc).ir_value(),
+                Int32(sA_stage).ir_value(),
+                # Int32(cute.arch.make_warp_uniform(smem_desc_start_b_lo)).ir_value(),
+                Int32(sB_base_addr_for_desc).ir_value(),
+                Int32(sB_stage).ir_value(),
+                Int32(not zero_init).ir_value(),
+                mask[0].ir_value(),
+                mask[1].ir_value(),
+                mask[2].ir_value(),
+                mask[3].ir_value(),
+            ],
+            "{\n\t"
+            ".reg .pred leader_thread;\n\t"
+            ".reg .pred p;\n\t"
+            ".reg .b32 idesc;\n\t"
+            ".reg .b32 tmem_acc;\n\t"
+            ".reg .b32 smem_desc_a_lo, smem_desc_b_lo;\n\t"
+            ".reg .b32 smem_desc_a_hi, smem_desc_b_hi;\n\t"
+            ".reg .b64 smem_desc_a, smem_desc_b;\n\t"
+            "elect.sync _|leader_thread, -1;\n\t"
+            f"mov.b32 idesc, {hex(idesc)};\n\t"
+            f"mov.b32 tmem_acc, {hex(acc_tmem_addr)};\n\t"
+            # "mov.b32 smem_desc_a_lo, $0;\n\t"
+            # f"add.u32 smem_desc_a_lo, $0, {hex(smem_desc_start_a_lo)};\n\t"
+            f"mad.lo.u32 smem_desc_a_lo, $1, {hex(sA_addr_offset_for_desc)}, $0;\n\t"
+            # "mov.b32 smem_desc_b_lo, $2;\n\t"
+            f"mad.lo.u32 smem_desc_b_lo, $3, {hex(sB_addr_offset_for_desc)}, $2;\n\t"
+            f"mov.b32 smem_desc_a_hi, {hex(smem_desc_a_hi)};\n\t"
+            f"mov.b32 smem_desc_b_hi, {hex(smem_desc_b_hi)};\n\t"
+            f"mov.b64 smem_desc_a, {{smem_desc_a_lo, smem_desc_a_hi}};\n\t"
+            f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t"
+            "setp.ne.b32 p, $4, 0;\n\t"
+            f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [tmem_acc], smem_desc_a, smem_desc_b, idesc, {{$5, $6, $7, $8}}, {pred_str};\n\t"
+            + "".join(
+                (
+                    f"add.u32 smem_desc_a_lo, smem_desc_a_lo, {hex(offset_a_diff[k - 1])};\n\t"
+                    f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t"
+                    f"mov.b64 smem_desc_a, {{smem_desc_a_lo, smem_desc_a_hi}};\n\t"
+                    f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t"
+                    f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [tmem_acc], smem_desc_a, smem_desc_b, idesc, {{$5, $6, $7, $8}}, 1;\n\t"
+                )
+                for k in range(1, cute.size(tCrA.shape[2]))
+            )
+            + "}\n",
+            "r,r,r,r,r,r,r,r,r",
+            has_side_effects=True,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    else:
+        llvm.inline_asm(
+            None,
+            [
+                # acc.iterator.toint().ir_value(),
+                Int32(tCrA[None, None, 0].iterator.toint()).ir_value(),
+                Int32(smem_desc_start_b_lo).ir_value(),
+                Int32(not zero_init).ir_value(),
+                mask[0].ir_value(),
+                mask[1].ir_value(),
+                mask[2].ir_value(),
+                mask[3].ir_value(),
+            ],
+            "{\n\t"
+            ".reg .pred leader_thread;\n\t"
+            ".reg .pred p;\n\t"
+            ".reg .b32 idesc;\n\t"
+            ".reg .b32 tmem_a;\n\t"
+            ".reg .b32 smem_desc_b_lo;\n\t"
+            ".reg .b32 smem_desc_b_hi;\n\t"
+            ".reg .b64 smem_desc_b;\n\t"
+            "elect.sync _|leader_thread, -1;\n\t"
+            f"mov.b32 idesc, {hex(idesc)};\n\t"
+            f"mov.b32 tmem_a, $1;\n\t"
+            f"mov.b32 smem_desc_b_lo, $2;\n\t"
+            f"mov.b32 smem_desc_b_hi, {hex(smem_desc_b_hi)};\n\t"
+            f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t"
+            "setp.ne.b32 p, $3, 0;\n\t"
+            f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [$0], [tmem_a], smem_desc_b, idesc, {{$4, $5, $6, $7}}, {pred_str};\n\t"
+            + "".join(
+                (
+                    f"add.u32 tmem_a, tmem_a, {hex(offset_a_diff[k - 1])};\n\t"
+                    f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t"
+                    f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t"
+                    f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [$0], [tmem_a], smem_desc_b, idesc, {{$4, $5, $6, $7}}, 1;\n\t"
+                )
+                for k in range(1, cute.size(tCrA.shape[2]))
+            )
+            + "}\n",
+            "r,r,r,r,r,r,r,r",
+            has_side_effects=True,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )

build/torch-universal/block_info.py

@@ -0,0 +1,104 @@
+# 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 .seqlen_info import SeqlenInfoQK
+
+
+@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: cutlass.Int32 = 0,
+        num_splits: cutlass.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 = (
+                cutlass.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):
+            m_block_min = max(
+                m_block_min,
+                (n_block * self.tile_n + seqlen_info.seqlen_q - seqlen_info.seqlen_k)
+                // self.tile_m,
+            )
+        return m_block_min, m_block_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))

build/torch-universal/block_sparse_utils.py

@@ -0,0 +1,419 @@
+"""
+Block-sparse runtime utilities for CUTE DSL kernels.
+
+This module contains runtime execution functions for block-sparse attention kernels.
+These utilities are used by CUTE DSL kernels to produce and consume block-sparse loads.
+"""
+
+from typing import Callable
+from functools import partial
+import cutlass
+import cutlass.cute as cute
+from cutlass import const_expr
+
+# Import data structures from block_sparsity
+from .block_sparsity import BlockSparseTensors
+
+
+@cute.jit
+def load_block_list(
+    block_indices: cute.Tensor,
+    block_count,
+    load_q_with_first: cutlass.Constexpr,
+    first_block_preloaded: cutlass.Constexpr,
+    kv_producer_state,
+    load_Q,
+    load_K,
+    load_V,
+    pipeline_k,
+    pipeline_v,
+    use_tma_q: cutlass.Constexpr,
+    tma_q_bytes: cutlass.Constexpr,
+    intra_wg_overlap: cutlass.Constexpr,
+):
+    """Iterate over the sparse blocks and load K, V (and Q) into the pipeline.
+    for the intra_wg_overlap case, we overlap the loads of K and V. And this
+    means we need to pipeline the last V load from the partial block case,
+    with the loads for the full blocks. Set first_block_preloaded when the
+    caller has already issued the first K load for the list.
+
+    Note:
+        we iterate along the block_n indices in reverse.
+
+    Returns:
+        Updated kv_producer_state after processing the block list.
+
+    """
+    if block_count > 0:
+        if const_expr(not intra_wg_overlap):
+            # Peel first iteration: the first block may need to load Q alongside K,
+            # Parameters are already Constexpr, so no need to wrap in const_expr()
+            n_block_first = block_indices[block_count - 1]
+            extra_tx = tma_q_bytes if const_expr(load_q_with_first) and const_expr(use_tma_q) else 0
+            pipeline_k.producer_acquire(kv_producer_state, extra_tx_count=extra_tx)
+
+            if const_expr(load_q_with_first and use_tma_q):
+                load_Q(tma_bar_ptr=pipeline_k.producer_get_barrier(kv_producer_state))
+
+            load_K(src_idx=n_block_first, producer_state=kv_producer_state)
+            pipeline_v.producer_acquire(kv_producer_state)
+            load_V(src_idx=n_block_first, producer_state=kv_producer_state)
+            kv_producer_state.advance()
+
+            for offset in cutlass.range(1, block_count):
+                n_block = block_indices[block_count - 1 - offset]
+                pipeline_k.producer_acquire(kv_producer_state)
+                load_K(src_idx=n_block, producer_state=kv_producer_state)
+                pipeline_v.producer_acquire(kv_producer_state)
+                load_V(src_idx=n_block, producer_state=kv_producer_state)
+                kv_producer_state.advance()
+        else:
+            n_block_first = block_indices[block_count - 1]
+            if const_expr(not first_block_preloaded):
+                extra_tx = (
+                    tma_q_bytes if const_expr(load_q_with_first) and const_expr(use_tma_q) else 0
+                )
+                pipeline_k.producer_acquire(kv_producer_state, extra_tx_count=extra_tx)
+
+                if const_expr(load_q_with_first and use_tma_q):
+                    load_Q(tma_bar_ptr=pipeline_k.producer_get_barrier(kv_producer_state))
+
+                load_K(src_idx=n_block_first, producer_state=kv_producer_state)
+
+            for idx in cutlass.range(block_count - 1, unroll=1):
+                n_block_prev = block_indices[block_count - 1 - idx]
+                n_block = block_indices[block_count - 2 - idx]
+                kv_producer_state_prev = kv_producer_state.clone()
+                kv_producer_state.advance()
+                pipeline_k.producer_acquire(kv_producer_state)
+                load_K(src_idx=n_block, producer_state=kv_producer_state)
+                pipeline_v.producer_acquire(kv_producer_state_prev)
+                load_V(src_idx=n_block_prev, producer_state=kv_producer_state_prev)
+
+    return kv_producer_state
+
+
+@cute.jit
+def finish_overlap_v_load(
+    block_indices: cute.Tensor,
+    block_count,
+    load_V,
+    pipeline_v,
+    kv_producer_state,
+):
+    """Load the final V block after overlapped K/V loads."""
+    if block_count > 0:
+        n_block_last = block_indices[0]
+        pipeline_v.producer_acquire(kv_producer_state)
+        load_V(src_idx=n_block_last, producer_state=kv_producer_state)
+        kv_producer_state.advance()
+
+    return kv_producer_state
+
+
+@cute.jit
+def produce_block_sparse_loads(
+    blocksparse_tensors: BlockSparseTensors,
+    batch_idx,
+    head_idx,
+    m_block,
+    kv_producer_state,
+    load_Q,
+    load_K,
+    load_V,
+    pipeline_k,
+    pipeline_v,
+    use_tma_q: cutlass.Constexpr,
+    tma_q_bytes: cutlass.Constexpr,
+    intra_wg_overlap: cutlass.Constexpr,
+):
+    """Iterate over the mask and full block lists for a single tile.
+
+    The masked (partial) list may leave the last V load pending when intra-warp-group
+    overlap is enabled. The first full block must consume that pending V while
+    issuing its own K load on the next pipeline stage.
+
+    In the intra-wg-overlap path, the last masked block leaves its V copy in flight
+    while we advance the producer state to start the next full K. Either the full list
+    overlaps that pending V load, or, if no full blocks exist, we explicitly drain it.
+
+    """
+
+    mask_block_cnt, mask_block_idx, full_block_cnt, full_block_idx = blocksparse_tensors
+
+    curr_mask_block_cnt = mask_block_cnt[batch_idx, head_idx, m_block]
+    curr_mask_block_idx = mask_block_idx[batch_idx, head_idx, m_block, None]
+    curr_full_block_cnt = full_block_cnt[batch_idx, head_idx, m_block]
+    curr_full_block_idx = full_block_idx[batch_idx, head_idx, m_block, None]
+
+    mask_empty = curr_mask_block_cnt == 0
+    full_empty = curr_full_block_cnt == 0
+
+    if mask_empty:
+        # No masked blocks: the full list owns the initial Q+K load.
+        kv_producer_state = load_block_list(
+            curr_full_block_idx,
+            curr_full_block_cnt,
+            load_q_with_first=True,
+            first_block_preloaded=False,
+            kv_producer_state=kv_producer_state,
+            load_Q=load_Q,
+            load_K=load_K,
+            load_V=load_V,
+            pipeline_k=pipeline_k,
+            pipeline_v=pipeline_v,
+            use_tma_q=use_tma_q,
+            tma_q_bytes=tma_q_bytes,
+            intra_wg_overlap=intra_wg_overlap,
+        )
+
+        if const_expr(intra_wg_overlap) and curr_full_block_cnt > 0:
+            kv_producer_state = finish_overlap_v_load(
+                curr_full_block_idx,
+                curr_full_block_cnt,
+                load_V,
+                pipeline_v,
+                kv_producer_state,
+            )
+    else:
+        # Masked blocks present: load Q together with the first masked K so consumers can
+        # start immediately. When overlap is disabled this fully drains the list.
+        kv_producer_state = load_block_list(
+            curr_mask_block_idx,
+            curr_mask_block_cnt,
+            load_q_with_first=True,
+            first_block_preloaded=False,
+            kv_producer_state=kv_producer_state,
+            load_Q=load_Q,
+            load_K=load_K,
+            load_V=load_V,
+            pipeline_k=pipeline_k,
+            pipeline_v=pipeline_v,
+            use_tma_q=use_tma_q,
+            tma_q_bytes=tma_q_bytes,
+            intra_wg_overlap=intra_wg_overlap,
+        )
+
+        if full_empty:
+            if const_expr(intra_wg_overlap):
+                kv_producer_state = finish_overlap_v_load(
+                    curr_mask_block_idx,
+                    curr_mask_block_cnt,
+                    load_V,
+                    pipeline_v,
+                    kv_producer_state,
+                )
+        else:
+            if const_expr(intra_wg_overlap):
+                # Bridge the masked list to the full list by overlapping the pending masked V
+                # with the first full K load.
+                n_block_mask_last = curr_mask_block_idx[0]
+                n_block_full_first = curr_full_block_idx[curr_full_block_cnt - 1]
+                kv_producer_state_prev = kv_producer_state.clone()
+                kv_producer_state.advance()
+                pipeline_k.producer_acquire(kv_producer_state)
+                load_K(src_idx=n_block_full_first, producer_state=kv_producer_state)
+                pipeline_v.producer_acquire(kv_producer_state_prev)
+                load_V(src_idx=n_block_mask_last, producer_state=kv_producer_state_prev)
+
+                kv_producer_state = load_block_list(
+                    curr_full_block_idx,
+                    curr_full_block_cnt,
+                    load_q_with_first=False,
+                    first_block_preloaded=True,
+                    kv_producer_state=kv_producer_state,
+                    load_Q=load_Q,
+                    load_K=load_K,
+                    load_V=load_V,
+                    pipeline_k=pipeline_k,
+                    pipeline_v=pipeline_v,
+                    use_tma_q=use_tma_q,
+                    tma_q_bytes=tma_q_bytes,
+                    intra_wg_overlap=intra_wg_overlap,
+                )
+
+                kv_producer_state = finish_overlap_v_load(
+                    curr_full_block_idx,
+                    curr_full_block_cnt,
+                    load_V,
+                    pipeline_v,
+                    kv_producer_state,
+                )
+            else:
+                # Non-overlap path with both lists: run the full list normally (skipping the Q
+                # reload because the masked list already issued it).
+                kv_producer_state = load_block_list(
+                    curr_full_block_idx,
+                    curr_full_block_cnt,
+                    load_q_with_first=False,
+                    first_block_preloaded=False,
+                    kv_producer_state=kv_producer_state,
+                    load_Q=load_Q,
+                    load_K=load_K,
+                    load_V=load_V,
+                    pipeline_k=pipeline_k,
+                    pipeline_v=pipeline_v,
+                    use_tma_q=use_tma_q,
+                    tma_q_bytes=tma_q_bytes,
+                    intra_wg_overlap=intra_wg_overlap,
+                )
+
+    return kv_producer_state
+
+
+@cute.jit
+def consume_block_sparse_loads(
+    blocksparse_tensors: BlockSparseTensors,
+    batch_idx,
+    head_idx,
+    m_block,
+    kv_consumer_state,
+    mma_pv_fn,
+    mma_one_n_block,
+    process_first_half_block,
+    process_last_half_block,
+    mask_fn,
+    score_mod_fn,
+    O_should_accumulate,
+    mask_mod,
+    intra_wg_overlap: cutlass.Constexpr,
+    warp_scheduler_barrier_sync: Callable,
+    warp_scheduler_barrier_arrive: Callable,
+):
+    """Consume the mask and full block lists for a single tile on the consumer side.
+
+    Mirrors `produce_block_sparse_loads` so that the consumer pipeline
+    """
+
+    mask_block_cnt, mask_block_idx, full_block_cnt, full_block_idx = blocksparse_tensors
+
+    curr_mask_block_cnt = mask_block_cnt[batch_idx, head_idx, m_block]
+    curr_mask_block_idx = mask_block_idx[batch_idx, head_idx, m_block, None]
+    curr_full_block_cnt = full_block_cnt[batch_idx, head_idx, m_block]
+    curr_full_block_idx = full_block_idx[batch_idx, head_idx, m_block, None]
+
+    processed_any = curr_mask_block_cnt + curr_full_block_cnt > 0
+
+    if const_expr(not intra_wg_overlap):
+        if curr_mask_block_cnt > 0:
+            mask_n_block = curr_mask_block_idx[curr_mask_block_cnt - 1]
+            warp_scheduler_barrier_sync()
+            kv_consumer_state = mma_one_n_block(
+                kv_consumer_state,
+                n_block=mask_n_block,
+                mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate),
+                mask_fn=partial(mask_fn, mask_mod=mask_mod, mask_seqlen=True),
+                is_first_n_block=True,
+            )
+            O_should_accumulate = True
+            for i in cutlass.range(1, curr_mask_block_cnt):
+                mask_n_block = curr_mask_block_idx[curr_mask_block_cnt - 1 - i]
+                kv_consumer_state = mma_one_n_block(
+                    kv_consumer_state,
+                    n_block=mask_n_block,
+                    mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate),
+                    mask_fn=partial(mask_fn, mask_mod=mask_mod, mask_seqlen=False),
+                    is_first_n_block=False,
+                )
+                O_should_accumulate = True
+            if curr_full_block_cnt == 0:
+                warp_scheduler_barrier_arrive()
+
+        if curr_full_block_cnt > 0:
+            full_n_block = curr_full_block_idx[curr_full_block_cnt - 1]
+            if curr_mask_block_cnt == 0:
+                warp_scheduler_barrier_sync()
+                kv_consumer_state = mma_one_n_block(
+                    kv_consumer_state,
+                    n_block=full_n_block,
+                    mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate),
+                    mask_fn=partial(mask_fn, mask_seqlen=True),
+                    is_first_n_block=True,
+                )
+                O_should_accumulate = True
+                for i in cutlass.range(1, curr_full_block_cnt):
+                    full_n_block = curr_full_block_idx[curr_full_block_cnt - 1 - i]
+                    kv_consumer_state = mma_one_n_block(
+                        kv_consumer_state,
+                        n_block=full_n_block,
+                        mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate),
+                        mask_fn=partial(mask_fn, mask_seqlen=False),
+                        is_first_n_block=False,
+                    )
+                    O_should_accumulate = True
+            else:
+                kv_consumer_state = mma_one_n_block(
+                    kv_consumer_state,
+                    n_block=full_n_block,
+                    mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate),
+                    mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=True),
+                    is_first_n_block=False,
+                )
+                O_should_accumulate = True
+                for i in cutlass.range(1, curr_full_block_cnt):
+                    full_n_block = curr_full_block_idx[curr_full_block_cnt - 1 - i]
+                    kv_consumer_state = mma_one_n_block(
+                        kv_consumer_state,
+                        n_block=full_n_block,
+                        mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate),
+                        mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=False),
+                        is_first_n_block=False,
+                    )
+                    O_should_accumulate = True
+            warp_scheduler_barrier_arrive()
+    else:
+        if curr_mask_block_cnt > 0:
+            mask_n_block = curr_mask_block_idx[curr_mask_block_cnt - 1]
+            kv_consumer_state = process_first_half_block(
+                n_block=mask_n_block,
+                kv_consumer_state=kv_consumer_state,
+                mask_fn=partial(mask_fn, mask_mod=mask_mod),
+                score_mod_fn=score_mod_fn,
+                is_first_block=True,
+            )
+            for i in cutlass.range(1, curr_mask_block_cnt):
+                mask_n_block = curr_mask_block_idx[curr_mask_block_cnt - 1 - i]
+                kv_consumer_state = mma_one_n_block(
+                    kv_consumer_state,
+                    n_block=mask_n_block,
+                    mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate),
+                    mask_fn=partial(mask_fn, mask_mod=mask_mod, mask_seqlen=False),
+                )
+                O_should_accumulate = True
+
+        if curr_full_block_cnt > 0:
+            full_n_block = curr_full_block_idx[curr_full_block_cnt - 1]
+            if curr_mask_block_cnt == 0:
+                kv_consumer_state = process_first_half_block(
+                    n_block=full_n_block,
+                    kv_consumer_state=kv_consumer_state,
+                    mask_fn=partial(mask_fn, mask_mod=None),
+                    score_mod_fn=score_mod_fn,
+                    is_first_block=True,
+                )
+            else:
+                kv_consumer_state = mma_one_n_block(
+                    kv_consumer_state,
+                    n_block=full_n_block,
+                    mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate),
+                    mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=True),
+                )
+                O_should_accumulate = True
+            for i in cutlass.range(1, curr_full_block_cnt):
+                full_n_block = curr_full_block_idx[curr_full_block_cnt - 1 - i]
+                kv_consumer_state = mma_one_n_block(
+                    kv_consumer_state,
+                    n_block=full_n_block,
+                    mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate),
+                    mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=False),
+                )
+                O_should_accumulate = True
+
+        if curr_mask_block_cnt + curr_full_block_cnt > 0:
+            kv_consumer_state = process_last_half_block(
+                kv_consumer_state=kv_consumer_state,
+                zero_init=not O_should_accumulate,
+            )
+            O_should_accumulate = True
+
+    return kv_consumer_state, O_should_accumulate, processed_any

build/torch-universal/block_sparsity.py

@@ -0,0 +1,650 @@
+"""
+Computes block-sparse attention masks for Flex Attention.
+
+This utility generates block sparsity patterns based on common attention masking
+strategies (e.g., causal, sliding window). The resulting tensors define which
+blocks are fully computed, which are partially computed (requiring a mask), and
+which are skipped entirely. This is a temporary solution intended to be replaced
+by a more robust preprocessing kernel in the future.
+"""
+
+from typing import Tuple, Optional, Callable, List, NamedTuple
+import torch
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+
+
+# placeholder
+Config = type("Config", (), {})
+
+
+class BlockSparseTensors(NamedTuple):
+    mask_block_cnt: cute.Tensor
+    mask_block_idx: cute.Tensor
+    full_block_cnt: Optional[cute.Tensor]
+    full_block_idx: Optional[cute.Tensor]
+
+    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: Optional[torch.Tensor] = None
+    full_block_idx: Optional[torch.Tensor] = None
+
+
+def _expand_sparsity_tensor(
+    tensor: torch.Tensor,
+    expected_shape: Tuple[int, ...],
+    tensor_name: str,
+) -> 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:
+        raise ValueError(
+            f"{tensor_name} with shape {tensor.shape} cannot be expanded to expected shape {expected_shape}."
+        )
+    return tensor.expand(*expected_shape).contiguous()
+
+
+def _check_and_expand_block(
+    name: str,
+    cnt: Optional[torch.Tensor],
+    idx: Optional[torch.Tensor],
+    expected_count_shape: Tuple[int, int, int],
+    expected_index_shape: Tuple[int, int, int, int],
+) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
+    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")
+    expanded_idx = _expand_sparsity_tensor(idx, expected_index_shape, f"{name}_block_idx")
+    return expanded_cnt, expanded_idx
+
+
+def normalize_block_sparse_tensors(
+    tensors: BlockSparseTensorsTorch,
+    *,
+    expected_count_shape: Tuple[int, int, int],
+    expected_index_shape: Tuple[int, int, int, int],
+) -> 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,
+    )
+    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,
+    )
+    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,
+    )
+
+
+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 to_cute_block_sparse_tensors(tensors: BlockSparseTensorsTorch) -> Optional[BlockSparseTensors]:
+    if not is_block_sparsity_enabled(tensors):
+        return None
+
+    mask_block_cnt_tensor = from_dlpack(
+        tensors.mask_block_cnt.detach(), assumed_align=4
+    ).mark_layout_dynamic(leading_dim=2)
+    mask_block_idx_tensor = from_dlpack(
+        tensors.mask_block_idx.detach(), assumed_align=4
+    ).mark_layout_dynamic(leading_dim=3)
+    full_block_cnt_tensor = (
+        from_dlpack(tensors.full_block_cnt.detach(), assumed_align=4).mark_layout_dynamic(
+            leading_dim=2
+        )
+        if tensors.full_block_cnt is not None
+        else None
+    )
+    full_block_idx_tensor = (
+        from_dlpack(tensors.full_block_idx.detach(), assumed_align=4).mark_layout_dynamic(
+            leading_dim=3
+        )
+        if tensors.full_block_idx is not None
+        else None
+    )
+
+    return BlockSparseTensors(
+        mask_block_cnt_tensor,
+        mask_block_idx_tensor,
+        full_block_cnt_tensor,
+        full_block_idx_tensor,
+    )
+
+
+def compute_block_sparsity(
+    config: Config,
+    mask_mod_flex: Optional[Callable],
+    device: str,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    aux_tensors: Optional[List[torch.Tensor]] = None,
+) -> Tuple[
+    Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]
+]:
+    """
+    Computes block sparsity tensors from a given masking function.
+
+    This function serves as the main entry point for generating block-sparse masks.
+    It dispatches to specialized handlers for variable-length and fixed-length
+    sequences.
+
+    Args:
+        config: A configuration object containing model and tiling parameters.
+        mask_mod_flex: The mask function for generic flex attention patterns.
+        device: The device to create tensors on (e.g., 'cuda').
+        cu_seqlens_q: Cumulative sequence lengths for Q (for varlen).
+        cu_seqlens_k: Cumulative sequence lengths for K (for varlen).
+        aux_tensors: A list of auxiliary tensors, e.g., for document masking.
+
+    Returns:
+        A tuple of four tensors:
+        - `full_block_cnt`: (batch, nheads, n_blocks_q) - Count of full n blocks per m block.
+        - `full_block_idx`: (batch, nheads, n_blocks_q, max_n_blocks) - Indices of full n blocks.
+        - `mask_block_cnt`: (batch, nheads, n_blocks_q) - Count of partial n blocks per m block.
+        - `mask_block_idx`: (batch, nheads, n_blocks_q, max_n_blocks) - Indices of partial n blocks.
+        Returns (None, None, None, None) if masking is disabled.
+    """
+    if not config.use_mask_mod or mask_mod_flex is None:
+        return None, None, None, None
+
+    if cu_seqlens_q is not None:
+        # Handle variable-length sequences
+        return _compute_varlen_sparsity(config, mask_mod_flex, device, cu_seqlens_q, cu_seqlens_k)
+    else:
+        # Handle fixed-length sequences
+        return _compute_sparsity(config, device, aux_tensors)
+
+
+## ---------------------------------------------------------------------------
+## Fixed-Length Sequence Kernels
+## ---------------------------------------------------------------------------
+
+
+def _compute_sparsity(
+    config: Config, device: str, aux_tensors: Optional[List[torch.Tensor]]
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Computes block sparsity for fixed-length sequences."""
+    n_blocks_q = (config.seqlen_q + config.tile_m - 1) // config.tile_m
+    n_blocks_k = (config.seqlen_k + config.tile_n - 1) // config.tile_n
+
+    # Pre-allocate output tensors
+    full_block_cnt = torch.zeros(
+        (config.batch_size, config.nheads, n_blocks_q), device=device, dtype=torch.int32
+    )
+    mask_block_cnt = torch.zeros(
+        (config.batch_size, config.nheads, n_blocks_q), device=device, dtype=torch.int32
+    )
+    full_block_idx = torch.zeros(
+        (config.batch_size, config.nheads, n_blocks_q, n_blocks_k), device=device, dtype=torch.int32
+    )
+    mask_block_idx = torch.zeros(
+        (config.batch_size, config.nheads, n_blocks_q, n_blocks_k), device=device, dtype=torch.int32
+    )
+
+    # --- Identity Mask ---
+    # All blocks are fully computed.
+    if config.mask_mod_name == "identity":
+        k_blocks = torch.arange(n_blocks_k, device=device)
+        for q_block_idx in range(n_blocks_q):
+            full_block_cnt[:, :, q_block_idx] = n_blocks_k
+            full_block_idx[:, :, q_block_idx, :n_blocks_k] = k_blocks
+
+    # --- Identity Partial Mask ---
+    # All blocks are partially computed (masked).
+    elif config.mask_mod_name == "identity_partial":
+        k_blocks = torch.arange(n_blocks_k, device=device)
+        for q_block_idx in range(n_blocks_q):
+            mask_block_cnt[:, :, q_block_idx] = n_blocks_k
+            mask_block_idx[:, :, q_block_idx, :n_blocks_k] = k_blocks
+
+    # --- Block Causal Mask ---
+    elif config.mask_mod_name == "block_causal":
+        k_blocks = torch.arange(n_blocks_k, device=device)
+        for q_block_idx in range(n_blocks_q):
+            causal_indices = k_blocks[k_blocks <= q_block_idx]
+            num_causal_indices = len(causal_indices)
+            if num_causal_indices > 0:
+                full_block_cnt[:, :, q_block_idx] = num_causal_indices
+                full_block_idx[:, :, q_block_idx, :num_causal_indices] = causal_indices
+
+    # --- Causal and Sliding Window Masks ---
+    elif config.mask_mod_name in ["causal", "sliding_window"]:
+        q_block_indices = torch.arange(n_blocks_q, device=device)
+        k_block_indices = torch.arange(n_blocks_k, device=device)
+
+        q_starts = q_block_indices * config.tile_m
+        q_ends = torch.minimum(
+            (q_block_indices + 1) * config.tile_m, torch.tensor(config.seqlen_q, device=device)
+        )
+        k_starts = k_block_indices * config.tile_n
+        k_ends = torch.minimum(
+            (k_block_indices + 1) * config.tile_n, torch.tensor(config.seqlen_k, device=device)
+        )
+
+        # Expand dims for broadcasting: (n_blocks_q, 1) and (1, n_blocks_k)
+        q_starts, q_ends = q_starts.unsqueeze(1), q_ends.unsqueeze(1)
+        k_starts, k_ends = k_starts.unsqueeze(0), k_ends.unsqueeze(0)
+
+        offset = config.seqlen_k - config.seqlen_q
+
+        if config.mask_mod_name == "causal":
+            is_full = (k_ends - 1) <= (q_starts + offset)
+            # min(k_pos) <= max(q_pos) AND not is_full.
+            is_partial = (k_starts <= (q_ends - 1 + offset)) & ~is_full
+
+        else:  # sliding_window
+            window_size = getattr(config, "window_size", 1024)
+            is_full = (k_ends - 1 <= q_starts + offset) & (
+                k_starts >= q_ends - 1 + offset - (window_size - 1)
+            )
+            # A block is EMPTY if no (q, k) pairs satisfy the constraint.
+            is_empty = (k_starts > q_ends - 1 + offset) | (
+                k_ends - 1 < q_starts + offset - (window_size - 1)
+            )
+            # A block is PARTIAL if it's not empty and not full.
+            is_partial = ~is_empty & ~is_full
+
+        # Populate indices based on the computed block classifications
+        for q_block_idx in range(n_blocks_q):
+            full_indices = k_block_indices[is_full[q_block_idx]]
+            if len(full_indices) > 0:
+                full_block_cnt[:, :, q_block_idx] = len(full_indices)
+                full_block_idx[:, :, q_block_idx, : len(full_indices)] = full_indices
+
+            partial_indices = k_block_indices[is_partial[q_block_idx]]
+            if len(partial_indices) > 0:
+                mask_block_cnt[:, :, q_block_idx] = len(partial_indices)
+                mask_block_idx[:, :, q_block_idx, : len(partial_indices)] = partial_indices
+
+    elif config.mask_mod_name == "document":
+        raise NotImplementedError("Block sparsity for document masking not yet implemented")
+
+    return full_block_cnt, full_block_idx, mask_block_cnt, mask_block_idx
+
+
+## ---------------------------------------------------------------------------
+## Variable-Length Sequence Kernels
+## ---------------------------------------------------------------------------
+
+
+def _compute_varlen_sparsity(
+    config: Config,
+    mask_mod_flex: Callable,
+    device: str,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Computes block sparsity for variable-length sequences."""
+    assert cu_seqlens_k is not None, "cu_seqlens_k is required for varlen attention"
+    assert cu_seqlens_q.shape[0] == config.batch_size + 1
+    assert cu_seqlens_k.shape[0] == config.batch_size + 1
+
+    # In varlen, each sequence can have a different number of Q blocks.
+    # We pad up to the maximum number of Q blocks in the batch.
+    max_m_blocks = 0
+    for seq_idx in range(config.batch_size):
+        seq_len_q = (cu_seqlens_q[seq_idx + 1] - cu_seqlens_q[seq_idx]).item()
+        n_blocks_q = (seq_len_q + config.tile_m - 1) // config.tile_m
+        max_m_blocks = max(max_m_blocks, n_blocks_q)
+
+    # The number of K blocks is determined by the total length of all sequences.
+    total_k_len = cu_seqlens_k[-1].item()
+    max_n_blocks = (total_k_len + config.tile_n - 1) // config.tile_n
+
+    # Pre-allocate padded output tensors
+    full_block_cnt = torch.zeros(
+        (config.batch_size, config.nheads, max_m_blocks), device=device, dtype=torch.int32
+    )
+    mask_block_cnt = torch.zeros(
+        (config.batch_size, config.nheads, max_m_blocks), device=device, dtype=torch.int32
+    )
+    full_block_idx = torch.zeros(
+        (config.batch_size, config.nheads, max_m_blocks, max_n_blocks),
+        device=device,
+        dtype=torch.int32,
+    )
+    mask_block_idx = torch.zeros(
+        (config.batch_size, config.nheads, max_m_blocks, max_n_blocks),
+        device=device,
+        dtype=torch.int32,
+    )
+
+    # Process each sequence in the batch individually
+    for seq_idx in range(config.batch_size):
+        seq_start_q = cu_seqlens_q[seq_idx].item()
+        seq_end_q = cu_seqlens_q[seq_idx + 1].item()
+        seq_len_q = seq_end_q - seq_start_q
+
+        seq_start_k = cu_seqlens_k[seq_idx].item()
+        seq_end_k = cu_seqlens_k[seq_idx + 1].item()
+        seq_len_k = seq_end_k - seq_start_k
+
+        n_blocks_q = (seq_len_q + config.tile_m - 1) // config.tile_m
+        n_blocks_k = (seq_len_k + config.tile_n - 1) // config.tile_n
+
+        # Global block indices are relative to the start of the entire batch tensor
+        first_m_block_global = seq_start_q // config.tile_m
+        first_n_block_global = seq_start_k // config.tile_n
+
+        common_args = {
+            "full_block_cnt": full_block_cnt,
+            "full_block_idx": full_block_idx,
+            "mask_block_cnt": mask_block_cnt,
+            "mask_block_idx": mask_block_idx,
+            "seq_idx": seq_idx,
+            "n_blocks_q": n_blocks_q,
+            "n_blocks_k": n_blocks_k,
+            "seq_start_q": seq_start_q,
+            "seq_end_q": seq_end_q,
+            "seq_start_k": seq_start_k,
+            "seq_end_k": seq_end_k,
+            "first_n_block_global": first_n_block_global,
+            "tile_m": config.tile_m,
+            "tile_n": config.tile_n,
+            "device": device,
+        }
+
+        if config.mask_mod_name == "causal":
+            _compute_causal_varlen_blocks(**common_args)
+        elif config.mask_mod_name == "sliding_window":
+            window_size = getattr(config, "window_size", 1024)
+            _compute_sliding_window_varlen_blocks(**common_args, window_size=window_size)
+        elif config.mask_mod_name == "identity":
+            _compute_identity_varlen_blocks(
+                full_block_cnt,
+                full_block_idx,
+                seq_idx,
+                n_blocks_q,
+                n_blocks_k,
+                first_n_block_global,
+                device,
+            )
+        else:
+            # Generic case relies on sampling the user-provided mask function
+            _compute_generic_varlen_blocks(
+                **common_args,
+                mask_mod_flex=mask_mod_flex,
+                seq_len_q=seq_len_q,
+                seq_len_k=seq_len_k,
+                num_heads=config.nheads,
+                nheads_kv=config.nheads_kv,
+            )
+
+    return full_block_cnt, full_block_idx, mask_block_cnt, mask_block_idx
+
+
+def _classify_varlen_block(
+    m_local: int,
+    n_local: int,
+    seq_start_q: int,
+    seq_end_q: int,
+    seq_start_k: int,
+    seq_end_k: int,
+    tile_m: int,
+    tile_n: int,
+    is_full_fn: Callable,
+    is_partial_fn: Callable,
+) -> Tuple[bool, bool]:
+    """Helper to classify a varlen block as full, partial, or empty."""
+    m_start_global = seq_start_q + m_local * tile_m
+    m_end_global = min(seq_start_q + (m_local + 1) * tile_m, seq_end_q)
+    n_start_global = seq_start_k + n_local * tile_n
+    n_end_global = min(seq_start_k + (n_local + 1) * tile_n, seq_end_k)
+
+    # Use sequence-local coordinates for the logical check
+    m_start_local = m_start_global - seq_start_q
+    m_end_local = m_end_global - seq_start_q
+    n_start_local = n_start_global - seq_start_k
+    n_end_local = n_end_global - seq_start_k
+
+    is_full = is_full_fn(m_start_local, m_end_local, n_start_local, n_end_local)
+    is_partial = (
+        is_partial_fn(m_start_local, m_end_local, n_start_local, n_end_local) and not is_full
+    )
+
+    # Any block that touches the sequence boundary is partial because it requires masking.
+    at_boundary = (m_end_global > seq_end_q) or (n_end_global > seq_end_k)
+
+    return is_full and not at_boundary, is_partial or (is_full and at_boundary)
+
+
+def _compute_causal_varlen_blocks(
+    full_block_cnt,
+    full_block_idx,
+    mask_block_cnt,
+    mask_block_idx,
+    seq_idx,
+    n_blocks_q,
+    n_blocks_k,
+    seq_start_q,
+    seq_end_q,
+    seq_start_k,
+    seq_end_k,
+    first_n_block_global,
+    tile_m,
+    tile_n,
+    device,
+    **kwargs,
+):
+    """Computes causal block sparsity for a single varlen sequence."""
+    is_full_fn = lambda m_start, m_end, n_start, n_end: (m_start >= n_end - 1)
+    is_partial_fn = lambda m_start, m_end, n_start, n_end: (m_end - 1 >= n_start)
+
+    for m_local in range(n_blocks_q):
+        full_blocks, partial_blocks = [], []
+        for n_local in range(n_blocks_k):
+            is_full, is_partial = _classify_varlen_block(
+                m_local,
+                n_local,
+                seq_start_q,
+                seq_end_q,
+                seq_start_k,
+                seq_end_k,
+                tile_m,
+                tile_n,
+                is_full_fn,
+                is_partial_fn,
+            )
+            n_block_global = first_n_block_global + n_local
+            if is_full:
+                full_blocks.append(n_block_global)
+            elif is_partial:
+                partial_blocks.append(n_block_global)
+
+        if full_blocks:
+            full_block_cnt[seq_idx, :, m_local] = len(full_blocks)
+            full_block_idx[seq_idx, :, m_local, : len(full_blocks)] = torch.tensor(
+                full_blocks, device=device
+            )
+        if partial_blocks:
+            mask_block_cnt[seq_idx, :, m_local] = len(partial_blocks)
+            mask_block_idx[seq_idx, :, m_local, : len(partial_blocks)] = torch.tensor(
+                partial_blocks, device=device
+            )
+
+
+def _compute_sliding_window_varlen_blocks(
+    full_block_cnt,
+    full_block_idx,
+    mask_block_cnt,
+    mask_block_idx,
+    seq_idx,
+    n_blocks_q,
+    n_blocks_k,
+    seq_start_q,
+    seq_end_q,
+    seq_start_k,
+    seq_end_k,
+    first_n_block_global,
+    tile_m,
+    tile_n,
+    window_size,
+    device,
+    **kwargs,
+):
+    """Computes sliding window block sparsity for a single varlen sequence."""
+    is_full_fn = lambda m_start, m_end, n_start, n_end: (n_end - 1 <= m_start) and (
+        n_start >= m_start - window_size + 1
+    )
+    is_partial_fn = lambda m_start, m_end, n_start, n_end: not (
+        (n_start > m_end - 1) or (n_end - 1 < m_start - window_size + 1)
+    )
+
+    for m_local in range(n_blocks_q):
+        full_blocks, partial_blocks = [], []
+        for n_local in range(n_blocks_k):
+            is_full, is_partial = _classify_varlen_block(
+                m_local,
+                n_local,
+                seq_start_q,
+                seq_end_q,
+                seq_start_k,
+                seq_end_k,
+                tile_m,
+                tile_n,
+                is_full_fn,
+                is_partial_fn,
+            )
+            n_block_global = first_n_block_global + n_local
+            if is_full:
+                full_blocks.append(n_block_global)
+            elif is_partial:
+                partial_blocks.append(n_block_global)
+
+        if full_blocks:
+            full_block_cnt[seq_idx, :, m_local] = len(full_blocks)
+            full_block_idx[seq_idx, :, m_local, : len(full_blocks)] = torch.tensor(
+                full_blocks, device=device
+            )
+        if partial_blocks:
+            mask_block_cnt[seq_idx, :, m_local] = len(partial_blocks)
+            mask_block_idx[seq_idx, :, m_local, : len(partial_blocks)] = torch.tensor(
+                partial_blocks, device=device
+            )
+
+
+def _compute_identity_varlen_blocks(
+    full_block_cnt,
+    full_block_idx,
+    seq_idx,
+    n_blocks_q,
+    n_blocks_k,
+    first_n_block_global,
+    device,
+    **kwargs,
+):
+    """Computes identity (all-attend) block sparsity for a single varlen sequence."""
+    n_blocks_global = torch.arange(
+        first_n_block_global, first_n_block_global + n_blocks_k, device=device, dtype=torch.int32
+    )
+    for m_local in range(n_blocks_q):
+        full_block_cnt[seq_idx, :, m_local] = n_blocks_k
+        full_block_idx[seq_idx, :, m_local, :n_blocks_k] = n_blocks_global
+
+
+def _compute_generic_varlen_blocks(
+    full_block_cnt,
+    full_block_idx,
+    mask_block_cnt,
+    mask_block_idx,
+    mask_mod_flex,
+    seq_idx,
+    num_heads,
+    n_blocks_q,
+    n_blocks_k,
+    seq_len_q,
+    seq_len_k,
+    first_n_block_global,
+    tile_m,
+    tile_n,
+    nheads_kv,
+    device,
+    **kwargs,
+):
+    """Generic sampling-based block classification for a varlen sequence."""
+    qhead_per_kvhead = num_heads // nheads_kv
+
+    for h_q in range(num_heads):
+        h_kv = h_q // qhead_per_kvhead
+        for m_local in range(n_blocks_q):
+            m_start_local = m_local * tile_m
+            m_end_local = min((m_local + 1) * tile_m, seq_len_q)
+
+            full_blocks, partial_blocks = [], []
+            for n_local in range(n_blocks_k):
+                n_start_local = n_local * tile_n
+                n_end_local = min((n_local + 1) * tile_n, seq_len_k)
+
+                # Sample points within the block (corners and center) to classify it.
+                # Coordinates are sequence-local, as required by mask_mod_flex.
+                sample_positions = [
+                    (m_start_local, n_start_local),
+                    (m_start_local, n_end_local - 1),
+                    (m_end_local - 1, n_start_local),
+                    (m_end_local - 1, n_end_local - 1),
+                    ((m_start_local + m_end_local) // 2, (n_start_local + n_end_local) // 2),
+                ]
+
+                unmasked_count = sum(
+                    1
+                    for q_pos, k_pos in sample_positions
+                    if mask_mod_flex(seq_idx, h_q, q_pos, k_pos, seq_len_q, seq_len_k)
+                )
+
+                n_block_global = first_n_block_global + n_local
+                if unmasked_count == len(sample_positions):  # All samples unmasked -> full
+                    full_blocks.append(n_block_global)
+                elif unmasked_count > 0:  # Some unmasked -> partial
+                    partial_blocks.append(n_block_global)
+
+            if full_blocks:
+                full_block_cnt[seq_idx, h_q, m_local] = len(full_blocks)
+                full_block_idx[seq_idx, h_q, m_local, : len(full_blocks)] = torch.tensor(
+                    full_blocks, device=device
+                )
+            if partial_blocks:
+                mask_block_cnt[seq_idx, h_q, m_local] = len(partial_blocks)
+                mask_block_idx[seq_idx, h_q, m_local, : len(partial_blocks)] = torch.tensor(
+                    partial_blocks, device=device
+                )

build/torch-universal/copy_utils.py

@@ -0,0 +1,340 @@
+# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
+
+import math
+from typing import Optional, Type, Tuple, Callable
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Int32, Boolean, const_expr
+from cutlass.cute.nvgpu import cpasync
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass.cutlass_dsl import T, dsl_user_op
+from cutlass._mlir.dialects import llvm
+import cutlass.pipeline
+
+
+@dsl_user_op
+def cvt_copy(
+    atom: cute.CopyAtom,
+    src: cute.Tensor,
+    dst: cute.Tensor,
+    *,
+    pred: Optional[cute.Tensor] = None,
+    loc=None,
+    ip=None,
+    **kwargs,
+) -> None:
+    assert isinstance(src.iterator, cute.Pointer) and src.memspace == cute.AddressSpace.rmem
+    if const_expr(src.element_type != dst.element_type):
+        src_cvt = cute.make_fragment_like(src, dst.element_type, loc=loc, ip=ip)
+        src_cvt.store(src.load().to(dst.element_type))
+        src = src_cvt
+    cute.copy(atom, src, dst, pred=pred, loc=loc, ip=ip, **kwargs)
+
+
+@dsl_user_op
+def load_s2r(src: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor:
+    dst = cute.make_fragment_like(src, src.element_type, loc=loc, ip=ip)
+    cute.autovec_copy(src, dst, loc=loc, ip=ip)
+    return dst
+
+
+@dsl_user_op
+def get_copy_atom(
+    dtype: Type[cutlass.Numeric], num_copy_elems: int, is_async: bool = False, *, loc=None, ip=None
+) -> cute.CopyAtom:
+    num_copy_bits = const_expr(min(128, num_copy_elems * dtype.width))
+    copy_op = cpasync.CopyG2SOp() if is_async else cute.nvgpu.CopyUniversalOp()
+    return cute.make_copy_atom(copy_op, dtype, num_bits_per_copy=num_copy_bits)
+
+
+@dsl_user_op
+def make_tmem_copy(
+    tmem_copy_atom: cute.CopyAtom, num_wg: int = 1, *, loc=None, ip=None
+) -> cute.CopyAtom:
+    num_dp, num_bits, num_rep, _ = sm100_utils.get_tmem_copy_properties(tmem_copy_atom)
+    assert num_dp == 32
+    assert num_bits == 32
+    tiler_mn = (cute.make_layout((128 * num_rep * num_wg // 32, 32), stride=(32, 1)),)
+    layout_tv = cute.make_layout(
+        ((32, 4, num_wg), (num_rep, 32)), stride=((0, 1, 4 * num_rep), (4, 4 * num_rep * num_wg))
+    )
+    return cute.make_tiled_copy(tmem_copy_atom, layout_tv, tiler_mn)
+
+
+@dsl_user_op
+def copy(
+    src: cute.Tensor,
+    dst: cute.Tensor,
+    *,
+    pred: Optional[cute.Tensor] = None,
+    num_copy_elems: int = 1,
+    is_async: bool = False,
+    loc=None,
+    ip=None,
+    **kwargs,
+) -> None:
+    copy_atom = get_copy_atom(src.element_type, num_copy_elems, is_async)
+    cute.copy(copy_atom, src, dst, pred=pred, loc=loc, ip=ip, **kwargs)
+
+
+def tiled_copy_1d(
+    dtype: Type[cutlass.Numeric], num_threads: int, num_copy_elems: int = 1, is_async: bool = False
+) -> cute.TiledCopy:
+    num_copy_bits = num_copy_elems * dtype.width
+    copy_op = cpasync.CopyG2SOp() if is_async else cute.nvgpu.CopyUniversalOp()
+    copy_atom = cute.make_copy_atom(copy_op, dtype, num_bits_per_copy=num_copy_bits)
+    thr_layout = cute.make_layout(num_threads)
+    val_layout = cute.make_layout(num_copy_elems)
+    return cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout)
+
+
+def tiled_copy_2d(
+    dtype: Type[cutlass.Numeric], major_mode_size: int, num_threads: int, is_async: bool = False
+) -> cute.TiledCopy:
+    num_copy_bits = math.gcd(major_mode_size, 128 // dtype.width) * dtype.width
+    copy_elems = num_copy_bits // dtype.width
+    copy_op = cpasync.CopyG2SOp() if is_async else cute.nvgpu.CopyUniversalOp()
+    copy_atom = cute.make_copy_atom(copy_op, dtype, num_bits_per_copy=num_copy_bits)
+    gmem_threads_per_row = major_mode_size // copy_elems
+    assert num_threads % gmem_threads_per_row == 0
+    thr_layout = cute.make_ordered_layout(
+        (num_threads // gmem_threads_per_row, gmem_threads_per_row),
+        order=(1, 0),
+    )
+    val_layout = cute.make_layout((1, copy_elems))
+    return cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout)
+
+
+@dsl_user_op
+def atomic_add_fp32x4(
+    a: Float32, b: Float32, c: Float32, d: 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),
+            Float32(b).ir_value(loc=loc, ip=ip),
+            Float32(c).ir_value(loc=loc, ip=ip),
+            Float32(d).ir_value(loc=loc, ip=ip),
+        ],
+        # [gmem_ptr_i64, Float32(a).ir_value(loc=loc, ip=ip), cache_hint.ir_value()],
+        "{\n\t"
+        # ".reg .b128 abcd;\n\t"
+        # "mov.b128 abcd, {$1, $2, $3, $4};\n\t"
+        ".reg .v4 .f32 abcd;\n\t"
+        # "mov.b128 abcd, {$1, $2, $3, $4};\n\t"
+        "mov.f32 abcd.x, $1;\n\t"
+        "mov.f32 abcd.y, $2;\n\t"
+        "mov.f32 abcd.z, $3;\n\t"
+        "mov.f32 abcd.w, $4;\n\t"
+        "red.global.add.v4.f32 [$0], abcd;\n\t"
+        # "red.global.add.L2::cache_hint.v4.f32 [$0], abcd, 0x14F0000000000000;\n\t"
+        "}\n",
+        # "red.global.add.L2::cache_hint.f32 [$0], $1, 0x12F0000000000000;",
+        # "red.global.add.L2::cache_hint.f32 [$0], $1, $2;",
+        "l,f,f,f,f",
+        # "l,f,l",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+@dsl_user_op
+def set_block_rank(
+    smem_ptr: cute.Pointer, peer_cta_rank_in_cluster: Int32, *, loc=None, ip=None
+) -> Int32:
+    """Map the given smem pointer to the address at another CTA rank in the cluster."""
+    smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
+    return Int32(
+        llvm.inline_asm(
+            T.i32(),
+            [smem_ptr_i32, peer_cta_rank_in_cluster.ir_value()],
+            "mapa.shared::cluster.u32 $0, $1, $2;",
+            "=r,r,r",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
+@dsl_user_op
+def store_shared_remote_fp32x4(
+    a: Float32,
+    b: Float32,
+    c: Float32,
+    d: Float32,
+    smem_ptr: cute.Pointer,
+    mbar_ptr: cute.Pointer,
+    peer_cta_rank_in_cluster: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    remote_smem_ptr_i32 = set_block_rank(
+        smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip
+    ).ir_value()
+    remote_mbar_ptr_i32 = set_block_rank(
+        mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip
+    ).ir_value()
+    llvm.inline_asm(
+        None,
+        [
+            remote_smem_ptr_i32,
+            remote_mbar_ptr_i32,
+            Float32(a).ir_value(loc=loc, ip=ip),
+            Float32(b).ir_value(loc=loc, ip=ip),
+            Float32(c).ir_value(loc=loc, ip=ip),
+            Float32(d).ir_value(loc=loc, ip=ip),
+        ],
+        "{\n\t"
+        ".reg .v4 .f32 abcd;\n\t"
+        "mov.f32 abcd.x, $2;\n\t"
+        "mov.f32 abcd.y, $3;\n\t"
+        "mov.f32 abcd.z, $4;\n\t"
+        "mov.f32 abcd.w, $5;\n\t"
+        "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.f32 [$0], abcd, [$1];\n\t"
+        "}\n",
+        "r,r,f,f,f,f",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+@dsl_user_op
+def cpasync_bulk_g2s(
+    gmem_ptr: cute.Pointer,
+    smem_ptr: cute.Pointer,
+    tma_bar_ptr: cute.Pointer,
+    size: int | Int32,
+    *,
+    loc=None,
+    ip=None,
+):
+    gmem_ptr_i64 = gmem_ptr.toint(loc=loc, ip=ip).ir_value()
+    smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
+    mbar_ptr_i32 = tma_bar_ptr.toint(loc=loc, ip=ip).ir_value()
+    llvm.inline_asm(
+        None,
+        [gmem_ptr_i64, smem_ptr_i32, mbar_ptr_i32, Int32(size).ir_value()],
+        "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [$1], [$0], $3, [$2];",
+        "l,r,r,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+@dsl_user_op
+def cpasync_reduce_bulk_add_f32(
+    smem_ptr: cute.Pointer,
+    gmem_ptr: cute.Pointer,
+    store_bytes: int | Int32,
+    *,
+    loc=None,
+    ip=None,
+):
+    smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
+    # cache_hint = cutlass.Int64(0x14F0000000000000)  # EVICT_LAST
+    llvm.inline_asm(
+        None,
+        [gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes).ir_value()],
+        "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32 [$0], [$1], $2;",
+        "l,r,r",
+        # [gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes).ir_value(), cache_hint.ir_value()],
+        # "cp.reduce.async.bulk.global.shared::cta.bulk_group.L2::cache_hint.add.f32 [$0], [$1], $2, $3;",
+        # "l,r,r,l",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+def cpasync_bulk_get_copy_fn(
+    src_tensor: cute.Tensor,
+    dst_tensor: cute.Tensor,
+    single_stage: bool = False,
+    **kwargs,
+) -> Callable:
+    # src_is_smem = const_expr(
+    #     isinstance(src_tensor.iterator, cute.Pointer)
+    #     and src_tensor.memspace == cute.AddressSpace.smem
+    # )
+    group_rank_src = const_expr(cute.rank(src_tensor) - (1 if not single_stage else 0))
+    group_rank_dst = const_expr(cute.rank(dst_tensor) - (1 if not single_stage else 0))
+    # ((atom_v, rest_v), STAGE), ((atom_v, rest_v), RestK)
+    src = cute.group_modes(src_tensor, 0, group_rank_src)
+    dst = cute.group_modes(dst_tensor, 0, group_rank_dst)
+
+    def copy_bulk(src_idx, dst_idx, **new_kwargs):
+        size = const_expr(cute.size(src.shape[:-1]) * src.element_type.width // 8)
+        cpasync_bulk_g2s(
+            src[None, src_idx].iterator,
+            dst[None, dst_idx].iterator,
+            size=size,
+            **new_kwargs,
+            **kwargs
+        )
+
+    def copy_bulk_single_stage(**new_kwargs):
+        size = const_expr(cute.size(src.shape) * src.element_type.width // 8)
+        cpasync_bulk_g2s(src.iterator, dst.iterator, size=size, **new_kwargs, **kwargs)
+
+    return copy_bulk if const_expr(not single_stage) else copy_bulk_single_stage
+
+
+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,
+    **kwargs,
+) -> Callable:
+    src_is_smem = const_expr(
+        isinstance(src_tensor.iterator, cute.Pointer)
+        and src_tensor.memspace == cute.AddressSpace.smem
+    )
+    smem_tensor, gmem_tensor = (src_tensor, dst_tensor) if src_is_smem else (dst_tensor, src_tensor)
+    group_rank_smem = const_expr(cute.rank(smem_tensor) - (1 if not single_stage else 0))
+    group_rank_gmem = const_expr(cute.rank(gmem_tensor) - (1 if not single_stage else 0))
+    # ((atom_v, rest_v), STAGE), ((atom_v, rest_v), RestK)
+    s, g = cpasync.tma_partition(
+        atom,
+        cta_coord,
+        cta_layout,
+        cute.group_modes(smem_tensor, 0, group_rank_smem),
+        cute.group_modes(gmem_tensor, 0, group_rank_gmem),
+    )
+    if const_expr(filter_zeros):
+        s = cute.filter_zeros(s)
+        g = cute.filter_zeros(g)
+    src, dst = (s, g) if src_is_smem else (g, s)
+
+    def copy_tma(src_idx, dst_idx, **new_kwargs):
+        cute.copy(atom, src[None, src_idx], dst[None, dst_idx], **new_kwargs, **kwargs)
+
+    def copy_tma_single_stage(**new_kwargs):
+        cute.copy(atom, src, dst, **new_kwargs, **kwargs)
+
+    return (copy_tma if const_expr(not single_stage) else copy_tma_single_stage), s, g
+
+
+def tma_producer_copy_fn(copy: Callable, pipeline: cutlass.pipeline.PipelineAsync):
+    def copy_fn(src_idx, producer_state: cutlass.pipeline.PipelineState, **new_kwargs):
+        copy(
+            src_idx=src_idx,
+            dst_idx=producer_state.index,
+            tma_bar_ptr=pipeline.producer_get_barrier(producer_state),
+            **new_kwargs,
+        )
+
+    return copy_fn

build/torch-universal/cute_dsl_utils.py

@@ -0,0 +1,124 @@
+# Copyright (c) 2025, Tri Dao.
+
+import os
+import pathlib
+from typing import Tuple
+from functools import partial, lru_cache
+from dataclasses import dataclass, fields
+
+import torch
+
+try:
+    from triton.tools.disasm import extract
+except ImportError:
+    extract = None
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.base_dsl.typing import JitArgument
+from cutlass.cutlass_dsl import NumericMeta
+
+
+StaticTypes = (cutlass.Constexpr, NumericMeta, int, bool, str, float, type(None))
+
+
+load_cubin_module_data_og = cutlass.base_dsl.runtime.cuda.load_cubin_module_data
+cute_compile_og = cute.compile
+
+
+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)
+
+
+@dataclass
+class ParamsBase:
+    def __extract_mlir_values__(self):
+        all_fields = [getattr(self, field.name) for field in fields(self)]
+        non_constexpr_fields = [f for f in all_fields if not isinstance(f, StaticTypes)]
+        values, self._values_pos = [], []
+        for obj in non_constexpr_fields:
+            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):
+        all_fields = {field.name: getattr(self, field.name) for field in fields(self)}
+        constexpr_fields = {n: f for n, f in all_fields.items() if isinstance(f, StaticTypes)}
+        non_constexpr_fields = {
+            n: f for n, f in all_fields.items() if not isinstance(f, StaticTypes)
+        }
+        for (name, field), n_items in zip(non_constexpr_fields.items(), self._values_pos):
+            non_constexpr_fields[name] = cutlass.new_from_mlir_values(field, values[:n_items])
+            values = values[n_items:]
+        return self.__class__(**non_constexpr_fields, **constexpr_fields)
+
+
+@dataclass
+class ArgumentsBase(JitArgument):
+    def __c_pointers__(self):
+        all_fields = [getattr(self, field.name) for field in fields(self)]
+        non_constexpr_fields = [f for f in all_fields if not isinstance(f, StaticTypes)]
+        c_ptrs = []
+        for obj in non_constexpr_fields:
+            if hasattr(obj, "__c_pointers__"):
+                c_ptrs.extend(obj.__c_pointers__())
+        return c_ptrs
+
+    def __get_mlir_types__(self):
+        all_fields = [getattr(self, field.name) for field in fields(self)]
+        non_constexpr_fields = [f for f in all_fields if not isinstance(f, StaticTypes)]
+        types, self._values_pos = [], []
+        for obj in non_constexpr_fields:
+            if hasattr(obj, "__get_mlir_types__"):
+                obj_types = obj.__get_mlir_types__()
+                types.extend(obj_types)
+                self._values_pos.append(len(obj_types))
+            else:
+                self._values_pos.append(0)
+        return types
+
+    def __new_from_mlir_values__(self, values):
+        all_fields = {field.name: getattr(self, field.name) for field in fields(self)}
+        constexpr_fields = {n: f for n, f in all_fields.items() if isinstance(f, StaticTypes)}
+        non_constexpr_fields = {
+            n: f for n, f in all_fields.items() if not isinstance(f, StaticTypes)
+        }
+        for (name, field), n_items in zip(non_constexpr_fields.items(), self._values_pos):
+            non_constexpr_fields[name] = cutlass.new_from_mlir_values(field, values[:n_items])
+            values = values[n_items:]
+        return self.__class__(**non_constexpr_fields, **constexpr_fields)
+
+
+def 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

build/torch-universal/fast_math.py

@@ -0,0 +1,97 @@
+# Copyright (c) 2025, Tri Dao.
+
+from typing import Tuple
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32, Uint32
+from cutlass.cutlass_dsl import T, dsl_user_op
+from cutlass._mlir.dialects import llvm
+
+
+@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
+
+
+def find_log2(x: Int32) -> Int32:
+    a: Int32 = Int32(31 - clz(x))
+    return a + ((x & (x - 1)) != 0)  # Round up, add 1 if not a power of 2.
+
+
+@dsl_user_op
+def umulhi(a: Int32, b: Int32, *, loc=None, ip=None) -> Uint32:
+    return Uint32(
+        llvm.inline_asm(
+            T.i32(),
+            [Int32(a).ir_value(loc=loc, ip=ip), Int32(b).ir_value(loc=loc, ip=ip)],
+            "mul.hi.u32 $0, $1, $2;",
+            "=r,r,r",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
+class FastDivmod:
+    def __init__(
+        self, divisor: Int32, multipler: Uint32, shift_right: Uint32, *, loc=None, ip=None
+    ):
+        self.divisor = divisor
+        self.multiplier = multipler
+        self.shift_right = shift_right
+        self._loc = loc
+
+    # called by host
+    @staticmethod
+    def create(divisor: Int32, *, loc=None, ip=None) -> "FastDivmod":
+        """Construct the FastDivmod object, in host code.
+        This precomputes some values based on the divisor and is computationally expensive.
+        """
+        p = Uint32(31 + find_log2(divisor))
+        divisor_u32 = Uint32(divisor)
+        multiplier = Uint32(((cutlass.Uint64(1) << p) + divisor_u32 - 1) // divisor_u32)
+        shift_right = Uint32(p - 32)
+        return FastDivmod(divisor, multiplier, shift_right, loc=loc, ip=ip)
+
+    @cute.jit
+    def div(self, dividend: Int32) -> Int32:
+        return (
+            Int32(umulhi(dividend, self.multiplier) >> self.shift_right)
+            if self.divisor != 1
+            else dividend
+        )
+
+    def divmod(self, dividend: Int32) -> Tuple[Int32, Int32]:
+        quotient = self.div(dividend)
+        remainder = dividend - quotient * self.divisor
+        return quotient, remainder
+
+    def __extract_mlir_values__(self):
+        values, self._values_pos = [], []
+        for obj in [self.divisor, self.multiplier, self.shift_right]:
+            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.divisor, self.multiplier, self.shift_right], self._values_pos
+        ):
+            obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items]))
+            values = values[n_items:]
+        return FastDivmod(*(tuple(obj_list)), loc=self._loc)

build/torch-universal/flash_attn_cute/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+from pathlib import Path
+from types import ModuleType
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch-universal/flash_bwd.py

@@ -0,0 +1,1261 @@
+# 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/mainloop_bwd_sm80.hpp
+# from Cutlass C++ to Cute-DSL.
+import math
+from types import SimpleNamespace
+from typing import Type, Callable, Optional
+from functools import partial
+
+import cuda.bindings.driver as cuda
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.nvgpu import cpasync, warp
+from cutlass import Float32, Int32
+import cutlass.utils as utils_basic
+
+from . import ampere_helpers as sm80_utils
+from . import utils
+from .mask import AttentionMask
+from .seqlen_info import SeqlenInfoQK
+from .tile_scheduler import ParamsBase, SingleTileScheduler, SingleTileVarlenScheduler, TileSchedulerArguments
+
+
+class FlashAttentionBackwardSm80:
+    def __init__(
+        self,
+        dtype: Type[cutlass.Numeric],
+        head_dim: int,
+        head_dim_v: Optional[int] = None,
+        qhead_per_kvhead: int = 1,
+        m_block_size: int = 64,
+        n_block_size: int = 128,
+        num_stages_Q: int = 2,
+        num_stages_dO: int = 2,
+        num_threads: int = 256,
+        pack_gqa: bool = False,
+        is_causal: bool = False,
+        SdP_swapAB: bool = False,
+        dKV_swapAB: bool = False,
+        dQ_swapAB: bool = False,
+        AtomLayoutMSdP: int = 1,
+        AtomLayoutNdKV: int = 8,
+        AtomLayoutMdQ: int = 1,
+        V_in_regs: bool = False,
+    ):
+        """Initializes the configuration for a flash attention v2 kernel.
+
+        All contiguous dimensions must be at least 16 bytes aligned which indicates the head dimension
+        should be a multiple of 8.
+
+        :param head_dim: head dimension
+        :type head_dim: int
+        :param m_block_size: m block size
+        :type m_block_size: int
+        :param n_block_size: n block size
+        :type n_block_size: int
+        :param num_threads: number of threads
+        :type num_threads: int
+        :param is_causal: is causal
+        """
+        self.dtype = dtype
+        # padding head_dim to a multiple of 16 as k_block_size
+        hdim_multiple_of = 32
+        self.head_dim_padded = 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.head_dim_v_padded = 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.head_dim_padded
+        self.check_hdim_v_oob = head_dim_v != self.head_dim_v_padded
+        self.qhead_per_kvhead = qhead_per_kvhead
+        self.m_block_size = m_block_size
+        self.n_block_size = n_block_size
+        self.num_threads = num_threads
+        self.pack_gqa = pack_gqa
+        self.is_causal = is_causal
+        self.num_stages_Q = num_stages_Q
+        self.num_stages_dO = num_stages_dO
+        self.SdP_swapAB = SdP_swapAB
+        self.dKV_swapAB = dKV_swapAB
+        self.dQ_swapAB = dQ_swapAB
+        self.AtomLayoutMSdP = AtomLayoutMSdP
+        self.AtomLayoutNdKV = AtomLayoutNdKV
+        self.AtomLayoutMdQ = AtomLayoutMdQ
+        num_mma_warps = self.num_threads // cute.arch.WARP_SIZE
+        self.Mma_dKV_is_RS = AtomLayoutMSdP == 1 and AtomLayoutNdKV == num_mma_warps and SdP_swapAB and not dKV_swapAB
+        self.V_in_regs = V_in_regs
+        self.share_QV_smem = V_in_regs
+
+    @staticmethod
+    def can_implement(
+        dtype, head_dim, head_dim_v, m_block_size, n_block_size, num_stages_Q, num_stages_dO,
+        num_threads, is_causal,
+        V_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 m_block_size: m block size
+        :type m_block_size: int
+        :param n_block_size: n block size
+        :type n_block_size: 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 n_block_size % 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 = m_block_size * head_dim * num_stages_Q * 2
+        smem_usage_dO = m_block_size * head_dim_v * num_stages_dO * 2
+        smem_usage_K = n_block_size * head_dim * 2
+        smem_usage_V = n_block_size * head_dim_v * 2
+        smem_usage_QV = (smem_usage_Q + smem_usage_V) if not V_in_regs else max(smem_usage_Q, smem_usage_V)
+        smem_usage = smem_usage_QV + smem_usage_dO + smem_usage_K
+        smem_capacity = utils_basic.get_smem_capacity_in_bytes("sm_80")
+        if smem_usage > smem_capacity:
+            return False
+        return True
+
+    def _check_type(
+        self,
+        mQ_type: Type[cutlass.Numeric],
+        mK_type: Type[cutlass.Numeric],
+        mV_type: Type[cutlass.Numeric],
+        mdO_type: Type[cutlass.Numeric],
+        mLSE_type: Type[cutlass.Numeric],
+        mdPsum_type: Type[cutlass.Numeric],
+        mdQaccum_type: Type[cutlass.Numeric],
+        mdK_type: Type[cutlass.Numeric],
+        mdV_type: Type[cutlass.Numeric],
+        mCuSeqlensQ_type: Type[cutlass.Numeric] | None,
+        mCuSeqlensK_type: Type[cutlass.Numeric] | None,
+        mSeqUsedQ_type: Type[cutlass.Numeric] | None,
+        mSeqUsedK_type: Type[cutlass.Numeric] | None,
+    ):
+        if cutlass.const_expr(not (mQ_type == mK_type == mV_type == mdO_type)):
+            raise TypeError("All tensors must have the same data type")
+        if cutlass.const_expr(self.qhead_per_kvhead == 1):
+            if cutlass.const_expr(not (mdK_type == mdV_type == mQ_type)):
+                raise TypeError("mdK and mdV tensors must have the same data type as mQ")
+        else:
+            if cutlass.const_expr(not (mdK_type == mdV_type == cutlass.Float32)):
+                raise TypeError("mdKaccum and mdVaccum tensors must have the data type Float32")
+        if cutlass.const_expr(not mQ_type in [cutlass.Float16, cutlass.BFloat16]):
+            raise TypeError("Only Float16 or BFloat16 is supported")
+        if cutlass.const_expr(not mLSE_type in [cutlass.Float32]):
+            raise TypeError("LSE tensor must be Float32")
+        if cutlass.const_expr(not mdPsum_type in [cutlass.Float32]):
+            raise TypeError("dPsum tensor must be Float32")
+        if cutlass.const_expr(not mdQaccum_type in [cutlass.Float32]):
+            raise TypeError("dQaccum tensor must be Float32")
+        if cutlass.const_expr(mCuSeqlensQ_type not in [None, cutlass.Int32]):
+            raise TypeError("cuSeqlensQ tensor must be Int32")
+        if cutlass.const_expr(mCuSeqlensK_type not in [None, cutlass.Int32]):
+            raise TypeError("cuSeqlensK tensor must be Int32")
+        if cutlass.const_expr(mSeqUsedQ_type not in [None, cutlass.Int32]):
+            raise TypeError("SeqUsedQ tensor must be Int32")
+        if cutlass.const_expr(mSeqUsedK_type not in [None, cutlass.Int32]):
+            raise TypeError("SeqUsedK tensor must be Int32")
+        assert mQ_type == self.dtype
+
+    def _setup_attributes(self):
+        # ///////////////////////////////////////////////////////////////////////////////
+        # Shared memory layout: Q/K/V
+        # ///////////////////////////////////////////////////////////////////////////////
+        sQ_layout_atom = sm80_utils.get_smem_layout_atom(self.dtype, self.head_dim_padded)
+        self.sQ_layout = cute.tile_to_shape(
+            sQ_layout_atom, (self.m_block_size, self.head_dim_padded, self.num_stages_Q), (0, 1, 2),
+        )
+        sK_layout_atom = sQ_layout_atom
+        self.sK_layout = cute.tile_to_shape(
+            sK_layout_atom, (self.n_block_size, self.head_dim_padded), (0, 1),
+        )
+        sV_layout_atom = sm80_utils.get_smem_layout_atom(self.dtype, self.head_dim_v_padded)
+        self.sV_layout = cute.tile_to_shape(
+            sV_layout_atom, (self.n_block_size, self.head_dim_v_padded), (0, 1),
+        )
+        sdO_layout_atom = sV_layout_atom
+        self.sdO_layout = cute.tile_to_shape(
+            sdO_layout_atom, (self.m_block_size, self.head_dim_v_padded, self.num_stages_dO), (0, 1, 2),
+        )
+        # TODO: do we set swizzle to be 3 here explicitly?
+        sPdS_layout_atom = sm80_utils.get_smem_layout_atom(self.dtype, self.n_block_size)
+        self.sPdS_layout = cute.tile_to_shape(
+            sPdS_layout_atom, (self.m_block_size, self.n_block_size), (0, 1),
+        )
+        # We set stride to be multiple of 64 so that if ShuffleLSE, even if threads read from sLSE but out of bounds,
+        # it's still a valid smem address.
+        self.sLSE_layout = cute.make_layout(
+            (self.m_block_size, self.num_stages_Q),
+            stride=(1, cute.round_up(self.m_block_size, 64)),
+        )
+        sLSEMma_layout = cute.make_layout(
+            (self.m_block_size, self.n_block_size, self.num_stages_Q),
+            stride=(1, 0, cute.round_up(self.m_block_size, 64)),
+        )
+        sLSEMma_layout_transposed = cute.make_layout(
+            (self.n_block_size, self.m_block_size, self.num_stages_Q),
+            stride=(0, 1, cute.round_up(self.m_block_size, 64)),
+        )
+        self.sLSEMma_layout = sLSEMma_layout if not self.SdP_swapAB else sLSEMma_layout_transposed
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        # 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,
+        )
+        # tQK_layout: thread layout for QK load
+        tQK_shape_dim_1 = sQ_layout_atom.outer.shape[1] // async_copy_elems
+        assert self.num_threads % tQK_shape_dim_1 == 0, "num_threads must be divisible by tQK_shape_dim_1"
+        tQK_layout = cute.make_ordered_layout(
+            (self.num_threads // tQK_shape_dim_1, tQK_shape_dim_1), order=(1, 0),
+        )
+        # Do we need to check if we overshot kBlockM when we load Q?
+        self.is_even_m_smem_q = self.m_block_size % tQK_layout.shape[0] == 0
+        # Do we need to check if we overshot kBlockN when we load K?
+        self.is_even_n_smem_k = self.n_block_size % tQK_layout.shape[0] == 0
+        tVdO_shape_dim_1 = sV_layout_atom.outer.shape[1] // async_copy_elems
+        assert self.num_threads % tVdO_shape_dim_1 == 0, "num_threads must be divisible by tVdO_shape_dim_1"
+        tVdO_layout = cute.make_ordered_layout(
+            (self.num_threads // tVdO_shape_dim_1, tVdO_shape_dim_1), order=(1, 0),
+        )
+        # Do we need to check if we overshot kBlockN when we load V?
+        self.is_even_n_smem_v = self.n_block_size % tVdO_layout.shape[0] == 0
+        self.is_even_m_smem_do = self.m_block_size % tVdO_layout.shape[0] == 0
+
+        # Value layouts for copies
+        vQKVdO_layout = cute.make_layout((1, async_copy_elems))
+
+        # gmem_tiled_copy_QK: tiled copy for QK load
+        self.gmem_tiled_copy_QK = cute.make_tiled_copy_tv(atom_async_copy, tQK_layout, vQKVdO_layout)
+        self.gmem_tiled_copy_VdO = cute.make_tiled_copy_tv(atom_async_copy, tVdO_layout, vQKVdO_layout)
+        self.gmem_tiled_copy_dK = cute.make_tiled_copy_tv(atom_universal_copy, tQK_layout, vQKVdO_layout)
+        self.gmem_tiled_copy_dV = cute.make_tiled_copy_tv(atom_universal_copy, tVdO_layout, vQKVdO_layout)
+        async_copy_elems_accum = universal_copy_bits // cutlass.Float32.width
+
+        # I think we wouldn't require this with smarter padding
+        if cutlass.const_expr(not self.varlen_q):
+            async_copy_elems_accum = universal_copy_bits // cutlass.Float32.width
+            atom_async_copy_accum = cute.make_copy_atom(
+                cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
+                cutlass.Float32,
+                num_bits_per_copy=universal_copy_bits,
+            )
+        else:
+            async_copy_elems_accum = 1
+            atom_async_copy_accum = cute.make_copy_atom(
+                cute.nvgpu.CopyUniversalOp(),
+                cutlass.Float32,
+                num_bits_per_copy=cutlass.Float32.width,
+            )
+        self.gmem_tiled_copy_LSE = cute.make_tiled_copy_tv(
+            atom_async_copy_accum,
+            cute.make_layout(self.num_threads),
+            cute.make_layout(async_copy_elems_accum),
+        )
+        self.gmem_tiled_copy_dQaccum = cute.make_tiled_copy_tv(
+            cute.make_copy_atom(
+                cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=cutlass.Float32.width
+            ),
+            cute.make_layout(self.num_threads),
+            cute.make_layout(1)
+        )
+        if cutlass.const_expr(self.qhead_per_kvhead > 1):
+            self.gmem_tiled_copy_dK = self.gmem_tiled_copy_dQaccum
+            self.gmem_tiled_copy_dV = self.gmem_tiled_copy_dQaccum
+
+    def _get_tiled_mma(self):
+        num_mma_warps = self.num_threads // 32
+        AtomLayoutSdP = (self.AtomLayoutMSdP, num_mma_warps // self.AtomLayoutMSdP, 1) if cutlass.const_expr(not self.SdP_swapAB) else (num_mma_warps // self.AtomLayoutMSdP, self.AtomLayoutMSdP, 1)
+        tiled_mma_sdp = cute.make_tiled_mma(
+            warp.MmaF16BF16Op(self.dtype, cutlass.Float32, (16, 8, 16)),
+            AtomLayoutSdP,
+            permutation_mnk=(AtomLayoutSdP[0] * 16, AtomLayoutSdP[1] * 16, 16),
+        )
+        AtomLayoutdKV = (self.AtomLayoutNdKV, num_mma_warps // self.AtomLayoutNdKV, 1) if cutlass.const_expr(not self.dKV_swapAB) else (num_mma_warps // self.AtomLayoutNdKV, self.AtomLayoutNdKV, 1)
+        tiled_mma_dkv = cute.make_tiled_mma(
+            warp.MmaF16BF16Op(self.dtype, cutlass.Float32, (16, 8, 16)),
+            AtomLayoutdKV,
+            permutation_mnk=(AtomLayoutdKV[0] * 16, AtomLayoutdKV[1] * 16, 16),
+        )
+        AtomLayoutdQ = (self.AtomLayoutMdQ, num_mma_warps // self.AtomLayoutMdQ, 1) if cutlass.const_expr(not self.dQ_swapAB) else (num_mma_warps // self.AtomLayoutMdQ, self.AtomLayoutMdQ, 1)
+        tiled_mma_dq = cute.make_tiled_mma(
+            warp.MmaF16BF16Op(self.dtype, cutlass.Float32, (16, 8, 16)),
+            AtomLayoutdQ,
+            permutation_mnk=(AtomLayoutdQ[0] * 16, AtomLayoutdQ[1] * 16, 16),
+        )
+        return tiled_mma_sdp, tiled_mma_dkv, tiled_mma_dq
+
+    def _get_shared_storage_cls(self):
+        sQ_struct, sK_struct, sV_struct, sdO_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, self.sdO_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]
+        sLSE_struct, sdPsum_struct = [
+            cute.struct.Align[cute.struct.MemRange[cutlass.Float32, cute.cosize(layout)], 128]
+            for layout in (self.sLSE_layout, self.sLSE_layout)
+        ]
+        sP_struct, sdS_struct = [
+            cute.struct.Align[cute.struct.MemRange[self.dtype, cute.cosize(layout)], 128]
+            for layout in (self.sPdS_layout, self.sPdS_layout)
+        ]
+
+        @cute.struct
+        class SharedStorageSeparateQV:
+            sK: sK_struct
+            sV: sV_struct
+            sQ: sQ_struct
+            sdO: sdO_struct
+            sLSE: sLSE_struct
+            sdPsum: sdPsum_struct
+            sP: sP_struct
+            sdS: sdS_struct
+            # TODO: the case where there's no sP
+
+        @cute.struct
+        class SharedStorageSharedQV:
+            sK: sK_struct
+            sV: sV_struct
+            sQ: sQV_struct
+            sdO: sdO_struct
+            sLSE: sLSE_struct
+            sdPsum: sdPsum_struct
+            sP: sP_struct
+            sdS: sdS_struct
+
+        return SharedStorageSeparateQV if cutlass.const_expr(not self.share_QV_smem) else SharedStorageSharedQV
+
+    @cute.jit
+    def __call__(
+        self,
+        mQ: cute.Tensor,
+        mK: cute.Tensor,
+        mV: cute.Tensor,
+        mdO: cute.Tensor,
+        mLSE: cute.Tensor,
+        mdPsum: cute.Tensor,
+        mdQaccum: cute.Tensor,
+        mdK: cute.Tensor,
+        mdV: cute.Tensor,
+        softmax_scale: cutlass.Float32,
+        stream: cuda.CUstream,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        mSeqUsedQ: Optional[cute.Tensor] = None,
+        mSeqUsedK: Optional[cute.Tensor] = None,
+        softcap: Float32 | float | None = None,
+        window_size_left: Int32 | int | None = None,
+        window_size_right: Int32 | int | None = None,
+        mdQ_semaphore: Optional[cute.Tensor] = None,
+    ):
+        assert mdQ_semaphore is None, "semaphore not supported yet"
+        # Get the data type and check if it is fp16 or bf16
+        self._check_type(*(t.element_type if t is not None else None
+                           for t in (mQ, mK, mV, mdO, mLSE, mdPsum, mdQaccum, mdK, mdV, mCuSeqlensQ, mCuSeqlensK, mSeqUsedQ, mSeqUsedK)))
+        # Assume all strides are divisible by 128 bits except the last stride
+        new_stride = lambda t: (*(cute.assume(s, divby=128 // t.element_type.width) for s in t.stride[:-1]), t.stride[-1])
+        mQ, mK, mV, mdO, mLSE, mdPsum, mdQaccum, mdK, mdV = [cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t))) if t is not None else None for t in (mQ, mK, mV, mdO, mLSE, mdPsum, mdQaccum, mdK, mdV)]
+        self.varlen_q = (mCuSeqlensQ is not None)
+        self._setup_attributes()
+        SharedStorage = self._get_shared_storage_cls()
+        tiled_mma_sdp, tiled_mma_dkv, tiled_mma_dq = self._get_tiled_mma()
+
+        num_head = mQ.shape[1] if cutlass.const_expr(mCuSeqlensQ is not None) else mQ.shape[2]
+
+        if cutlass.const_expr(mCuSeqlensK is not None):
+            TileScheduler = SingleTileVarlenScheduler
+            num_batch = mCuSeqlensK.shape[0] - 1
+        else:
+            TileScheduler = SingleTileScheduler
+            num_batch = mK.shape[0]
+
+        # Uses seqlen k, etc. since main bwd kernel's blocks are over n 
+        tile_sched_args = TileSchedulerArguments(
+            num_block=cute.ceil_div(mK.shape[1], self.n_block_size),
+            num_head=num_head,
+            num_batch=num_batch,
+            num_splits=1,
+            seqlen_k=0,
+            headdim=mK.shape[2],
+            headdim_v=mV.shape[2],
+            total_q=mK.shape[0],
+            tile_shape_mn=(self.n_block_size, self.m_block_size),
+            qhead_per_kvhead_packgqa=self.qhead_per_kvhead if cutlass.const_expr(self.pack_gqa) else 1,
+            mCuSeqlensQ=mCuSeqlensK,
+            mSeqUsedQ=mSeqUsedK,
+        )
+        
+        tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
+        grid_dim = TileScheduler.get_grid_shape(tile_sched_params)
+
+        softmax_scale_log2 = softmax_scale * math.log2(math.e)
+        self.kernel(
+            mQ,
+            mK,
+            mV,
+            mdO,
+            mLSE,
+            mdPsum,
+            mdQaccum,
+            mdK,
+            mdV,
+            mCuSeqlensQ,
+            mCuSeqlensK,
+            mSeqUsedQ,
+            mSeqUsedK,
+            softmax_scale,
+            softmax_scale_log2,
+            self.sQ_layout,
+            self.sK_layout,
+            self.sV_layout,
+            self.sdO_layout,
+            self.sPdS_layout,
+            self.sLSE_layout,
+            self.sLSEMma_layout,
+            self.gmem_tiled_copy_QK,
+            self.gmem_tiled_copy_VdO,
+            self.gmem_tiled_copy_dK,
+            self.gmem_tiled_copy_dV,
+            self.gmem_tiled_copy_LSE,
+            self.gmem_tiled_copy_dQaccum,
+            tiled_mma_sdp,
+            tiled_mma_dkv,
+            tiled_mma_dq,
+            SharedStorage,
+            tile_sched_params,
+            TileScheduler,
+        ).launch(
+            grid=grid_dim,
+            block=[self.num_threads, 1, 1],
+            smem=SharedStorage.size_in_bytes(),
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        mQ: cute.Tensor,
+        mK: cute.Tensor,
+        mV: cute.Tensor,
+        mdO: cute.Tensor,
+        mLSE: cute.Tensor,
+        mdPsum: cute.Tensor,
+        mdQaccum: cute.Tensor,
+        mdK: cute.Tensor,
+        mdV: cute.Tensor,
+        mCuSeqlensQ: Optional[cute.Tensor],
+        mCuSeqlensK: Optional[cute.Tensor],
+        mSeqUsedQ: Optional[cute.Tensor],
+        mSeqUsedK: Optional[cute.Tensor],
+        softmax_scale: cutlass.Float32,
+        softmax_scale_log2: cutlass.Float32,
+        sQ_layout: cute.ComposedLayout,
+        sK_layout: cute.ComposedLayout,
+        sV_layout: cute.ComposedLayout,
+        sdO_layout: cute.ComposedLayout,
+        sPdS_layout: cute.ComposedLayout,
+        sLSE_layout: cute.Layout,
+        sLSEMma_layout: cute.Layout,
+        gmem_tiled_copy_QK: cute.TiledCopy,
+        gmem_tiled_copy_VdO: cute.TiledCopy,
+        gmem_tiled_copy_dK: cute.TiledCopy,
+        gmem_tiled_copy_dV: cute.TiledCopy,
+        gmem_tiled_copy_LSE: cute.TiledCopy,
+        gmem_tiled_copy_dQaccum: cute.TiledCopy,
+        tiled_mma_sdp: cute.TiledMma,
+        tiled_mma_dkv: cute.TiledMma,
+        tiled_mma_dq: cute.TiledMma,
+        SharedStorage: cutlass.Constexpr,
+        tile_sched_params: ParamsBase,
+        TileScheduler: cutlass.Constexpr[Callable],
+    ):
+        # Thread index, block index
+        tidx, _, _ = cute.arch.thread_idx()
+
+        tile_scheduler = TileScheduler.create(tile_sched_params)
+        work_tile = tile_scheduler.initial_work_tile_info()
+
+        n_block, head_idx, batch_idx, _ = work_tile.tile_idx
+
+        if work_tile.is_valid_tile:
+            seqlen = SeqlenInfoQK.create(batch_idx, mQ.shape[1], mK.shape[1], mCuSeqlensQ=mCuSeqlensQ, mCuSeqlensK=mCuSeqlensK, mSeqUsedQ=mSeqUsedQ, mSeqUsedK=mSeqUsedK)
+
+            m_block_max = cute.ceil_div(seqlen.seqlen_q, self.m_block_size)
+            m_block_min = 0
+            if cutlass.const_expr(self.is_causal):
+                m_block_min = max(
+                    (n_block * self.n_block_size + seqlen.seqlen_q - seqlen.seqlen_k) // self.m_block_size,
+                    m_block_min,
+                )
+            # TODO: return early if m_block_max == 0
+
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Get the appropriate tiles for this thread block.
+            # ///////////////////////////////////////////////////////////////////////////////
+            blkQ_shape = (self.m_block_size, self.head_dim_padded)
+            blkK_shape = (self.n_block_size, self.head_dim_padded)
+            blkV_shape = (self.n_block_size, self.head_dim_v_padded)
+            blkdO_shape = (self.m_block_size, self.head_dim_v_padded)
+
+            if cutlass.const_expr(not seqlen.has_cu_seqlens_q):
+                mQ_cur = mQ[batch_idx, None, head_idx, None]
+                mLSE_cur = mLSE[batch_idx, head_idx, None]
+                mdO_cur = mdO[batch_idx, None, head_idx, None]
+                mdPsum_cur = mdPsum[batch_idx, head_idx, None]
+                mdQaccum_cur = mdQaccum[batch_idx, head_idx, None]
+            else:
+                padded_offset_q = seqlen.offset_q + batch_idx * self.m_block_size
+                mQ_cur = cute.domain_offset((seqlen.offset_q, 0), mQ[None, head_idx, None])
+                mLSE_cur = cute.domain_offset((padded_offset_q,), mLSE[head_idx, None])
+                mdO_cur = cute.domain_offset((seqlen.offset_q, 0), mdO[None, head_idx, None])
+                mdPsum_cur = cute.domain_offset((padded_offset_q,), mdPsum[head_idx, None])
+                mdQaccum_cur = cute.domain_offset((padded_offset_q * self.head_dim_padded,), mdQaccum[head_idx, None])
+            head_idx_kv = head_idx // self.qhead_per_kvhead if cutlass.const_expr(not self.pack_gqa) else head_idx
+
+            if cutlass.const_expr(not seqlen.has_cu_seqlens_k):
+                mK_cur, mV_cur = [t[batch_idx, None, head_idx_kv, None] for t in (mK, mV)]
+            else:
+                mK_cur, mV_cur = [cute.domain_offset((seqlen.offset_k, 0), t[None, head_idx_kv, None]) for t in (mK, mV)]
+
+            # (m_block_size, head_dim, m_block)
+            gQ = cute.local_tile(mQ_cur, blkQ_shape, (None, 0))
+            # (n_block_size, head_dim)
+            gK = cute.local_tile(mK_cur, blkK_shape, (n_block, 0))
+            # (n_block_size, head_dim_v)
+            gV = cute.local_tile(mV_cur, blkV_shape, (n_block, 0))
+            # (m_block_size, head_dim_v, m_block)
+            gdO = cute.local_tile(mdO_cur, blkdO_shape, (None, 0))
+            gLSE = cute.local_tile(mLSE_cur, (self.m_block_size,), (None,))
+            gdPsum = cute.local_tile(mdPsum_cur, (self.m_block_size,), (None,))
+            gdQaccum = cute.local_tile(mdQaccum_cur, (self.m_block_size * self.head_dim_padded,), (None,))
+
+            # ///////////////////////////////////////////////////////////////////////////////
+            # 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 cutlass.const_expr(not self.share_QV_smem):
+                sV = storage.sV.get_tensor(sV_layout)
+            else:
+                sV = cute.make_tensor(cute.recast_ptr(sQ.iterator, dtype=self.dtype), sV_layout)
+            sdO = storage.sdO.get_tensor(sdO_layout)
+            sP = storage.sP.get_tensor(sPdS_layout)
+            sdS = storage.sdS.get_tensor(sPdS_layout)
+            sLSE = storage.sLSE.get_tensor(sLSE_layout)
+            sdPsum = storage.sdPsum.get_tensor(sLSE_layout)
+            sLSEMma = storage.sLSE.get_tensor(sLSEMma_layout)
+            sdPsumMma = storage.sdPsum.get_tensor(sLSEMma_layout)
+
+            # Transpose view of tensors for tiled mma
+            sQt, sdOt, sKt, sPt, sdSt = [utils.transpose_view(t) for t in (sQ, sdO, sK, sP, sdS)]
+
+            gmem_thr_copy_QK = gmem_tiled_copy_QK.get_slice(tidx)
+            gmem_thr_copy_VdO = gmem_tiled_copy_VdO.get_slice(tidx)
+            gmem_thr_copy_lse = gmem_tiled_copy_LSE.get_slice(tidx)
+            gmem_thr_copy_dQaccum = gmem_tiled_copy_dQaccum.get_slice(tidx)
+            # (CPY_Atom, CPY_M, CPY_K, m_block)
+            tQgQ = gmem_thr_copy_QK.partition_S(gQ)
+            tQsQ = gmem_thr_copy_QK.partition_D(sQ)
+            # (CPY_Atom, CPY_N, CPY_K)
+            tKgK = gmem_thr_copy_QK.partition_S(gK)
+            tKsK = gmem_thr_copy_QK.partition_D(sK)
+            # (CPY_Atom, CPY_N, CPY_K)
+            tVgV = gmem_thr_copy_VdO.partition_S(gV)
+            tVsV = gmem_thr_copy_VdO.partition_D(sV)
+            # (CPY_Atom, CPY_M, CPY_K, m_block)
+            tdOgdO = gmem_thr_copy_VdO.partition_S(gdO)
+            tdOsdO = gmem_thr_copy_VdO.partition_D(sdO)
+            tLSEgLSE = gmem_thr_copy_lse.partition_S(gLSE)
+            tLSEsLSE = gmem_thr_copy_lse.partition_D(sLSE)
+            tLSEgdPsum = gmem_thr_copy_lse.partition_S(gdPsum)
+            tLSEsdPsum = gmem_thr_copy_lse.partition_D(sdPsum)
+            tdQgdQaccum = gmem_thr_copy_dQaccum.partition_S(gdQaccum)
+
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Tile MMA compute thread partitions and allocate accumulators
+            # ///////////////////////////////////////////////////////////////////////////////
+            thr_mma_sdp = tiled_mma_sdp.get_slice(tidx)
+            thr_mma_dkv = tiled_mma_dkv.get_slice(tidx)
+            thr_mma_dq = tiled_mma_dq.get_slice(tidx)
+            acc_shape_dK = thr_mma_dkv.partition_shape_C((self.n_block_size, self.head_dim_padded))
+            acc_shape_dV = thr_mma_dkv.partition_shape_C((self.n_block_size, self.head_dim_v_padded))
+            acc_dK = cute.make_fragment(acc_shape_dK, cutlass.Float32)
+            acc_dV = cute.make_fragment(acc_shape_dV, cutlass.Float32)
+            acc_dK.fill(0.0)
+            acc_dV.fill(0.0)
+
+            tSrQ = utils.mma_make_fragment_A(sQ[None, None, 0], thr_mma_sdp, swapAB=self.SdP_swapAB)
+            tSrK = utils.mma_make_fragment_B(sK, thr_mma_sdp, swapAB=self.SdP_swapAB)
+            tdPrdO = utils.mma_make_fragment_A(sdO[None, None, 0], thr_mma_sdp, swapAB=self.SdP_swapAB)
+            tdPrV = utils.mma_make_fragment_B(sV, thr_mma_sdp, swapAB=self.SdP_swapAB)
+            tdVrP = utils.mma_make_fragment_A(sPt, thr_mma_dkv, swapAB=self.dKV_swapAB)
+            tdVrdO = utils.mma_make_fragment_B(sdOt[None, None, 0], thr_mma_dkv, swapAB=self.dKV_swapAB)
+            tdKrdS = utils.mma_make_fragment_A(sdSt, thr_mma_dkv, swapAB=self.dKV_swapAB)
+            tdKrQ = utils.mma_make_fragment_B(sQt[None, None, 0], thr_mma_dkv, swapAB=self.dKV_swapAB)
+            tdQrdS = utils.mma_make_fragment_A(sdS, thr_mma_dq, swapAB=self.dQ_swapAB)
+            tdQrK = utils.mma_make_fragment_B(sKt, thr_mma_dq, swapAB=self.dQ_swapAB)
+
+            LSEslice = (None, 0, None) if cutlass.const_expr(not self.SdP_swapAB) else (0, None, None)
+            tSsLSEMma = utils.make_acc_tensor_mn_view(thr_mma_sdp.partition_C(sLSEMma))[LSEslice]
+            tSsdPsumMma = utils.make_acc_tensor_mn_view(thr_mma_sdp.partition_C(sdPsumMma))[LSEslice]
+
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Smem copy atom tiling
+            # ///////////////////////////////////////////////////////////////////////////////
+            smem_copy_atom = cute.make_copy_atom(
+                warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), self.dtype,
+            )
+            smem_copy_atom_transposed = cute.make_copy_atom(
+                warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4), self.dtype,
+            )
+            smem_thr_copy_QdO = utils.make_tiled_copy_A(
+                smem_copy_atom, tiled_mma_sdp, swapAB=self.SdP_swapAB
+            ).get_slice(tidx)
+            smem_thr_copy_KV = utils.make_tiled_copy_B(
+                smem_copy_atom, tiled_mma_sdp, swapAB=self.SdP_swapAB
+            ).get_slice(tidx)
+            # TODO: should this be smem_copy_atom_transposed?
+            smem_thr_copy_PdSt = utils.make_tiled_copy_A(
+                smem_copy_atom_transposed, tiled_mma_dkv, swapAB=self.dKV_swapAB
+            ).get_slice(tidx)
+            smem_thr_copy_QdOt = utils.make_tiled_copy_B(
+                smem_copy_atom_transposed, tiled_mma_dkv, swapAB=self.dKV_swapAB
+            ).get_slice(tidx)
+            smem_thr_copy_dS = utils.make_tiled_copy_A(
+                smem_copy_atom, tiled_mma_dq, swapAB=self.dQ_swapAB
+            ).get_slice(tidx)
+            smem_thr_copy_Kt = utils.make_tiled_copy_B(
+                smem_copy_atom_transposed, tiled_mma_dq, swapAB=self.dQ_swapAB
+            ).get_slice(tidx)
+            # TODO: what's the number of bits? What if SdP_swapAB
+            r2s_thr_copy_PdS = cute.make_tiled_copy_C(
+                cute.make_copy_atom(
+                    cute.nvgpu.CopyUniversalOp(), self.dtype, num_bits_per_copy=2 * self.dtype.width
+                ),
+                tiled_mma_sdp,
+            ).get_slice(tidx)
+
+            tSsQ = smem_thr_copy_QdO.partition_S(sQ)
+            tdPsdO = smem_thr_copy_QdO.partition_S(sdO)
+            tSsK = smem_thr_copy_KV.partition_S(sK)
+            tdPsV = smem_thr_copy_KV.partition_S(sV)
+            tdVsPt = smem_thr_copy_PdSt.partition_S(sPt)
+            tdKsdSt = smem_thr_copy_PdSt.partition_S(sdSt)
+            tdVsdOt = smem_thr_copy_QdOt.partition_S(sdOt)
+            tdKsQt = smem_thr_copy_QdOt.partition_S(sQt)
+            tdQsdS = smem_thr_copy_dS.partition_S(sdS)
+            tdQsKt = smem_thr_copy_Kt.partition_S(sKt)
+            tPsP = r2s_thr_copy_PdS.partition_D(sP)
+            tdSsdS = r2s_thr_copy_PdS.partition_D(sdS)
+
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Predicate: Mark indices that need to copy when problem_shape isn't a multiple
+            # of tile_shape
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Construct identity layout for KV
+            cQ = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded))
+            tQcQ = gmem_thr_copy_QK.partition_S(cQ)
+            t0QcQ = gmem_thr_copy_QK.get_slice(0).partition_S(cQ)
+            if cutlass.const_expr(self.head_dim_padded == self.head_dim_v_padded):
+                tdOcdO = tQcQ
+                t0dOcdO = t0QcQ
+            else:
+                cdO = cute.make_identity_tensor((self.m_block_size, self.head_dim_v_padded))
+                tdOcdO = gmem_thr_copy_VdO.partition_S(cdO)
+                t0dOcdO = gmem_thr_copy_VdO.get_slice(0).partition_S(cdO)
+            cLSE = cute.make_identity_tensor((self.m_block_size,))
+            tLSEcLSE = gmem_thr_copy_lse.partition_S(cLSE)
+
+            # 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.
+
+            d_head = mQ.shape[cute.rank(mQ) - 1]
+            d_head_v = mdO.shape[cute.rank(mdO) - 1]
+
+            tQpQ = utils.predicate_k(tQcQ, limit=d_head)
+            if cutlass.const_expr(self.same_hdim_kv):
+                tdOpdO = tQpQ
+            else:
+                tdOpdO = utils.predicate_k(tdOcdO, limit=d_head_v)
+
+            # group parameters for compute_one_m_block
+            mma_params = SimpleNamespace(
+                thr_mma_sdp=thr_mma_sdp, thr_mma_dkv=thr_mma_dkv, thr_mma_dq=thr_mma_dq,
+                tSrQ=tSrQ, tSrK=tSrK, tdPrdO=tdPrdO, tdPrV=tdPrV,
+                tdVrP=tdVrP, tdVrdO=tdVrdO, tdKrdS=tdKrdS, tdKrQ=tdKrQ,
+                tdQrdS=tdQrdS, tdQrK=tdQrK,
+                acc_dK=acc_dK, acc_dV=acc_dV,
+            )
+            smem_copy_params = SimpleNamespace(
+                smem_thr_copy_QdO=smem_thr_copy_QdO,
+                smem_thr_copy_KV=smem_thr_copy_KV,
+                smem_thr_copy_PdSt=smem_thr_copy_PdSt,
+                smem_thr_copy_QdOt=smem_thr_copy_QdOt,
+                smem_thr_copy_dS=smem_thr_copy_dS,
+                smem_thr_copy_Kt=smem_thr_copy_Kt,
+                r2s_thr_copy_PdS=r2s_thr_copy_PdS,
+                tSsQ=tSsQ, tSsK=tSsK, tdPsdO=tdPsdO, tdPsV=tdPsV,
+                tSsLSEMma=tSsLSEMma, tSsdPsumMma=tSsdPsumMma,
+                tPsP=tPsP, tdSsdS=tdSsdS,
+                tdVsPt=tdVsPt, tdVsdOt=tdVsdOt, tdKsdSt=tdKsdSt, tdKsQt=tdKsQt,
+                tdQsdS=tdQsdS, tdQsKt=tdQsKt,
+            )
+            gmem_copy_params = SimpleNamespace(
+                gmem_thr_copy_dQaccum=gmem_thr_copy_dQaccum, tdQgdQaccum=tdQgdQaccum
+            )
+            load_Q_LSE = partial(
+                self.load_Q_LSE, gmem_tiled_copy_QK, gmem_tiled_copy_LSE,
+                tQgQ, tQsQ, tQcQ, t0QcQ, tQpQ,
+                tLSEgLSE, tLSEsLSE, tLSEcLSE, seqlen=seqlen.seqlen_q
+            )
+            load_dO_dPsum = partial(
+                self.load_dO_dPsum, gmem_tiled_copy_VdO, gmem_tiled_copy_LSE,
+                tdOgdO, tdOsdO, tdOcdO, t0dOcdO, tdOpdO,
+                tLSEgdPsum, tLSEsdPsum, tLSEcLSE, seqlen=seqlen.seqlen_q
+            )
+            compute_one_m_block = partial(
+                self.compute_one_m_block, mma_params=mma_params,
+                smem_copy_params=smem_copy_params, gmem_copy_params=gmem_copy_params,
+                load_Q_LSE=load_Q_LSE, load_dO_dPsum=load_dO_dPsum,
+                m_block_max=m_block_max,
+                softmax_scale_log2=softmax_scale_log2,
+            )
+
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Prologue
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Start async loads of the last mn-tile, where we take care of the mn residue
+            self.load_V(gmem_thr_copy_VdO, tVgV, tVsV, n_block, seqlen=seqlen.seqlen_k,
+                        headdim=d_head_v)
+            if cutlass.const_expr(self.V_in_regs):
+                cute.arch.cp_async_commit_group()
+            self.load_K(gmem_thr_copy_QK, tKgK, tKsK, n_block, seqlen=seqlen.seqlen_k,
+                        headdim=d_head)
+            cute.arch.cp_async_commit_group()
+
+            if cutlass.const_expr(self.V_in_regs):
+                cute.arch.cp_async_wait_group(1)
+                cute.arch.barrier()
+                tdPrV_copy_view = smem_thr_copy_KV.retile(tdPrV)
+                cute.copy(smem_thr_copy_KV, tdPsV, tdPrV_copy_view)
+                # Sync to avoid loading Q to smem_q, which overlaps with smem_v
+                cute.arch.barrier()
+
+            m_block = m_block_min
+            assert self.num_stages_Q >= self.num_stages_dO
+            for stage in cutlass.range_constexpr(self.num_stages_Q):
+                if cutlass.const_expr(self.num_stages_Q == 1 or stage < self.num_stages_Q - 1):
+                    if stage == 0 or m_block + stage < m_block_max:
+                        load_Q_LSE(m_block + stage, smem_pipe_write_q=stage)
+                    cute.arch.cp_async_commit_group()
+                if cutlass.const_expr(stage < self.num_stages_dO):
+                    if stage == 0 or m_block + stage < m_block_max:
+                        load_dO_dPsum(m_block + stage, smem_pipe_write_q=stage)
+                    cute.arch.cp_async_commit_group()
+
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Mainloop
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Start processing of the first n-block.
+            mask = AttentionMask(self.m_block_size, self.n_block_size, seqlen.seqlen_q, seqlen.seqlen_k)
+            mask_fn = partial(
+                mask.apply_mask, n_block=n_block, thr_mma=thr_mma_sdp,
+                mask_seqlen=True, mask_causal=self.is_causal
+            )
+            smem_pipe_read_q = cutlass.Int32(0)
+            smem_pipe_read_do = cutlass.Int32(0)
+            smem_pipe_write_q = cutlass.Int32(self.num_stages_Q - 1)
+            smem_pipe_write_do = cutlass.Int32(0)
+            for m_tile in cutlass.range(m_block_min, m_block_max, unroll=1):
+                compute_one_m_block(
+                    m_tile, smem_pipe_read_q, smem_pipe_read_do, smem_pipe_write_q, smem_pipe_write_do,
+                    mask_fn=mask_fn,
+                )
+                smem_pipe_read_q = self.advance_pipeline(smem_pipe_read_q, self.num_stages_Q)
+                smem_pipe_read_do = self.advance_pipeline(smem_pipe_read_do, self.num_stages_dO)
+                smem_pipe_write_q = self.advance_pipeline(smem_pipe_write_q, self.num_stages_Q)
+                smem_pipe_write_do = self.advance_pipeline(smem_pipe_write_do, self.num_stages_dO)
+
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Epilogue
+            # ///////////////////////////////////////////////////////////////////////////////
+            # If GQA, we scale dK in the postprocessing kernel instead
+            if cutlass.const_expr(self.qhead_per_kvhead == 1):
+                acc_dK.store(acc_dK.load() * softmax_scale)
+            # reuse sK and sV data iterator
+            sdK = cute.make_tensor(sK.iterator, sK_layout)
+            sdV = cute.make_tensor(sV.iterator, sV_layout)
+            self.epilogue(
+                acc_dK, acc_dV, mdK, mdV, sdK, sdV,
+                gmem_tiled_copy_dK, gmem_tiled_copy_dV, tiled_mma_dkv,
+                tidx, n_block, head_idx, batch_idx, seqlen, d_head, d_head_v
+            )
+
+    @cute.jit
+    def compute_one_m_block(
+        self,
+        m_block: cutlass.Int32,
+        smem_pipe_read_q: cutlass.Int32,
+        smem_pipe_read_do: cutlass.Int32,
+        smem_pipe_write_q: cutlass.Int32,
+        smem_pipe_write_do: cutlass.Int32,
+        mma_params: SimpleNamespace,
+        smem_copy_params: SimpleNamespace,
+        gmem_copy_params: SimpleNamespace,
+        load_Q_LSE: Callable,
+        load_dO_dPsum: Callable,
+        m_block_max: cutlass.Int32,
+        softmax_scale_log2: cutlass.Float32,
+        mask_fn: Optional[Callable] = None,
+    ):
+        def load_Q_next():
+            m_block_next = m_block + (self.num_stages_Q - 1 if cutlass.const_expr(self.num_stages_Q > 1) else 1)
+            if m_block_next < m_block_max:
+                load_Q_LSE(m_block_next, smem_pipe_write_q)
+            cute.arch.cp_async_commit_group()
+
+        def load_dO_next():
+            if m_block + self.num_stages_dO < m_block_max:
+                load_dO_dPsum(m_block + self.num_stages_dO, smem_pipe_write_do)
+            cute.arch.cp_async_commit_group()
+
+        # MMA S
+        acc_shape_SdP = mma_params.thr_mma_sdp.partition_shape_C(
+            (self.m_block_size, self.n_block_size) if cutlass.const_expr(not self.SdP_swapAB) else (self.n_block_size, self.m_block_size)
+        )
+        acc_S = cute.make_fragment(acc_shape_SdP, cutlass.Float32)
+        acc_S.fill(0.0)
+        cute.arch.cp_async_wait_group(1 if cutlass.const_expr(self.num_stages_Q > 1) else 0)
+        cute.arch.barrier()
+        sm80_utils.gemm(
+            mma_params.thr_mma_sdp, acc_S, mma_params.tSrQ, mma_params.tSrK,
+            smem_copy_params.tSsQ[None, None, None, smem_pipe_read_q if cutlass.const_expr(self.num_stages_Q > 1) else 0],
+            smem_copy_params.tSsK,
+            smem_copy_params.smem_thr_copy_QdO, smem_copy_params.smem_thr_copy_KV,
+            swap_AB=self.SdP_swapAB,
+        )
+        tLSErLSE = cute.make_fragment_like(smem_copy_params.tSsLSEMma[None, 0])
+        cute.autovec_copy(
+            smem_copy_params.tSsLSEMma[None, smem_pipe_read_q if cutlass.const_expr(self.num_stages_Q > 1) else 0], tLSErLSE
+        )
+        if cutlass.const_expr(mask_fn is not None):
+            mask_fn(acc_S, m_block=m_block)
+        acc_S_mn = utils.make_acc_tensor_mn_view(acc_S)
+        bidx = 0
+        # if cute.arch.thread_idx()[0] == 0 and cute.arch.block_idx()[0] == bidx: cute.print_tensor(acc_S_mn)
+        # if cute.arch.thread_idx()[0] == 0 and cute.arch.block_idx()[0] == 1: cute.print_tensor(tLSErLSE)
+        assert cute.size(acc_S_mn, mode=[0]) == cute.size(tLSErLSE)
+        for r in cutlass.range(cute.size(acc_S_mn, mode=[0]), unroll_full=True):
+            acc_S_mn[r, None].store(utils.exp2f(acc_S_mn[r, None].load() * softmax_scale_log2 - tLSErLSE[r]))
+        # if cute.arch.thread_idx()[0] == 0 and cute.arch.block_idx()[0] == bidx: cute.print_tensor(acc_S_mn)
+
+        # MMA dP
+        acc_dP = cute.make_fragment(acc_shape_SdP, cutlass.Float32)
+        acc_dP.fill(0.0)
+        cute.arch.cp_async_wait_group(1 if cutlass.const_expr(self.num_stages_dO > 1) else 0)
+        cute.arch.barrier()
+        sm80_utils.gemm(
+            mma_params.thr_mma_sdp, acc_dP, mma_params.tdPrdO, mma_params.tdPrV,
+            smem_copy_params.tdPsdO[None, None, None, smem_pipe_read_do if cutlass.const_expr(self.num_stages_dO > 1) else 0],
+            smem_copy_params.tdPsV,
+            smem_copy_params.smem_thr_copy_QdO, smem_copy_params.smem_thr_copy_KV,
+            hook_fn=load_Q_next if cutlass.const_expr(self.num_stages_Q > 1) else None,
+            swap_AB=self.SdP_swapAB,
+        )
+        tLSErdPsum = cute.make_fragment_like(smem_copy_params.tSsdPsumMma[None, 0])
+        cute.autovec_copy(
+            smem_copy_params.tSsdPsumMma[None, smem_pipe_read_do if cutlass.const_expr(self.num_stages_dO > 1) else 0], tLSErdPsum
+        )
+        acc_dP_mn = utils.make_acc_tensor_mn_view(acc_dP)
+        # if cute.arch.thread_idx()[0] == 0 and cute.arch.block_idx()[0] == bidx: cute.print_tensor(acc_dP_mn)
+        assert cute.size(acc_dP_mn, mode=[0]) == cute.size(tLSErdPsum)
+        for r in cutlass.range(cute.size(acc_dP_mn, mode=[0]), unroll_full=True):
+            acc_dP_mn[r, None].store(acc_S_mn[r, None].load() * (acc_dP_mn[r, None].load() - tLSErdPsum[r]))
+        # if cute.arch.thread_idx()[0] == 0 and cute.arch.block_idx()[0] == bidx: cute.print_tensor(acc_dP_mn)
+        rP = cute.make_fragment_like(acc_S, self.dtype)
+        rP.store(acc_S.load().to(self.dtype))
+        if cutlass.const_expr(not self.Mma_dKV_is_RS):
+            tPrP = smem_copy_params.r2s_thr_copy_PdS.retile(rP)  # ((Atom,AtomNum), MMA_N, MMA_N)
+            cute.copy(smem_copy_params.r2s_thr_copy_PdS, tPrP, smem_copy_params.tPsP)
+        rdS = cute.make_fragment_like(acc_dP, self.dtype)
+        rdS.store(acc_dP.load().to(self.dtype))
+        if cutlass.const_expr(not self.Mma_dKV_is_RS):
+            cute.arch.barrier()  # Make sure P is written
+        # For hdim 64, It's faster to write to smem_dS first before the dV gemm
+        if cutlass.const_expr(not self.Mma_dKV_is_RS):
+            tdSrdS = smem_copy_params.r2s_thr_copy_PdS.retile(rdS)
+            cute.copy(smem_copy_params.r2s_thr_copy_PdS, tdSrdS, smem_copy_params.tdSsdS)
+        if cutlass.const_expr(self.Mma_dKV_is_RS):
+            tdVrP = cute.make_tensor(rP.iterator, utils.convert_layout_acc_frgA(rP.layout))
+        else:
+            tdVrP = mma_params.tdVrP
+
+        # MMA dK
+        sm80_utils.gemm(
+            mma_params.thr_mma_dkv, mma_params.acc_dV, tdVrP, mma_params.tdVrdO,
+            smem_copy_params.tdVsPt,
+            smem_copy_params.tdVsdOt[None, None, None, smem_pipe_read_do if cutlass.const_expr(self.num_stages_dO > 1) else 0],
+            smem_copy_params.smem_thr_copy_PdSt, smem_copy_params.smem_thr_copy_QdOt,
+            A_in_regs=self.Mma_dKV_is_RS,
+            swap_AB=self.dKV_swapAB,
+        )
+        # if cute.arch.thread_idx()[0] == 0 and cute.arch.block_idx()[0] == bidx: cute.print_tensor(mma_params.acc_dV)
+        cute.arch.barrier()  # Make sure dS is written
+
+        # MMA dQ
+        def dQ_mma(hook_fn):
+            acc_shape_dQ = mma_params.thr_mma_dq.partition_shape_C(
+                (self.m_block_size, self.head_dim_padded) if cutlass.const_expr(not self.dQ_swapAB) else (self.head_dim_padded, self.m_block_size)
+            )
+            acc_dQ = cute.make_fragment(acc_shape_dQ, cutlass.Float32)
+            acc_dQ.fill(0.0)
+            sm80_utils.gemm(
+                mma_params.thr_mma_dq, acc_dQ, mma_params.tdQrdS, mma_params.tdQrK,
+                smem_copy_params.tdQsdS, smem_copy_params.tdQsKt,
+                smem_copy_params.smem_thr_copy_dS, smem_copy_params.smem_thr_copy_Kt,
+                swap_AB=self.dQ_swapAB,
+                hook_fn=hook_fn
+            )
+            # ((1, 1), num_elements)
+            acc_dQ_atomic = gmem_copy_params.gmem_thr_copy_dQaccum.retile(acc_dQ)
+            tdQgdQaccum_atomic = gmem_copy_params.tdQgdQaccum[None, None, m_block]
+            assert cute.size(acc_dQ_atomic) == cute.size(tdQgdQaccum_atomic)
+            for i in cutlass.range(cute.size(acc_dQ_atomic), unroll_full=True):
+                utils.atomic_add_fp32(acc_dQ_atomic[i], utils.elem_pointer(tdQgdQaccum_atomic, i))
+                # utils.atomic_add_fp32(acc_dQ[i], tdQgdQaccum_atomic.iterator + i * tdQgdQaccum_atomic.stride[1])
+            # if cute.arch.thread_idx()[0] == 64 and cute.arch.block_idx()[0] == bidx: cute.print_tensor(acc_dQ)
+
+        # If num_stages_Q == 1, we want to do Mma_dK first so we can start loading Q for the next iteration
+        if cutlass.const_expr(self.num_stages_Q > 1):
+            dQ_mma(load_dO_next)
+
+        # MMA dK
+        if cutlass.const_expr(self.Mma_dKV_is_RS):
+            tdKrdS = cute.make_tensor(rdS.iterator, utils.convert_layout_acc_frgA(rdS.layout))
+        else:
+            tdKrdS = mma_params.tdKrdS
+        sm80_utils.gemm(
+            mma_params.thr_mma_dkv, mma_params.acc_dK, tdKrdS, mma_params.tdKrQ,
+            smem_copy_params.tdKsdSt,
+            smem_copy_params.tdKsQt[None, None, None, smem_pipe_read_q if cutlass.const_expr(self.num_stages_Q > 1) else 0],
+            smem_copy_params.smem_thr_copy_PdSt, smem_copy_params.smem_thr_copy_QdOt,
+            A_in_regs=self.Mma_dKV_is_RS,
+            swap_AB=self.dKV_swapAB,
+            hook_fn=load_dO_next if cutlass.const_expr(self.num_stages_Q == 1) else None,
+        )
+        # if cute.arch.thread_idx()[0] == 0: cute.print_tensor(mma_params.acc_dK)
+        if cutlass.const_expr(self.num_stages_Q == 1):
+            cute.arch.barrier()
+            dQ_mma(load_Q_next)
+
+    @cute.jit
+    def epilogue(
+        self,
+        acc_dK: cute.Tensor,
+        acc_dV: cute.Tensor,
+        mdK: cute.Tensor,
+        mdV: cute.Tensor,
+        sdK: cute.Tensor,
+        sdV: cute.Tensor,
+        gmem_tiled_copy_dK: cute.TiledCopy,
+        gmem_tiled_copy_dV: cute.TiledCopy,
+        tiled_mma: cute.TiledMma,
+        tidx: cutlass.Int32,
+        n_block: cutlass.Int32,
+        num_head: cutlass.Int32,
+        batch_size: cutlass.Int32,
+        seqlen: SeqlenInfoQK,
+        d_head: cutlass.Int32, 
+        d_head_v: cutlass.Int32
+    ):
+        rdV = cute.make_fragment_like(acc_dV, self.dtype)
+        rdV.store(acc_dV.load().to(self.dtype))
+        rdK = cute.make_fragment_like(acc_dK, self.dtype)
+        rdK.store(acc_dK.load().to(self.dtype))
+        gmem_thr_copy_dK = gmem_tiled_copy_dK.get_slice(tidx)
+        gmem_thr_copy_dV = gmem_tiled_copy_dV.get_slice(tidx)
+
+        batch_idx = batch_size
+        head_idx_kv = num_head // self.qhead_per_kvhead if cutlass.const_expr(not self.pack_gqa) else num_head
+
+        if cutlass.const_expr(self.qhead_per_kvhead == 1):
+            # Make sure all threads have finished reading K and V, otherwise we get racy dQ
+            # because smem_q could be changed.
+            cute.arch.barrier()
+            # smem copy atom for dKV
+            smem_copy_atom_dKV = cute.make_copy_atom(
+                cute.nvgpu.CopyUniversalOp(), self.dtype, num_bits_per_copy=2 * self.dtype.width
+            )
+            smem_thr_copy_dKV = cute.make_tiled_copy_C(smem_copy_atom_dKV, tiled_mma).get_slice(tidx)
+            taccdVrdV = smem_thr_copy_dKV.retile(rdV)
+            taccdKrdK = smem_thr_copy_dKV.retile(rdK)
+            taccdVsdV = smem_thr_copy_dKV.partition_D(sdV)
+            taccdKsdK = smem_thr_copy_dKV.partition_D(sdK)
+            # copy acc O from rmem to smem with the smem copy atom
+            cute.copy(smem_copy_atom_dKV, taccdVrdV, taccdVsdV)
+            cute.copy(smem_copy_atom_dKV, taccdKrdK, taccdKsdK)
+
+
+            if cutlass.const_expr(not seqlen.has_cu_seqlens_k):
+                mdK_cur, mdV_cur = [t[batch_idx, None, head_idx_kv, None] for t in (mdK, mdV)]
+            else:
+                mdK_cur, mdV_cur = [cute.domain_offset((seqlen.offset_k, 0), t[None, head_idx_kv, None]) for t in (mdK, mdV)]
+
+            blkdK_shape = (self.n_block_size, self.head_dim_padded)
+            blkdV_shape = (self.n_block_size, self.head_dim_v_padded)
+            gdK = cute.local_tile(mdK_cur, blkdK_shape, (n_block, 0))
+            gdV = cute.local_tile(mdV_cur, blkdV_shape, (n_block, 0))
+            tdKsdK = gmem_thr_copy_dK.partition_S(sdK)
+            tdKgdK = gmem_thr_copy_dK.partition_D(gdK)
+            tdVsdV = gmem_thr_copy_dV.partition_S(sdV)
+            tdVgdV = gmem_thr_copy_dV.partition_D(gdV)
+            tdKrdK = cute.make_fragment_like(tdKgdK, self.dtype)
+            tdVrdV = cute.make_fragment_like(tdVgdV, self.dtype)
+            # sync before all smem stores are done.
+            cute.arch.barrier()
+            # load acc dK and dV from smem to rmem for wider vectorization
+            # Need to check OOB when reading from smem if kBlockN isn't evenly tiled
+            # TODO
+            cute.autovec_copy(tdKsdK, tdKrdK)
+            cute.autovec_copy(tdVsdV, tdVrdV)
+
+            cdK = cute.make_identity_tensor((self.n_block_size, self.head_dim_padded))
+            tdKcdK = gmem_thr_copy_dK.partition_S(cdK)
+            t0dKcdK = gmem_tiled_copy_dK.get_slice(0).partition_S(cdK)
+            if cutlass.const_expr(self.head_dim_padded == self.head_dim_v_padded):
+                tdVcdV = tdKcdK
+                t0dVcdV = t0dKcdK
+            else:
+                cdV = cute.make_identity_tensor((self.n_block_size, self.head_dim_v_padded))
+                tdVcdV = gmem_thr_copy_dV.partition_S(cdV)
+                t0dVcdV = gmem_tiled_copy_dV.get_slice(0).partition_S(cdV)
+            tdKpdK = utils.predicate_k(tdKcdK, limit=d_head)
+            if cutlass.const_expr(self.same_hdim_kv):
+                tdVpdV = tdKpdK
+            else:
+                tdVpdV = utils.predicate_k(tdVcdV, limit=d_head_v)
+            # copy acc dK and acc_dV from rmem to gmem
+            for rest_m in cutlass.range_constexpr(cute.size(tdKrdK.shape[1])):
+                if t0dKcdK[0, rest_m, 0][0] < seqlen.seqlen_k - n_block * self.n_block_size - tdKcdK[0][0]:
+                    cute.copy(
+                        gmem_tiled_copy_dK,
+                        tdKrdK[None, rest_m, None],
+                        tdKgdK[None, rest_m, None],
+                        pred=tdKpdK[None, rest_m, None] if cutlass.const_expr(self.check_hdim_oob) else None,
+                    )
+            for rest_m in cutlass.range_constexpr(cute.size(tdVrdV.shape[1])):
+                if t0dVcdV[0, rest_m, 0][0] < seqlen.seqlen_k - n_block * self.n_block_size - tdVcdV[0][0]:
+                    cute.copy(
+                        gmem_tiled_copy_dV,
+                        tdVrdV[None, rest_m, None],
+                        tdVgdV[None, rest_m, None],
+                        pred=tdVpdV[None, rest_m, None] if cutlass.const_expr(self.check_hdim_v_oob) else None,
+                    )
+
+        else:  # qhead_per_kvhead > 1, do atomic add
+            # For Sm90, we need to sync to avoid racy writes to smem_q
+            # For Sm80, we don't need to sync since we're not touching smem
+            head_idx_kv = num_head // self.qhead_per_kvhead if cutlass.const_expr(not self.pack_gqa) else num_head
+
+            if cutlass.const_expr(not seqlen.has_cu_seqlens_k):
+                mdK_cur, mdV_cur = [t[batch_idx, head_idx_kv, None] for t in (mdK, mdV)]
+            else:
+                padded_offset_k = seqlen.offset_k + batch_idx * self.n_block_size
+                mdK_cur = cute.domain_offset((padded_offset_k * self.head_dim_padded,), mdK[head_idx_kv, None])
+                mdV_cur = cute.domain_offset((padded_offset_k * self.head_dim_v_padded,), mdV[head_idx_kv, None])
+
+            gdV = cute.local_tile(mdV_cur, (self.n_block_size * self.head_dim_v_padded,), (n_block,))
+            gdK = cute.local_tile(mdK_cur, (self.n_block_size * self.head_dim_padded,), (n_block,))
+            tdVgdVaccum = gmem_thr_copy_dV.partition_S(gdV)
+            tdKgdKaccum = gmem_thr_copy_dK.partition_S(gdK)
+            acc_dV_atomic = gmem_thr_copy_dV.retile(acc_dV)
+            acc_dK_atomic = gmem_thr_copy_dK.retile(acc_dK)
+            assert cute.size(acc_dV_atomic) == cute.size(tdVgdVaccum)
+            assert cute.size(acc_dK_atomic) == cute.size(tdKgdKaccum)
+            for i in cutlass.range(cute.size(acc_dV_atomic), unroll_full=True):
+                utils.atomic_add_fp32(acc_dV_atomic[i], utils.elem_pointer(tdVgdVaccum, i))
+            for i in cutlass.range(cute.size(acc_dK_atomic), unroll_full=True):
+                utils.atomic_add_fp32(acc_dK_atomic[i], utils.elem_pointer(tdKgdKaccum, i))
+
+    @cute.jit
+    def advance_pipeline(self, pipeline_index, num_stages: cutlass.Constexpr):
+        return pipeline_index + 1 if pipeline_index < num_stages - 1 else 0
+
+    @cute.jit
+    def load_K(
+        self,
+        gmem_thr_copy: cute.TiledCopy,
+        tKgK: cute.Tensor,
+        tKsK: cute.Tensor,
+        block: cutlass.Int32,
+        seqlen: cutlass.Int32,
+        headdim: cutlass.Int32,
+    ):
+        cK = cute.make_identity_tensor((self.n_block_size, self.head_dim_padded))
+        tKcK = gmem_thr_copy.partition_S(cK)
+        t0KcK = gmem_thr_copy.get_slice(0).partition_S(cK)
+        tKpK = utils.predicate_k(tKcK, limit=headdim)
+        for n in cutlass.range_constexpr(cute.size(tKsK.shape[1])):
+            # If kBlockN doesn't evenly divide the tiled copy, only the last `n` needs to be checked
+            if self.is_even_n_smem_k or n < cute.size(tKsK.shape[1]) - 1 or tKcK[0, n, 0][0] < self.n_block_size:
+                # 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.
+                predicate_n = t0KcK[0, n, 0][0] < seqlen - block * self.n_block_size - tKcK[0][0]
+                predicate = cute.make_fragment_like(tKpK[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] = (tKpK[i, n, k] if cutlass.const_expr(self.check_hdim_oob) else True) and predicate_n
+                cute.copy(
+                    gmem_thr_copy, tKgK[None, n, None], tKsK[None, n, None], pred=predicate,
+                )
+            # We need to clear the sK smem tiles since we'll use sKt for mma_dq
+
+    @cute.jit
+    def load_V(
+        self,
+        gmem_thr_copy: cute.TiledCopy,
+        tVgV: cute.Tensor,
+        tVsV: cute.Tensor,
+        block: cutlass.Int32,
+        seqlen: cutlass.Int32,
+        headdim: cutlass.Int32,
+    ):
+        cV = cute.make_identity_tensor((self.n_block_size, self.head_dim_v_padded))
+        tVcV = gmem_thr_copy.partition_S(cV)
+        t0VcV = gmem_thr_copy.get_slice(0).partition_S(cV)
+        tVpV = utils.predicate_k(tVcV, limit=headdim)
+        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 self.is_even_n_smem_v or n < cute.size(tVsV.shape[1]) - 1 or tVcV[0, n, 0][0] < self.n_block_size:
+                # Instead of using tVcV, we using t0VcV and subtract the offset from the limit
+                # (seqlen - block * kBlockN). This is because the entries of t0VcV are known at compile time.
+                predicate_n = t0VcV[0, n, 0][0] < seqlen - block * self.n_block_size - tVcV[0][0]
+                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 cutlass.const_expr(self.check_hdim_oob) else True) and predicate_n
+                cute.copy(
+                    gmem_thr_copy, tVgV[None, n, None], tVsV[None, n, None], pred=predicate,
+                )
+
+    @cute.jit
+    def load_Q_LSE(
+        self,
+        gmem_tiled_copy_Q: cute.TiledCopy,
+        gmem_tiled_copy_LSE: cute.TiledCopy,
+        tQgQ: cute.Tensor,
+        tQsQ: cute.Tensor,
+        tQcQ: cute.Tensor,
+        t0QcQ: cute.Tensor,
+        tQpQ: cute.Tensor,
+        tLSEgLSE: cute.Tensor,
+        tLSEsLSE: cute.Tensor,
+        tLSEcLSE: cute.Tensor,
+        block: cutlass.Int32,
+        smem_pipe_write_q: cutlass.Int32,
+        seqlen: cutlass.Int32,
+    ):
+        for m in cutlass.range_constexpr(cute.size(tQsQ.shape[1])):
+            # If kBlockM doesn't evenly divide the tiled copy, only the last `m` needs to be checked
+            if self.is_even_m_smem_q or m < cute.size(tQsQ.shape[1]) - 1 or tQcQ[0, m, 0][0] < self.m_block_size:
+                # 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.
+                predicate_m = t0QcQ[0, m, 0][0] < seqlen - block * self.m_block_size - tQcQ[0][0]
+                predicate = cute.make_fragment_like(tQpQ[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] = (tQpQ[i, m, k] if cutlass.const_expr(self.check_hdim_oob) else True) and predicate_m
+                cute.copy(
+                    gmem_tiled_copy_Q,
+                    tQgQ[None, m, None, block],
+                    tQsQ[None, m, None, smem_pipe_write_q if cutlass.const_expr(self.num_stages_Q) > 1 else 0],
+                    pred=predicate,
+                )
+            # We need to clear the sQ smem tiles since we'll use sQt for mma_dK
+        # We made sure LSE length is padded so we read `kBlockM` elements so that all
+        # elements in sLSE are filled. Without this we might have uninitialized sLSE values.
+        for m in cutlass.range_constexpr(cute.size(tLSEsLSE.shape[1])):
+            if tLSEcLSE[0, m][0] < self.m_block_size:
+                cute.copy(
+                    gmem_tiled_copy_LSE,
+                    tLSEgLSE[None, m, block],
+                    tLSEsLSE[None, m, smem_pipe_write_q if cutlass.const_expr(self.num_stages_Q > 1) else 0],
+                )
+
+    @cute.jit
+    def load_dO_dPsum(
+        self,
+        gmem_tiled_copy_dO: cute.TiledCopy,
+        gmem_tiled_copy_dPsum: cute.TiledCopy,
+        tdOgdO: cute.Tensor,
+        tdOsdO: cute.Tensor,
+        tdOcdO: cute.Tensor,
+        t0dOcdO: cute.Tensor,
+        tdOpdO: cute.Tensor,
+        tdPsumgdPsum: cute.Tensor,
+        tdPsumsdPsum: cute.Tensor,
+        tdPsumcdPsum: cute.Tensor,
+        block: cutlass.Int32,
+        smem_pipe_write_q: cutlass.Int32,
+        seqlen: cutlass.Int32,
+    ):
+        for m in cutlass.range_constexpr(cute.size(tdOsdO.shape[1])):
+            # If kBlockM doesn't evenly divide the tiled copy, only the last `m` needs to be checked
+            if self.is_even_m_smem_do or m < cute.size(tdOsdO.shape[1]) - 1 or tdOcdO[0, m, 0][0] < self.m_block_size:
+                # Instead of using tdOcdO, we using t0dOcdO and subtract the offset from the limit
+                # (seqlen - block * kBlockM). This is because the entries of t0dOcdO are known at compile time.
+                predicate_m = t0dOcdO[0, m, 0][0] < seqlen - block * self.m_block_size - tdOcdO[0][0]
+                predicate = cute.make_fragment_like(tdOpdO[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] = (tdOpdO[i, m, k] if cutlass.const_expr(self.check_hdim_oob) else True) and predicate_m
+                cute.copy(
+                    gmem_tiled_copy_dO,
+                    tdOgdO[None, m, None, block],
+                    tdOsdO[None, m, None, smem_pipe_write_q if cutlass.const_expr(self.num_stages_dO > 1) else 0],
+                    pred=predicate,
+                )
+            # We need to clear the sQ smem tiles since we'll use sQt for mma_dK
+        # We made sure LSE length is padded so we read `kBlockM` elements so that all
+        # elements in sLSE are filled. Without this we might have uninitialized sLSE values.
+        for m in cutlass.range_constexpr(cute.size(tdPsumgdPsum.shape[1])):
+            if tdPsumcdPsum[0, m][0] < self.m_block_size:
+                cute.copy(
+                    gmem_tiled_copy_dPsum,
+                    tdPsumgdPsum[None, m, block],
+                    tdPsumsdPsum[None, m, smem_pipe_write_q if cutlass.const_expr(self.num_stages_dO > 1) else 0],
+                )

build/torch-universal/flash_bwd_postprocess.py

@@ -0,0 +1,729 @@
+# 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_bwd_postprocess_kernel.h
+# from Cutlass C++ to Cute-DSL.
+import math
+from typing import Callable, Optional, Type, Literal
+
+import cuda.bindings.driver as cuda
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils.hopper_helpers as sm90_utils_basic
+import cutlass.utils.blackwell_helpers as sm100_utils_basic
+from cutlass.cute.nvgpu import cpasync, warp, warpgroup
+from cutlass import Float32, const_expr
+from cutlass.utils import LayoutEnum
+
+from . import utils
+from . import copy_utils
+from . import ampere_helpers as sm80_utils
+from . import hopper_helpers as sm90_utils
+from .seqlen_info import SeqlenInfoQK
+import cutlass.cute.nvgpu.tcgen05 as tcgen05
+from .tile_scheduler import (
+    ParamsBase,
+    SingleTileScheduler,
+    SingleTileVarlenScheduler,
+    TileSchedulerArguments,
+)
+
+
+class FlashAttentionBackwardPostprocess:
+    def __init__(
+        self,
+        dtype: Type[cutlass.Numeric],
+        head_dim: int,
+        arch: Literal[80, 90, 100],
+        tile_m: int = 128,
+        num_threads: int = 256,
+        AtomLayoutMdQ: int = 1,
+        dQ_swapAB: bool = False,
+    ):
+        """
+        :param head_dim: head dimension
+        :type head_dim: int
+        :param tile_m: m block size
+        :type tile_m: int
+        """
+        self.dtype = dtype
+        self.tile_m = tile_m
+        assert arch in [80, 90, 100], (
+            "Only Ampere (80), Hopper (90), and Blackwell (100) are supported"
+        )
+        self.arch = arch
+        # padding head_dim to a multiple of 32 as k_block_size
+        hdim_multiple_of = 32
+        self.tile_hdim = int(math.ceil(head_dim / hdim_multiple_of) * hdim_multiple_of)
+        self.check_hdim_oob = head_dim != self.tile_hdim
+        self.num_threads = num_threads
+        self.AtomLayoutMdQ = AtomLayoutMdQ
+        self.dQ_swapAB = dQ_swapAB
+
+    @staticmethod
+    def can_implement(dtype, head_dim, tile_m, num_threads) -> 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
+
+        :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 num_threads % 32 != 0:
+            return False
+        return True
+
+    def _get_tiled_mma(self):
+        if const_expr(self.arch == 80):
+            num_mma_warps = self.num_threads // 32
+            atom_layout_dQ = (
+                (self.AtomLayoutMdQ, num_mma_warps // self.AtomLayoutMdQ, 1)
+                if const_expr(not self.dQ_swapAB)
+                else (num_mma_warps // self.AtomLayoutMdQ, self.AtomLayoutMdQ, 1)
+            )
+            tiled_mma = cute.make_tiled_mma(
+                warp.MmaF16BF16Op(self.dtype, Float32, (16, 8, 16)),
+                atom_layout_dQ,
+                permutation_mnk=(atom_layout_dQ[0] * 16, atom_layout_dQ[1] * 16, 16),
+            )
+        elif const_expr(self.arch == 90):
+            num_mma_warp_groups = self.num_threads // 128
+            atom_layout_dQ = (self.AtomLayoutMdQ, num_mma_warp_groups // self.AtomLayoutMdQ)
+            tiler_mn_dQ = (self.tile_m // atom_layout_dQ[0], self.tile_hdim // atom_layout_dQ[1])
+            tiled_mma = sm90_utils_basic.make_trivial_tiled_mma(
+                self.dtype,
+                self.dtype,
+                warpgroup.OperandMajorMode.K,  # These don't matter, we only care about the accum
+                warpgroup.OperandMajorMode.K,
+                Float32,
+                atom_layout_mnk=(atom_layout_dQ if not self.dQ_swapAB else atom_layout_dQ[::-1])
+                + (1,),
+                tiler_mn=tiler_mn_dQ if not self.dQ_swapAB else tiler_mn_dQ[::-1],
+            )
+        else:
+            cta_group = tcgen05.CtaGroup.ONE
+            tiled_mma = sm100_utils_basic.make_trivial_tiled_mma(
+                self.dtype,
+                tcgen05.OperandMajorMode.MN,  # dS_major_mode
+                tcgen05.OperandMajorMode.MN,  # Kt_major_mode
+                Float32,
+                cta_group,
+                (self.tile_m, self.tile_hdim),
+            )
+        if const_expr(self.arch in [80, 90]):
+            assert self.num_threads == tiled_mma.size
+        return tiled_mma
+
+    def _setup_attributes(self):
+        # ///////////////////////////////////////////////////////////////////////////////
+        # GMEM Tiled copy:
+        # ///////////////////////////////////////////////////////////////////////////////
+        # Thread layouts for copies
+        universal_copy_bits = 128
+        async_copy_elems_accum = universal_copy_bits // Float32.width
+        atom_async_copy_accum = cute.make_copy_atom(
+            cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
+            Float32,
+            num_bits_per_copy=universal_copy_bits,
+        )
+        # We don't do bound checking for the gmem -> smem load so we just assert here.
+        assert (self.tile_m * self.tile_hdim // async_copy_elems_accum) % self.num_threads == 0
+        self.g2s_tiled_copy_dQaccum = cute.make_tiled_copy_tv(
+            atom_async_copy_accum,
+            cute.make_layout(self.num_threads),
+            cute.make_layout(async_copy_elems_accum),
+        )
+        num_s2r_copy_elems = 1 if const_expr(self.arch == 80) else 4
+        if const_expr(self.arch == 80):
+            self.s2r_tiled_copy_dQaccum = copy_utils.tiled_copy_1d(
+                Float32, self.num_threads, num_s2r_copy_elems
+            )
+            self.sdQaccum_layout = cute.make_layout(self.tile_m * self.tile_hdim)
+        elif const_expr(self.arch == 90):
+            num_threads_per_warp_group = 128
+            num_mma_warp_groups = self.num_threads // 128
+            self.s2r_tiled_copy_dQaccum = cute.make_tiled_copy_tv(
+                cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128),
+                cute.make_layout((num_threads_per_warp_group, num_mma_warp_groups)),  # thr_layout
+                cute.make_layout(128 // Float32.width),  # val_layout
+            )
+            self.sdQaccum_layout = cute.make_layout(
+                (self.tile_m * self.tile_hdim // num_mma_warp_groups, num_mma_warp_groups)
+            )
+        else:
+            self.dQ_reduce_ncol = 32
+            dQaccum_reduce_stage = self.tile_hdim // self.dQ_reduce_ncol
+            assert self.num_threads == 128  # TODO: currently hard-coded
+            self.s2r_tiled_copy_dQaccum = copy_utils.tiled_copy_1d(
+                Float32, self.num_threads, num_s2r_copy_elems
+            )
+            self.sdQaccum_layout = cute.make_layout(
+                (self.tile_m * self.tile_hdim // dQaccum_reduce_stage, dQaccum_reduce_stage)
+            )
+
+        self.gmem_tiled_copy_dQ = copy_utils.tiled_copy_2d(
+            self.dtype, self.tile_hdim, self.num_threads
+        )
+        # ///////////////////////////////////////////////////////////////////////////////
+        # Shared memory layout: dQ
+        # ///////////////////////////////////////////////////////////////////////////////
+        # We can't just use kHeadDim here. E.g. if MMA shape is 64 x 96 but split across 2 WGs,
+        # then setting kBlockKSmem to 32 will cause "Static shape_div failure".
+        # We want to treat it as 64 x 48, so kBlockKSmem should be 16.
+        mma_shape_n = self.tiled_mma.get_tile_size(1)
+        if const_expr(self.arch == 80):
+            sdQ_layout_atom = sm80_utils.get_smem_layout_atom(self.dtype, mma_shape_n)
+            self.sdQ_layout = cute.tile_to_shape(
+                sdQ_layout_atom, (self.tile_m, self.tile_hdim), (0, 1)
+            )
+        elif const_expr(self.arch == 90):
+            self.sdQ_layout = sm90_utils.make_smem_layout(
+                self.dtype, LayoutEnum.ROW_MAJOR, (self.tile_m, self.tile_hdim)
+            )
+        else:
+            # TODO: this is hard-coded for hdim 128
+            self.sdQ_layout = sm100_utils_basic.make_smem_layout_epi(
+                self.dtype, LayoutEnum.ROW_MAJOR, (self.tile_m, self.tile_hdim), 1
+            )
+
+    @cute.jit
+    def __call__(
+        self,
+        mdQaccum: cute.Tensor,
+        mdQ: cute.Tensor,
+        scale: cutlass.Float32,
+        mCuSeqlensQ: Optional[cute.Tensor],
+        mSeqUsedQ: Optional[cute.Tensor],
+        stream: cuda.CUstream,
+    ):
+        # Get the data type and check if it is fp16 or bf16
+        if const_expr(mdQ.element_type not in [cutlass.Float16, cutlass.BFloat16]):
+            raise TypeError("Only Float16 or BFloat16 is supported")
+        if const_expr(mdQaccum is not None):
+            if const_expr(mdQaccum.element_type not in [cutlass.Float32]):
+                raise TypeError("dQaccum tensor must be Float32")
+
+        # Assume all strides are divisible by 128 bits except the last stride
+        new_stride = lambda t: (
+            *(cute.assume(s, divby=128 // t.element_type.width) for s in t.stride[:-1]),
+            t.stride[-1],
+        )
+        mdQaccum, mdQ = [
+            cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t)))
+            for t in (mdQaccum, mdQ)
+        ]
+
+        self.tiled_mma = self._get_tiled_mma()
+        self._setup_attributes()
+
+        smem_size = max(
+            cute.size_in_bytes(cutlass.Float32, self.sdQaccum_layout),
+            cute.size_in_bytes(self.dtype, self.sdQ_layout),
+        )
+
+        if const_expr(mCuSeqlensQ is not None):
+            TileScheduler = SingleTileVarlenScheduler
+            num_head = mdQ.shape[1]
+            num_batch = mCuSeqlensQ.shape[0] - 1
+        else:
+            TileScheduler = SingleTileScheduler
+            num_head = mdQ.shape[2]
+            num_batch = mdQ.shape[0]
+
+        tile_sched_args = TileSchedulerArguments(
+            num_block=cute.ceil_div(mdQ.shape[1], self.tile_m),
+            num_head=num_head,
+            num_batch=num_batch,
+            num_splits=1,
+            seqlen_k=0,
+            headdim=mdQ.shape[2],
+            headdim_v=0,
+            total_q=mdQ.shape[0],
+            tile_shape_mn=(self.tile_m, 1),
+            mCuSeqlensQ=mCuSeqlensQ,
+            mSeqUsedQ=mSeqUsedQ,
+        )
+
+        tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
+        grid_dim = TileScheduler.get_grid_shape(tile_sched_params)
+
+        # grid_dim: (m_block, num_head, batch_size)
+        self.kernel(
+            mdQaccum,
+            mdQ,
+            mCuSeqlensQ,
+            mSeqUsedQ,
+            scale,
+            self.tiled_mma,
+            self.dQ_swapAB,
+            self.sdQaccum_layout,
+            self.sdQ_layout,
+            self.g2s_tiled_copy_dQaccum,
+            self.s2r_tiled_copy_dQaccum,
+            self.gmem_tiled_copy_dQ,
+            tile_sched_params,
+            TileScheduler,
+        ).launch(
+            grid=grid_dim,
+            block=[self.num_threads, 1, 1],
+            smem=smem_size,
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        mdQaccum: cute.Tensor,
+        mdQ: cute.Tensor,
+        mCuSeqlensQ: Optional[cute.Tensor],
+        mSeqUsedQ: Optional[cute.Tensor],
+        scale: cutlass.Float32,
+        tiled_mma: cute.TiledMma,
+        dQ_swapAB: cutlass.Constexpr,
+        sdQaccum_layout: cute.Layout,
+        sdQ_layout: cute.ComposedLayout,
+        g2s_tiled_copy_dQaccum: cute.TiledCopy,
+        s2r_tiled_copy_dQaccum: cute.TiledCopy,
+        gmem_tiled_copy_dQ: cute.TiledCopy,
+        tile_sched_params: ParamsBase,
+        TileScheduler: cutlass.Constexpr[Callable],
+    ):
+        # ///////////////////////////////////////////////////////////////////////////////
+        # Get shared memory buffer
+        # ///////////////////////////////////////////////////////////////////////////////
+        smem = cutlass.utils.SmemAllocator()
+        sdQaccum = smem.allocate_tensor(cutlass.Float32, sdQaccum_layout, byte_alignment=1024)
+        sdQaccum_flat = cute.make_tensor(sdQaccum.iterator, cute.make_layout(cute.size(sdQaccum)))
+        if const_expr(self.arch in [80, 90]):
+            sdQ = cute.make_tensor(cute.recast_ptr(sdQaccum.iterator, dtype=self.dtype), sdQ_layout)
+        else:
+            # extra stage dimension
+            sdQ = cute.make_tensor(
+                cute.recast_ptr(sdQaccum.iterator, sdQ_layout.inner, dtype=self.dtype),
+                sdQ_layout.outer,
+            )[None, None, 0]
+        sdQt = utils.transpose_view(sdQ)
+
+        # 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
+
+        if work_tile.is_valid_tile:
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Get the appropriate tiles for this thread block.
+            # ///////////////////////////////////////////////////////////////////////////////
+
+            seqlen = SeqlenInfoQK.create(
+                batch_size,
+                mdQ.shape[1],
+                0,
+                mCuSeqlensQ=mCuSeqlensQ,
+                mCuSeqlensK=None,
+                mSeqUsedQ=mSeqUsedQ,
+                mSeqUsedK=None,
+            )
+            if const_expr(not seqlen.has_cu_seqlens_q):
+                mdQ_cur = mdQ[batch_size, None, num_head, None]
+                mdQaccum_cur = mdQaccum[batch_size, num_head, None]
+                head_dim = mdQ.shape[3]
+            else:
+                padded_offset_q = seqlen.offset_q + batch_size * self.tile_m
+                mdQ_cur = cute.domain_offset((seqlen.offset_q, 0), mdQ[None, num_head, None])
+                mdQaccum_cur = cute.domain_offset(
+                    (padded_offset_q * self.tile_hdim,), mdQaccum[num_head, None]
+                )
+                head_dim = mdQ.shape[2]
+
+                # HACK: Compiler doesn't seem to recognize that padding
+                # by padded_offset_q * self.tile_hdim keeps alignment
+                # since statically divisible by 4
+
+                mdQaccum_cur_ptr = cute.make_ptr(
+                    dtype=mdQaccum_cur.element_type,
+                    value=mdQaccum_cur.iterator.toint(),
+                    mem_space=mdQaccum_cur.iterator.memspace,
+                    assumed_align=mdQaccum.iterator.alignment,
+                )
+                mdQaccum_cur = cute.make_tensor(mdQaccum_cur_ptr, mdQaccum_cur.layout)
+
+            gdQaccum = cute.local_tile(mdQaccum_cur, (self.tile_m * self.tile_hdim,), (m_block,))
+            gdQ = cute.local_tile(mdQ_cur, (self.tile_m, self.tile_hdim), (m_block, 0))
+
+            seqlen_q = seqlen.seqlen_q
+            seqlen_q_rounded = cute.round_up(seqlen_q, self.tile_m)
+
+            # Step 1: load dQaccum from gmem to smem
+            g2s_thr_copy_dQaccum = g2s_tiled_copy_dQaccum.get_slice(tidx)
+            tdQgdQaccum = g2s_thr_copy_dQaccum.partition_S(gdQaccum)
+            tdQsdQaccumg2s = g2s_thr_copy_dQaccum.partition_D(sdQaccum_flat)
+            cute.copy(g2s_tiled_copy_dQaccum, tdQgdQaccum, tdQsdQaccumg2s)
+            cute.arch.cp_async_commit_group()
+            cute.arch.cp_async_wait_group(0)
+            cute.arch.barrier()
+
+            # Step 2: load dQ from smem to rmem
+            s2r_thr_copy_dQaccum = s2r_tiled_copy_dQaccum.get_slice(tidx)
+            tdQsdQaccum = s2r_thr_copy_dQaccum.partition_S(sdQaccum)
+            tile_shape = (self.tile_m, self.tile_hdim)
+            acc = None
+            tiled_copy_t2r = None
+            if const_expr(self.arch in [80, 90]):
+                acc_shape = tiled_mma.partition_shape_C(
+                    tile_shape if const_expr(not dQ_swapAB) else tile_shape[::-1]
+                )
+                acc = cute.make_fragment(acc_shape, cutlass.Float32)
+                assert cute.size(acc) == cute.size(tdQsdQaccum)
+            else:
+                thr_mma = tiled_mma.get_slice(0)  # 1-CTA
+                dQacc_shape = tiled_mma.partition_shape_C((self.tile_m, self.tile_hdim))
+                tdQtdQ = tiled_mma.make_fragment_C(dQacc_shape)
+                tdQcdQ = thr_mma.partition_C(
+                    cute.make_identity_tensor((self.tile_m, self.tile_hdim))
+                )
+                tmem_load_atom = cute.make_copy_atom(
+                    tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(self.dQ_reduce_ncol)), Float32
+                )
+                tiled_copy_t2r = tcgen05.make_tmem_copy(tmem_load_atom, tdQtdQ)
+                thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
+                tdQrdQ_t2r_shape = thr_copy_t2r.partition_D(tdQcdQ).shape
+                acc = cute.make_fragment(tdQrdQ_t2r_shape, Float32)
+            tdQrdQaccum = cute.make_tensor(acc.iterator, cute.make_layout(tdQsdQaccum.shape))
+            cute.autovec_copy(tdQsdQaccum, tdQrdQaccum)
+            # Convert tdQrdQaccum from fp32 to fp16/bf16
+            rdQ = cute.make_fragment_like(acc, self.dtype)
+            rdQ.store((acc.load() * scale).to(self.dtype))
+
+            # Step 3: Copy dQ from register to smem
+            cute.arch.barrier()  # make sure all threads have finished loading dQaccum
+            if const_expr(self.arch in [80, 90]):
+                copy_atom_r2s_dQ = utils.get_smem_store_atom(
+                    self.arch, self.dtype, transpose=self.dQ_swapAB
+                )
+                tiled_copy_r2s_dQ = cute.make_tiled_copy_C(copy_atom_r2s_dQ, tiled_mma)
+            else:
+                # copy_atom_r2s_dQ = sm100_utils_basic.get_smem_store_op(
+                #     LayoutEnum.ROW_MAJOR, self.dtype, Float32, tiled_copy_t2r,
+                # )
+                # tiled_copy_r2s_dQ = cute.make_tiled_copy_D(copy_atom_r2s_dQ, tiled_copy_t2r)
+                thr_layout_r2s_dQ = cute.make_layout((self.num_threads, 1))  # 128 threads
+                val_layout_r2s_dQ = cute.make_layout((1, 128 // self.dtype.width))
+                copy_atom_r2s_dQ = cute.make_copy_atom(
+                    cute.nvgpu.CopyUniversalOp(),
+                    self.dtype,
+                    num_bits_per_copy=128,
+                )
+                tiled_copy_r2s_dQ = cute.make_tiled_copy_tv(
+                    copy_atom_r2s_dQ, thr_layout_r2s_dQ, val_layout_r2s_dQ
+                )
+            thr_copy_r2s_dQ = tiled_copy_r2s_dQ.get_slice(tidx)
+            cdQ = cute.make_identity_tensor((self.tile_m, self.tile_hdim))
+            if const_expr(self.arch in [80, 90]):
+                taccdQrdQ = thr_copy_r2s_dQ.retile(rdQ)
+            else:
+                taccdQcdQ_shape = thr_copy_r2s_dQ.partition_S(cdQ).shape
+                taccdQrdQ = cute.make_tensor(rdQ.iterator, taccdQcdQ_shape)
+            taccdQsdQ = thr_copy_r2s_dQ.partition_D(sdQ if const_expr(not self.dQ_swapAB) else sdQt)
+            cute.copy(thr_copy_r2s_dQ, taccdQrdQ, taccdQsdQ)
+
+            # Step 4: Copy dQ from smem to register to prepare for coalesced write to gmem
+            cute.arch.barrier()  # make sure all smem stores are done
+            gmem_thr_copy_dQ = gmem_tiled_copy_dQ.get_slice(tidx)
+            tdQgdQ = gmem_thr_copy_dQ.partition_S(gdQ)
+            tdQsdQ = gmem_thr_copy_dQ.partition_D(sdQ)
+            tdQrdQ = cute.make_fragment_like(tdQsdQ, self.dtype)
+            # TODO: check OOB when reading from smem if kBlockM isn't evenly tiled
+            cute.autovec_copy(tdQsdQ, tdQrdQ)
+
+            # Step 5: Copy dQ from register to gmem
+            tdQcdQ = gmem_thr_copy_dQ.partition_S(cdQ)
+            tdQpdQ = utils.predicate_k(tdQcdQ, limit=head_dim)
+            for rest_m in cutlass.range(cute.size(tdQrdQ.shape[1]), unroll_full=True):
+                if tdQcdQ[0, rest_m, 0][0] < seqlen_q - m_block * self.tile_m:
+                    cute.copy(
+                        gmem_tiled_copy_dQ,
+                        tdQrdQ[None, rest_m, None],
+                        tdQgdQ[None, rest_m, None],
+                        pred=tdQpdQ[None, rest_m, None],
+                    )
+
+
+class FlashAttentionBackwardPostprocess_sm100(FlashAttentionBackwardPostprocess):
+    def __init__(
+        self,
+        dtype: Type[cutlass.Numeric],
+        head_dim: int,
+        tile_m: int = 128,
+        num_threads: int = 256,
+        AtomLayoutMdQ: int = 1,
+        dQ_swapAB: bool = False,
+    ):
+        super().__init__(
+            dtype=dtype,
+            head_dim=head_dim,
+            arch=90,  # tmp dummy placement for now
+            tile_m=tile_m,
+            num_threads=num_threads,
+            AtomLayoutMdQ=AtomLayoutMdQ,
+            dQ_swapAB=dQ_swapAB,
+        )
+
+    def _setup_attributes(self):
+        self.num_stages = self.tile_hdim // 32  # 2 for D=64, 4 for D=128
+
+        self.sdQaccum_layout = cute.make_layout(
+            shape=(self.tile_m * 32, 2), stride=(1, self.tile_m * 32)
+        )
+        self.epi_tile_q = (self.tile_m, self.tile_hdim)
+        self.sdQ_layout = sm100_utils_basic.make_smem_layout_epi(
+            self.dtype,
+            LayoutEnum.ROW_MAJOR,
+            self.epi_tile_q,
+            1,
+        )
+
+    @cute.jit
+    def __call__(
+        self,
+        mdQaccum: cute.Tensor,
+        mdQ: cute.Tensor,
+        scale: cutlass.Float32,
+        stream: cuda.CUstream,
+    ):
+        # Assume all strides are divisible by 128 bits except the last stride
+        new_stride = lambda t: (
+            *(cute.assume(s, divby=128 // t.element_type.width) for s in t.stride[:-1]),
+            t.stride[-1],
+        )
+        mdQaccum, mdQ = [
+            cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t)))
+            for t in (mdQaccum, mdQ)
+        ]
+        # (b, h, s*d) -> (s*d, h, b)
+        mdQaccum = cute.make_tensor(mdQaccum.iterator, cute.select(mdQaccum.layout, mode=[2, 1, 0]))
+        # (b, s, h, d) -> (s, d, h, b)
+        mdQ = cute.make_tensor(mdQ.iterator, cute.select(mdQ.layout, mode=[1, 3, 2, 0]))
+
+        self._setup_attributes()
+
+        grid_dim = [
+            cute.ceil_div(mdQ.shape[0], self.tile_m),
+            cute.size(mdQ.shape[2]),
+            cute.size(mdQ.shape[3]),
+        ]
+
+        cta_group = tcgen05.CtaGroup.ONE
+        self.mma_tiler_dsk = (self.tile_m, self.tile_hdim)
+
+        dS_major_mode = tcgen05.OperandMajorMode.MN
+        kt_major_mode_dsq = tcgen05.OperandMajorMode.MN
+
+        tiled_mma_dsk = sm100_utils_basic.make_trivial_tiled_mma(
+            cutlass.BFloat16,
+            dS_major_mode,
+            kt_major_mode_dsq,
+            cutlass.Float32,
+            cta_group,
+            self.mma_tiler_dsk,
+        )
+
+        dQ_cta_v_layout = cute.composition(cute.make_identity_layout(mdQ.shape), self.mma_tiler_dsk)
+        tma_store_op = cpasync.CopyBulkTensorTileS2GOp()
+        tma_atom_dQ, tma_tensor_dQ = cute.nvgpu.cpasync.make_tiled_tma_atom(
+            tma_store_op,
+            mdQ,
+            cute.select(self.sdQ_layout, mode=[0, 1]),
+            dQ_cta_v_layout,
+        )
+
+        buffer_align_bytes = 1024
+
+        @cute.struct
+        class SharedStorage:
+            sdQaccum: cute.struct.Align[
+                cute.struct.MemRange[cutlass.Float32, cute.cosize(self.sdQaccum_layout)],
+                128,
+            ]
+
+            sdQ: cute.struct.Align[
+                cute.struct.MemRange[self.dtype, cute.cosize(self.sdQ_layout)],
+                buffer_align_bytes,
+            ]
+
+        self.shared_storage = SharedStorage
+
+        self.kernel(
+            mdQaccum,
+            tma_tensor_dQ,
+            tma_atom_dQ,
+            self.sdQaccum_layout,
+            self.sdQ_layout,
+            tiled_mma_dsk,
+            scale,
+        ).launch(
+            grid=grid_dim,
+            block=[self.num_threads, 1, 1],
+            smem=self.shared_storage.size_in_bytes(),
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        mdQaccum: cute.Tensor,
+        mdQ: cute.Tensor,
+        tma_atom_dQ: cute.CopyAtom,
+        sdQaccum_layout: cute.Layout,
+        sdQ_layout: cute.ComposedLayout,
+        tiled_mma_dsk: cute.TiledMma,
+        scale: cutlass.Float32,
+    ):
+        tidx = cute.arch.thread_idx()[0]
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+        m_block, head_idx, batch_idx = cute.arch.block_idx()
+
+        # SMEM
+        smem = cutlass.utils.SmemAllocator()
+        storage = smem.allocate(self.shared_storage)
+        swz128 = cute.make_swizzle(3, 4, 3)
+        sdQaccum = storage.sdQaccum.get_tensor(sdQaccum_layout, swizzle=swz128)
+
+        sdQ = storage.sdQ.get_tensor(sdQ_layout.outer, swizzle=sdQ_layout.inner)
+
+        mdQaccum_cur = mdQaccum[None, head_idx, batch_idx]
+        mdQ_cur = mdQ[None, None, head_idx, batch_idx]
+
+        thr_mma_dsk = tiled_mma_dsk.get_slice(tidx)
+        dQacc_shape = thr_mma_dsk.partition_shape_C(self.mma_tiler_dsk[:2])
+        tdQtdQ = thr_mma_dsk.make_fragment_C(dQacc_shape)
+        tdQtdQ = cute.make_tensor(tdQtdQ.iterator, tdQtdQ.layout)
+
+        tmem_ld_atom = cute.make_copy_atom(
+            tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), cutlass.Float32
+        )
+        tiled_tmem_ld = tcgen05.make_tmem_copy(tmem_ld_atom, tdQtdQ)
+        thr_tmem_ld = tiled_tmem_ld.get_slice(tidx)
+
+        cdQ = cute.make_identity_tensor((self.mma_tiler_dsk[0], self.mma_tiler_dsk[1]))
+        tdQcdQ = thr_mma_dsk.partition_C(cdQ)
+        tdQcdQ_tensor = cute.make_tensor(tdQcdQ.iterator, tdQcdQ.layout)
+        tdQrdQ = thr_tmem_ld.partition_D(tdQcdQ_tensor)
+
+        gdQaccum = cute.local_tile(mdQaccum_cur, (self.tile_m * self.tile_hdim,), (m_block,))
+
+        num_reduce_warps = 4
+        num_reduce_threads = cute.arch.WARP_SIZE * num_reduce_warps
+
+        atom_universal_copy = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=128
+        )
+        tiler_mn, layout_tv = cute.make_layout_tv(
+            thr_layout=cute.make_layout(shape=num_reduce_threads, stride=1),
+            val_layout=cute.make_layout(shape=4, stride=1),
+        )
+        G2S_tiled_copy_dQaccum = cute.make_tiled_copy(
+            atom_universal_copy, layout_tv=layout_tv, tiler_mn=tiler_mn
+        )
+
+        smem_thr_copy_g2s = G2S_tiled_copy_dQaccum.get_slice(tidx)
+
+        # S->R
+        tdQrdQ_t2r = cute.make_fragment(tdQrdQ.shape, cutlass.Float32)
+        tiled_smem_store_s2r = cute.make_tiled_copy(
+            atom_universal_copy, layout_tv=layout_tv, tiler_mn=tiler_mn
+        )
+
+        s2r_thr_copy_dQaccum = tiled_smem_store_s2r.get_slice(tidx)
+        tdQsdQ_s2r = s2r_thr_copy_dQaccum.partition_S(sdQaccum)
+        tdQrdQ_s2r = cute.make_tensor(tdQrdQ_t2r.iterator, tdQrdQ_t2r.shape)
+
+        # R->S
+        smem_copy_atom = sm100_utils_basic.get_smem_store_op(
+            LayoutEnum.ROW_MAJOR, self.dtype, cutlass.Float32, tiled_tmem_ld
+        )
+        tiled_smem_store_r2s = cute.make_tiled_copy(
+            smem_copy_atom,
+            layout_tv=tiled_tmem_ld.layout_dst_tv_tiled,
+            tiler_mn=tiled_tmem_ld.tiler_mn,
+        )
+        tdQsdQ_r2s = thr_tmem_ld.partition_D(thr_mma_dsk.partition_C(sdQ))
+        tdQrdQ_r2s = cute.make_fragment(tdQsdQ_r2s.shape, self.dtype)
+
+        num_stages = cute.size(tdQrdQ_t2r, mode=[1])
+        for stage in cutlass.range_constexpr(num_stages):
+            # G->S
+            gdQaccum_stage = cute.local_tile(
+                gdQaccum,
+                (self.tile_m * 32,),
+                (stage,),
+            )
+
+            gdQaccum_layout_g2s = cute.make_layout(shape=(self.tile_m * 32, 1), stride=(1, 0))
+            gdQaccum_stage_g2s = cute.make_tensor(
+                cute.recast_ptr(gdQaccum_stage.iterator, swizzle_=swz128), gdQaccum_layout_g2s
+            )
+
+            tdQgdQ = smem_thr_copy_g2s.partition_S(gdQaccum_stage_g2s)
+            tdQsdQ = smem_thr_copy_g2s.partition_D(sdQaccum)
+
+            cute.copy(smem_thr_copy_g2s, tdQgdQ[None, None, 0], tdQsdQ[None, None, 0])
+
+            cute.arch.fence_proxy(
+                cute.arch.ProxyKind.async_shared, space=cute.arch.SharedSpace.shared_cta
+            )
+            cute.arch.barrier(barrier_id=6, number_of_threads=num_reduce_threads)
+
+            # S -> R
+            tdQrdQ_s2r_cpy = tdQrdQ_s2r[None, stage, None, None]
+            tdQsdQ_s2r_p = tdQsdQ_s2r[None, None, 0]
+            tdQrdQ_r2s_cpy = cute.make_tensor(
+                tdQrdQ_s2r_cpy.iterator, cute.make_layout(tdQsdQ_s2r_p.shape)
+            )
+
+            cute.copy(s2r_thr_copy_dQaccum, tdQsdQ_s2r_p, tdQrdQ_r2s_cpy)
+
+            cute.arch.fence_proxy(
+                cute.arch.ProxyKind.async_shared, space=cute.arch.SharedSpace.shared_cta
+            )
+            cute.arch.barrier(barrier_id=7, number_of_threads=num_reduce_threads)
+
+            # R->S
+            tdQrdQ_r2s_cpy = cute.make_tensor(
+                cute.recast_ptr(tdQrdQ_r2s_cpy.iterator),
+                tdQrdQ_r2s[((None, 0), stage, 0, 0, 0)].shape,
+            )
+            dQ_vec = tdQrdQ_r2s_cpy.load() * scale
+            tdQrdQ_r2s[((None, 0), stage, 0, 0, 0)].store(dQ_vec.to(self.dtype))
+
+        cute.copy(
+            tiled_smem_store_r2s,
+            tdQrdQ_r2s[None, None, None, None, 0],
+            tdQsdQ_r2s[None, None, None, None, 0],
+        )
+        cute.arch.fence_proxy(
+            cute.arch.ProxyKind.async_shared, space=cute.arch.SharedSpace.shared_cta
+        )
+        cute.arch.barrier(barrier_id=8, number_of_threads=num_reduce_threads)
+
+        # S-> G
+        gdQ = cute.local_tile(mdQ_cur, (self.tile_m, self.tile_hdim), (None, 0))
+        tdQsdQ, tdQgdQ = cpasync.tma_partition(
+            tma_atom_dQ,
+            0,
+            cute.make_layout(1),
+            cute.group_modes(sdQ, 0, 2),
+            cute.group_modes(gdQ, 0, 2),
+        )
+
+        cute.copy(tma_atom_dQ, tdQsdQ[None, 0], tdQgdQ[None, m_block])

build/torch-universal/flash_bwd_preprocess.py

@@ -0,0 +1,363 @@
+# 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_bwd_preprocess_kernel.h
+# from Cutlass C++ to Cute-DSL.
+import math
+import operator
+from typing import Callable, Type, Optional
+
+import cuda.bindings.driver as cuda
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32
+
+from . import utils
+from . import copy_utils
+from .seqlen_info import SeqlenInfoQK
+from .tile_scheduler import (
+    ParamsBase,
+    SingleTileScheduler,
+    SingleTileVarlenScheduler,
+    TileSchedulerArguments,
+)
+
+
+class FlashAttentionBackwardPreprocess:
+    def __init__(
+        self,
+        dtype: Type[cutlass.Numeric],
+        head_dim: int,
+        m_block_size: int = 128,
+        num_threads: int = 128,
+    ):
+        """
+        All contiguous dimensions must be at least 16 bytes aligned which indicates the head dimension
+        should be a multiple of 8.
+
+        :param head_dim: head dimension
+        :type head_dim: int
+        :param m_block_size: m block size
+        :type m_block_size: int
+        :param num_threads: number of threads
+        :type num_threads: int
+        """
+        self.dtype = dtype
+        self.m_block_size = m_block_size
+        # padding head_dim to a multiple of 32 as k_block_size
+        hdim_multiple_of = 32
+        self.head_dim_padded = int(math.ceil(head_dim / hdim_multiple_of) * hdim_multiple_of)
+        self.check_hdim_oob = head_dim != self.head_dim_padded
+        self.num_threads = num_threads
+
+    @staticmethod
+    def can_implement(dtype, head_dim, m_block_size, num_threads) -> 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 m_block_size: m block size
+        :type m_block_size: int
+        :param num_threads: number of threads
+        :type num_threads: int
+
+        :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 num_threads % 32 != 0:
+            return False
+        if num_threads < m_block_size:  # For multiplying lse with log2
+            return False
+        return True
+
+    def _setup_attributes(self):
+        # ///////////////////////////////////////////////////////////////////////////////
+        # GMEM Tiled copy:
+        # ///////////////////////////////////////////////////////////////////////////////
+        # Thread layouts for copies
+        # We want kBlockKGmem to be a power of 2 so that when we do the summing,
+        # it's just between threads in the same warp
+        gmem_k_block_size = (
+            128
+            if self.head_dim_padded % 128 == 0
+            else (
+                64
+                if self.head_dim_padded % 64 == 0
+                else (32 if self.head_dim_padded % 32 == 0 else 16)
+            )
+        )
+        self.gmem_tiled_copy_O = copy_utils.tiled_copy_2d(
+            self.dtype, gmem_k_block_size, self.num_threads
+        )
+        universal_copy_bits = 128
+        num_copy_elems_dQaccum = universal_copy_bits // Float32.width
+        assert (
+            self.m_block_size * self.head_dim_padded // num_copy_elems_dQaccum
+        ) % self.num_threads == 0
+        self.gmem_tiled_copy_dQaccum = copy_utils.tiled_copy_1d(
+            Float32, self.num_threads, num_copy_elems_dQaccum
+        )
+
+    @cute.jit
+    def __call__(
+        self,
+        mO: cute.Tensor,
+        mdO: cute.Tensor,
+        mdPsum: cute.Tensor,
+        mLSE: Optional[cute.Tensor],
+        mLSElog2: Optional[cute.Tensor],
+        mdQaccum: Optional[cute.Tensor],
+        mCuSeqlensQ: Optional[cute.Tensor],
+        mSeqUsedQ: Optional[cute.Tensor],
+        stream: cuda.CUstream,
+    ):
+        # Get the data type and check if it is fp16 or bf16
+        if cutlass.const_expr(not (mO.element_type == mdO.element_type)):
+            raise TypeError("All tensors must have the same data type")
+        if cutlass.const_expr(mO.element_type not in [cutlass.Float16, cutlass.BFloat16]):
+            raise TypeError("Only Float16 or BFloat16 is supported")
+        if cutlass.const_expr(mdPsum.element_type not in [Float32]):
+            raise TypeError("dPsum tensor must be Float32")
+        if cutlass.const_expr(mdQaccum is not None):
+            if cutlass.const_expr(mdQaccum.element_type not in [Float32]):
+                raise TypeError("dQaccum tensor must be Float32")
+        if cutlass.const_expr(mLSE is not None):
+            assert mLSElog2 is not None, "If mLSE is provided, mLSElog2 must also be provided"
+            if cutlass.const_expr(mLSE.element_type not in [Float32]):
+                raise TypeError("LSE tensor must be Float32")
+            if cutlass.const_expr(mLSElog2.element_type not in [Float32]):
+                raise TypeError("LSElog2 tensor must be Float32")
+
+        # Assume all strides are divisible by 128 bits except the last stride
+        new_stride = lambda t: (
+            *(cute.assume(s, divby=128 // t.element_type.width) for s in t.stride[:-1]),
+            t.stride[-1],
+        )
+        mO, mdO, mdQaccum = [
+            cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t)))
+            if t is not None
+            else None
+            for t in (mO, mdO, mdQaccum)
+        ]
+
+        self._setup_attributes()
+
+        if cutlass.const_expr(mCuSeqlensQ is not None):
+            TileScheduler = SingleTileVarlenScheduler
+            num_head = mO.shape[1]
+            num_batch = mCuSeqlensQ.shape[0] - 1
+        else:
+            TileScheduler = SingleTileScheduler
+            num_head = mO.shape[2]
+            num_batch = mO.shape[0]
+
+        tile_sched_args = TileSchedulerArguments(
+            num_block=cute.ceil_div(mO.shape[1], self.m_block_size),
+            num_head=num_head,
+            num_batch=num_batch,
+            num_splits=1,
+            seqlen_k=0,
+            headdim=0,
+            headdim_v=mO.shape[2],
+            total_q=mO.shape[0],
+            tile_shape_mn=(self.m_block_size, 1),
+            mCuSeqlensQ=mCuSeqlensQ,
+            mSeqUsedQ=mSeqUsedQ,
+        )
+
+        tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
+        grid_dim = TileScheduler.get_grid_shape(tile_sched_params)
+
+        self.kernel(
+            mO,
+            mdO,
+            mdPsum,
+            mLSE,
+            mLSElog2,
+            mdQaccum,
+            mCuSeqlensQ,
+            mSeqUsedQ,
+            self.gmem_tiled_copy_O,
+            self.gmem_tiled_copy_dQaccum,
+            tile_sched_params,
+            TileScheduler,
+        ).launch(
+            grid=grid_dim,
+            block=[self.num_threads, 1, 1],
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        mO: cute.Tensor,
+        mdO: cute.Tensor,
+        mdPsum: cute.Tensor,
+        mLSE: Optional[cute.Tensor],
+        mLSElog2: Optional[cute.Tensor],
+        mdQaccum: Optional[cute.Tensor],
+        mCuSeqlensQ: Optional[cute.Tensor],
+        mSeqUsedQ: Optional[cute.Tensor],
+        gmem_tiled_copy_O: cute.TiledCopy,
+        gmem_tiled_copy_dQaccum: cute.TiledCopy,
+        tile_sched_params: ParamsBase,
+        TileScheduler: cutlass.Constexpr[Callable],
+    ):
+        # 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
+
+        if work_tile.is_valid_tile:
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Get the appropriate tiles for this thread block.
+            # ///////////////////////////////////////////////////////////////////////////////
+            seqlen = SeqlenInfoQK.create(
+                batch_size,
+                mO.shape[1],
+                0,
+                mCuSeqlensQ=mCuSeqlensQ,
+                mCuSeqlensK=None,
+                mSeqUsedQ=mSeqUsedQ,
+                mSeqUsedK=None,
+            )
+
+            if cutlass.const_expr(not seqlen.has_cu_seqlens_q):
+                mO_cur = mO[batch_size, None, num_head, None]
+                mdO_cur = mdO[batch_size, None, num_head, None]
+                mdPsum_cur = mdPsum[batch_size, num_head, None]
+                headdim_v = mO.shape[3]
+            else:
+                mO_cur = cute.domain_offset((seqlen.offset_q, 0), mO[None, num_head, None])
+                mdO_cur = cute.domain_offset((seqlen.offset_q, 0), mdO[None, num_head, None])
+
+                padded_offset_q = seqlen.offset_q + batch_size * self.m_block_size
+                mdPsum_cur = cute.domain_offset((padded_offset_q,), mdPsum[num_head, None])
+                headdim_v = mO.shape[2]
+
+            blkOdO_shape = (self.m_block_size, self.head_dim_padded)
+            # (m_block_size, head_dim)
+            gO = cute.local_tile(mO_cur, blkOdO_shape, (m_block, 0))
+            gdO = cute.local_tile(mdO_cur, blkOdO_shape, (m_block, 0))
+
+            gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx)
+            # (CPY_Atom, CPY_M, CPY_K)
+            tOgO = gmem_thr_copy_O.partition_S(gO)
+            tOgdO = gmem_thr_copy_O.partition_S(gdO)
+
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Predicate: Mark indices that need to copy when problem_shape isn't a multiple
+            # of tile_shape
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Construct identity layout for KV
+            cO = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded))
+            tOcO = gmem_thr_copy_O.partition_S(cO)
+            t0OcO = gmem_thr_copy_O.get_slice(0).partition_S(cO)
+            tOpO = utils.predicate_k(tOcO, limit=headdim_v)
+            tOpdO = utils.predicate_k(tOcO, limit=headdim_v)
+
+            seqlen_q = seqlen.seqlen_q
+            seqlen_q_rounded = cute.round_up(seqlen_q, self.m_block_size)
+
+            if cutlass.const_expr(mLSE is not None):
+                if cutlass.const_expr(not seqlen.has_cu_seqlens_q):
+                    mLSE_cur = mLSE[batch_size, num_head, None]
+                else:
+                    mLSE_cur = cute.domain_offset((seqlen.offset_q,), mLSE[num_head, None])
+
+                gLSE = cute.local_tile(mLSE_cur, (self.m_block_size,), (m_block,))
+                lse = Float32.inf
+                if tidx < seqlen_q - m_block * self.m_block_size:
+                    lse = gLSE[tidx]
+
+            tOrO = cute.make_fragment_like(tOgO)
+            tOrdO = cute.make_fragment_like(tOgdO)
+            assert cute.size(tOgO, mode=[0]) == cute.size(tOgdO, mode=[0])
+            assert cute.size(tOgO, mode=[1]) == cute.size(tOgdO, mode=[1])
+            assert cute.size(tOgO, mode=[2]) == cute.size(tOgdO, mode=[2])
+            for m in cutlass.range(cute.size(tOrO.shape[1]), unroll_full=True):
+                # Instead of using tOcO, we using t0OcO and subtract the offset from the limit
+                # (seqlen_q - m_block * kBlockM). This is because the entries of t0OcO are known at compile time.
+                if t0OcO[0, m, 0][0] < seqlen_q - m_block * self.m_block_size - tOcO[0][0]:
+                    cute.copy(
+                        gmem_thr_copy_O,
+                        tOgO[None, m, None],
+                        tOrO[None, m, None],
+                        pred=tOpO[None, m, None]
+                        if cutlass.const_expr(self.check_hdim_oob)
+                        else None,
+                    )
+                    cute.copy(
+                        gmem_thr_copy_O,
+                        tOgdO[None, m, None],
+                        tOrdO[None, m, None],
+                        pred=tOpdO[None, m, None]
+                        if cutlass.const_expr(self.check_hdim_oob)
+                        else None,
+                    )
+            # Sum across the "k" dimension
+            dpsum = (tOrO.load().to(Float32) * tOrdO.load().to(Float32)).reduce(
+                cute.ReductionOp.ADD, init_val=0.0, reduction_profile=(0, None, 1)
+            )
+            threads_per_row = gmem_tiled_copy_O.layout_src_tv_tiled[0].shape[0]
+            assert cute.arch.WARP_SIZE % threads_per_row == 0
+            dpsum = utils.warp_reduce(dpsum, operator.add, width=threads_per_row)
+            dP_sum = cute.make_fragment(cute.size(tOrO, mode=[1]), Float32)
+            dP_sum.store(dpsum)
+
+            # Write dPsum from rmem -> gmem
+            gdPsum = cute.local_tile(mdPsum_cur, (self.m_block_size,), (m_block,))
+            # Only the thread corresponding to column 0 writes out the dPsum to gmem
+            if tOcO[0, 0, 0][1] == 0:
+                for m in cutlass.range(cute.size(dP_sum), unroll_full=True):
+                    row = tOcO[0, m, 0][0]
+                    gdPsum[row] = dP_sum[m] if row < seqlen_q - m_block * self.m_block_size else 0.0
+
+            # Clear dQaccum
+            if cutlass.const_expr(mdQaccum is not None):
+                if cutlass.const_expr(not seqlen.has_cu_seqlens_q):
+                    mdQaccum_cur = mdQaccum[batch_size, num_head, None]
+                else:
+                    padded_offset_q = seqlen.offset_q + batch_size * self.m_block_size
+                    mdQaccum_cur = cute.domain_offset(
+                        (padded_offset_q * self.head_dim_padded,), mdQaccum[num_head, None]
+                    )
+
+                    # HACK: Compiler doesn't seem to recognize that padding
+                    # by padded_offset_q * self.head_dim_padded keeps alignment
+                    # since statically divisible by 4
+
+                    mdQaccum_cur_ptr = cute.make_ptr(
+                        dtype=mdQaccum_cur.element_type,
+                        value=mdQaccum_cur.iterator.toint(),
+                        mem_space=mdQaccum_cur.iterator.memspace,
+                        assumed_align=mdQaccum.iterator.alignment,
+                    )
+                    mdQaccum_cur = cute.make_tensor(mdQaccum_cur_ptr, mdQaccum_cur.layout)
+
+                blkdQaccum_shape = (self.m_block_size * self.head_dim_padded,)
+                gdQaccum = cute.local_tile(mdQaccum_cur, blkdQaccum_shape, (m_block,))
+                gmem_thr_copy_dQaccum = gmem_tiled_copy_dQaccum.get_slice(tidx)
+                tdQgdQaccum = gmem_thr_copy_dQaccum.partition_S(gdQaccum)
+                zero = cute.make_fragment_like(tdQgdQaccum)
+                zero.fill(0.0)
+                cute.copy(gmem_tiled_copy_dQaccum, zero, tdQgdQaccum)
+
+            if cutlass.const_expr(mLSE is not None):
+                if cutlass.const_expr(not seqlen.has_cu_seqlens_q):
+                    mLSElog2_cur = mLSElog2[batch_size, num_head, None]
+                else:
+                    padded_offset_q = seqlen.offset_q + batch_size * self.m_block_size
+                    mLSElog2_cur = cute.domain_offset((padded_offset_q,), mLSElog2[num_head, None])
+
+                gLSElog2 = cute.local_tile(mLSElog2_cur, (self.m_block_size,), (m_block,))
+                LOG2_E = math.log2(math.e)
+                if tidx < seqlen_q_rounded - m_block * self.m_block_size:
+                    gLSElog2[tidx] = lse * LOG2_E if lse != -Float32.inf else 0.0

build/torch-universal/flash_bwd_sm90.py

@@ -0,0 +1,1244 @@
+import math
+from typing import Callable, Optional, Type
+from functools import partial
+
+import cuda.bindings.driver as cuda
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils.hopper_helpers as sm90_utils_basic
+from cutlass.cute.nvgpu import cpasync, warpgroup
+from cutlass.cute.arch import ProxyKind, SharedSpace
+from cutlass import Float32, Int32, Boolean, const_expr
+from cutlass.utils import LayoutEnum
+
+from . import hopper_helpers as sm90_utils
+from . import utils
+from . import copy_utils
+from .hopper_helpers import gemm_zero_init, gemm_w_idx
+from .mask import AttentionMask
+from .seqlen_info import SeqlenInfoQK
+from .block_info import BlockInfo
+from . import pipeline
+from .tile_scheduler import TileSchedulerArguments, SingleTileScheduler, ParamsBase
+from .named_barrier import NamedBarrierFwd, NamedBarrierBwd
+
+
+def mma_partition_fragment_AB(
+    thr_mma: cute.core.ThrMma, sA: Optional[cute.Tensor], sB: Optional[cute.Tensor], swap_AB: bool
+):
+    if const_expr(not swap_AB):
+        return (
+            thr_mma.make_fragment_A(thr_mma.partition_A(sA)) if sA is not None else None,
+            thr_mma.make_fragment_B(thr_mma.partition_B(sB)) if sB is not None else None,
+        )
+    else:
+        return (
+            thr_mma.make_fragment_B(thr_mma.partition_B(sA)) if sA is not None else None,
+            thr_mma.make_fragment_A(thr_mma.partition_A(sB)) if sB is not None else None,
+        )
+
+
+class FlashAttentionBackwardSm90:
+    arch = 90
+
+    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,
+        tile_m: int = 64,
+        tile_n: int = 128,
+        Q_stage: int = 2,
+        dO_stage: int = 2,
+        PdS_stage: int = 2,
+        SdP_swapAB: bool = False,
+        dKV_swapAB: bool = False,
+        dQ_swapAB: bool = False,
+        AtomLayoutMSdP: int = 1,
+        AtomLayoutNdKV: int = 2,
+        AtomLayoutMdQ: int = 1,
+        num_threads: int = 384,
+        V_in_regs: bool = False,
+    ):
+        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 = False
+        self.tile_m = tile_m
+        self.tile_n = tile_n
+        self.num_threads = num_threads
+        self.Q_stage = Q_stage
+        self.dO_stage = dO_stage
+        self.PdS_stage = PdS_stage
+        assert self.dO_stage in [1, self.Q_stage]
+        assert self.PdS_stage in [1, self.Q_stage]
+        self.SdP_swapAB = SdP_swapAB
+        self.dKV_swapAB = dKV_swapAB
+        self.dQ_swapAB = dQ_swapAB
+        self.AtomLayoutMSdP = AtomLayoutMSdP
+        self.AtomLayoutNdKV = AtomLayoutNdKV
+        self.AtomLayoutMdQ = AtomLayoutMdQ
+        self.num_mma_warp_groups = (self.num_threads // 128) - 1
+        self.mma_dkv_is_rs = (
+            AtomLayoutMSdP == 1
+            and AtomLayoutNdKV == self.num_mma_warp_groups
+            and SdP_swapAB
+            and not dKV_swapAB
+        )
+        self.V_in_regs = V_in_regs
+        # These are tuned for speed
+        # Do we keep the LSE and dPsum in each thread, or split them across 8 threads that share
+        # them and then shuffle to get the value whenever we need? This can reduce register
+        # pressure when SdP_swapAB, where each thread needs to keep statistics for (kBlockM / 4)
+        # rows. If !SdP_swapAB, each thread only needs to keep statistics for 2 rows.
+        # TODO: impl these for hdim 64
+        self.shuffle_LSE = self.SdP_swapAB and self.tile_hdim <= 64
+        self.shuffle_dPsum = self.SdP_swapAB and self.tile_hdim <= 64
+
+    @staticmethod
+    def can_implement(
+        dtype,
+        head_dim,
+        head_dim_v,
+        tile_m,
+        tile_n,
+        Q_stage,
+        num_threads,
+        V_in_regs=False,
+    ) -> 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
+        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],
+        mdO_type: Type[cutlass.Numeric],
+        mLSE_type: Type[cutlass.Numeric],
+        mdPsum_type: Type[cutlass.Numeric],
+        mdQaccum_type: Type[cutlass.Numeric],
+        mdK_type: Type[cutlass.Numeric],
+        mdV_type: Type[cutlass.Numeric],
+    ):
+        # Get the data type and check if it is fp16 or bf16
+        if const_expr(not (mQ_type == mK_type == mV_type == mdO_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 [Float32]):
+            raise TypeError("LSE tensor must be Float32")
+        if const_expr(mdPsum_type not in [Float32]):
+            raise TypeError("dPsum tensor must be Float32")
+        if const_expr(mdQaccum_type not in [Float32]):
+            raise TypeError("dQaccum tensor must be Float32")
+        if const_expr(self.qhead_per_kvhead == 1):
+            if const_expr(not (mdK_type == mdV_type == mQ_type)):
+                raise TypeError("mdK and mdV tensors must have the same data type as mQ")
+        else:
+            if const_expr(not (mdK_type == mdV_type == Float32)):
+                raise TypeError("mdKaccum and mdVaccum tensors must have the data type Float32")
+        assert mQ_type == self.dtype
+
+    def _setup_attributes(self):
+        self.sQ_layout, self.sK_layout, self.sV_layout, self.sdO_layout, self.sPdS_layout = [
+            sm90_utils.make_smem_layout(self.dtype, LayoutEnum.ROW_MAJOR, shape, stage)
+            for shape, stage in [
+                ((self.tile_m, self.tile_hdim), self.Q_stage),
+                ((self.tile_n, self.tile_hdim), None),
+                ((self.tile_n, self.tile_hdimv), None),
+                ((self.tile_m, self.tile_hdimv), self.dO_stage),
+                ((self.tile_m, self.tile_n), self.PdS_stage),
+            ]
+        ]
+        self.sdQaccum_layout = cute.make_layout(
+            (self.tile_m * self.tile_hdim // self.num_mma_warp_groups, self.num_mma_warp_groups)
+        )
+        # dQaccum R->S
+        self.r2s_tiled_copy_dQaccum = cute.make_tiled_copy_tv(
+            cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128),
+            # thr_layout
+            cute.make_layout((self.num_threads_per_warp_group, self.num_mma_warp_groups)),
+            cute.make_layout(128 // Float32.width),  # val_layout
+        )
+
+    def _get_tiled_mma(self):
+        # S = Q @ K.T, dP = dO @ V.T
+        atom_layout_SdP = (self.AtomLayoutMSdP, self.num_mma_warp_groups // self.AtomLayoutMSdP)
+        tiler_mn_SdP = (self.tile_m // atom_layout_SdP[0], self.tile_n // atom_layout_SdP[1])
+        tiled_mma_SdP = sm90_utils_basic.make_trivial_tiled_mma(
+            self.dtype,
+            self.dtype,
+            warpgroup.OperandMajorMode.K,
+            warpgroup.OperandMajorMode.K,
+            Float32,
+            atom_layout_mnk=(atom_layout_SdP if not self.SdP_swapAB else atom_layout_SdP[::-1])
+            + (1,),
+            tiler_mn=tiler_mn_SdP if not self.SdP_swapAB else tiler_mn_SdP[::-1],
+        )
+        # dV = P.T @ dO, dK = dS.T @ Q
+        atom_layout_dKV = (self.AtomLayoutNdKV, self.num_mma_warp_groups // self.AtomLayoutNdKV)
+        tiler_mn_dK = (self.tile_n // atom_layout_dKV[0], self.tile_hdim // atom_layout_dKV[1])
+        tiler_mn_dV = (self.tile_n // atom_layout_dKV[0], self.tile_hdimv // atom_layout_dKV[1])
+        tiled_mma_dK, tiled_mma_dV = [
+            sm90_utils_basic.make_trivial_tiled_mma(
+                self.dtype,
+                self.dtype,
+                warpgroup.OperandMajorMode.MN
+                if not self.mma_dkv_is_rs
+                else warpgroup.OperandMajorMode.K,
+                warpgroup.OperandMajorMode.MN,
+                Float32,
+                atom_layout_mnk=(atom_layout_dKV if not self.dKV_swapAB else atom_layout_dKV[::-1])
+                + (1,),
+                tiler_mn=tiler_mn_d if not self.dKV_swapAB else tiler_mn_d[::-1],
+                a_source=warpgroup.OperandSource.RMEM
+                if self.mma_dkv_is_rs
+                else warpgroup.OperandSource.SMEM,
+            )
+            for tiler_mn_d in (tiler_mn_dK, tiler_mn_dV)
+        ]
+        # dQ = dS @ K
+        atom_layout_dQ = (self.AtomLayoutMdQ, self.num_mma_warp_groups // self.AtomLayoutMdQ)
+        tiler_mn_dQ = (self.tile_m // atom_layout_dQ[0], self.tile_hdim // atom_layout_dQ[1])
+        tiled_mma_dQ = sm90_utils_basic.make_trivial_tiled_mma(
+            self.dtype,
+            self.dtype,
+            warpgroup.OperandMajorMode.K if not self.dQ_swapAB else warpgroup.OperandMajorMode.MN,
+            warpgroup.OperandMajorMode.MN if not self.dQ_swapAB else warpgroup.OperandMajorMode.K,
+            Float32,
+            atom_layout_mnk=(atom_layout_dQ if not self.dQ_swapAB else atom_layout_dQ[::-1]) + (1,),
+            tiler_mn=tiler_mn_dQ if not self.dQ_swapAB else tiler_mn_dQ[::-1],
+        )
+        return tiled_mma_SdP, tiled_mma_dK, tiled_mma_dV, tiled_mma_dQ
+
+    def _get_shared_storage_cls(self):
+        sQ_alignment = sK_alignment = sV_alighment = sdQaccum_alignment = sdO_alignment = 1024
+
+        sQ_struct, sK_struct, sV_struct, sdO_struct, sdQaccum_struct = [
+            cute.struct.Align[cute.struct.MemRange[type, cute.cosize(layout)], alignment]
+            for (layout, type, alignment) in [
+                (self.sQ_layout, self.dtype, sQ_alignment),
+                (self.sK_layout, self.dtype, sK_alignment),
+                (self.sV_layout, self.dtype, sV_alighment),
+                (self.sdO_layout, self.dtype, sdO_alignment),
+                (self.sdQaccum_layout, Float32, sdQaccum_alignment),
+            ]
+        ]
+
+        cosize_sdS = cute.cosize(self.sPdS_layout)
+        cosize_sP = cute.cosize(self.sPdS_layout) if const_expr(not self.mma_dkv_is_rs) else 0
+        sLSE_struct = cute.struct.Align[
+            cute.struct.MemRange[Float32, cute.round_up(self.tile_m, 64) * self.Q_stage], 128
+        ]
+        sdPsum_struct = cute.struct.Align[
+            cute.struct.MemRange[Float32, cute.round_up(self.tile_m, 64) * self.dO_stage], 128
+        ]
+
+        @cute.struct
+        class SharedStorageQKV:
+            mbar_ptr_Q: cute.struct.MemRange[cutlass.Int64, self.Q_stage * 2]
+            mbar_ptr_dO: cute.struct.MemRange[cutlass.Int64, self.dO_stage * 2]
+            sLSE: sLSE_struct
+            sdPsum: sdPsum_struct
+            sQ: sQ_struct
+            sV: sV_struct
+            sK: sK_struct
+            sdO: sdO_struct
+            sP: cute.struct.Align[cute.struct.MemRange[self.dtype, cosize_sP], 1024]
+            sdS: cute.struct.Align[cute.struct.MemRange[self.dtype, cosize_sdS], 1024]
+            sdQaccum: sdQaccum_struct
+
+        return SharedStorageQKV
+
+    @cute.jit
+    def __call__(
+        self,
+        mQ: cute.Tensor,
+        mK: cute.Tensor,
+        mV: cute.Tensor,
+        mdO: cute.Tensor,
+        mLSE: cute.Tensor,
+        mdPsum: cute.Tensor,
+        mdQaccum: cute.Tensor,
+        mdK: cute.Tensor,
+        mdV: cute.Tensor,
+        softmax_scale: Float32,
+        stream: cuda.CUstream,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        mSeqUsedQ: Optional[cute.Tensor] = None,
+        mSeqUsedK: Optional[cute.Tensor] = None,
+        softcap: Float32 | float | None = None,
+        window_size_left: Int32 | int | None = None,
+        window_size_right: Int32 | int | None = None,
+    ):
+        self._check_type(
+            *(
+                t.element_type if t is not None else None
+                for t in (mQ, mK, mV, mdO, mLSE, mdPsum, mdQaccum, mdK, mdV)
+            )
+        )
+
+        # Assume all strides are divisible by 128 bits except the last stride
+        new_stride = lambda t: (
+            *(cute.assume(s, divby=128 // t.element_type.width) for s in t.stride[:-1]),
+            t.stride[-1],
+        )
+        mQ, mK, mV, mdO, mLSE, mdPsum, mdQaccum, mdK, mdV = [
+            cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t)))
+            if t is not None
+            else None
+            for t in (mQ, mK, mV, mdO, mLSE, mdPsum, mdQaccum, mdK, mdV)
+        ]
+
+        layout_transpose = [1, 3, 2, 0]  # (b, s, n, h) --> (s, h, n, b)
+        mQ, mK, mV, mdK, mdV, mdO = [
+            utils.select(t, layout_transpose) for t in (mQ, mK, mV, mdK, mdV, mdO)
+        ]
+        LSE_dPsum_dQaccum_transpose = [2, 1, 0]  # (b, n, s) -> (s, n, b)
+        mLSE, mdPsum, mdQaccum = [
+            utils.select(t, LSE_dPsum_dQaccum_transpose) for t in (mLSE, mdPsum, mdQaccum)
+        ]
+
+        tiled_mma_SdP, tiled_mma_dK, tiled_mma_dV, tiled_mma_dQ = self._get_tiled_mma()
+
+        self.num_mma_threads = tiled_mma_SdP.size
+        assert self.num_mma_threads + 128 == self.num_threads
+
+        self.num_threads_per_warp_group = 128
+        self.num_producer_threads = 32
+
+        self.num_mma_regs = 240
+        self.num_producer_regs = 24
+        # self.num_mma_regs = 232
+        # self.num_producer_regs = 40
+
+        self._setup_attributes()
+        SharedStorage = self._get_shared_storage_cls()
+
+        self.tma_copy_bytes = {
+            name: cute.size_in_bytes(mX.element_type, cute.select(layout, mode=[0, 1]))
+            for name, mX, layout in [
+                ("Q", mQ, self.sQ_layout),
+                ("K", mK, self.sK_layout),
+                ("V", mV, self.sV_layout),
+                ("dO", mdO, self.sdO_layout),
+            ]
+        }
+        self.tma_copy_bytes["LSE"] = self.tile_m * Float32.width // 8
+        self.tma_copy_bytes["dPsum"] = self.tile_m * Float32.width // 8
+        self.tma_copy_bytes["dQ"] = (
+            self.tile_m * self.tile_hdim * Float32.width // 8 // self.num_mma_warp_groups
+        )
+
+        tma_atom_Q, tma_tensor_Q = cpasync.make_tiled_tma_atom(
+            cpasync.CopyBulkTensorTileG2SOp(),
+            mQ,
+            cute.select(self.sQ_layout, mode=[0, 1]),
+            (self.tile_m, self.tile_hdim),
+        )
+        tma_atom_K, tma_tensor_K = cpasync.make_tiled_tma_atom(
+            cpasync.CopyBulkTensorTileG2SOp(),
+            mK,
+            cute.select(self.sK_layout, mode=[0, 1]),
+            (self.tile_n, self.tile_hdim),
+        )
+        tma_atom_V, tma_tensor_V = cpasync.make_tiled_tma_atom(
+            cpasync.CopyBulkTensorTileG2SOp(),
+            mV,
+            cute.select(self.sV_layout, mode=[0, 1]),
+            (self.tile_n, self.tile_hdimv),
+        )
+        tma_atom_dO, tma_tensor_dO = cpasync.make_tiled_tma_atom(
+            cpasync.CopyBulkTensorTileG2SOp(),
+            mdO,
+            cute.select(self.sdO_layout, mode=[0, 1]),
+            (self.tile_m, self.tile_hdimv),
+        )
+        tma_atom_dK, tma_tensor_dK = cpasync.make_tiled_tma_atom(
+            cpasync.CopyBulkTensorTileS2GOp(),
+            mdK,
+            cute.select(self.sK_layout, mode=[0, 1]),
+            (self.tile_n, self.tile_hdim),
+        )
+        tma_atom_dV, tma_tensor_dV = cpasync.make_tiled_tma_atom(
+            cpasync.CopyBulkTensorTileS2GOp(),
+            mdV,
+            cute.select(self.sV_layout, mode=[0, 1]),
+            (self.tile_n, self.tile_hdimv),
+        )
+
+        TileScheduler = SingleTileScheduler
+        tile_sched_args = TileSchedulerArguments(
+            cute.ceil_div(cute.size(mK.shape[0]), self.tile_n),
+            cute.size(mK.shape[2]),
+            cute.size(mK.shape[3]),
+            1,  # num_splits
+            cute.size(mK.shape[0]),
+            mQ.shape[1],
+            mV.shape[1],
+            total_q=cute.size(mQ.shape[0]) * cute.size(mQ.shape[3]),
+            tile_shape_mn=(self.tile_m, self.tile_n),
+            mCuSeqlensQ=None,
+            mSeqUsedQ=None,
+            qhead_per_kvhead_packgqa=1,
+            element_size=self.dtype.width // 8,
+            is_persistent=False,
+            lpt=False,
+        )
+
+        tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
+        grid_dim = TileScheduler.get_grid_shape(tile_sched_params)
+
+        LOG2_E = math.log2(math.e)
+        softmax_scale_log2 = softmax_scale * LOG2_E
+
+        self.kernel(
+            tma_tensor_Q,
+            tma_tensor_K,
+            tma_tensor_V,
+            tma_tensor_dO,
+            tma_tensor_dK,
+            tma_tensor_dV,
+            tma_atom_Q,
+            tma_atom_K,
+            tma_atom_V,
+            tma_atom_dO,
+            tma_atom_dK,
+            tma_atom_dV,
+            mLSE,
+            mdPsum,
+            mdQaccum,
+            self.sQ_layout,
+            self.sK_layout,
+            self.sV_layout,
+            self.sPdS_layout,
+            self.sdO_layout,
+            self.sdQaccum_layout,
+            self.r2s_tiled_copy_dQaccum,
+            tiled_mma_SdP,
+            tiled_mma_dK,
+            tiled_mma_dV,
+            tiled_mma_dQ,
+            softmax_scale_log2,
+            softmax_scale,
+            tile_sched_params,
+            TileScheduler,
+            SharedStorage,
+        ).launch(
+            grid=grid_dim,
+            block=[self.num_threads, 1, 1],
+            smem=SharedStorage.size_in_bytes(),
+            stream=stream,
+            min_blocks_per_mp=1,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        mQ: cute.Tensor,
+        mK: cute.Tensor,
+        mV: cute.Tensor,
+        mdO: cute.Tensor,
+        mdK: cute.Tensor,
+        mdV: cute.Tensor,
+        tma_atom_Q: cute.CopyAtom,
+        tma_atom_K: cute.CopyAtom,
+        tma_atom_V: cute.CopyAtom,
+        tma_atom_dO: cute.CopyAtom,
+        tma_atom_dK: cute.CopyAtom,
+        tma_atom_dV: cute.CopyAtom,
+        mLSE: cute.Tensor,
+        mdPsum: cute.Tensor,
+        mdQaccum: cute.Tensor,
+        sQ_layout: cute.ComposedLayout,
+        sK_layout: cute.ComposedLayout,
+        sV_layout: cute.ComposedLayout,
+        sPdS_layout: cute.ComposedLayout,
+        sdO_layout: cute.ComposedLayout,
+        sdQaccum_layout: cute.Layout,
+        r2s_tiled_copy_dQaccum: cute.TiledCopy,
+        tiled_mma_SdP: cute.TiledMma,
+        tiled_mma_dK: cute.TiledMma,
+        tiled_mma_dV: cute.TiledMma,
+        tiled_mma_dQ: cute.TiledMma,
+        softmax_scale_log2,
+        softmax_scale,
+        tile_sched_params: ParamsBase,
+        TileScheduler: cutlass.Constexpr[Callable],
+        SharedStorage: cutlass.Constexpr[Callable],
+    ):
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+
+        # prefetch TMA descriptors
+        if warp_idx == 0:
+            cpasync.prefetch_descriptor(tma_atom_Q)
+            cpasync.prefetch_descriptor(tma_atom_K)
+            cpasync.prefetch_descriptor(tma_atom_V)
+            cpasync.prefetch_descriptor(tma_atom_dO)
+
+        smem = cutlass.utils.SmemAllocator()
+        storage = smem.allocate(SharedStorage)
+
+        pipeline_producer_group = cutlass.pipeline.CooperativeGroup(cutlass.pipeline.Agent.Thread)
+        pipeline_consumer_group = cutlass.pipeline.CooperativeGroup(
+            cutlass.pipeline.Agent.Thread, self.num_mma_threads // self.num_threads_per_warp_group
+        )
+        pipeline_Q = pipeline.PipelineTmaAsync.create(
+            barrier_storage=storage.mbar_ptr_Q.data_ptr(),
+            num_stages=self.Q_stage,
+            producer_group=pipeline_producer_group,
+            consumer_group=pipeline_consumer_group,
+            tx_count=self.tma_copy_bytes["Q"] + self.tma_copy_bytes["LSE"],
+            init_wait=False,
+        )
+        pipeline_dO = pipeline.PipelineTmaAsync.create(
+            barrier_storage=storage.mbar_ptr_dO.data_ptr(),
+            num_stages=self.dO_stage,
+            producer_group=pipeline_producer_group,
+            consumer_group=pipeline_consumer_group,
+            tx_count=self.tma_copy_bytes["dO"] + self.tma_copy_bytes["dPsum"],
+            init_wait=True,
+        )
+
+        sQ = storage.sQ.get_tensor(sQ_layout.outer, swizzle=sQ_layout.inner)
+        sdO = storage.sdO.get_tensor(sdO_layout.outer, swizzle=sdO_layout.inner)
+        sK = storage.sK.get_tensor(sK_layout.outer, swizzle=sK_layout.inner)
+        sV = storage.sV.get_tensor(sV_layout.outer, swizzle=sV_layout.inner)
+        sP = None
+        if const_expr(not self.mma_dkv_is_rs):
+            sP = storage.sP.get_tensor(sPdS_layout.outer, swizzle=sPdS_layout.inner)
+        sdS = storage.sdS.get_tensor(sPdS_layout.outer, swizzle=sPdS_layout.inner)
+        sLSE = storage.sLSE.get_tensor(
+            cute.make_layout(
+                (self.tile_m, self.Q_stage),
+                stride=(1, cute.round_up(self.tile_m, 64)),
+            )
+        )
+        sdPsum = storage.sdPsum.get_tensor(
+            cute.make_layout(
+                (self.tile_m, self.dO_stage),
+                stride=(1, cute.round_up(self.tile_m, 64)),
+            )
+        )
+        sdQaccum = storage.sdQaccum.get_tensor(sdQaccum_layout)
+
+        block_info = BlockInfo(
+            self.tile_m,
+            self.tile_n,
+            self.is_causal,
+            self.is_local,
+            False,  # is_split_kv
+            None,
+            None,
+            qhead_per_kvhead_packgqa=1,
+        )
+        SeqlenInfoCls = partial(
+            SeqlenInfoQK.create,
+            seqlen_q_static=mQ.shape[0],
+            seqlen_k_static=mK.shape[0],
+            mCuSeqlensQ=None,
+            mCuSeqlensK=None,
+            mSeqUsedQ=None,
+            mSeqUsedK=None,
+        )
+        AttentionMaskCls = partial(
+            AttentionMask,
+            self.tile_m,
+            self.tile_n,
+            window_size_left=None,
+            window_size_right=None,
+        )
+        TileSchedulerCls = partial(TileScheduler.create, tile_sched_params)
+
+        if warp_idx < 4:
+            cute.arch.warpgroup_reg_dealloc(self.num_producer_regs)
+            if warp_idx == 0:
+                self.load(
+                    mQ,
+                    mK,
+                    mV,
+                    mdO,
+                    mLSE,
+                    mdPsum,
+                    sQ,
+                    sK,
+                    sV,
+                    sdO,
+                    sLSE,
+                    sdPsum,
+                    tma_atom_Q,
+                    tma_atom_K,
+                    tma_atom_V,
+                    tma_atom_dO,
+                    pipeline_Q,
+                    pipeline_dO,
+                    block_info,
+                    SeqlenInfoCls,
+                    TileSchedulerCls,
+                )
+            if warp_idx == 1:
+                for warp_group_idx in cutlass.range(self.num_mma_warp_groups):
+                    cute.arch.barrier_arrive(
+                        barrier_id=int(NamedBarrierBwd.dQEmptyWG0) + warp_group_idx,
+                        number_of_threads=self.num_threads_per_warp_group + cute.arch.WARP_SIZE,
+                    )
+                self.dQaccum_store(mdQaccum, sdQaccum, block_info, TileSchedulerCls, SeqlenInfoCls)
+        else:
+            cute.arch.warpgroup_reg_alloc(self.num_mma_regs)
+            tidx, _, _ = cute.arch.thread_idx()
+            tidx = tidx - 128
+            self.mma(
+                tiled_mma_SdP,
+                tiled_mma_dK,
+                tiled_mma_dV,
+                tiled_mma_dQ,
+                mdK,
+                mdV,
+                mdQaccum,
+                sQ,
+                sK,
+                sV,
+                sdO,
+                sP,
+                sdS,
+                sLSE,
+                sdPsum,
+                sdQaccum,
+                pipeline_Q,
+                pipeline_dO,
+                tidx,
+                tma_atom_dK,
+                tma_atom_dV,
+                r2s_tiled_copy_dQaccum,
+                softmax_scale_log2,
+                softmax_scale,
+                block_info,
+                SeqlenInfoCls,
+                AttentionMaskCls,
+                TileSchedulerCls,
+            )
+
+    @cute.jit
+    def load(
+        self,
+        mQ: cute.Tensor,
+        mK: cute.Tensor,
+        mV: cute.Tensor,
+        mdO: cute.Tensor,
+        mLSE: cute.Tensor,
+        mdPsum: cute.Tensor,
+        sQ: cute.Tensor,
+        sK: cute.Tensor,
+        sV: cute.Tensor,
+        sdO: cute.Tensor,
+        sLSE: cute.Tensor,
+        sdPsum: cute.Tensor,
+        tma_atom_Q: cute.CopyAtom,
+        tma_atom_K: cute.CopyAtom,
+        tma_atom_V: cute.CopyAtom,
+        tma_atom_dO: cute.CopyAtom,
+        pipeline_Q: cutlass.pipeline.PipelineAsync,
+        pipeline_dO: cutlass.pipeline.PipelineAsync,
+        block_info: BlockInfo,
+        SeqlenInfoCls: Callable,
+        TileSchedulerCls: Callable,
+    ):
+        warp_idx_in_wg = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4
+
+        if warp_idx_in_wg == 0:
+            producer_state_Q = cutlass.pipeline.make_pipeline_state(
+                cutlass.pipeline.PipelineUserType.Producer, self.Q_stage
+            )
+            producer_state_dO = cutlass.pipeline.make_pipeline_state(
+                cutlass.pipeline.PipelineUserType.Producer, self.dO_stage
+            )
+            tile_scheduler = TileSchedulerCls()
+            work_tile = tile_scheduler.initial_work_tile_info()
+            while work_tile.is_valid_tile:
+                n_block, head_idx, batch_idx, _ = work_tile.tile_idx
+                seqlen = SeqlenInfoCls(batch_idx)
+                mK_cur = mK[None, None, head_idx, batch_idx]
+                gK = cute.local_tile(mK_cur, (self.tile_n, self.tile_hdim), (n_block, 0))
+                mV_cur = mV[None, None, head_idx, batch_idx]
+                gV = cute.local_tile(mV_cur, (self.tile_n, self.tile_hdimv), (n_block, 0))
+
+                mQ_cur = mQ[None, None, head_idx, batch_idx]
+                gQ = cute.local_tile(mQ_cur, (self.tile_m, self.tile_hdim), (None, 0))
+                mdO_cur = mdO[None, None, head_idx, batch_idx]
+                gdO = cute.local_tile(mdO_cur, (self.tile_m, self.tile_hdimv), (None, 0))
+                mLSE_cur = mLSE[None, head_idx, batch_idx]
+                gLSE = cute.local_tile(mLSE_cur, (self.tile_m,), (None,))
+                mdPsum_cur = mdPsum[None, head_idx, batch_idx]
+                gdPsum = cute.local_tile(mdPsum_cur, (self.tile_m,), (None,))
+
+                load_K, _, _ = copy_utils.tma_get_copy_fn(
+                    tma_atom_K, 0, cute.make_layout(1), gK, sK, single_stage=True
+                )
+                load_V, _, _ = copy_utils.tma_get_copy_fn(
+                    tma_atom_V, 0, cute.make_layout(1), gV, sV, single_stage=True
+                )
+                load_Q, _, _ = copy_utils.tma_get_copy_fn(
+                    tma_atom_Q, 0, cute.make_layout(1), gQ, sQ
+                )
+                load_Q = copy_utils.tma_producer_copy_fn(load_Q, pipeline_Q)
+                load_dO, _, _ = copy_utils.tma_get_copy_fn(
+                    tma_atom_dO, 0, cute.make_layout(1), gdO, sdO
+                )
+                load_dO = copy_utils.tma_producer_copy_fn(load_dO, pipeline_dO)
+                load_LSE = copy_utils.cpasync_bulk_get_copy_fn(gLSE, sLSE)
+                load_LSE = copy_utils.tma_producer_copy_fn(load_LSE, pipeline_Q)
+                load_dPsum = copy_utils.cpasync_bulk_get_copy_fn(gdPsum, sdPsum)
+                load_dPsum = copy_utils.tma_producer_copy_fn(load_dPsum, pipeline_dO)
+
+                m_block_min, m_block_max = block_info.get_m_block_min_max(seqlen, n_block)
+                # First iteration: load K together w Q & LSE, then V together w dO & dPsum
+                m_block = m_block_min
+                pipeline_Q.producer_acquire(
+                    producer_state_Q, extra_tx_count=self.tma_copy_bytes["K"]
+                )
+                load_K(tma_bar_ptr=pipeline_Q.producer_get_barrier(producer_state_Q))
+                load_Q(m_block, producer_state=producer_state_Q)
+                # cp.async.bulk is using ptx, so we need to elect one thread to do it
+                with cute.arch.elect_one():
+                    load_LSE(m_block, producer_state=producer_state_Q)
+                producer_state_dO_cur = (
+                    producer_state_dO
+                    if const_expr(self.Q_stage != self.dO_stage)
+                    else producer_state_Q
+                )
+                pipeline_dO.producer_acquire(
+                    producer_state_dO_cur, extra_tx_count=self.tma_copy_bytes["V"]
+                )
+                load_V(tma_bar_ptr=pipeline_dO.producer_get_barrier(producer_state_dO_cur))
+                load_dO(m_block, producer_state=producer_state_dO_cur)
+                with cute.arch.elect_one():
+                    load_dPsum(m_block, producer_state=producer_state_dO_cur)
+                producer_state_Q.advance()
+                producer_state_dO.advance()
+                # Subsequent iterations: load Q & LSE, then dO & dPsum
+                for m_block in cutlass.range(m_block_min + 1, m_block_max, unroll=1):
+                    pipeline_Q.producer_acquire(producer_state_Q)
+                    load_Q(m_block, producer_state=producer_state_Q)
+                    # cp.async.bulk is using ptx, so we need to elect one thread to do it
+                    with cute.arch.elect_one():
+                        load_LSE(m_block, producer_state=producer_state_Q)
+                    producer_state_dO_cur = (
+                        producer_state_dO
+                        if const_expr(self.Q_stage != self.dO_stage)
+                        else producer_state_Q
+                    )
+                    pipeline_dO.producer_acquire(producer_state_dO_cur)
+                    load_dO(m_block, producer_state=producer_state_dO_cur)
+                    with cute.arch.elect_one():
+                        load_dPsum(m_block, producer_state=producer_state_dO_cur)
+                    producer_state_Q.advance()
+                    producer_state_dO.advance()
+
+                tile_scheduler.prefetch_next_work()
+                tile_scheduler.advance_to_next_work()
+                work_tile = tile_scheduler.get_current_work()
+
+    @cute.jit
+    def mma(
+        self,
+        tiled_mma_SdP: cute.TiledMma,
+        tiled_mma_dK: cute.TiledMma,
+        tiled_mma_dV: cute.TiledMma,
+        tiled_mma_dQ: cute.TiledMma,
+        mdK: cute.Tensor,
+        mdV: cute.Tensor,
+        mdQaccum: cute.Tensor,
+        sQ: cute.Tensor,
+        sK: cute.Tensor,
+        sV: cute.Tensor,
+        sdO: cute.Tensor,
+        sP: Optional[cute.Tensor],
+        sdS: cute.Tensor,
+        sLSE: cute.Tensor,
+        sdPsum: cute.Tensor,
+        sdQaccum: cute.Tensor,
+        pipeline_Q: cutlass.pipeline.PipelineAsync,
+        pipeline_dO: cutlass.pipeline.PipelineAsync,
+        tidx: Int32,
+        tma_atom_dK: cute.CopyAtom,
+        tma_atom_dV: cute.CopyAtom,
+        r2s_tiled_copy_dQaccum: cute.TiledCopy,
+        softmax_scale_log2: Float32,
+        softmax_scale: Float32,
+        block_info: BlockInfo,
+        SeqlenInfoCls: Callable,
+        AttentionMaskCls: Callable,
+        TileSchedulerCls: Callable,
+    ):
+        warp_group_idx = cute.arch.make_warp_uniform(tidx // self.num_threads_per_warp_group)
+        warp_group_thread_layout = cute.make_layout(
+            self.num_mma_warp_groups, stride=self.num_threads_per_warp_group
+        )
+        thr_mma_SdP = tiled_mma_SdP.get_slice(tidx)
+        wg_mma_SdP = tiled_mma_SdP.get_slice(warp_group_thread_layout(warp_group_idx))
+        wg_mma_dK = tiled_mma_dK.get_slice(warp_group_thread_layout(warp_group_idx))
+        wg_mma_dV = tiled_mma_dV.get_slice(warp_group_thread_layout(warp_group_idx))
+        wg_mma_dQ = tiled_mma_dQ.get_slice(warp_group_thread_layout(warp_group_idx))
+        # S = Q @ K.T
+        tSrQ, tSrK = mma_partition_fragment_AB(wg_mma_SdP, sQ, sK, self.SdP_swapAB)
+        # dP = dO @ V.T
+        tdPrdO, tdPrV = mma_partition_fragment_AB(wg_mma_SdP, sdO, sV, self.SdP_swapAB)
+        # dV += P.T @ dO
+        sPt = utils.transpose_view(sP) if sP is not None else None
+        sdOt = utils.transpose_view(sdO)
+        tdVrPt, tdVrdOt = mma_partition_fragment_AB(wg_mma_dV, sPt, sdOt, self.dKV_swapAB)
+        # dK += dS.T @ Q
+        sdSt = utils.transpose_view(sdS)
+        sQt = utils.transpose_view(sQ)
+        tdKrdSt, tdKrQt = mma_partition_fragment_AB(wg_mma_dK, sdSt, sQt, self.dKV_swapAB)
+        # dQ = dS @ K
+        sKt = utils.transpose_view(sK)
+        tdQrdS, tdQrKt = mma_partition_fragment_AB(wg_mma_dQ, sdS, sKt, self.dQ_swapAB)
+
+        # Smem copy atom tiling
+        smem_copy_atom_PdS = utils.get_smem_store_atom(
+            self.arch, self.dtype, transpose=self.SdP_swapAB
+        )
+        smem_thr_copy_PdS = cute.make_tiled_copy_C(smem_copy_atom_PdS, tiled_mma_SdP).get_slice(
+            tidx
+        )
+        tPsP = None
+        if const_expr(sP is not None):
+            tPsP = smem_thr_copy_PdS.partition_D(sP if const_expr(not self.SdP_swapAB) else sPt)
+        tdSsdS = smem_thr_copy_PdS.partition_D(sdS if const_expr(not self.SdP_swapAB) else sdSt)
+
+        sLSE_mma = cute.make_tensor(
+            sLSE.iterator,
+            cute.make_layout(
+                (self.tile_m, self.tile_n, self.Q_stage),
+                stride=(1, 0, cute.round_up(self.tile_m, 64)),
+            ),
+        )
+        sdPsum_mma = cute.make_tensor(
+            sdPsum.iterator,
+            cute.make_layout(
+                (self.tile_m, self.tile_n, self.dO_stage),
+                stride=(1, 0, cute.round_up(self.tile_m, 64)),
+            ),
+        )
+        if const_expr(self.SdP_swapAB):
+            sLSE_mma = utils.transpose_view(sLSE_mma)
+            sdPsum_mma = utils.transpose_view(sdPsum_mma)
+        LSEslice = (None, 0, None) if const_expr(not self.SdP_swapAB) else (0, None, None)
+        tLSEsLSE = utils.make_acc_tensor_mn_view(thr_mma_SdP.partition_C(sLSE_mma))[LSEslice]
+        tLSEsdPsum = utils.make_acc_tensor_mn_view(thr_mma_SdP.partition_C(sdPsum_mma))[LSEslice]
+
+        smem_thr_copy_dQaccum = r2s_tiled_copy_dQaccum.get_slice(tidx)
+        tdQsdQaccum = smem_thr_copy_dQaccum.partition_D(sdQaccum)
+
+        dV_shape = (self.tile_n, self.tile_hdimv)
+        acc_dV = cute.make_fragment(
+            tiled_mma_dV.partition_shape_C(dV_shape if not self.dKV_swapAB else dV_shape[::-1]),
+            Float32,
+        )
+        dK_shape = (self.tile_n, self.tile_hdim)
+        acc_dK = cute.make_fragment(
+            tiled_mma_dK.partition_shape_C(dK_shape if not self.dKV_swapAB else dK_shape[::-1]),
+            Float32,
+        )
+
+        mma_qk_fn = partial(
+            gemm_zero_init,
+            tiled_mma_SdP,
+            (self.tile_m, self.tile_n),
+            tSrQ,
+            tSrK,
+            swap_AB=self.SdP_swapAB,
+        )
+        mma_dov_fn = partial(
+            gemm_zero_init,
+            tiled_mma_SdP,
+            (self.tile_m, self.tile_n),
+            tdPrdO,
+            tdPrV,
+            swap_AB=self.SdP_swapAB,
+        )
+        if const_expr(not self.mma_dkv_is_rs):
+            mma_pdo_fn = partial(
+                gemm_w_idx, tiled_mma_dV, acc_dV, tdVrPt, tdVrdOt, swap_AB=self.dKV_swapAB
+            )
+            mma_dsq_fn = partial(
+                gemm_w_idx, tiled_mma_dK, acc_dK, tdKrdSt, tdKrQt, swap_AB=self.dKV_swapAB
+            )
+        else:
+            assert not self.dKV_swapAB
+            mma_pdo_fn = partial(gemm_w_idx, tiled_mma_dV, acc_dV, tCrB=tdVrdOt)
+            mma_dsq_fn = partial(gemm_w_idx, tiled_mma_dK, acc_dK, tCrB=tdKrQt)
+        mma_dsk_fn = partial(
+            gemm_zero_init,
+            tiled_mma_dQ,
+            (self.tile_m, self.tile_hdim),
+            tdQrdS,
+            tdQrKt,
+            swap_AB=self.dQ_swapAB,
+        )
+
+        mma_one_m_block_all = partial(
+            self.mma_one_m_block,
+            warp_group_idx=warp_group_idx,
+            mma_qk_fn=mma_qk_fn,
+            mma_dov_fn=mma_dov_fn,
+            mma_pdo_fn=mma_pdo_fn,
+            mma_dsq_fn=mma_dsq_fn,
+            mma_dsk_fn=mma_dsk_fn,
+            pipeline_Q=pipeline_Q,
+            pipeline_dO=pipeline_dO,
+            tLSEsLSE=tLSEsLSE,
+            tLSEsdPsum=tLSEsdPsum,
+            tPsP=tPsP,
+            tdSsdS=tdSsdS,
+            tdQsdQaccum=tdQsdQaccum,
+            smem_thr_copy_PdS=smem_thr_copy_PdS,
+            smem_thr_copy_dQaccum=smem_thr_copy_dQaccum,
+            softmax_scale_log2=softmax_scale_log2,
+            # acc_dV=acc_dV,
+            # acc_dK=acc_dK,
+        )
+
+        consumer_state_Q = cutlass.pipeline.make_pipeline_state(
+            cutlass.pipeline.PipelineUserType.Consumer, self.Q_stage
+        )
+        consumer_state_dO = cutlass.pipeline.make_pipeline_state(
+            cutlass.pipeline.PipelineUserType.Consumer, self.dO_stage
+        )
+        tile_scheduler = TileSchedulerCls()
+        work_tile = tile_scheduler.initial_work_tile_info()
+        while work_tile.is_valid_tile:
+            n_block, head_idx, batch_idx, _ = work_tile.tile_idx
+            seqlen = SeqlenInfoCls(batch_idx)
+            mask = AttentionMaskCls(seqlen.seqlen_q, seqlen.seqlen_k)
+            mask_fn = partial(
+                mask.apply_mask,
+                batch_idx=None,
+                head_idx=None,
+                n_block=n_block,
+                thr_mma=thr_mma_SdP,
+                mask_seqlen=True,
+                mask_causal=self.is_causal,
+                mask_local=self.is_local,
+            )
+            m_block_min, m_block_max = block_info.get_m_block_min_max(seqlen, n_block)
+            # if cute.arch.thread_idx()[0] % 32 == 0: cute.printf("tidx = {}, m_block_min = {}, m_block_max = {}", cute.arch.thread_idx()[0], m_block_min, m_block_max)
+            dKV_accumulate = False
+            for m_block in cutlass.range(m_block_min, m_block_max, unroll=1):
+                consumer_state_Q, consumer_state_dO = mma_one_m_block_all(
+                    m_block,
+                    consumer_state_Q,
+                    consumer_state_dO,
+                    mask_fn=mask_fn,
+                    dKV_accumulate=dKV_accumulate,
+                )
+                dKV_accumulate = True
+
+            # if cute.arch.thread_idx()[0] == 128: cute.print_tensor(acc_dV)
+            # scale dK
+            acc_dK.store(acc_dK.load() * softmax_scale)
+            self.epilogue_dKV(
+                acc_dV,
+                mdV,
+                sV,
+                acc_dK,
+                mdK,
+                sK,
+                seqlen,
+                tma_atom_dK,
+                tma_atom_dV,
+                tiled_mma_dK,
+                tiled_mma_dV,
+                tidx,
+                n_block,
+                head_idx,
+                batch_idx,
+            )
+
+            tile_scheduler.advance_to_next_work()
+            work_tile = tile_scheduler.get_current_work()
+
+    @cute.jit
+    def mma_one_m_block(
+        self,
+        m_block: Int32,
+        consumer_state_Q: cutlass.pipeline.PipelineState | pipeline.PipelineStateSimple,
+        consumer_state_dO: cutlass.pipeline.PipelineState | pipeline.PipelineStateSimple,
+        warp_group_idx: Int32,
+        mma_qk_fn: Callable,
+        mma_dov_fn: Callable,
+        mma_pdo_fn: Callable,
+        mma_dsq_fn: Callable,
+        mma_dsk_fn: Callable,
+        pipeline_Q: cutlass.pipeline.PipelineAsync,
+        pipeline_dO: cutlass.pipeline.PipelineAsync,
+        tLSEsLSE: cute.Tensor,
+        tLSEsdPsum: cute.Tensor,
+        tPsP: Optional[cute.Tensor],
+        tdSsdS: Optional[cute.Tensor],
+        tdQsdQaccum: cute.Tensor,
+        smem_thr_copy_PdS: cute.TiledCopy,
+        smem_thr_copy_dQaccum: cute.TiledCopy,
+        softmax_scale_log2: Float32,
+        mask_fn: Optional[Callable] = None,
+        # acc_dV,
+        # acc_dK,
+        dKV_accumulate: Boolean = True,
+    ):
+        consumer_state_dO_cur = (
+            consumer_state_dO if const_expr(self.Q_stage == self.dO_stage) else consumer_state_Q
+        )
+        smem_idx_Q = consumer_state_Q.index
+        smem_idx_dO = consumer_state_dO_cur.index if const_expr(self.dO_stage > 1) else 0
+        smem_idx_PdS = smem_idx_Q if const_expr(self.PdS_stage > 1) else 0
+        # (1) [GEMM 1] S = Q @ K^T
+        pipeline_Q.consumer_wait(consumer_state_Q, pipeline_Q.consumer_try_wait(consumer_state_Q))
+        acc_S = mma_qk_fn(A_idx=smem_idx_Q, wg_wait=-1)
+        tLSErLSE = copy_utils.load_s2r(tLSEsLSE[None, smem_idx_Q])
+        # (2) [GEMM 2] dP = dO @ V.T
+        pipeline_dO.consumer_wait(
+            consumer_state_dO_cur, pipeline_dO.consumer_try_wait(consumer_state_dO_cur)
+        )
+        acc_dP = mma_dov_fn(A_idx=smem_idx_Q, wg_wait=1)
+        # (3) [Pointwise 1] P = exp(S - LSE)
+        if cutlass.const_expr(mask_fn is not None):
+            mask_fn(acc_S, m_block=m_block)
+        acc_S_mn = utils.make_acc_tensor_mn_view(acc_S, transpose=self.SdP_swapAB)
+        # if cute.arch.thread_idx()[0] == 256: cute.print_tensor(acc_S_mn)
+        for r in cutlass.range_constexpr(cute.size(acc_S_mn, mode=[0])):
+            for c in cutlass.range(cute.size(acc_S_mn, mode=[1]), unroll_full=True):
+                acc_S_mn[r, c] = cute.math.exp2(
+                    acc_S_mn[r, c] * softmax_scale_log2 - tLSErLSE[r], fastmath=True
+                )
+        # if cute.arch.thread_idx()[0] == 128: cute.print_tensor(acc_S_mn)
+        tLSErdPsum = copy_utils.load_s2r(tLSEsdPsum[None, smem_idx_dO])
+
+        # Convert P from f32 -> f16
+        tdVrP = utils.cvt_f16(utils.make_acc_tensor_frgA_view(acc_S), self.dtype)
+        # R2S for P
+        if const_expr(not self.mma_dkv_is_rs):
+            # sync to ensure P has already been used in the previous iteration before overwriting
+            if const_expr(self.PdS_stage == 1):
+                cute.arch.barrier(
+                    barrier_id=int(NamedBarrierBwd.PdS), number_of_threads=self.num_mma_threads
+                )
+            tPrP = smem_thr_copy_PdS.retile(tdVrP)
+            cute.copy(smem_thr_copy_PdS, tPrP, tPsP[None, None, None, smem_idx_PdS])
+
+        # (4) [Pointwise 2] dS = P*(dP-dPsum)
+        warpgroup.wait_group(0)
+        acc_dP_mn = utils.make_acc_tensor_mn_view(acc_dP, transpose=self.SdP_swapAB)
+        # if cute.arch.thread_idx()[0] == 128: cute.print_tensor(acc_dP_mn)
+        for r in cutlass.range_constexpr(cute.size(acc_dP_mn, mode=[0])):
+            for c in cutlass.range(cute.size(acc_dP_mn, mode=[1]), unroll_full=True):
+                acc_dP_mn[r, c] = acc_S_mn[r, c] * (acc_dP_mn[r, c] - tLSErdPsum[r])
+        # if cute.arch.thread_idx()[0] == 128: cute.print_tensor(acc_dP_mn)
+        # Convert dS from f32 -> f16
+        tdKrdS = utils.cvt_f16(utils.make_acc_tensor_frgA_view(acc_dP), self.dtype)
+
+        # If there's double buffering on dS, we don't need to sync here.
+        # Otherwise we might have WG1 writing to dS before WG2 is done reading from it during MmadQ.
+        # But because both WGs have to sync at the end of the loop and double buffering,
+        # this race condition is not possible.
+        # This sync is to ensure (1) P is written in case of !mma_dkv_is_rs and
+        # (2) dS is already read by the Mma in the previous iteration in case of mma_dkv_is_rs.
+        if const_expr(not self.mma_dkv_is_rs or (self.PdS_stage == 1 and self.mma_dkv_is_rs)):
+            cute.arch.fence_proxy(ProxyKind.async_shared, space=SharedSpace.shared_cta)
+            cute.arch.barrier(
+                barrier_id=int(NamedBarrierBwd.PdS), number_of_threads=self.num_mma_threads
+            )
+
+        # R2S for dS
+        tdSrdS = smem_thr_copy_PdS.retile(tdKrdS)
+        cute.copy(smem_thr_copy_PdS, tdSrdS, tdSsdS[None, None, None, smem_idx_PdS])
+
+        # (5) [GEMM 3] dV += P.T @ dO
+        if const_expr(not self.mma_dkv_is_rs):
+            mma_pdo_fn(
+                A_idx=smem_idx_PdS, B_idx=smem_idx_dO, zero_init=not dKV_accumulate, wg_wait=-1
+            )
+        else:
+            mma_pdo_fn(tCrA=tdVrP, B_idx=smem_idx_dO, zero_init=not dKV_accumulate, wg_wait=-1)
+
+        # smem fence to make sure sdS is written before it's read by WGMMA
+        cute.arch.fence_proxy(ProxyKind.async_shared, space=SharedSpace.shared_cta)
+        cute.arch.barrier(
+            barrier_id=int(NamedBarrierBwd.PdS), number_of_threads=self.num_mma_threads
+        )
+        # (6) [GEMM 4] dQ = dS @ K
+        acc_dQ = mma_dsk_fn(A_idx=smem_idx_PdS, wg_wait=1)
+        # if cute.arch.thread_idx()[0] == 128: cute.print_tensor(acc_dV)
+        pipeline_dO.consumer_release(consumer_state_dO_cur)  # release dO as dV mma is done
+
+        # (7) [GEMM 5] dK += dS.T @ Q
+        if const_expr(not self.mma_dkv_is_rs):
+            mma_dsq_fn(
+                A_idx=smem_idx_PdS, B_idx=smem_idx_Q, zero_init=not dKV_accumulate, wg_wait=1
+            )
+        else:
+            mma_dsq_fn(tCrA=tdKrdS, B_idx=smem_idx_Q, zero_init=not dKV_accumulate, wg_wait=1)
+        # if cute.arch.thread_idx()[0] == 128: cute.print_tensor(acc_dQ)
+
+        cute.arch.barrier(
+            barrier_id=int(NamedBarrierBwd.dQEmptyWG0) + warp_group_idx,
+            number_of_threads=self.num_threads_per_warp_group + cute.arch.WARP_SIZE,
+        )
+        tdQrdQaccum_flat = cute.make_tensor(acc_dQ.iterator, cute.make_layout(tdQsdQaccum.shape))
+        cute.autovec_copy(tdQrdQaccum_flat, tdQsdQaccum)
+        cute.arch.fence_proxy(ProxyKind.async_shared, space=SharedSpace.shared_cta)
+        cute.arch.barrier_arrive(
+            barrier_id=int(NamedBarrierBwd.dQFullWG0) + warp_group_idx,
+            number_of_threads=self.num_threads_per_warp_group + cute.arch.WARP_SIZE,
+        )
+
+        warpgroup.wait_group(0)
+        # if cute.arch.thread_idx()[0] == 128: cute.print_tensor(acc_dK)
+        pipeline_Q.consumer_release(consumer_state_Q)
+        # if cute.arch.thread_idx()[0] % 32 == 0: cute.printf("tidx = {}, m_block = {}, after pipeline_Q consumer release", cute.arch.thread_idx()[0], m_block)
+
+        consumer_state_Q.advance()
+        consumer_state_dO.advance()
+        return consumer_state_Q, consumer_state_dO
+
+    @cute.jit
+    def epilogue_dKV(
+        self,
+        acc_dV: cute.Tensor,
+        mdV: cute.Tensor,
+        sV: cute.Tensor,
+        acc_dK: cute.Tensor,
+        mdK: cute.Tensor,
+        sK: cute.Tensor,
+        seqlen: SeqlenInfoQK,
+        tma_atom_dK: cute.CopyAtom,
+        tma_atom_dV: cute.CopyAtom,
+        tiled_mma_dK: cute.TiledMma,
+        tiled_mma_dV: cute.TiledMma,
+        tidx: Int32,
+        n_block: Int32,
+        head_idx: Int32,
+        batch_idx: Int32,
+    ):
+        rdV = cute.make_fragment_like(acc_dV, self.dtype)
+        rdV.store(acc_dV.load().to(self.dtype))
+        rdK = utils.cvt_f16(acc_dK, self.dtype)
+
+        cute.arch.barrier(
+            barrier_id=int(NamedBarrierFwd.Epilogue), number_of_threads=self.num_mma_threads
+        )
+
+        smem_copy_atom_dKV = cute.make_copy_atom(
+            cute.nvgpu.warp.StMatrix8x8x16bOp(transpose=self.dKV_swapAB, num_matrices=4),
+            self.dtype,
+        )
+        smem_thr_copy_dK = cute.make_tiled_copy_C(smem_copy_atom_dKV, tiled_mma_dK).get_slice(tidx)
+        smem_thr_copy_dV = cute.make_tiled_copy_C(smem_copy_atom_dKV, tiled_mma_dV).get_slice(tidx)
+        mdV_cur = mdV[None, None, head_idx, batch_idx]
+        mdK_cur = mdK[None, None, head_idx, batch_idx]
+        gdK = cute.local_tile(mdK_cur, (self.tile_n, self.tile_hdim), (n_block, 0))
+        gdV = cute.local_tile(mdV_cur, (self.tile_n, self.tile_hdimv), (n_block, 0))
+        store_dK, _, _ = copy_utils.tma_get_copy_fn(
+            tma_atom_dK, 0, cute.make_layout(1), sK, gdK, single_stage=True
+        )
+        store_dV, _, _ = copy_utils.tma_get_copy_fn(
+            tma_atom_dV, 0, cute.make_layout(1), sV, gdV, single_stage=True
+        )
+
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+        # rmem -> smem
+        taccdVrdV = smem_thr_copy_dV.retile(rdV)
+        sdV = sV if const_expr(not self.dKV_swapAB) else utils.transpose_view(sV)  # reuse sV SMEM
+        taccdVsdV = smem_thr_copy_dV.partition_D(sdV)
+        cute.copy(smem_copy_atom_dKV, taccdVrdV, taccdVsdV)
+        # ensure smem writes are visible to TMA
+        cute.arch.fence_proxy(ProxyKind.async_shared, space=SharedSpace.shared_cta)
+        cute.arch.barrier(
+            barrier_id=int(NamedBarrierFwd.Epilogue), number_of_threads=self.num_mma_threads
+        )
+        if warp_idx == 4:
+            store_dV()
+        taccdKrdK = smem_thr_copy_dK.retile(rdK)
+        sdK = sK if const_expr(not self.dKV_swapAB) else utils.transpose_view(sK)  # reuse sK SMEM
+        taccdKsdK = smem_thr_copy_dK.partition_D(sdK)  # reuse sK SMEM
+        cute.copy(smem_copy_atom_dKV, taccdKrdK, taccdKsdK)
+        # ensure smem writes are visible to TMA
+        cute.arch.fence_proxy(ProxyKind.async_shared, space=SharedSpace.shared_cta)
+        cute.arch.barrier(
+            barrier_id=int(NamedBarrierFwd.Epilogue), number_of_threads=self.num_mma_threads
+        )
+        # smem -> gmem
+        if warp_idx == 4:
+            store_dK()
+            cute.arch.cp_async_bulk_commit_group()
+            cute.arch.cp_async_bulk_wait_group(0, read=True)
+
+    @cute.jit
+    def dQaccum_store(
+        self,
+        mdQaccum: cute.Tensor,
+        sdQaccum: cute.Tensor,
+        block_info: BlockInfo,
+        TileSchedulerCls: cutlass.Constexpr[Callable],
+        SeqlenInfoCls: cutlass.Constexpr[Callable],
+    ):
+        tile_scheduler = TileSchedulerCls()
+        work_tile = tile_scheduler.initial_work_tile_info()
+        while work_tile.is_valid_tile:
+            n_block, head_idx, batch_idx, _ = work_tile.tile_idx
+            seqlen = SeqlenInfoCls(batch_idx)
+            mdQaccum_cur = mdQaccum[None, head_idx, batch_idx]
+            gdQaccum_ = cute.local_tile(mdQaccum_cur, (self.tile_m * self.tile_hdim,), (None,))
+            # (M * K / WG, WG, _)
+            gdQaccum = cute.flat_divide(
+                gdQaccum_, (self.tile_m * self.tile_hdim // self.num_mma_warp_groups,)
+            )
+            m_block_min, m_block_max = block_info.get_m_block_min_max(seqlen, n_block)
+            for m_block in cutlass.range(m_block_min, m_block_max, unroll=1):
+                for warp_group_idx in cutlass.range_constexpr(self.num_mma_warp_groups):
+                    cute.arch.barrier(
+                        barrier_id=int(NamedBarrierBwd.dQFullWG0) + warp_group_idx,
+                        number_of_threads=self.num_threads_per_warp_group + cute.arch.WARP_SIZE,
+                    )
+                    with cute.arch.elect_one():
+                        copy_utils.cpasync_reduce_bulk_add_f32(
+                            sdQaccum[None, warp_group_idx].iterator,
+                            gdQaccum[None, warp_group_idx, m_block].iterator,
+                            self.tma_copy_bytes["dQ"],
+                        )
+                    cute.arch.cp_async_bulk_commit_group()
+                for warp_group_idx in cutlass.range_constexpr(self.num_mma_warp_groups):
+                    cute.arch.cp_async_bulk_wait_group(
+                        self.num_mma_warp_groups - 1 - warp_group_idx, read=True
+                    )
+                    cute.arch.barrier_arrive(
+                        barrier_id=int(NamedBarrierBwd.dQEmptyWG0) + warp_group_idx,
+                        number_of_threads=self.num_threads_per_warp_group + cute.arch.WARP_SIZE,
+                    )
+            tile_scheduler.advance_to_next_work()
+            work_tile = tile_scheduler.get_current_work()

build/torch-universal/flash_fwd_combine.py

@@ -0,0 +1,644 @@
+# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+# A reimplementation of https://github.com/Dao-AILab/flash-attention/blob/main/hopper/flash_fwd_combine_kernel.h
+# from Cutlass C++ to Cute-DSL.
+import math
+import operator
+from typing import Type, Optional
+from functools import partial
+
+import cuda.bindings.driver as cuda
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.nvgpu import cpasync
+from cutlass import Float32, Int32, const_expr
+
+from . import utils
+from .fast_math import FastDivmod
+from .seqlen_info import SeqlenInfo
+
+
+class FlashAttentionForwardCombine:
+    def __init__(
+        self,
+        dtype: Type[cutlass.Numeric],
+        dtype_partial: Type[cutlass.Numeric],
+        head_dim: int,
+        m_block_size: int = 8,
+        k_block_size: int = 64,
+        log_max_splits: int = 4,
+        num_threads: int = 256,
+        stages: int = 4,
+    ):
+        """
+        Forward combine kernel for split attention computation.
+
+        :param dtype: output data type
+        :param dtype_partial: partial accumulation data type
+        :param head_dim: head dimension
+        :param m_block_size: m block size
+        :param k_block_size: k block size
+        :param log_max_splits: log2 of maximum splits
+        :param num_threads: number of threads
+        :param varlen: whether using variable length sequences
+        :param stages: number of pipeline stages
+        """
+        self.dtype = dtype
+        self.dtype_partial = dtype_partial
+        self.head_dim = head_dim
+        self.m_block_size = m_block_size
+        self.k_block_size = k_block_size
+        self.max_splits = 1 << log_max_splits
+        self.num_threads = num_threads
+        self.is_even_k = head_dim % k_block_size == 0
+        self.stages = stages
+
+    @staticmethod
+    def can_implement(
+        dtype, dtype_partial, head_dim, m_block_size, k_block_size,
+        log_max_splits, num_threads,
+    ) -> bool:
+        """Check if the kernel can be implemented with the given parameters."""
+        if dtype not in [cutlass.Float16, cutlass.BFloat16, cutlass.Float32]:
+            return False
+        if dtype_partial not in [cutlass.Float16, cutlass.BFloat16, Float32]:
+            return False
+        if head_dim % 8 != 0:
+            return False
+        if num_threads % 32 != 0:
+            return False
+        if m_block_size % 8 != 0:
+            return False
+        max_splits = 1 << log_max_splits
+        if max_splits > 256:
+            return False
+        if (m_block_size * max_splits) % num_threads != 0:
+            return False
+        return True
+
+    def _setup_attributes(self):
+        # GMEM copy setup for O partial
+        universal_copy_bits = 128
+        async_copy_elems = universal_copy_bits // self.dtype_partial.width
+        assert self.k_block_size % async_copy_elems == 0
+
+        k_block_gmem = (
+            128 if self.k_block_size % 128 == 0 else
+            (64 if self.k_block_size % 64 == 0 else 32)
+        )
+        gmem_threads_per_row = k_block_gmem // async_copy_elems
+        assert self.num_threads % gmem_threads_per_row == 0
+
+        # Async copy atom for O partial load
+        atom_async_copy_partial = cute.make_copy_atom(
+            cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
+            self.dtype_partial,
+            num_bits_per_copy=universal_copy_bits,
+        )
+        tOpartial_layout = cute.make_ordered_layout(
+            (self.num_threads // gmem_threads_per_row, gmem_threads_per_row),
+            order=(1, 0),
+        )
+        vOpartial_layout = cute.make_layout((1, async_copy_elems))  # 4 vals per load
+        self.gmem_tiled_copy_O_partial = cute.make_tiled_copy_tv(
+            atom_async_copy_partial, tOpartial_layout, vOpartial_layout
+        )
+
+        # GMEM copy setup for final O (use universal copy for store)
+        atom_universal_copy = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            self.dtype,
+            num_bits_per_copy=async_copy_elems * self.dtype.width,
+        )
+        self.gmem_tiled_copy_O = cute.make_tiled_copy_tv(
+            atom_universal_copy, tOpartial_layout, vOpartial_layout  # 4 vals per store
+        )
+
+        # LSE copy setup with async copy (alignment = 1)
+        lse_copy_bits = Float32.width  # 1 element per copy, width is in bits
+        m_block_smem = (
+            128 if self.m_block_size % 128 == 0 else
+            (64 if self.m_block_size % 64 == 0 else
+             (32 if self.m_block_size % 32 == 0 else
+              (16 if self.m_block_size % 16 == 0 else 8)))
+        )
+        gmem_threads_per_row_lse = m_block_smem
+        assert self.num_threads % gmem_threads_per_row_lse == 0
+
+        # Async copy atom for LSE load
+        atom_async_copy_lse = cute.make_copy_atom(
+            cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.ALWAYS),
+            Float32,
+            num_bits_per_copy=lse_copy_bits,
+        )
+        tLSE_layout = cute.make_ordered_layout(
+            (self.num_threads // gmem_threads_per_row_lse, gmem_threads_per_row_lse),
+            order=(1, 0),
+        )
+        vLSE_layout = cute.make_layout(1)
+        self.gmem_tiled_copy_LSE = cute.make_tiled_copy_tv(
+            atom_async_copy_lse, tLSE_layout, vLSE_layout
+        )
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        # Shared memory
+        # ///////////////////////////////////////////////////////////////////////////////
+
+        # Shared memory to register copy for LSE
+        self.smem_threads_per_col_lse = self.num_threads // m_block_smem
+        assert 32 % self.smem_threads_per_col_lse == 0  # Must divide warp size
+
+        s2r_layout_atom_lse = cute.make_ordered_layout(
+            (self.smem_threads_per_col_lse, self.num_threads // self.smem_threads_per_col_lse),
+            order=(0, 1),
+        )
+        self.s2r_tiled_copy_LSE = cute.make_tiled_copy_tv(
+            cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Float32),
+            s2r_layout_atom_lse,
+            cute.make_layout(1),
+        )
+
+        # LSE shared memory layout with swizzling to avoid bank conflicts
+        # This works for kBlockMSmem = 8, 16, 32, 64, 128, no bank conflicts
+        if const_expr(m_block_smem == 8):
+            smem_lse_swizzle = cute.make_swizzle(5, 0, 5)
+        elif const_expr(m_block_smem == 16):
+            smem_lse_swizzle = cute.make_swizzle(4, 0, 4)
+        else:
+            smem_lse_swizzle = cute.make_swizzle(3, 2, 3)
+        smem_layout_atom_lse = cute.make_composed_layout(
+            smem_lse_swizzle,
+            0,
+            cute.make_ordered_layout((8, m_block_smem), order=(1, 0))
+        )
+        self.smem_layout_lse = cute.tile_to_shape(
+            smem_layout_atom_lse, (self.max_splits, self.m_block_size), (0, 1)
+        )
+
+        # O partial shared memory layout (simple layout for pipeline stages)
+        self.smem_layout_o = cute.make_ordered_layout(
+            (self.m_block_size, self.k_block_size, self.stages),
+            order=(1, 0, 2)
+        )
+
+
+    @cute.jit
+    def __call__(
+        self,
+        mO_partial: cute.Tensor,
+        mLSE_partial: cute.Tensor,
+        mO: cute.Tensor,
+        mLSE: Optional[cute.Tensor] = None,
+        cu_seqlens: Optional[cute.Tensor] = None,
+        seqused: Optional[cute.Tensor] = None,
+        num_splits_dynamic_ptr: Optional[cute.Tensor] = None,
+        semaphore_to_reset: Optional[cute.Tensor] = None,
+        stream: cuda.CUstream = None,
+    ):
+        # Type checking
+        if const_expr(not (mO_partial.element_type == self.dtype_partial)):
+            raise TypeError("O partial tensor must match dtype_partial")
+        if const_expr(not (mO.element_type == self.dtype)):
+            raise TypeError("O tensor must match dtype")
+        if const_expr(not mLSE_partial.element_type in [Float32]):
+            raise TypeError("LSE partial tensor must be Float32")
+        if const_expr(mLSE is not None and not mLSE.element_type in [Float32]):
+            raise TypeError("LSE tensor must be Float32")
+
+        # Shape validation - input tensors are in user format, need to be converted to kernel format
+        if const_expr(len(mO_partial.shape) not in [4, 5]):
+            raise ValueError("O partial tensor must have 4 or 5 dimensions: (num_splits, batch, seqlen, nheads, headdim) or (num_splits, total_q, nheads, headdim)")
+        if const_expr(len(mLSE_partial.shape) not in [3, 4]):
+            raise ValueError("LSE partial tensor must have 3 or 4 dimensions: (num_splits, batch, seqlen, nheads) or (num_splits, total_q, nheads)")
+        if const_expr(len(mO.shape) not in [3, 4]):
+            raise ValueError("O tensor must have 3 or 4 dimensions: (batch, seqlen, nheads, headdim) or (total_q, nheads, headdim)")
+        if const_expr(mLSE is not None and len(mLSE.shape) not in [2, 3]):
+            raise ValueError("LSE tensor must have 2 or 3 dimensions: (batch, seqlen, nheads) or (total_q, nheads)")
+
+        # Assume all strides are divisible by 128 bits except the last stride
+        new_stride = lambda t: (*(cute.assume(s, divby=128 // t.element_type.width) for s in t.stride[:-1]), t.stride[-1])
+        mO_partial, mO = [cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t))) for t in (mO_partial, mO)]
+        # (num_splits, b, seqlen, h, d) -> (seqlen, d, num_splits, h, b)
+        # or (num_splits, total_q, h, d) -> (total_q, d, num_splits, h)
+        O_partial_layout_transpose = [2, 4, 0, 3, 1] if const_expr(cu_seqlens is None) else [1, 3, 0, 2]
+        # (b, seqlen, h, d) -> (seqlen, d, h, b) or (total_q, h, d) -> (total_q, d, h)
+        mO_partial = cute.make_tensor(mO_partial.iterator, cute.select(mO_partial.layout, mode=O_partial_layout_transpose))
+        O_layout_transpose = [1, 3, 2, 0] if const_expr(cu_seqlens is None) else [0, 2, 1]
+        mO = cute.make_tensor(mO.iterator, cute.select(mO.layout, mode=O_layout_transpose))
+        # (num_splits, b, seqlen, h) -> (seqlen, num_splits, h, b)
+        # or (num_splits, total_q, h) -> (total_q, num_splits, h)
+        LSE_partial_layout_transpose = [2, 0, 3, 1] if const_expr(cu_seqlens is None) else [1, 0, 2]
+        mLSE_partial = cute.make_tensor(mLSE_partial.iterator, cute.select(mLSE_partial.layout, mode=LSE_partial_layout_transpose))
+        # (b, seqlen, h) -> (seqlen, h, b) or (total_q, h) -> (total_q, h)
+        LSE_layout_transpose = [1, 2, 0] if const_expr(cu_seqlens is None) else [0, 1]
+        mLSE = cute.make_tensor(mLSE.iterator, cute.select(mLSE.layout, mode=LSE_layout_transpose)) if mLSE is not None else None
+
+        # Determine if we have variable length sequences
+        varlen = const_expr(cu_seqlens is not None or seqused is not None)
+
+        self._setup_attributes()
+
+        @cute.struct
+        class SharedStorage:
+            sLSE: cute.struct.Align[
+                cute.struct.MemRange[Float32, cute.cosize(self.smem_layout_lse)], 128
+            ]
+            sMaxValidSplit: cute.struct.Align[
+                cute.struct.MemRange[Int32, self.m_block_size], 128
+            ]
+            sO: cute.struct.Align[
+                cute.struct.MemRange[self.dtype_partial, cute.cosize(self.smem_layout_o)], 128
+            ]
+
+        smem_size = SharedStorage.size_in_bytes()
+
+        # Grid dimensions: (ceil_div(seqlen, m_block), ceil_div(head_dim, k_block), num_head * batch)
+        seqlen = mO_partial.shape[0]
+        num_head = mO_partial.shape[3]
+        batch_size = mO_partial.shape[4] if const_expr(cu_seqlens is None) else Int32(cu_seqlens.shape[0] - 1)
+
+        # Create FastDivmod objects for efficient division
+        seqlen_divmod = FastDivmod.create(seqlen)
+        head_divmod = FastDivmod.create(num_head)
+
+        grid_dim = (
+            cute.ceil_div(seqlen * num_head, self.m_block_size),
+            cute.ceil_div(self.head_dim, self.k_block_size),
+            batch_size,
+        )
+
+        self.kernel(
+            mO_partial,
+            mLSE_partial,
+            mO,
+            mLSE,
+            cu_seqlens,
+            seqused,
+            num_splits_dynamic_ptr,
+            semaphore_to_reset,
+            SharedStorage,
+            self.smem_layout_lse,
+            self.smem_layout_o,
+            self.gmem_tiled_copy_O_partial,
+            self.gmem_tiled_copy_O,
+            self.gmem_tiled_copy_LSE,
+            self.s2r_tiled_copy_LSE,
+            seqlen_divmod,
+            head_divmod,
+            varlen,
+        ).launch(
+            grid=grid_dim,
+            block=[self.num_threads, 1, 1],
+            smem=smem_size,
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        mO_partial: cute.Tensor,
+        mLSE_partial: cute.Tensor,
+        mO: cute.Tensor,
+        mLSE: Optional[cute.Tensor],
+        cu_seqlens: Optional[cute.Tensor],
+        seqused: Optional[cute.Tensor],
+        num_splits_dynamic_ptr: Optional[cute.Tensor],
+        semaphore_to_reset: Optional[cute.Tensor],
+        SharedStorage: cutlass.Constexpr,
+        smem_layout_lse: cute.Layout | cute.ComposedLayout,
+        smem_layout_o: cute.Layout,
+        gmem_tiled_copy_O_partial: cute.TiledCopy,
+        gmem_tiled_copy_O: cute.TiledCopy,
+        gmem_tiled_copy_LSE: cute.TiledCopy,
+        s2r_tiled_copy_LSE: cute.TiledCopy,
+        seqlen_divmod: FastDivmod,
+        head_divmod: FastDivmod,
+        varlen: cutlass.Constexpr[bool],
+    ):
+        # Thread and block indices
+        tidx, _, _ = cute.arch.thread_idx()
+        m_block, k_block, batch_idx = cute.arch.block_idx()
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        # Get shared memory buffer
+        # ///////////////////////////////////////////////////////////////////////////////
+        smem = cutlass.utils.SmemAllocator()
+        storage = smem.allocate(SharedStorage)
+        sLSE = storage.sLSE.get_tensor(smem_layout_lse)
+        sMaxValidSplit = storage.sMaxValidSplit.get_tensor((self.m_block_size,))
+        sO = storage.sO.get_tensor(smem_layout_o)
+
+        # Handle semaphore reset
+        if const_expr(semaphore_to_reset is not None):
+            if (tidx == 0 and m_block == cute.arch.grid_dim()[0] - 1 and
+                k_block == cute.arch.grid_dim()[1] - 1 and
+                batch_idx == cute.arch.grid_dim()[2] - 1):
+                semaphore_to_reset[0] = 0
+
+        # Get number of splits
+        num_splits = (
+            num_splits_dynamic_ptr[batch_idx] if const_expr(num_splits_dynamic_ptr is not None)
+            else mLSE_partial.shape[1]
+        )
+        # Handle variable length sequences using SeqlenInfo
+        seqlen_info = SeqlenInfo.create(
+            batch_idx=batch_idx,
+            seqlen_static=mO_partial.shape[0],
+            cu_seqlens=cu_seqlens,
+            seqused=seqused
+        )
+        seqlen, offset = seqlen_info.seqlen, seqlen_info.offset
+
+        # Extract number of heads (head index will be determined dynamically)
+        num_head = mO_partial.shape[3]
+        max_idx = seqlen * num_head
+
+        # Early exit for single split if dynamic
+        if (const_expr(num_splits_dynamic_ptr is None) or num_splits > 1) and (const_expr(not varlen) or m_block * self.m_block_size < max_idx):
+
+            # ===============================
+            # Step 1: Load LSE_partial from gmem to shared memory
+            # ===============================
+
+            if const_expr(cu_seqlens is None):
+                # mLSE_partial_cur = mLSE_partial[None, None, None, batch_idx]
+                mLSE_partial_cur = utils.coord_offset_i64(mLSE_partial, batch_idx, dim=3)
+            else:
+                # mLSE_partial_cur = cute.domain_offset((offset, 0, 0), mLSE_partial)
+                mLSE_partial_cur = utils.domain_offset_i64((offset, 0, 0), mLSE_partial)
+            mLSE_partial_copy = cute.tiled_divide(mLSE_partial_cur, (1,))
+
+            gmem_thr_copy_LSE = gmem_tiled_copy_LSE.get_slice(tidx)
+            tLSEsLSE = gmem_thr_copy_LSE.partition_D(sLSE)
+
+            # Create identity tensor for coordinate tracking
+            cLSE = cute.make_identity_tensor((self.max_splits, self.m_block_size))
+            tLSEcLSE = gmem_thr_copy_LSE.partition_S(cLSE)
+
+            # Load LSE partial values
+            for m in cutlass.range(cute.size(tLSEcLSE, mode=[2]), unroll_full=True):
+                mi = tLSEcLSE[0, 0, m][1]  # Get m coordinate
+                idx = m_block * self.m_block_size + mi
+                if idx < max_idx:
+                    # Calculate actual sequence position and head using FastDivmod
+                    if const_expr(not varlen):
+                        head_idx, m_idx = seqlen_divmod.divmod(idx)
+                    else:
+                        head_idx = idx // seqlen
+                        m_idx = idx - head_idx * seqlen
+                    mLSE_partial_cur_copy = mLSE_partial_copy[None, m_idx, None, head_idx]
+                    for s in cutlass.range(cute.size(tLSEcLSE, mode=[1]), unroll_full=True):
+                        si = tLSEcLSE[0, s, 0][0]  # Get split coordinate
+                        if si < num_splits:
+                            cute.copy(gmem_thr_copy_LSE, mLSE_partial_cur_copy[None, si], tLSEsLSE[None, s, m])
+                        else:
+                            tLSEsLSE[None, s, m].fill(-Float32.inf)
+                # Don't need to zero out the rest of the LSEs, as we will not write the output to gmem
+            cute.arch.cp_async_commit_group()
+
+            # ===============================
+            # Step 2: Load O_partial for pipeline stages
+            # ===============================
+
+            gmem_thr_copy_O_partial = gmem_tiled_copy_O_partial.get_slice(tidx)
+            cO = cute.make_identity_tensor((self.m_block_size, self.k_block_size))
+            tOcO = gmem_thr_copy_O_partial.partition_D(cO)
+            tOsO_partial = gmem_thr_copy_O_partial.partition_D(sO)
+            if const_expr(cu_seqlens is None):
+                # mO_partial_cur = mO_partial[None, None, None, None, batch_idx]
+                mO_partial_cur = utils.coord_offset_i64(mO_partial, batch_idx, dim=4)
+            else:
+                # mO_partial_cur = cute.domain_offset((offset, 0, 0, 0), mO_partial)
+                mO_partial_cur = utils.domain_offset_i64((offset, 0, 0, 0), mO_partial)
+
+            # Precompute these values to avoid recomputing them in the loop
+            num_rows = const_expr(cute.size(tOcO, mode=[1]))
+            tOmidx = cute.make_fragment(num_rows, cutlass.Int32)
+            tOhidx = cute.make_fragment(num_rows, cutlass.Int32)
+            tOrOptr = cute.make_fragment(num_rows, cutlass.Int64)
+            for m in cutlass.range(num_rows, unroll_full=True):
+                mi = tOcO[0, m, 0][0]  # m coordinate
+                idx = m_block * self.m_block_size + mi
+                if const_expr(not varlen):
+                    tOhidx[m], tOmidx[m] = seqlen_divmod.divmod(idx)
+                else:
+                    tOhidx[m] = idx // seqlen
+                    tOmidx[m] = idx - tOhidx[m] * seqlen
+                tOrOptr[m] = utils.elem_pointer_i64(mO_partial_cur, (tOmidx[m], k_block * self.k_block_size, 0, tOhidx[m])).toint()
+                if idx >= max_idx:
+                    tOhidx[m] = -1
+
+            tOpO = cute.make_fragment(cute.size(tOcO, [2]), cutlass.Boolean)
+            if const_expr(not self.is_even_k):
+                for k in cutlass.range(cute.size(tOpO), unroll_full=True):
+                    tOpO[k] = tOcO[0, 0, k][1] < mO_partial.shape[1] - k_block * self.k_block_size
+            # if cute.arch.thread_idx()[0] == 0 and k_block == 1: cute.print_tensor(tOpO)
+
+            load_O_partial = partial(
+                self.load_O_partial,
+                gmem_tiled_copy_O_partial,
+                tOrOptr,
+                tOsO_partial,
+                tOhidx,
+                tOpO,
+                tOcO,
+                mO_partial_cur.layout,
+            )
+
+            # Load first few stages of O_partial
+            for stage in cutlass.range(self.stages - 1, unroll_full=True):
+                if stage < num_splits:
+                    load_O_partial(stage, stage)
+                cute.arch.cp_async_commit_group()
+
+            # ===============================
+            # Step 3: Load and transpose LSE from smem to registers
+            # ===============================
+
+            # Wait for LSE and initial O partial stages to complete
+            cute.arch.cp_async_wait_group(self.stages - 1)
+            cute.arch.sync_threads()
+            # if cute.arch.thread_idx()[0] == 0:
+            #     # cute.print_tensor(sLSE)
+            #     for i in range(64):
+            #         cute.printf("sLSE[%d, 0] = %f", i, sLSE[i, 0])
+            # cute.arch.sync_threads()
+
+            s2r_thr_copy_LSE = s2r_tiled_copy_LSE.get_slice(tidx)
+            ts2rsLSE = s2r_thr_copy_LSE.partition_S(sLSE)
+            ts2rrLSE = cute.make_fragment_like(ts2rsLSE)
+            cute.copy(s2r_tiled_copy_LSE, ts2rsLSE, ts2rrLSE)
+
+            # ===============================
+            # Step 4: Compute final LSE along split dimension
+            # ===============================
+
+            lse_sum = cute.make_fragment(cute.size(ts2rrLSE, mode=[2]), Float32)
+            ts2rcLSE = s2r_thr_copy_LSE.partition_D(cLSE)
+            # We compute the max valid split for each row to short-circuit the computation later
+            max_valid_split = cute.make_fragment(cute.size(ts2rrLSE, mode=[2]), Int32)
+            assert cute.size(ts2rrLSE, mode=[0]) == 1
+            # Compute max, scales, and final LSE for each row
+            for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True):
+                # Find max LSE value across splits
+                threads_per_col = const_expr(self.smem_threads_per_col_lse)
+                lse_max = utils.warp_reduce(
+                    ts2rrLSE[None, None, m].load().reduce(cute.ReductionOp.MAX, init_val=-Float32.inf, reduction_profile=0),
+                    op=cute.arch.fmax,
+                    width=threads_per_col,
+                )
+                # if cute.arch.thread_idx()[0] == 0: cute.printf(lse_max)
+                # Find max valid split index
+                max_valid_idx = -1
+                for s in cutlass.range(cute.size(ts2rrLSE, mode=[1]), unroll_full=True):
+                    if ts2rrLSE[0, s, m] != -Float32.inf:
+                        max_valid_idx = ts2rcLSE[0, s, 0][0]  # Get split coordinate
+                # if cute.arch.thread_idx()[0] < 32: cute.printf(max_valid_idx)
+                max_valid_split[m] = utils.warp_reduce(max_valid_idx, max, width=threads_per_col)
+                # Compute exp scales and sum
+                lse_max_cur = 0.0 if lse_max == -Float32.inf else lse_max  # In case all local LSEs are -inf
+                LOG2_E = math.log2(math.e)
+                lse_sum_cur = 0.0
+                for s in cutlass.range(cute.size(ts2rrLSE, mode=[1]), unroll_full=True):
+                    scale = utils.exp2f(ts2rrLSE[0, s, m] * LOG2_E - (lse_max_cur * LOG2_E))
+                    lse_sum_cur += scale
+                    ts2rrLSE[0, s, m] = scale  # Store scale for later use
+                lse_sum_cur = utils.warp_reduce(lse_sum_cur, operator.add, width=threads_per_col)
+                lse_sum[m] = utils.logf(lse_sum_cur) + lse_max
+                # Normalize scales
+                inv_sum = 0.0 if (lse_sum_cur == 0.0 or lse_sum_cur != lse_sum_cur) else 1.0 / lse_sum_cur
+                ts2rrLSE[None, None, m].store(ts2rrLSE[None, None, m].load() * inv_sum)
+            # Store the scales exp(lse - lse_logsum) back to smem
+            cute.copy(s2r_tiled_copy_LSE, ts2rrLSE, ts2rsLSE)
+
+            # Store max valid split to smem
+            for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True):
+                if ts2rcLSE[0, 0, m][0] == 0:  # Only thread responsible for s=0 writes
+                    mi = ts2rcLSE[0, 0, m][1]
+                    if mi < self.m_block_size:
+                        sMaxValidSplit[mi] = max_valid_split[m]
+
+            # ===============================
+            # Step 5: Store final LSE to gmem
+            # ===============================
+
+            if const_expr(mLSE is not None):
+                if const_expr(cu_seqlens is None):
+                    # mLSE_cur = mLSE[None, None, batch_idx]
+                    mLSE_cur = utils.coord_offset_i64(mLSE, batch_idx, dim=2)
+                else:
+                    # mLSE_cur = cute.domain_offset((offset, 0), mLSE)
+                    mLSE_cur = utils.domain_offset_i64((offset, 0), mLSE)
+                if k_block == 0:  # Only first k_block writes LSE when mLSE is provided
+                    for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True):
+                        if ts2rcLSE[0, 0, m][0] == 0:  # Only thread responsible for s=0 writes
+                            mi = ts2rcLSE[0, 0, m][1]
+                            idx = m_block * self.m_block_size + mi
+                            if idx < max_idx:
+                                if const_expr(not varlen):
+                                    head_idx, m_idx = seqlen_divmod.divmod(idx)
+                                else:
+                                    head_idx = idx // seqlen
+                                    m_idx = idx - head_idx * seqlen
+                                mLSE_cur[m_idx, head_idx] = lse_sum[m]
+
+            # ===============================
+            # Step 6: Read O_partial and accumulate final O
+            # ===============================
+
+            cute.arch.sync_threads()
+
+            # Get max valid split for this thread
+            thr_max_valid_split = sMaxValidSplit[tOcO[0, 0, 0][0]]
+            for m in cutlass.range(1, cute.size(tOcO, mode=[1])):
+                thr_max_valid_split = max(thr_max_valid_split, sMaxValidSplit[tOcO[0, m, 0][0]])
+
+            tOrO_partial = cute.make_fragment_like(tOsO_partial[None, None, None, 0])
+            tOrO = cute.make_fragment_like(tOrO_partial, Float32)
+            tOrO.fill(0.0)
+
+            stage_load = self.stages - 1
+            stage_compute = 0
+
+            # Main accumulation loop
+            for s in cutlass.range(thr_max_valid_split + 1, unroll=4):
+                # Get scales for this split
+                scale = cute.make_fragment(num_rows, Float32)
+                for m in cutlass.range(num_rows, unroll_full=True):
+                    scale[m] = sLSE[s, tOcO[0, m, 0][0]]  # Get scale from smem
+
+                # Load next stage if needed
+                split_to_load = s + self.stages - 1
+                if split_to_load <= thr_max_valid_split:
+                    load_O_partial(split_to_load, stage_load)
+                cute.arch.cp_async_commit_group()
+                stage_load = 0 if stage_load == self.stages - 1 else stage_load + 1
+
+                # Wait for the current stage to be ready
+                cute.arch.cp_async_wait_group(self.stages - 1)
+                # We don't need __syncthreads() because each thread is just reading its own data from smem
+                # Copy from smem to registers
+                cute.autovec_copy(tOsO_partial[None, None, None, stage_compute], tOrO_partial)
+                stage_compute = 0 if stage_compute == self.stages - 1 else stage_compute + 1
+
+                # Accumulate scaled partial results
+                for m in cutlass.range(num_rows, unroll_full=True):
+                    if tOhidx[m] >= 0 and scale[m] > 0.0:
+                        tOrO[None, m, None].store(tOrO[None, m, None].load() + scale[m] * tOrO_partial[None, m, None].load().to(Float32))
+
+            # ===============================
+            # Step 7: Write final O to gmem
+            # ===============================
+
+            rO = cute.make_fragment_like(tOrO, self.dtype)
+            rO.store(tOrO.load().to(self.dtype))
+            if const_expr(cu_seqlens is None):
+                # mO_cur = mO[None, None, None, batch_idx]
+                mO_cur = utils.coord_offset_i64(mO, batch_idx, dim=3)
+            else:
+                # mO_cur = cute.domain_offset((offset, 0, 0), mO)
+                mO_cur = utils.domain_offset_i64((offset, 0, 0), mO)
+            mO_cur = utils.domain_offset_aligned((0, k_block * self.k_block_size, 0), mO_cur)
+            elems_per_store = const_expr(cute.size(gmem_tiled_copy_O.layout_tv_tiled[1]))
+            # mO_cur_copy = cute.tiled_divide(mO_cur, (1, elems_per_store,))
+            gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx)
+            # Write final results
+            for m in cutlass.range(num_rows, unroll_full=True):
+                if tOhidx[m] >= 0:
+                    mO_cur_copy = cute.tiled_divide(mO_cur[tOmidx[m], None, tOhidx[m]], (elems_per_store,))
+                    for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True):
+                        k_idx = tOcO[0, 0, k][1] // elems_per_store
+                        if const_expr(self.is_even_k) or tOpO[k]:
+                            cute.copy(gmem_thr_copy_O, rO[None, m, k], mO_cur_copy[None, k_idx])
+
+    @cute.jit
+    def load_O_partial(
+        self,
+        gmem_tiled_copy_O_partial: cute.TiledCopy,
+        tOrOptr: cute.Tensor,
+        tOsO_partial: cute.Tensor,
+        tOhidx: cute.Tensor,
+        tOpO: cute.Tensor,
+        tOcO: cute.Tensor,
+        mO_cur_partial_layout: cute.Layout,
+        split: Int32,
+        stage: Int32,
+    ) -> None:
+        elems_per_load = const_expr(cute.size(gmem_tiled_copy_O_partial.layout_tv_tiled[1]))
+        tOsO_partial_cur = tOsO_partial[None, None, None, stage]
+        for m in cutlass.range(cute.size(tOcO, [1]), unroll_full=True):
+            if tOhidx[m] >= 0:
+                o_gmem_ptr = cute.make_ptr(
+                    tOsO_partial.element_type, tOrOptr[m], cute.AddressSpace.gmem, assumed_align=16
+                )
+                mO_partial_cur = cute.make_tensor(o_gmem_ptr, cute.slice_(mO_cur_partial_layout, (0, None, None, 0)))
+                mO_partial_cur_copy = cute.tiled_divide(mO_partial_cur, (elems_per_load,))
+                for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True):
+                    k_idx = tOcO[0, 0, k][1] // elems_per_load
+                    if const_expr(self.is_even_k) or tOpO[k]:
+                        cute.copy(
+                            gmem_tiled_copy_O_partial,
+                            # mO_partial_cur_copy[None, k_idx, split],
+                            utils.coord_offset_i64(mO_partial_cur_copy, split, dim=2)[None, k_idx],
+                            tOsO_partial_cur[None, m, k]
+                        )

build/torch-universal/hopper_helpers.py

@@ -0,0 +1,102 @@
+# Copyright (c) 2025, Tri Dao.
+from typing import Type, Union, Optional
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32, Float32, Boolean, const_expr
+from cutlass.cute.nvgpu import warpgroup
+from cutlass._mlir.dialects import llvm
+from cutlass.cutlass_dsl import Numeric, dsl_user_op
+from cutlass.utils import LayoutEnum
+import cutlass.utils.hopper_helpers as sm90_utils_og
+
+
+@cute.jit
+def gemm(
+    tiled_mma: cute.TiledMma,
+    acc: cute.Tensor,
+    tCrA: cute.Tensor,
+    tCrB: cute.Tensor,
+    zero_init: cutlass.Constexpr[bool] = False,
+    wg_wait: cutlass.Constexpr[int] = 0,
+    # A_in_regs: cutlass.Constexpr[bool] = False,
+    swap_AB: cutlass.Constexpr[bool] = False,
+) -> None:
+    if const_expr(swap_AB):
+        gemm(tiled_mma, acc, tCrB, tCrA, zero_init=zero_init, wg_wait=wg_wait, swap_AB=False)
+    else:
+        warpgroup.fence()
+        # We make a new mma_atom since we'll be modifying its attribute (accumulate).
+        # Otherwise the compiler complains "operand #0 does not dominate this use"
+        mma_atom = cute.make_mma_atom(tiled_mma.op)
+        mma_atom.set(warpgroup.Field.ACCUMULATE, not zero_init)
+        for k in cutlass.range_constexpr(cute.size(tCrA.shape[2])):
+            cute.gemm(mma_atom, acc, tCrA[None, None, k], tCrB[None, None, k], acc)
+            mma_atom.set(warpgroup.Field.ACCUMULATE, True)
+        warpgroup.commit_group()
+        if const_expr(wg_wait >= 0):
+            warpgroup.wait_group(wg_wait)
+
+
+def gemm_zero_init(
+    tiled_mma: cute.TiledMma,
+    shape: cute.Shape,
+    tCrA: cute.Tensor,
+    tCrB: cute.Tensor,
+    A_idx: Optional[Int32] = None,
+    B_idx: Optional[Int32] = None,
+    wg_wait: int = -1,
+    swap_AB: bool = False,
+) -> cute.Tensor:
+    if const_expr(swap_AB):
+        return gemm_zero_init(
+            tiled_mma, shape[::-1], tCrB, tCrA, B_idx, A_idx, wg_wait, swap_AB=False
+        )
+    else:
+        acc = cute.make_fragment(tiled_mma.partition_shape_C(shape), Float32)
+        rA = tCrA if const_expr(A_idx is None) else tCrA[None, None, None, A_idx]
+        rB = tCrB if const_expr(B_idx is None) else tCrB[None, None, None, B_idx]
+        gemm(tiled_mma, acc, rA, rB, zero_init=True, wg_wait=wg_wait)
+        return acc
+
+
+def gemm_w_idx(
+    tiled_mma: cute.TiledMma,
+    acc: cute.Tensor,
+    tCrA: cute.Tensor,
+    tCrB: cute.Tensor,
+    zero_init: Boolean,
+    A_idx: Optional[Int32] = None,
+    B_idx: Optional[Int32] = None,
+    wg_wait: int = -1,
+    swap_AB: bool = False,
+) -> None:
+    if const_expr(swap_AB):
+        gemm_w_idx(tiled_mma, acc, tCrB, tCrA, zero_init, B_idx, A_idx, wg_wait, swap_AB=False)
+    else:
+        rA = tCrA if const_expr(A_idx is None) else tCrA[None, None, None, A_idx]
+        rB = tCrB if const_expr(B_idx is None) else tCrB[None, None, None, B_idx]
+        gemm(tiled_mma, acc, rA, rB, zero_init=zero_init, wg_wait=wg_wait)
+
+
+@dsl_user_op
+def make_smem_layout(
+    dtype: Type[Numeric],
+    layout: LayoutEnum,
+    shape: cute.Shape,
+    stage: Optional[int] = None,
+    *,
+    loc=None,
+    ip=None,
+) -> Union[cute.Layout, cute.ComposedLayout]:
+    major_mode_size = shape[1] if layout.is_n_major_c() else shape[0]
+    smem_layout_atom = warpgroup.make_smem_layout_atom(
+        sm90_utils_og.get_smem_layout_atom(layout, dtype, major_mode_size),
+        dtype,
+    )
+    order = (1, 0, 2) if const_expr(layout.is_m_major_c()) else (0, 1, 2)
+    smem_layout_staged = cute.tile_to_shape(
+        smem_layout_atom,
+        cute.append(shape, stage) if const_expr(stage is not None) else shape,
+        order=order if const_expr(stage is not None) else order[:2],
+    )
+    return smem_layout_staged

build/torch-universal/interface.py

@@ -0,0 +1,1483 @@
+# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+# [2025-07-04] Version in Cute-DSL, for Hopper and Blackwell. You'll need install nvidia-cutlass-dsl==4.2.0.
+# [2025-07-04] Version in Cute-DSL, for Hopper and Blackwell. You'll need install nvidia-cutlass-dsl==4.2.0.
+# [2025-07-04] Version in Cute-DSL, for Hopper and Blackwell. You'll need install nvidia-cutlass-dsl==4.2.0.
+
+# Supported features:
+# - BF16 & FP16 dtype
+# - noncausal & causal attention
+# - MHA, GQA, MQA
+# - hdim 64, 96, 128.
+# - (hdim_qk, hdim_v) = (192, 128) for Blackwell (i.e. DeepSeek shape)
+# - varlen
+# - sliding window
+# - bwd pass for Ampere (will also run on Hopper/Blackwell, but will be slow)
+
+# Features not supported yet:
+# - split (i.e. FlashDecoding)
+# - tuned block sizes
+# - paged KV
+# - append KV to existing KV cache
+# - FP8
+# - bwd pass optimized for Hopper/Blackwell
+
+import math
+from typing import Optional, Tuple, Callable
+
+import torch
+
+import cuda.bindings.driver as cuda
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+
+from . import utils
+from .flash_fwd import FlashAttentionForwardSm80, FlashAttentionForwardSm90
+from .flash_fwd_sm100 import FlashAttentionForwardSm100
+from .flash_bwd_preprocess import FlashAttentionBackwardPreprocess
+from .flash_bwd import FlashAttentionBackwardSm80
+from .flash_bwd_sm90 import FlashAttentionBackwardSm90
+from .flash_bwd_sm100 import FlashAttentionBackwardSm100
+from .flash_bwd_postprocess import FlashAttentionBackwardPostprocess
+from .flash_fwd_combine import FlashAttentionForwardCombine
+
+from .block_sparsity import (
+    BlockSparseTensorsTorch,
+    to_cute_block_sparse_tensors,
+    normalize_block_sparse_tensors,
+)
+
+def maybe_contiguous(x):
+    return x.contiguous() if x is not None and x.stride(-1) != 1 else x
+
+
+torch2cute_dtype_map = {
+    torch.float16: cutlass.Float16,
+    torch.bfloat16: cutlass.BFloat16,
+    torch.float32: cutlass.Float32,
+}
+
+
+def num_splits_heuristic(total_mblocks, num_SMs, num_n_blocks, max_splits):
+    # If num_n_blocks is too small, use 1 split. For example, we never split for hdim = 128 and seqlen_k = 512.
+    if num_n_blocks <= 4:
+        return 1
+
+    # NOTE: We should revisit this heuristic after persistence is supported for split KV.
+    # Sometimes, it's ideal to over-schedule splits for better efficiency.
+    return min(num_SMs // total_mblocks, max_splits, num_n_blocks)
+
+
+def _flash_attn_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    softcap: Optional[float] = None,
+    window_size_left: Optional[int] = None,
+    window_size_right: Optional[int] = None,
+    learnable_sink: Optional[torch.Tensor] = None,
+    # m_block_size: int = 128,
+    # n_block_size: int = 64,
+    # num_threads: int = 128,
+    m_block_size: int = 128,
+    n_block_size: int = 128,
+    num_threads: int = 384,
+    num_splits: int = 1,
+    pack_gqa: Optional[bool] = None,
+    _compute_capability: Optional[int] = None,
+    score_mod: Optional[Callable] = None,
+    mask_mod: Optional[Callable] = None,
+    block_sparse_tensors: Optional[BlockSparseTensorsTorch] = None,
+    return_lse: bool = False,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    aux_tensors: Optional[list[torch.Tensor]] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Forward pass for FlashAttention.
+
+    Args:
+        ...
+        score_mod: A callable that takes the attention scores and applies a modification.
+        return_lse: Whether to return the log softmax of the attention scores. If set to True will always calculate
+        out: Optional pre-allocated output tensor. If None, will be allocated internally.
+        lse: Optional pre-allocated log-sum-exp tensor. If None, will be allocated when needed.
+        aux_tensors: Some score_mods will want to read from global aux_tensors. This is how we thread them through to the inner kernel.
+    """
+    q, k, v = [maybe_contiguous(t) for t in (q, k, v)]
+    num_head, head_dim = q.shape[-2:]
+    if cu_seqlens_q is None:
+        batch_size, seqlen_q = q.shape[:2]
+        total_q = batch_size * seqlen_q
+    else:
+        batch_size = cu_seqlens_q.shape[0] - 1
+        seqlen_q = None
+        total_q = q.shape[0]
+    if page_table is not None:
+        assert cu_seqlens_k is None, "page_table is not supported with cu_seqlens_k"
+        assert page_table.dtype == torch.int32, "page_table must be int32"
+        assert page_table.stride(-1) == 1, "page_table must be contiguous in the last dimension"
+        max_num_pages_per_seq = page_table.shape[1]
+        assert page_table.shape == (batch_size, max_num_pages_per_seq)
+        num_pages, page_size = k.shape[:2]
+        seqlen_k = num_pages * page_size
+    else:
+        num_pages, page_size = None, None
+        seqlen_k = k.shape[-3]
+    num_head_kv = k.shape[-2]
+    head_dim_v = v.shape[-1]
+    if cu_seqlens_k is None:
+        if page_table is None:
+            assert k.shape == (batch_size, seqlen_k, num_head_kv, head_dim)
+            assert v.shape == (batch_size, seqlen_k, num_head_kv, head_dim_v)
+        else:
+            assert k.shape == (num_pages, page_size, num_head_kv, head_dim)
+            assert v.shape == (num_pages, page_size, num_head_kv, head_dim_v)
+    else:
+        assert k.shape == (seqlen_k, num_head_kv, head_dim)
+        assert v.shape == (seqlen_k, num_head_kv, head_dim_v)
+        assert cu_seqlens_k.shape == (batch_size + 1,), (
+            "cu_seqlens_k must have shape (batch_size + 1,)"
+        )
+
+    if cu_seqlens_q is not None:
+        assert cu_seqlens_q.shape == (batch_size + 1,), (
+            "cu_seqlens_q must have shape (batch_size + 1,)"
+        )
+    assert seqused_q is None or seqused_q.shape == (batch_size,), (
+        "seqused_q must have shape (batch_size,)"
+    )
+    assert seqused_k is None or seqused_k.shape == (batch_size,), (
+        "seqused_k must have shape (batch_size,)"
+    )
+    assert q.dtype in [torch.float16, torch.bfloat16], "inputs must be float16 or bfloat16"
+    assert q.dtype == k.dtype == v.dtype, "inputs must have the same dtype"
+    for t in [cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k]:
+        if t is not None:
+            assert t.dtype == torch.int32, (
+                "cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k must be int32"
+            )
+            assert t.stride(0) == 1, (
+                "cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k must be contiguous"
+            )
+    if learnable_sink is not None:
+        assert learnable_sink.shape == (num_head,)
+        assert learnable_sink.dtype == torch.bfloat16, "learnable_sink must be bfloat16"
+
+    assert all(
+        t is None or t.is_cuda
+        for t in (
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqused_q,
+            seqused_k,
+            page_table,
+            learnable_sink,
+        )
+    ), "inputs must be on CUDA device"
+    assert num_head % num_head_kv == 0, "num_head must be divisible by num_head_kv"
+    assert head_dim <= 256, "head_dim must be less than or equal to 256"
+    alignment = 16 // q.element_size()
+    assert head_dim % alignment == 0, f"head_dim must be divisible by {alignment}"
+    assert head_dim_v % alignment == 0, f"head_dim_v must be divisible by {alignment}"
+    if softmax_scale is None:
+        softmax_scale = 1.0 / math.sqrt(head_dim)
+    if softcap == 0.0:
+        softcap = None
+    qhead_per_kvhead = num_head // num_head_kv
+    if pack_gqa is None:
+        pack_gqa = qhead_per_kvhead > 1
+
+    out_torch_dtype = q.dtype
+    device = q.device
+    q_batch_seqlen_shape = (batch_size, seqlen_q) if cu_seqlens_q is None else (total_q,)
+    lse_shape = (batch_size, num_head, seqlen_q) if cu_seqlens_q is None else (num_head, total_q)
+    requires_grad = q.requires_grad or k.requires_grad or v.requires_grad
+
+    if out is None:
+        out = torch.empty(
+            *q_batch_seqlen_shape, num_head, head_dim_v, dtype=out_torch_dtype, device=device
+        )
+    else:
+        expected_out_shape = (*q_batch_seqlen_shape, num_head, head_dim_v)
+        assert out.shape == expected_out_shape, (
+            f"out tensor shape {out.shape} does not match expected shape {expected_out_shape}"
+        )
+        assert out.dtype == out_torch_dtype, (
+            f"out tensor dtype {out.dtype} does not match expected dtype {out_torch_dtype}"
+        )
+        assert out.device == device, (
+            f"out tensor device {out.device} does not match input device {device}"
+        )
+        assert out.is_cuda, "out tensor must be on CUDA device"
+
+    if lse is None:
+        lse = (
+            torch.empty(lse_shape, dtype=torch.float32, device=device)
+            if requires_grad or return_lse
+            else None
+        )
+    elif lse is not None:
+        assert lse.shape == lse_shape, (
+            f"lse tensor shape {lse.shape} does not match expected shape {lse_shape}"
+        )
+        assert lse.dtype == torch.float32, (
+            f"lse tensor dtype {lse.dtype} does not match expected dtype torch.float32"
+        )
+        assert lse.device == device, (
+            f"lse tensor device {lse.device} does not match input device {device}"
+        )
+        assert lse.is_cuda, "lse tensor must be on CUDA device"
+
+    dtype = torch2cute_dtype_map[q.dtype]
+    (
+        cu_seqlens_q_tensor,
+        cu_seqlens_k_tensor,
+        seqused_q_tensor,
+        seqused_k_tensor,
+        learnable_sink_tensor,
+    ) = [
+        from_dlpack(t.detach(), assumed_align=4).mark_layout_dynamic(leading_dim=0)
+        if t is not None
+        else None
+        for t in (cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, learnable_sink)
+    ]
+    page_table_tensor = (
+        from_dlpack(page_table.detach(), assumed_align=4).mark_layout_dynamic(leading_dim=1)
+        if page_table is not None
+        else None
+    )
+    sparse_tensors = None
+    if block_sparse_tensors is not None:
+        if seqlen_q is None:
+            raise ValueError("Block sparsity requires fixed-length sequences (seqlen_q must be known).")
+        expected_m_blocks = (seqlen_q + m_block_size - 1) // m_block_size
+        expected_n_blocks = (seqlen_k + n_block_size - 1) // n_block_size
+        block_sparse_tensors = normalize_block_sparse_tensors(
+            block_sparse_tensors,
+            expected_count_shape=(batch_size, num_head, expected_m_blocks),
+            expected_index_shape=(batch_size, num_head, expected_m_blocks, expected_n_blocks),
+        )
+        sparse_tensors = to_cute_block_sparse_tensors(block_sparse_tensors)
+
+    use_block_sparsity = sparse_tensors is not None
+
+    if mask_mod is None:
+        if causal:
+            window_size_right = 0
+        local = window_size_left is not None or window_size_right is not None
+        if window_size_left is not None or window_size_right is not None:
+            if window_size_left is None and window_size_right == 0:
+                causal, local = True, False
+            else:
+                causal, local = False, True
+    else:
+        causal, local = False, False
+
+    compute_capability = (
+        torch.cuda.get_device_capability()[0]
+        if _compute_capability is None
+        else _compute_capability
+    )
+    assert compute_capability in [9, 10], "Unsupported compute capability. Supported: 9.x, 10.x"
+    current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+
+    if compute_capability == 9:  # TODO: tune block size according to hdim.
+        if head_dim == head_dim_v == 128 and not causal and not local and not use_block_sparsity:
+            n_block_size = 192
+    if compute_capability == 10:
+        # TODO: fix the varlen case
+        if (
+            pack_gqa
+            and (128 % qhead_per_kvhead != 0)
+            or (cu_seqlens_q is not None or seqused_q is not None)
+        ):
+            pack_gqa = False
+        # TODO: fix GQA + SplitKV + non-varlen
+        if pack_gqa and num_splits != 1 and cu_seqlens_q is None:
+            pack_gqa = False
+
+    if num_splits < 1:
+        max_seqlen_k = seqlen_k if cu_seqlens_k is None else (cu_seqlens_k[1:] - cu_seqlens_k[:-1]).max().item()
+        max_seqlen_q = seqlen_q if cu_seqlens_q is None else (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).max().item()
+        seqlen_q_packgqa = max_seqlen_q * qhead_per_kvhead
+        seqlen_k_loaded = max_seqlen_k if not local else max(0, min(max_seqlen_k, window_size_right + window_size_left + 1 + m_block_size))
+        num_n_blocks = (seqlen_k_loaded + n_block_size - 1) // n_block_size
+        num_m_blocks = (seqlen_q_packgqa + m_block_size - 1) // m_block_size
+        total_mblocks = batch_size * num_head_kv * num_m_blocks
+        num_splits = num_splits_heuristic(
+            total_mblocks,
+            torch.cuda.get_device_properties(device).multi_processor_count,
+            num_n_blocks,
+            128,
+        )
+
+    is_split_kv = num_splits > 1
+    if is_split_kv:
+        out_partial = torch.empty(num_splits, *q_batch_seqlen_shape, num_head, head_dim_v, dtype=torch.float32, device=device)
+        lse_partial = torch.empty(num_splits, *lse_shape, dtype=torch.float32, device=device)
+
+    q_tensor, k_tensor, v_tensor, o_tensor = [
+        from_dlpack(t.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=t.ndim - 1)
+        for t in (q, k, v, out if not is_split_kv else out_partial)
+    ]
+    if is_split_kv:
+        lse_tensor = from_dlpack(lse_partial.detach(), assumed_align=4).mark_layout_dynamic(leading_dim=lse_partial.ndim - 1)
+    elif lse is not None:
+        lse_tensor = from_dlpack(lse.detach(), assumed_align=4).mark_layout_dynamic(leading_dim=lse.ndim - 1)
+    else:
+        lse_tensor = None 
+
+    # hash score and mask mods for compile cache
+    score_mod_hash = utils.hash_callable(score_mod) if score_mod is not None else False
+    mask_mod_hash = utils.hash_callable(mask_mod) if mask_mod is not None else False
+
+    if softcap is not None:
+        assert score_mod is None, "softcap and score_mod cannot be used together"
+        score_mod = utils.create_softcap_scoremod(softcap)
+
+    is_varlen = (
+        cu_seqlens_q is not None
+        or cu_seqlens_k is not None
+        or seqused_q is not None
+        or seqused_k is not None
+    )
+    if score_mod is not None:
+        if is_varlen:
+            raise NotImplementedError(
+                "score_mod with aux_tensors is not yet supported for varlen sequences. This will be fixed in a future PR."
+            )
+
+    if mask_mod is not None:
+        if not use_block_sparsity:
+            raise NotImplementedError(
+                "mask_mod requires the use of block sparsity. This will be fixed in a future PR."
+            )
+        if is_varlen:
+            raise NotImplementedError(
+                "mask_mod with aux_tensors is not yet supported for varlen sequences. This will be fixed in a future PR."
+            )
+        if pack_gqa:
+            raise NotImplementedError(
+                "mask_mod with aux_tensors is not yet supported with pack_gqa=True. This will be fixed in a future PR."
+            )
+
+    if use_block_sparsity:
+        if is_varlen:
+            raise NotImplementedError(
+                "Block sparsity is not yet supported for varlen sequences. This will be fixed in a future PR."
+            )
+        if pack_gqa:
+            raise NotImplementedError(
+                "Block sparsity is not yet supported with pack_gqa=True. This will be fixed in a future PR."
+            )
+
+    cute_aux_tensors = None
+    if aux_tensors is not None:
+        cute_aux_tensors = [from_dlpack(buf).mark_layout_dynamic() for buf in aux_tensors]
+
+    compile_key = (
+        dtype,
+        head_dim,
+        head_dim_v,
+        qhead_per_kvhead,
+        causal,
+        score_mod_hash,
+        mask_mod_hash,
+        use_block_sparsity,
+        len(aux_tensors) if aux_tensors is not None else 0,
+        lse is None,
+        cu_seqlens_q is None,
+        cu_seqlens_k is None,
+        seqused_q is None,
+        seqused_k is None,
+        page_table is not None,
+        window_size_left is not None,
+        window_size_right is not None,
+        learnable_sink is not None,
+        m_block_size,
+        n_block_size,
+        num_threads,
+        is_split_kv,
+        pack_gqa,
+        compute_capability,
+    )
+
+    if compile_key not in _flash_attn_fwd.compile_cache:
+        if compute_capability == 9:
+            assert page_table is None, "paged KV not supported on SM 9.0"
+            assert not is_split_kv, "SplitKV not supported on SM 9.0"
+            # fa_fwd = FlashAttentionForwardSm80(
+            fa_fwd = FlashAttentionForwardSm90(
+                dtype,
+                head_dim,
+                head_dim_v,
+                qhead_per_kvhead,
+                is_causal=causal,
+                is_local=local,
+                pack_gqa=pack_gqa,
+                tile_m=m_block_size,
+                tile_n=n_block_size,
+                # num_stages=1,
+                num_stages=2,
+                num_threads=num_threads,
+                Q_in_regs=False,
+                intra_wg_overlap=True,
+                mma_pv_is_rs=True,
+                mask_mod=mask_mod,
+                score_mod=score_mod,
+                has_aux_tensors=aux_tensors is not None,
+            )
+        elif compute_capability == 10:
+            assert page_size in [None, 128], (
+                "Only page_size=128 is supported for paged KV on SM 10.0"
+            )
+            if sparse_tensors is not None:
+                raise NotImplementedError("BlockSparsity not yet supported on SM 10.0")
+            fa_fwd = FlashAttentionForwardSm100(
+                head_dim,
+                head_dim_v,
+                qhead_per_kvhead=qhead_per_kvhead,
+                is_causal=causal,
+                is_local=local,
+                is_split_kv=is_split_kv,
+                pack_gqa=pack_gqa,
+                is_persistent=not causal
+                and not local
+                and cu_seqlens_q is None
+                and seqused_q is None
+                and not is_split_kv,
+                score_mod=score_mod,
+                has_aux_tensors=aux_tensors is not None,
+            )
+        else:
+            raise ValueError(
+                f"Unsupported compute capability: {compute_capability}. Supported: 9.x, 10.x"
+            )
+        # TODO: check @can_implement
+        _flash_attn_fwd.compile_cache[compile_key] = cute.compile(
+            fa_fwd,
+            q_tensor,
+            k_tensor,
+            v_tensor,
+            o_tensor,
+            lse_tensor,
+            softmax_scale,
+            current_stream,
+            cu_seqlens_q_tensor,
+            cu_seqlens_k_tensor,
+            seqused_q_tensor,
+            seqused_k_tensor,
+            page_table_tensor,
+            window_size_left,
+            window_size_right,
+            learnable_sink_tensor,
+            sparse_tensors,
+            cute_aux_tensors,
+        )
+    _flash_attn_fwd.compile_cache[compile_key](
+        q_tensor,
+        k_tensor,
+        v_tensor,
+        o_tensor,
+        lse_tensor,
+        softmax_scale,
+        current_stream,
+        cu_seqlens_q_tensor,
+        cu_seqlens_k_tensor,
+        seqused_q_tensor,
+        seqused_k_tensor,
+        page_table_tensor,
+        window_size_left,
+        window_size_right,
+        learnable_sink_tensor,
+        sparse_tensors,
+        cute_aux_tensors,
+    )
+    if is_split_kv:
+        _flash_attn_fwd_combine(
+            out_partial,
+            lse_partial.transpose(-1, -2),
+            out,
+            lse.transpose(-1, -2) if lse is not None else None,
+            cu_seqlens_q,
+            seqused_q,
+        )
+    return out, lse
+
+
+_flash_attn_fwd.compile_cache = {}
+
+
+def _flash_attn_bwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    dout: torch.Tensor,
+    lse: torch.Tensor,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    softcap: float = 0.0,
+    m_block_size: int = 64,
+    n_block_size: int = 128,
+    num_threads: int = 256,
+    pack_gqa: bool = False,
+    num_stages_Q: int = 2,
+    num_stages_dO: int = 2,
+    SdP_swapAB: bool = False,
+    dKV_swapAB: bool = False,
+    dQ_swapAB: bool = False,
+    AtomLayoutMSdP: int = 2,
+    AtomLayoutNdKV: int = 2,
+    AtomLayoutMdQ: int = 2,
+    V_in_regs: bool = False,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    compute_capability = torch.cuda.get_device_capability()[0]
+    assert compute_capability in [9, 10], "Unsupported compute capability. Supported: 9.x, 10.x"
+
+    if compute_capability == 9:
+        m_block_size = 80 if not causal else 64
+        n_block_size = 128
+        num_stages_Q = 2
+        num_stages_dO = 2
+        num_stages_PdS = 2
+        SdP_swapAB = True
+        dKV_swapAB = False
+        dQ_swapAB = not causal
+        AtomLayoutMSdP = 1
+        AtomLayoutNdKV = 2
+        AtomLayoutMdQ = 1
+        cluster_size = 1
+    else:
+        m_block_size = 128
+        n_block_size = 128
+        dQ_swapAB = False
+        dKV_swapAB = False
+        AtomLayoutMdQ = 1
+        AtomLayoutNdKV = 1
+        # TODO: support cluster size 2
+        cluster_size = 1
+    q, k, v, out, dout, lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = [
+        maybe_contiguous(t)
+        for t in (q, k, v, out, dout, lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k)
+    ]
+    num_head, head_dim = q.shape[-2:]
+    if cu_seqlens_q is None:
+        batch_size, seqlen_q = q.shape[:2]
+        total_q = batch_size * seqlen_q
+    else:
+        batch_size = cu_seqlens_q.shape[0] - 1
+        seqlen_q = None
+        total_q = q.shape[0]
+
+    if cu_seqlens_k is None:
+        batch_size, seqlen_k = k.shape[:2]
+        total_k = batch_size * seqlen_k
+    else:
+        batch_size = cu_seqlens_k.shape[0] - 1
+        seqlen_k = None
+        total_k = k.shape[0]
+
+    num_head_kv = k.shape[-2]
+    head_dim_v = v.shape[-1]
+
+    if cu_seqlens_k is None:
+        assert k.shape == (batch_size, seqlen_k, num_head_kv, head_dim)
+        assert v.shape == (batch_size, seqlen_k, num_head_kv, head_dim_v)
+    else:
+        assert k.shape == (total_k, num_head_kv, head_dim)
+        assert v.shape == (total_k, num_head_kv, head_dim_v)
+        assert cu_seqlens_k.shape == (batch_size + 1,), (
+            "cu_seqlens_k must have shape (batch_size + 1,)"
+        )
+
+    if cu_seqlens_q is not None:
+        assert cu_seqlens_q.shape == (batch_size + 1,), (
+            "cu_seqlens_q must have shape (batch_size + 1,)"
+        )
+
+        assert out.shape == (total_q, num_head, head_dim_v)
+        assert dout.shape == (total_q, num_head, head_dim_v)
+        assert lse.shape == (num_head, total_q), "lse must have shape (num_head, total_q)"
+    else:
+        assert out.shape == (batch_size, seqlen_q, num_head, head_dim_v)
+        assert dout.shape == (batch_size, seqlen_q, num_head, head_dim_v)
+        assert lse.shape == (batch_size, num_head, seqlen_q), (
+            "lse must have shape (batch_size, num_head, seqlen_q)"
+        )
+
+    assert q.dtype in [torch.float16, torch.bfloat16], "inputs must be float16 or bfloat16"
+    assert q.dtype == k.dtype == v.dtype == out.dtype == dout.dtype, (
+        "inputs must have the same dtype"
+    )
+    for t in [cu_seqlens_q, cu_seqlens_k]:
+        if t is not None:
+            assert t.dtype == torch.int32, "cu_seqlens_q, cu_seqlens_k must be int32"
+    assert lse.dtype == torch.float32, "lse must be float32"
+    assert all(
+        t is None or t.is_cuda for t in (q, k, v, out, dout, lse, cu_seqlens_q, cu_seqlens_k)
+    ), "inputs must be on CUDA device"
+    assert num_head % num_head_kv == 0, "num_head must be divisible by num_head_kv"
+    assert head_dim <= 256, "head_dim must be less than or equal to 256"
+    alignment = 16 // q.element_size()
+    assert head_dim % alignment == 0, f"head_dim must be divisible by {alignment}"
+    assert head_dim_v % alignment == 0, f"head_dim_v must be divisible by {alignment}"
+    if softmax_scale is None:
+        softmax_scale = 1.0 / math.sqrt(head_dim)
+    qhead_per_kvhead = num_head // num_head_kv
+    if pack_gqa is None:
+        pack_gqa = qhead_per_kvhead > 1
+    if compute_capability == 10:
+        pack_gqa = False # override for now
+
+    device = q.device
+    # TODO: check if this is the right rounding
+    dq = torch.empty_like(q)
+    dk = torch.empty_like(k)
+    dv = torch.empty_like(v)
+
+    head_dim_rounded = (head_dim + 32 - 1) // 32 * 32
+
+    if cu_seqlens_q is None:
+        seqlen_q_rounded = (seqlen_q + m_block_size - 1) // m_block_size * m_block_size
+        dq_accum = torch.empty(
+            batch_size,
+            num_head,
+            seqlen_q_rounded * head_dim_rounded,
+            dtype=torch.float32,
+            device=device,
+        )
+        dpsum = torch.empty(
+            batch_size, num_head, seqlen_q_rounded, dtype=torch.float32, device=device
+        )
+        lse_log2 = torch.empty(
+            batch_size, num_head, seqlen_q_rounded, dtype=torch.float32, device=device
+        )
+    else:
+        total_q_rounded_padded = (
+            (total_q + cu_seqlens_q.shape[0] * m_block_size - 1) // m_block_size * m_block_size
+        )
+        dq_accum = torch.empty(
+            num_head, total_q_rounded_padded * head_dim_rounded, dtype=torch.float32, device=device
+        )
+        dpsum = torch.empty(num_head, total_q_rounded_padded, dtype=torch.float32, device=device)
+        lse_log2 = torch.empty(num_head, total_q_rounded_padded, dtype=torch.float32, device=device)
+
+    if qhead_per_kvhead > 1:
+        head_dim_v_rounded = (head_dim_v + 32 - 1) // 32 * 32
+        if cu_seqlens_k is None:
+            seqlen_k_rounded = (seqlen_k + n_block_size - 1) // n_block_size * n_block_size
+            num_n_blocks = seqlen_k_rounded // n_block_size
+            if cluster_size == 2 and num_n_blocks % cluster_size != 0:
+                seqlen_k_rounded = seqlen_k_rounded + n_block_size
+            dk_accum = torch.zeros(
+                batch_size,
+                num_head_kv,
+                seqlen_k_rounded * head_dim_rounded,
+                dtype=torch.float32,
+                device=device,
+            )
+            dv_accum = torch.zeros(
+                batch_size,
+                num_head_kv,
+                seqlen_k_rounded * head_dim_v_rounded,
+                dtype=torch.float32,
+                device=device,
+            )
+        else:
+            total_k_rounded_padded = (
+                (total_k + cu_seqlens_k.shape[0] * n_block_size - 1) // n_block_size * n_block_size
+            )
+            num_n_blocks = total_k_rounded_padded // n_block_size
+            if cluster_size == 2 and num_n_blocks % cluster_size != 0:
+                total_k_rounded_padded = total_k_rounded_padded + n_block_size
+            dk_accum = torch.zeros(
+                num_head_kv,
+                total_k_rounded_padded * head_dim_rounded,
+                dtype=torch.float32,
+                device=device,
+            )
+            dv_accum = torch.zeros(
+                num_head_kv,
+                total_k_rounded_padded * head_dim_v_rounded,
+                dtype=torch.float32,
+                device=device,
+            )
+
+    dtype = torch2cute_dtype_map[q.dtype]
+    q_tensor, k_tensor, v_tensor, o_tensor, do_tensor, dq_tensor, dk_tensor, dv_tensor = [
+        from_dlpack(t.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=t.ndim - 1)
+        for t in (q, k, v, out, dout, dq, dk, dv)
+    ]
+    lse_tensor = from_dlpack(lse.detach(), assumed_align=4).mark_layout_dynamic(
+        leading_dim=lse.ndim - 1
+    )
+    dq_accum_tensor, dpsum_tensor, lse_log2_tensor = [
+        from_dlpack(t.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=t.ndim - 1)
+        for t in (dq_accum, dpsum, lse_log2)
+    ]
+    if qhead_per_kvhead > 1:
+        dk_accum_tensor, dv_accum_tensor = [
+            from_dlpack(t.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=t.ndim - 1)
+            for t in (dk_accum, dv_accum)
+        ]
+    cu_seqlens_q_tensor, cu_seqlens_k_tensor, seqused_q_tensor, seqused_k_tensor = [
+        from_dlpack(t.detach(), assumed_align=4).mark_layout_dynamic(leading_dim=t.ndim - 1)
+        if t is not None
+        else None
+        for t in (cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k)
+    ]
+    current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+
+    # Preprocess kernel: compute (o * dout).sum(dim=-1), lse * log2_e, and zero out dq_accum.
+    compile_key_pre = (compute_capability, dtype, head_dim_v, m_block_size, num_threads)
+    if compile_key_pre not in _flash_attn_bwd.compile_cache_pre:
+        fa_bwd_pre = FlashAttentionBackwardPreprocess(
+            dtype,
+            head_dim_v,
+            m_block_size,
+            num_threads=num_threads,
+        )
+        # TODO: check @can_implement
+        _flash_attn_bwd.compile_cache_pre[compile_key_pre] = cute.compile(
+            fa_bwd_pre,
+            o_tensor,
+            do_tensor,
+            dpsum_tensor,
+            lse_tensor,
+            lse_log2_tensor,
+            dq_accum_tensor,
+            cu_seqlens_q_tensor,
+            seqused_q_tensor,
+            current_stream,
+        )
+    _flash_attn_bwd.compile_cache_pre[compile_key_pre](
+        o_tensor,
+        do_tensor,
+        dpsum_tensor,
+        lse_tensor,
+        lse_log2_tensor,
+        dq_accum_tensor,
+        cu_seqlens_q_tensor,
+        seqused_q_tensor,
+        current_stream,
+    )
+
+    # Backward kernel: compute dk, dv, dq_accum.
+    if compute_capability == 9:
+        compile_key = (
+            compute_capability,
+            dtype,
+            head_dim,
+            head_dim_v,
+            qhead_per_kvhead,
+            causal,
+            softcap != 0.0,
+            m_block_size,
+            n_block_size,
+            num_threads,
+            pack_gqa,
+            num_stages_Q,
+            num_stages_dO,
+            SdP_swapAB,
+            dKV_swapAB,
+            dQ_swapAB,
+            AtomLayoutMSdP,
+            AtomLayoutNdKV,
+            AtomLayoutMdQ,
+            V_in_regs,
+        )
+    else:
+        compile_key = (
+            compute_capability,
+            dtype,
+            head_dim,
+            head_dim_v,
+            qhead_per_kvhead,
+            causal,
+            softcap != 0.0,
+            m_block_size,
+            n_block_size,
+            num_threads,
+            pack_gqa,
+            cluster_size,
+        )
+    num_threads = 384
+    if compile_key not in _flash_attn_bwd.compile_cache:
+        fa_bwd_sm80 = FlashAttentionBackwardSm80(
+            dtype,
+            head_dim,
+            head_dim_v,
+            qhead_per_kvhead,
+            m_block_size,
+            n_block_size,
+            num_stages_Q,
+            num_stages_dO,
+            num_threads,
+            pack_gqa,
+            causal,
+            SdP_swapAB,
+            dKV_swapAB,
+            dQ_swapAB,
+            AtomLayoutMSdP,
+            AtomLayoutNdKV,
+            AtomLayoutMdQ,
+            V_in_regs=V_in_regs,
+        )
+        if compute_capability == 9:
+            fa_bwd_obj = FlashAttentionBackwardSm90(
+                dtype,
+                head_dim,
+                head_dim_v,
+                qhead_per_kvhead,
+                causal,
+                m_block_size,
+                n_block_size,
+                num_stages_Q,
+                num_stages_dO,
+                num_stages_PdS,
+                SdP_swapAB,
+                dKV_swapAB,
+                dQ_swapAB,
+                AtomLayoutMSdP,
+                AtomLayoutNdKV,
+                AtomLayoutMdQ,
+                num_threads,
+                V_in_regs=V_in_regs,
+            )
+        else:
+            fa_bwd_obj = FlashAttentionBackwardSm100(
+                head_dim,
+                head_dim_v,
+                is_causal=causal,
+                qhead_per_kvhead=qhead_per_kvhead,
+                # tile_m=m_block_size,
+                # tile_n=n_block_size,
+                cluster_size=cluster_size,
+                # cluster_size=1,
+            )
+        # TODO: check @can_implement
+        _flash_attn_bwd.compile_cache[compile_key] = cute.compile(
+            fa_bwd_obj,
+            q_tensor,
+            k_tensor,
+            v_tensor,
+            do_tensor,
+            lse_log2_tensor,
+            dpsum_tensor,
+            dq_accum_tensor,
+            dk_tensor if qhead_per_kvhead == 1 else dk_accum_tensor,
+            dv_tensor if qhead_per_kvhead == 1 else dv_accum_tensor,
+            softmax_scale,
+            current_stream,
+            cu_seqlens_q_tensor,
+            cu_seqlens_k_tensor,
+            seqused_q_tensor,
+            seqused_k_tensor,
+        )
+    _flash_attn_bwd.compile_cache[compile_key](
+        q_tensor,
+        k_tensor,
+        v_tensor,
+        do_tensor,
+        lse_log2_tensor,
+        dpsum_tensor,
+        dq_accum_tensor,
+        dk_tensor if qhead_per_kvhead == 1 else dk_accum_tensor,
+        dv_tensor if qhead_per_kvhead == 1 else dv_accum_tensor,
+        softmax_scale,
+        current_stream,
+        cu_seqlens_q_tensor,
+        cu_seqlens_k_tensor,
+        seqused_q_tensor,
+        seqused_k_tensor,
+    )
+
+    num_threads = 256 if compute_capability == 9 else 128
+    # Postprocess kernel: convert dq_accum from float32 to dq in bf16/fp16
+    compile_key_post = (dtype, head_dim, m_block_size, num_threads, AtomLayoutMdQ, dQ_swapAB)
+    if compile_key_post not in _flash_attn_bwd.compile_cache_post:
+        arch = compute_capability * 10
+        fa_bwd_post = FlashAttentionBackwardPostprocess(
+            dtype, head_dim, arch, m_block_size, num_threads, AtomLayoutMdQ, dQ_swapAB
+        )
+        # TODO: check @can_implement
+        _flash_attn_bwd.compile_cache_post[compile_key_post] = cute.compile(
+            fa_bwd_post,
+            dq_accum_tensor,
+            dq_tensor,
+            softmax_scale,
+            cu_seqlens_q_tensor,
+            seqused_q_tensor,
+            current_stream,
+        )
+    _flash_attn_bwd.compile_cache_post[compile_key_post](
+        dq_accum_tensor,
+        dq_tensor,
+        softmax_scale,
+        cu_seqlens_q_tensor,
+        seqused_q_tensor,
+        current_stream,
+    )
+
+    if qhead_per_kvhead > 1:
+        # Postprocess kernel: convert dk_accum & dv_accum from float32 to bf16/fp16
+        compile_key_post = (dtype, head_dim, n_block_size, num_threads, AtomLayoutNdKV, dKV_swapAB)
+        if compile_key_post not in _flash_attn_bwd.compile_cache_post:
+            fa_bwd_post = FlashAttentionBackwardPostprocess(
+                dtype, head_dim, n_block_size, num_threads, AtomLayoutNdKV, dKV_swapAB
+            )
+            # TODO: check @can_implement
+            _flash_attn_bwd.compile_cache_post[compile_key_post] = cute.compile(
+                fa_bwd_post,
+                dk_accum_tensor,
+                dk_tensor,
+                softmax_scale,
+                cu_seqlens_k_tensor,
+                seqused_k_tensor,
+                current_stream,
+            )
+        _flash_attn_bwd.compile_cache_post[compile_key_post](
+            dk_accum_tensor,
+            dk_tensor,
+            softmax_scale,
+            cu_seqlens_k_tensor,
+            seqused_k_tensor,
+            current_stream,
+        )
+        compile_key_post = (
+            dtype,
+            head_dim_v,
+            n_block_size,
+            num_threads,
+            AtomLayoutNdKV,
+            dKV_swapAB,
+        )
+        if compile_key_post not in _flash_attn_bwd.compile_cache_post:
+            fa_bwd_post = FlashAttentionBackwardPostprocess(
+                dtype, head_dim_v, n_block_size, num_threads, AtomLayoutNdKV, dKV_swapAB
+            )
+            # TODO: check @can_implement
+            _flash_attn_bwd.compile_cache_post[compile_key_post] = cute.compile(
+                fa_bwd_post,
+                dv_accum_tensor,
+                dv_tensor,
+                cutlass.Float32(1.0),
+                cu_seqlens_k_tensor,
+                seqused_k_tensor,
+                current_stream,
+            )
+        _flash_attn_bwd.compile_cache_post[compile_key_post](
+            dv_accum_tensor,
+            dv_tensor,
+            cutlass.Float32(1.0),
+            cu_seqlens_k_tensor,
+            seqused_k_tensor,
+            current_stream,
+        )
+
+    return dq, dk, dv
+
+
+_flash_attn_bwd.compile_cache_pre = {}
+_flash_attn_bwd.compile_cache = {}
+_flash_attn_bwd.compile_cache_post = {}
+
+
+class FlashAttnFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        softmax_scale: Optional[float] = None,
+        causal: bool = False,
+        window_size: Tuple[Optional[int], Optional[int]] = (None, None),
+        learnable_sink: Optional[torch.Tensor] = None,
+        softcap: float = 0.0,
+        num_splits: int = 1,
+        pack_gqa: Optional[bool] = None,
+        mask_mod: Optional[Callable] = None,
+        full_block_cnt: Optional[torch.Tensor] = None,
+        full_block_idx: Optional[torch.Tensor] = None,
+        mask_block_cnt: Optional[torch.Tensor] = None,
+        mask_block_idx: Optional[torch.Tensor] = None,
+    ):
+        # Only create block sparse tensors if at least one block sparse parameter is provided
+        block_sparse_tensors = None
+        if any(t is not None for t in [full_block_cnt, full_block_idx, mask_block_cnt, mask_block_idx]):
+            block_sparse_tensors = BlockSparseTensorsTorch(
+                full_block_cnt=full_block_cnt,
+                full_block_idx=full_block_idx,
+                mask_block_cnt=mask_block_cnt,
+                mask_block_idx=mask_block_idx,
+            )
+        out, lse = _flash_attn_fwd(
+            q,
+            k,
+            v,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            learnable_sink=learnable_sink,
+            softcap=softcap,
+            num_splits=num_splits,
+            pack_gqa=pack_gqa,
+            mask_mod=mask_mod,
+            block_sparse_tensors=block_sparse_tensors
+        )
+        ctx.save_for_backward(q, k, v, out, lse)
+        ctx.softmax_scale = softmax_scale
+        ctx.causal = causal
+        ctx.window_size = window_size
+        ctx.softcap = softcap
+        return out, lse
+
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, lse = ctx.saved_tensors
+        dq, dk, dv = _flash_attn_bwd(
+            q,
+            k,
+            v,
+            out,
+            dout,
+            lse,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.softcap,
+        )
+        return dq, dk, dv, *((None,) * 20)  # Extra Nones is fine
+
+
+class FlashAttnVarlenFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        cu_seqlens_q: Optional[torch.Tensor],
+        cu_seqlens_k: Optional[torch.Tensor],
+        seqused_q: Optional[torch.Tensor] = None,
+        seqused_k: Optional[torch.Tensor] = None,
+        page_table: Optional[torch.Tensor] = None,
+        softmax_scale: Optional[float] = None,
+        causal: bool = False,
+        window_size: Tuple[Optional[int], Optional[int]] = (None, None),
+        learnable_sink: Optional[torch.Tensor] = None,
+        softcap: float = 0.0,
+        num_splits: int = 1,
+        pack_gqa: Optional[bool] = None,
+    ):
+        out, lse = _flash_attn_fwd(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqused_q,
+            seqused_k,
+            page_table=page_table,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            learnable_sink=learnable_sink,
+            softcap=softcap,
+            num_splits=num_splits,
+            pack_gqa=pack_gqa,
+        )
+        ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k)
+        ctx.softmax_scale = softmax_scale
+        ctx.causal = causal
+        ctx.window_size = window_size
+        ctx.softcap = softcap
+        return out, lse
+
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors
+        assert seqused_q == seqused_k == None
+        assert ctx.softcap == 0.0
+        dq, dk, dv = _flash_attn_bwd(
+            q,
+            k,
+            v,
+            out,
+            dout,
+            lse,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.softcap,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            seqused_q=seqused_q,
+            seqused_k=seqused_k,
+        )
+
+        return dq, dk, dv, *((None,) * 20)
+
+
+def flash_attn_func(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    window_size: Tuple[Optional[int], Optional[int]] = (None, None),
+    learnable_sink: Optional[torch.Tensor] = None,
+    softcap: float = 0.0,
+    num_splits: int = 1,
+    pack_gqa: Optional[bool] = None,
+    mask_mod: Optional[Callable] = None,
+    full_block_cnt: Optional[torch.Tensor] = None,
+    full_block_idx: Optional[torch.Tensor] = None,
+    mask_block_cnt: Optional[torch.Tensor] = None,
+    mask_block_idx: Optional[torch.Tensor] = None,
+):
+    return FlashAttnFunc.apply(
+        q,
+        k,
+        v,
+        softmax_scale,
+        causal,
+        window_size,
+        learnable_sink,
+        softcap,
+        num_splits,
+        pack_gqa,
+        mask_mod,
+        full_block_cnt,
+        full_block_idx,
+        mask_block_cnt,
+        mask_block_idx,
+    )
+
+
+def flash_attn_varlen_func(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    window_size: Tuple[Optional[int], Optional[int]] = (None, None),
+    learnable_sink: Optional[torch.Tensor] = None,
+    softcap: float = 0.0,
+    num_splits: int = 1,
+    pack_gqa: Optional[bool] = None,
+):
+    return FlashAttnVarlenFunc.apply(
+        q,
+        k,
+        v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        seqused_q,
+        seqused_k,
+        page_table,
+        softmax_scale,
+        causal,
+        window_size,
+        learnable_sink,
+        softcap,
+        num_splits,
+        pack_gqa,
+    )
+
+
+def _flash_attn_fwd_combine(
+    out_partial: torch.Tensor,
+    lse_partial: torch.Tensor,
+    out: torch.Tensor,
+    lse: Optional[torch.Tensor] = None,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    seqused: Optional[torch.Tensor] = None,
+    num_splits_dynamic_ptr: Optional[torch.Tensor] = None,
+    semaphore_to_reset: Optional[torch.Tensor] = None,
+) -> None:
+    """Forward combine kernel for split attention computation.
+
+    Combines partial outputs and log-sum-exp values from multiple splits
+    of attention computation into final outputs.
+
+    Args:
+        out_partial: Partial outputs tensor (num_splits, batch, seqlen, nheads, headdim) or
+                                            (num_splits, total_q, nheads, headdim) if there's cu_seqlens
+        lse_partial: Partial LSE tensor (num_splits, batch, seqlen, nheads) or
+                                       (num_splits, total_q, nheads) if there's cu_seqlens
+        out: Output tensor (batch, seqlen, nheads, headdim) or (total_q, nheads, headdim) if there's cu_seqlens
+        lse: Output LSE tensor (batch, seqlen, nheads) or (total_q, nheads) if there's cu_seqlens.
+        cu_seqlens: Cumulative sequence lengths for variable length sequences
+        seqused: Used sequence lengths for each batch
+        num_splits_dynamic_ptr: Dynamic number of splits per batch
+        semaphore_to_reset: Semaphore for synchronization
+        k_block_size: Block size for head dimension
+
+    Returns:
+        None
+    """
+    # Input validation
+    assert out_partial.dim() in [4, 5], "out_partial must have 4 or 5 dimensions"
+    assert lse_partial.dim() in [3, 4], "lse_partial must have 3 or 4 dimensions"
+    assert out_partial.dtype in [torch.float16, torch.bfloat16, torch.float32], (
+        "out_partial must be fp16, bf16, or fp32"
+    )
+    assert lse_partial.dtype == torch.float32, "lse_partial must be fp32"
+    assert out_partial.is_cuda and lse_partial.is_cuda, "tensors must be on CUDA device"
+    assert out_partial.stride(-1) == 1, "out_partial must be contiguous in the last dimension"
+    assert lse_partial.stride(-2) == 1, "lse_partial must be contiguous in the seqlen dimension"
+    assert lse_partial.shape == out_partial.shape[:-1]
+
+    # Determine if this is variable length based on dimensions
+    is_varlen = out_partial.dim() == 4
+
+    # Validate output tensor shapes and types
+    assert out.shape == out_partial.shape[1:], "out shape mismatch"
+    if lse is not None:
+        assert lse.shape == lse_partial.shape[1:], "lse shape mismatch"
+        assert lse.dtype == torch.float32, "lse must be fp32"
+
+    # Validate optional tensors
+    for t, name in [
+        (cu_seqlens, "cu_seqlens"),
+        (seqused, "seqused"),
+        (num_splits_dynamic_ptr, "num_splits_dynamic_ptr"),
+    ]:
+        if t is not None:
+            assert t.dtype == torch.int32, f"{name} must be int32"
+            assert t.is_cuda, f"{name} must be on CUDA device"
+            assert t.is_contiguous(), f"{name} must be contiguous"
+
+    head_dim = out_partial.shape[-1]
+    num_splits = out_partial.shape[0]
+    assert num_splits <= 256
+    # If hdim is 96 or 192, it's faster to round them to 128 or 256 respectively
+    # so that kBlockM is smaller and we have more parallelism.
+    k_block_size = 64 if head_dim <= 64 else 128
+    # We want kBlockM to be as small as possible to maximize parallelism.
+    # E.g., if hdim is 64, we want kBlockM to be 16 so that we can use 256 threads, each reading 4 elements (floats).
+    m_block_size = 8 if k_block_size % 128 == 0 else (16 if k_block_size % 64 == 0 else 32)
+    log_max_splits = max(math.ceil(math.log2(num_splits)), 4)
+    if m_block_size == 8:
+        # If kBlockM == 8 then the minimum number of splits is 32.
+        # TODO: we can deal w this by using 128 threads instead
+        log_max_splits = max(log_max_splits, 5)
+
+    # Convert to cute tensors (using kernel-formatted tensors)
+    out_partial_tensor = from_dlpack(out_partial.detach(), assumed_align=16).mark_layout_dynamic(
+        leading_dim=4 if not is_varlen else 3
+    )
+    lse_partial_tensor = from_dlpack(lse_partial.detach(), assumed_align=4).mark_layout_dynamic(
+        leading_dim=lse_partial.ndim - 2
+    )
+    out_tensor = from_dlpack(out.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=3 if not is_varlen else 2)
+    lse_tensor = (
+        from_dlpack(lse.detach(), assumed_align=4).mark_layout_dynamic(leading_dim=lse.ndim - 2)
+        if lse is not None
+        else None
+    )
+
+    optional_tensors = [
+        from_dlpack(t.detach(), assumed_align=4).mark_layout_dynamic(leading_dim=0)
+        if t is not None
+        else None
+        for t in (cu_seqlens, seqused, num_splits_dynamic_ptr, semaphore_to_reset)
+    ]
+    cu_seqlens_tensor, seqused_tensor, num_splits_dynamic_tensor, semaphore_tensor = (
+        optional_tensors
+    )
+
+    current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+
+    # Create combine kernel configuration
+    dtype = torch2cute_dtype_map[out.dtype]
+    dtype_partial = torch2cute_dtype_map[out_partial.dtype]
+
+    compile_key = (
+        dtype,
+        dtype_partial,
+        head_dim,
+        m_block_size,
+        k_block_size,
+        log_max_splits,
+        cu_seqlens is not None,
+        seqused is not None,
+        lse is not None,
+    )
+
+    if compile_key not in _flash_attn_fwd_combine.compile_cache:
+        fa_combine = FlashAttentionForwardCombine(
+            dtype=dtype,
+            dtype_partial=dtype_partial,
+            head_dim=head_dim,
+            m_block_size=m_block_size,
+            k_block_size=k_block_size,
+            log_max_splits=log_max_splits,
+        )
+
+        # Check if implementation is supported
+        if not fa_combine.can_implement(
+            dtype,
+            dtype_partial,
+            head_dim,
+            m_block_size,
+            k_block_size,
+            log_max_splits,
+            num_threads=256,
+        ):
+            raise RuntimeError(
+                "FlashAttention combine kernel cannot be implemented with given parameters"
+            )
+
+        _flash_attn_fwd_combine.compile_cache[compile_key] = cute.compile(
+            fa_combine,
+            out_partial_tensor,
+            lse_partial_tensor,
+            out_tensor,
+            lse_tensor,
+            cu_seqlens_tensor,
+            seqused_tensor,
+            num_splits_dynamic_tensor,
+            semaphore_tensor,
+            current_stream,
+        )
+
+    _flash_attn_fwd_combine.compile_cache[compile_key](
+        out_partial_tensor,
+        lse_partial_tensor,
+        out_tensor,
+        lse_tensor,
+        cu_seqlens_tensor,
+        seqused_tensor,
+        num_splits_dynamic_tensor,
+        semaphore_tensor,
+        current_stream,
+    )
+
+
+_flash_attn_fwd_combine.compile_cache = {}
+
+
+def flash_attn_combine(
+    out_partial: torch.Tensor,
+    lse_partial: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    out_dtype: Optional[torch.dtype] = None,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    seqused: Optional[torch.Tensor] = None,
+    return_lse: bool = True,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """Flash Attention combine function for split attention computation.
+
+    Combines partial outputs and log-sum-exp values from multiple splits
+    of attention computation into final outputs. This is the main user-facing
+    interface for the combine kernel.
+
+    Args:
+        out_partial: Partial outputs tensor with shape:
+            - (num_splits, batch_size, seqlen, num_heads, head_size) for regular batched input
+            - (num_splits, total_q, num_heads, head_size) for variable length input
+        lse_partial: Partial LSE tensor with shape:
+            - (num_splits, batch_size, seqlen, num_heads) for regular batched input
+            - (num_splits, total_q, num_heads) for variable length input
+        out: Optional output tensor. If None, will be created automatically.
+        out_dtype: Optional output dtype. If None, will use fp16/bf16 based on input.
+        cu_seqlens: Cumulative sequence lengths for variable length sequences
+        seqused: Used sequence lengths for each batch
+        return_lse: Whether to return the combined LSE tensor. Default is True.
+
+    Returns:
+        Tuple of (out, lse) where:
+        - out: Combined output tensor with shape (batch_size, seqlen, num_heads, head_size)
+              or (total_q, num_heads, head_size) for varlen
+        - lse: Combined log-sum-exp tensor with shape (batch_size, seqlen, num_heads)
+              or (total_q, num_heads) for varlen. None if return_lse=False
+
+    Note:
+        This function expects the input tensors to be in the format produced by
+        split attention computation, where the first dimension is num_splits.
+        The permuting from user format to kernel format is now done inside the kernel.
+    """
+    # Input validation
+    assert out_partial.dim() in [4, 5], "out_partial must have 4 or 5 dimensions"
+    assert lse_partial.dim() in [3, 4], "lse_partial must have 3 or 4 dimensions"
+    assert out_partial.dtype == torch.float32, "out_partial must be fp32 (from accumulation)"
+    assert lse_partial.dtype == torch.float32, "lse_partial must be fp32"
+
+    # Determine if this is variable length based on dimensions
+    is_varlen = out_partial.dim() == 4
+
+    if is_varlen:
+        # Variable length: (num_splits, total_q, num_heads, head_size)
+        num_splits, total_q, num_heads, head_size = out_partial.shape
+        assert lse_partial.shape == (num_splits, total_q, num_heads), (
+            "lse_partial shape mismatch for varlen"
+        )
+        batch_size = 1  # Treat as single batch for varlen
+        seqlen = total_q
+    else:
+        # Regular batched: (num_splits, batch_size, seqlen, num_heads, head_size)
+        num_splits, batch_size, seqlen, num_heads, head_size = out_partial.shape
+        assert lse_partial.shape == (num_splits, batch_size, seqlen, num_heads), (
+            "lse_partial shape mismatch"
+        )
+
+    # Determine output dtype
+    if out_dtype is None:
+        out_dtype = out_partial.dtype
+
+    # Create output if not provided
+    device = out_partial.device
+    if out is None:
+        if is_varlen:
+            out = torch.empty(total_q, num_heads, head_size, dtype=out_dtype, device=device)
+        else:
+            out = torch.empty(
+                batch_size, seqlen, num_heads, head_size, dtype=out_dtype, device=device
+            )
+
+    # Create lse output only if requested
+    if return_lse:
+        if is_varlen:
+            lse = torch.empty(num_heads, total_q, dtype=torch.float32, device=device).transpose(
+                0, 1
+            )
+        else:
+            lse = torch.empty(
+                batch_size, num_heads, seqlen, dtype=torch.float32, device=device
+            ).transpose(1, 2)
+    else:
+        lse = None
+
+    _flash_attn_fwd_combine(
+        out_partial,
+        lse_partial,
+        out,
+        lse,
+        cu_seqlens,
+        seqused,
+    )
+    return out, lse

build/torch-universal/mask.py

@@ -0,0 +1,413 @@
+# Copyright (c) 2025, Tri Dao.
+
+from typing import Optional, Callable
+from dataclasses import dataclass
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Int32, const_expr
+
+from . import utils
+
+
+@cute.jit
+def mask_r2p(X: cute.Tensor, col_limit: Int32, arch: int = 90, rank1: bool = False) -> None:
+    # Bit manipulation, compiles down to the R2P instruction
+    # For sm100: we know that tScS_t2r[i][1] == i, for the particular tmem copy atom we're using.
+    # For sm90: instead of comparing limit to 0, 1, 8, 9, 16, 17, ...,
+    # we compare a transformed version of limit to 0, 1, 2, 3, 4, 5, ...
+    if const_expr(arch == 90):
+        col_limit_transformed = col_limit // 8 * 2 + min(col_limit % 8, 2)
+    else:
+        col_limit_transformed = col_limit
+    ncol = const_expr(cute.size(X.shape[cute.rank(X) - 1]) if not rank1 else cute.size(X.shape))
+    # Ideally we'd move by 32 instead of 24, but mask >> i isn't correct for i == 31
+    for s in cutlass.range_constexpr(cute.ceil_div(ncol, 24)):
+        # Don't need to clamp to 32 since the shr.u32 instruction does that already
+        col_limit_right_s = max(col_limit_transformed - s * 24, 0)
+        # 0 -> 0b00...00, 1 -> 0b00...01, ..., 31 -> 0b01...11, 32 -> 0b11...11
+        mask = (1 << col_limit_right_s) - 1
+        # This needs to be range_constexpr, o/w the compiler can't generate the R2P instruction
+        for i in cutlass.range_constexpr(min(24, ncol - s * 24)):
+            in_bound = cutlass.Boolean(mask & (1 << i))
+            c = s * 24 + i
+            if const_expr(rank1):
+                X[c] = X[c] if in_bound else -Float32.inf
+                # This is the equivalent of:
+                # X[s * 24 + i] = X[s * 24 + i] if col_limit_right_s <= i 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 mask_r2p_transposed(X: cute.Tensor, row_limit_top: Int32, num_rep: int) -> None:
+    # Bit manipulation, compiles down to the R2P instruction
+    # For sm100: we know that tScS_t2r[i][0] has the form 0, 1, ..., 31, 64, ..., 127
+    # or 0, 1, ..., 15, 32, ..., 47, 64, ...
+    # We compare a transformed version of limit to 0, 1, 2, 3, 4, 5, ...
+    # Here we hardcode for the case of 2 warp groups.
+    num_wg = 2
+    row_limit_top_transformed = row_limit_top // (num_rep * num_wg) * num_rep + min(
+        row_limit_top % (num_rep * num_wg), num_rep
+    )
+    ncol = cute.size(X.shape)
+    # Ideally we'd move by 32 instead of 24, but mask >> i isn't correct for i == 31
+    for s in cutlass.range_constexpr(cute.ceil_div(ncol, 24)):
+        row_limit_top_s = max(row_limit_top_transformed - s * 24, 0)
+        # 0 -> 0b00...00, 1 -> 0b00...01, ..., 31 -> 0b01...11, 32 -> 0b11...11
+        mask = (1 << row_limit_top_s) - 1
+        # This needs to be range_constexpr, o/w the compiler can't generate the R2P instruction
+        for i in cutlass.range_constexpr(min(24, ncol - s * 24)):
+            out_bound = cutlass.Boolean(mask & (1 << i))
+            c = s * 24 + i
+            X[c] = -Float32.inf if out_bound else X[c]
+            # tidx = cute.arch.thread_idx()[0] % 256
+            # if tidx == 128:
+            #     cute.printf("tidx = {}, s = {}, i = {}, row_limit_top = {}, row_limit_top_s = {}, mask = {}, out_bound = {}", tidx, s, i, row_limit_top, row_limit_top_s, mask, out_bound)
+
+
+@dataclass(frozen=True)
+class AttentionMask:
+    tile_m: cutlass.Constexpr[int]
+    tile_n: cutlass.Constexpr[int]
+    seqlen_q: Int32
+    seqlen_k: Int32
+    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
+
+    @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,
+    ) -> None:
+        assert not (mask_causal and mask_local), "mask_causal and mask_local cannot be both True"
+        acc_S_mn = utils.make_acc_tensor_mn_view(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 = utils.make_acc_tensor_mn_view(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 = utils.make_acc_tensor_mn_view(
+            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]
+        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):
+                # The compiler now choses not to use R2P
+                r2p = const_expr(False and 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:
+                    mask_r2p(acc_S_mn, seqlenk_col_limit, arch=90)
+
+        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]))
+            thr_col_offset = tScS_mn[0, 0][1]
+
+            for r in cutlass.range_constexpr(nrow):
+                global_row_idx = tScS_mn[r, 0][0] + m_block * self.tile_m
+
+                for col in cutlass.range_constexpr(ncol):
+                    col_idx_local = t0ScS_mn[0, col][1]
+                    # Convert to absolute column index
+                    global_col_idx = thr_col_offset + col_idx_local + n_block * self.tile_n
+
+                    batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32)
+                    head_idx_ssa = utils.scalar_to_ssa(head_idx, cutlass.Int32)
+                    q_idx_ssa = utils.scalar_to_ssa(
+                        tScS_mn[r, 0][0] + m_block * self.tile_m, cutlass.Int32
+                    )
+                    kv_idx_ssa = utils.scalar_to_ssa(
+                        thr_col_offset + t0ScS_mn[0, col][1] + n_block * self.tile_n,
+                        cutlass.Int32,
+                    )
+                    mask_value = mask_mod(
+                        batch_idx_ssa,
+                        head_idx_ssa,
+                        q_idx_ssa,
+                        kv_idx_ssa,
+                        aux_tensors,
+                    )
+                    cond = cutlass.Boolean(utils.ssa_to_scalar(mask_value))
+                    if const_expr(mask_seqlen):
+                        out_of_bounds = (global_row_idx >= 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:
+                            mask_r2p(acc_S_mn[r, None], col_limit_right, arch=90, 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
+                    )
+                    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
+                            if const_expr(mask_seqlen):
+                                col_limit_right = cutlass.min(col_limit_right, seqlenk_col_limit)
+                        else:
+                            col_limit_right = self.tile_n
+                        col_limit_left = (
+                            row_idx + local_row_offset_left
+                            if const_expr(self.window_size_left is not None)
+                            else 0
+                        )
+                        # if cute.arch.thread_idx()[0] == 128: cute.printf("n_block = {}, r = {}, row_idx = {}, causal_row_offset = {}, col_limit_right = {}, col_limit_left = {}", n_block, r, row_idx, causal_row_offset, col_limit_right, col_limit_left)
+                        # traverse column index.
+                        for c in cutlass.range(cute.size(tScS_mn.shape[1]), unroll_full=True):
+                            col_idx = t0ScS_mn[0, c][1]
+                            # only consider the column index, so the row index sets to 0.
+                            if col_idx >= col_limit_right or col_idx < col_limit_left:
+                                acc_S_mn[r, c] = -Float32.inf
+            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
+                            else col0 - causal_row_offset - self.window_size_right
+                        )
+                        # TODO: do we need col_limit_sink?
+                        row_limit_bot = col0 - causal_row_offset + self.window_size_left
+                        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,
+    ) -> 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_t2r = thr_tmem_load.partition_D(tScS)
+        seqlenk_col_limit = self.seqlen_k - n_block * self.tile_n
+        r2p = True
+        if const_expr(not mask_causal and not mask_local):
+            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(acc_S, seqlenk_col_limit, arch=100, rank1=True)
+        else:  # Causal or local
+            causal_row_offset = 1 + 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
+                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(acc_S, col_limit_right, arch=100, rank1=True)
+            else:
+                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
+                )
+                if const_expr(self.window_size_right is not None):
+                    col_limit_right = row_idx + local_row_offset_right
+                    if const_expr(mask_seqlen):
+                        col_limit_right = cutlass.min(col_limit_right, seqlenk_col_limit)
+                else:
+                    col_limit_right = self.tile_n
+                col_limit_left = (
+                    row_idx + local_row_offset_left
+                    if const_expr(self.window_size_left is not None)
+                    else 0
+                )
+                # 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]
+                    )
+
+    @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,
+    ) -> None:
+        """
+        Backward pass: mask S = K @ Q.T where n_block tiles seqlen_k and m_block tiles seqlen_q.
+        """
+        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
+        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):
+            if const_expr(mask_seqlen):
+                if t0ScS_t2r[0][COL] >= seqlenk_col_limit:
+                    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]
+            causal_row_offset = (
+                seqlenk_col_limit - self.seqlen_q + m_block * self.tile_m + thr_row_offset
+            )
+            if const_expr(mask_causal):
+                col0 = t0ScS_t2r[0][COL]
+                row_limit_top = col0 - causal_row_offset
+                # tidx = cute.arch.thread_idx()[0] % 256
+                # if tidx < 32:
+                #     cute.printf("tidx = {}, {} {}, {} {}, col0 = {}", tidx, tScS_t2r[0][0], tScS_t2r[0][1], tScS_t2r[1][0], tScS_t2r[1][1], col0)
+                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 t0ScS_t2r[0][COL] >= seqlenk_col_limit:
+                        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
+                    mask_r2p_transposed(acc_S, row_limit_top, num_rep)
+            else:
+                assert False, "Local masking isn't supported yet"

build/torch-universal/mask_definitions.py

@@ -0,0 +1,220 @@
+from typing import Callable, Optional
+
+import random
+import math
+
+import cutlass
+import cutlass.cute as cute
+import torch
+
+from . import utils
+
+
+MaskModCallable = Optional[
+    Callable[
+        [
+            "cute.TensorSSA",
+            "cute.TensorSSA",
+            "cute.TensorSSA",
+            "cute.TensorSSA",
+            "Optional[list]",
+        ],
+        "cute.TensorSSA",
+    ]
+]
+
+
+# Flex Attention mask functions (PyTorch signatures for reference implementation)
+def get_flex_causal_mask(offset: int):
+    def _flex_causal_mask(b, h, q_idx, kv_idx):
+        return kv_idx <= q_idx + offset
+
+    return _flex_causal_mask
+
+
+def get_flex_block_causal_mask(offset: int):
+    def _flex_block_causal_mask(b, h, q_idx, kv_idx):
+        return kv_idx <= q_idx + offset
+
+    return _flex_block_causal_mask
+
+
+def get_flex_sliding_window_mask(window_left: int, window_right: int, offset: int):
+    def _flex_sliding_window_mask(b, h, q_idx, kv_idx):
+        center = q_idx + offset
+        lower = center - window_left
+        upper = center + window_right
+        return (kv_idx >= lower) & (kv_idx <= upper)
+
+    return _flex_sliding_window_mask
+
+
+def flex_block_diagonal_mask(b, h, q_idx, kv_idx):
+    block_size = 64
+    return (q_idx // block_size) == (kv_idx // block_size)
+
+
+def flex_mini_causal_mask(b, h, q_idx, kv_idx):
+    return (q_idx % 128) >= (kv_idx % 128)
+
+
+def flex_document_mask(b, h, q_idx, kv_idx, doc_id):
+    return doc_id[b, h, q_idx] == doc_id[b, h, kv_idx]
+
+
+# CuTe versions for kernel compilation
+def get_cute_causal_mask(offset: int):
+    @cute.jit
+    def _cute_causal_mask(
+        batch: cute.TensorSSA,
+        head: cute.TensorSSA,
+        m_idx: cute.TensorSSA,
+        n_idx: cute.TensorSSA,
+        aux_tensors: None,
+    ) -> cute.TensorSSA:
+        offset_ssa = utils.scalar_to_ssa(offset, cutlass.Int32)
+        return n_idx <= (m_idx + offset_ssa)
+
+    return _cute_causal_mask
+
+
+def get_cute_block_causal_mask(offset: int):
+    @cute.jit
+    def _cute_block_causal_mask(
+        batch: cute.TensorSSA,
+        head: cute.TensorSSA,
+        m_idx: cute.TensorSSA,
+        n_idx: cute.TensorSSA,
+        aux_tensors: None,
+    ) -> cute.TensorSSA:
+        offset_ssa = utils.scalar_to_ssa(offset, cutlass.Int32)
+        return n_idx <= (m_idx + offset_ssa)
+
+    return _cute_block_causal_mask
+
+
+def get_cute_sliding_window_mask(window_left: int, window_right: int, offset: int):
+    @cute.jit
+    def _cute_sliding_window_mask(
+        batch: cute.TensorSSA,
+        head: cute.TensorSSA,
+        m_idx: cute.TensorSSA,
+        n_idx: cute.TensorSSA,
+        aux_tensors,
+    ) -> cute.TensorSSA:
+        offset_ssa = utils.scalar_to_ssa(offset, cutlass.Int32)
+        window_left_ssa = utils.scalar_to_ssa(window_left, cutlass.Int32)
+        window_right_ssa = utils.scalar_to_ssa(window_right, cutlass.Int32)
+        center = m_idx + offset_ssa
+        lower = center - window_left_ssa
+        upper = center + window_right_ssa
+        return (n_idx >= lower) & (n_idx <= upper)
+
+    return _cute_sliding_window_mask
+
+
+@cute.jit
+def cute_document_mask(
+    batch: cute.TensorSSA,
+    head: cute.TensorSSA,
+    m_idx: cute.TensorSSA,
+    n_idx: cute.TensorSSA,
+    aux_tensors: list,
+) -> cute.TensorSSA:
+    doc_id = aux_tensors[0]
+    m_doc = utils.scalar_to_ssa(doc_id[batch[0], head[0], m_idx[0]], cutlass.Int32)
+    n_doc = utils.scalar_to_ssa(doc_id[batch[0], head[0], n_idx[0]], cutlass.Int32)
+    return m_doc == n_doc
+
+
+@cute.jit
+def cute_block_diagonal_mask(
+    batch: cute.TensorSSA,
+    head: cute.TensorSSA,
+    m_idx: cute.TensorSSA,
+    n_idx: cute.TensorSSA,
+    aux_tensors,
+) -> cute.TensorSSA:
+    block_size_ssa = utils.scalar_to_ssa(64, cutlass.Int32)
+    return (m_idx // block_size_ssa) == (n_idx // block_size_ssa)
+
+
+@cute.jit
+def cute_mini_causal_mask(
+    batch: cute.TensorSSA,
+    head: cute.TensorSSA,
+    m_idx: cute.TensorSSA,
+    n_idx: cute.TensorSSA,
+    aux_tensors,
+) -> cute.TensorSSA:
+    tile_size_ssa = utils.scalar_to_ssa(128, cutlass.Int32)
+    m_mod = m_idx % tile_size_ssa
+    n_mod = n_idx % tile_size_ssa
+    return m_mod >= n_mod
+
+
+def random_doc_id_tensor(nheads, batch, seqlen_q, device="cpu"):
+    doc_ids_tensor = torch.zeros(batch, nheads, seqlen_q, dtype=torch.int32, device=device)
+    for b in range(batch):
+        for h in range(nheads):
+            N = seqlen_q
+            max_segments = max(1, math.ceil(math.sqrt(max(N // 4, 1))))
+            n = random.randint(1, max_segments)
+            n = min(n, N)
+            cuts = sorted(random.sample(range(1, N), n - 1))
+            lengths = [b - a for a, b in zip((0, *cuts), (*cuts, N))]
+
+            doc_ids = []
+            for i, length in enumerate(lengths):
+                doc_ids += [i for _ in range(length)]
+
+            doc_ids_tensor[b, h, :] = torch.tensor(doc_ids, dtype=torch.int32, device=device)
+    return doc_ids_tensor
+
+
+STATIC_MASKS = {
+    "block_diagonal": (cute_block_diagonal_mask, flex_block_diagonal_mask),
+    "mini_causal": (cute_mini_causal_mask, flex_mini_causal_mask),
+    "document": (cute_document_mask, flex_document_mask),
+}
+
+PARAMETERIZED_MASK_FACTORIES = {
+    "causal": (get_cute_causal_mask, get_flex_causal_mask),
+    "block_causal": (get_cute_block_causal_mask, get_flex_block_causal_mask),
+    "sliding_window": (get_cute_sliding_window_mask, get_flex_sliding_window_mask),
+}
+
+
+def get_mask_pair(mask_name, seqlen_q=None, seqlen_k=None, window_size=None):
+    """Get (cute_mask, flex_mask) pair for the given mask name.
+
+    For static masks, seqlen info is not needed.
+    For parameterized masks, seqlen_q and seqlen_k are required.
+    """
+    if mask_name in STATIC_MASKS:
+        return STATIC_MASKS[mask_name]
+
+    if mask_name not in PARAMETERIZED_MASK_FACTORIES:
+        raise ValueError(f"Unknown mask: {mask_name}")
+
+    if seqlen_q is None or seqlen_k is None:
+        raise ValueError(f"Parameterized mask '{mask_name}' requires seqlen_q and seqlen_k")
+
+    cute_factory, flex_factory = PARAMETERIZED_MASK_FACTORIES[mask_name]
+    offset = seqlen_k - seqlen_q
+
+    if mask_name == "sliding_window":
+        if window_size is None:
+            raise ValueError("sliding_window mask requires window_size parameter")
+        cute_mask = cute_factory(window_size, window_size, offset)
+        flex_mask = flex_factory(window_size, window_size, offset)
+    else:
+        cute_mask = cute_factory(offset)
+        flex_mask = flex_factory(offset)
+
+    return cute_mask, flex_mask
+
+
+if __name__ == "__main__":
+    doc_ids = random_doc_id_tensor(1, 2, 128)
+    print(f"{doc_ids = }")

build/torch-universal/metadata.json

@@ -0,0 +1 @@
+{"python-depends":["nvidia-cutlass-dsl"]}

build/torch-universal/mma_sm100_desc.py

@@ -0,0 +1,291 @@
+# Copyright (c) 2025, Tri Dao.
+# Ported Cutlass code from C++ to Python:
+# https://github.com/NVIDIA/cutlass/blob/main/include/cute/arch/mma_sm100_desc.hpp
+# https://github.com/NVIDIA/cutlass/blob/main/include/cute/atom/mma_traits_sm100.hpp
+
+from enum import IntEnum
+
+import cutlass
+import cutlass.cute as cute
+
+# ---------------------------------------------------------------------------
+# Enumerations that match the HW encodings (values MUST stay identical)
+# ---------------------------------------------------------------------------
+
+
+class Major(IntEnum):  # matrix “layout” in the ISA docs
+    K = 0
+    MN = 1
+
+
+class ScaleIn(IntEnum):  # negate flags
+    One = 0
+    Neg = 1
+
+
+class Saturate(IntEnum):
+    False_ = 0
+    True_ = 1
+
+
+class CFormat(IntEnum):  # 2-bit field (bits 4-5)
+    F16 = 0
+    F32 = 1
+    S32 = 2
+
+
+class F16F32Format(IntEnum):  # 3-bit field (A/B element type)
+    F16 = 0
+    BF16 = 1
+    TF32 = 2
+
+
+class S8Format(IntEnum):
+    UINT8 = 0
+    INT8 = 1
+
+
+class MXF8F6F4Format(IntEnum):
+    E4M3 = 0
+    E5M2 = 1
+    E2M3 = 3
+    E3M2 = 4
+    E2M1 = 5
+
+
+class MaxShift(IntEnum):
+    NoShift = 0
+    MaxShift8 = 1
+    MaxShift16 = 2
+    MaxShift32 = 3
+
+
+# ---------------------------------------------------------------------------
+# CUTLASS-type → encoding helpers
+# ---------------------------------------------------------------------------
+
+
+def to_UMMA_format(cutlass_type) -> int:
+    """
+    Map a CUTLASS scalar class to the 3-bit encoding for Matrix A/B.
+    """
+    if cutlass_type is cutlass.Int8:
+        return S8Format.INT8
+    # Unsigned 8-bit (if available in your CUTLASS build)
+    if cutlass_type is cutlass.Uint8:
+        return S8Format.UINT8
+    # FP-16 / BF-16
+    if cutlass_type is cutlass.Float16:
+        return F16F32Format.F16
+    if cutlass_type is cutlass.BFloat16:
+        return F16F32Format.BF16
+    # TensorFloat-32 (8-bit exponent, 10-bit mantissa packed in 19 bits)
+    if cutlass_type is cutlass.TFloat32:
+        return F16F32Format.TF32
+    # Float-8 / Float-6 / Float-4 – add whenever CUTLASS exposes them
+    if cutlass_type is cutlass.FloatE4M3FN:
+        return MXF8F6F4Format.E4M3
+    if cutlass_type is cutlass.FloatE5M2:
+        return MXF8F6F4Format.E5M2
+    raise TypeError(f"Unsupported CUTLASS scalar type for A/B: {cutlass_type!r}")
+
+
+def to_C_format(cutlass_type) -> int:
+    """
+    Map a CUTLASS scalar class to the 2-bit accumulator encoding.
+    """
+    if cutlass_type is cutlass.Float16:
+        return CFormat.F16
+    if cutlass_type is cutlass.Float32:
+        return CFormat.F32
+    if cutlass_type is cutlass.Int32:
+        return CFormat.S32
+    raise TypeError(f"Unsupported CUTLASS scalar type for accumulator: {cutlass_type!r}")
+
+
+# ---------------------------------------------------------------------------
+# The constructor – accepts only CUTLASS scalar classes
+# ---------------------------------------------------------------------------
+
+
+def make_instr_desc(
+    a_type,  # CUTLASS scalar class, e.g. cutlass.Int8
+    b_type,
+    c_type,
+    M: int,  # 64, 128 or 256
+    N: int,  # 8 … 256 (multiple of 8)
+    a_major: Major,
+    b_major: Major,
+    a_neg: ScaleIn = ScaleIn.One,
+    b_neg: ScaleIn = ScaleIn.One,
+    c_sat: Saturate = Saturate.False_,
+    is_sparse: bool = False,
+    max_shift: MaxShift = MaxShift.NoShift,
+) -> int:
+    """
+    Build the 32-bit instruction descriptor for Blackwell MMA.
+    All matrix/accumulator **types must be CUTLASS scalar classes** –
+    passing integers is forbidden.
+    """
+    # --- encode element formats -------------------------------------------------
+    a_fmt = int(to_UMMA_format(a_type))
+    b_fmt = int(to_UMMA_format(b_type))
+    c_fmt = int(to_C_format(c_type))
+
+    # --- range checks on M/N -----------------------------------------------------
+    if M not in (64, 128, 256):
+        raise ValueError("M must be 64, 128 or 256")
+    if N < 8 or N > 256 or (N & 7):
+        raise ValueError("N must be a multiple of 8 in the range 8…256")
+
+    m_dim = M >> 4  # 5-bit field
+    n_dim = N >> 3  # 6-bit field
+
+    # fmt: off
+    # --- pack the bit-fields -----------------------------------------------------
+    desc = 0
+    desc |= (0                 & 0x3) << 0        # sparse_id2 (always 0 here)
+    desc |= (int(is_sparse)    & 0x1) << 2        # sparse_flag
+    desc |= (int(c_sat)        & 0x1) << 3        # saturate
+    desc |= (c_fmt             & 0x3) << 4        # c_format
+    desc |= (a_fmt             & 0x7) << 7        # a_format
+    desc |= (b_fmt             & 0x7) << 10       # b_format
+    desc |= (int(a_neg)        & 0x1) << 13       # a_negate
+    desc |= (int(b_neg)        & 0x1) << 14       # b_negate
+    desc |= (int(a_major)      & 0x1) << 15       # a_major
+    desc |= (int(b_major)      & 0x1) << 16       # b_major
+    desc |= (n_dim             & 0x3F) << 17      # n_dim (6 bits)
+    desc |= (m_dim             & 0x1F) << 24      # m_dim (5 bits)
+    desc |= (int(max_shift)    & 0x3) << 30       # max_shift (2 bits)
+    # fmt: on
+
+    return desc & 0xFFFF_FFFF  # ensure 32-bit result
+
+
+def mma_op_to_idesc(op: cute.nvgpu.tcgen05.mma.MmaOp):
+    return make_instr_desc(
+        op.a_dtype,
+        op.b_dtype,
+        op.acc_dtype,
+        op.shape_mnk[0],
+        op.shape_mnk[1],
+        Major.K if op.a_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K else Major.MN,
+        Major.K if op.b_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K else Major.MN,
+    )
+
+
+class LayoutType(IntEnum):  # occupies the top-3 bits [61:64)
+    SWIZZLE_NONE = 0  # (a.k.a. “INTERLEAVE” in older docs)
+    SWIZZLE_128B_BASE32B = 1
+    SWIZZLE_128B = 2
+    SWIZZLE_64B = 4
+    SWIZZLE_32B = 6
+    # values 3,5,7 are reserved / illegal for UMMA
+
+
+# ---------------------------------------------------------------------------
+#  Helpers – figure out the SWIZZLE_* family from the tensor layout
+# ---------------------------------------------------------------------------
+
+
+def _layout_type(swizzle: cute.Swizzle) -> LayoutType:
+    # No idea what the right way to get B, M, S is – so we're just parsing it from the __str__
+    # Swizzle string has the form "S<B,M,S>"
+    swz_str = str(swizzle)
+    inside = swz_str[swz_str.index("<") + 1 : swz_str.index(">")]  # '3,4,3'
+    B, M, S = [int(x) for x in inside.split(",")]  # [3, 4, 3]
+
+    if M == 4:  # Swizzle<*,4,3>
+        if S != 3:
+            raise ValueError("Unexpected swizzle shift – want S==3 for M==4")
+        return {
+            0: LayoutType.SWIZZLE_NONE,
+            1: LayoutType.SWIZZLE_32B,
+            2: LayoutType.SWIZZLE_64B,
+            3: LayoutType.SWIZZLE_128B,
+        }[B]  # KeyError ⇒ invalid B→ raise
+    if M == 5:  # Swizzle<2,5,2> (the only legal triple for M==5)
+        if (B, S) != (2, 2):
+            raise ValueError("Only Swizzle<2,5,2> supported for 128B_BASE32B")
+        return LayoutType.SWIZZLE_128B_BASE32B
+
+    # Any other (M,B,S) triple is not a UMMA-legal shared-memory layout
+    raise ValueError("Unsupported swizzle triple for UMMA smem descriptor")
+
+
+def make_smem_desc_base(layout: cute.Layout, swizzle: cute.Swizzle, major: Major) -> int:
+    """
+    Convert a 2-D *shared-memory* Cute layout into the Blackwell 64-bit
+    smem-descriptor, without the smem start address.
+    layout must correspond to layout of an uint128 tensor.
+    """
+    # ------------------------------------------------------------------ meta
+    layout_type = _layout_type(swizzle)  # resolve SWIZZLE_* family
+
+    VERSION = 1  # bits 46–47
+    LBO_MODE = 0  # bit  52
+    BASE_OFFSET = 0  # bits 49–51   (CUTLASS always 0)
+
+    # ---------------------------------------------------------- strides  (units: uint128_t = 16 B)
+    swizzle_atom_mn_size = {
+        LayoutType.SWIZZLE_NONE: 1,
+        LayoutType.SWIZZLE_32B: 2,
+        LayoutType.SWIZZLE_64B: 4,
+        LayoutType.SWIZZLE_128B: 8,
+        LayoutType.SWIZZLE_128B_BASE32B: 8,
+    }[layout_type]
+
+    if major is Major.MN:
+        swizzle_atom_k_size = 4 if layout_type is LayoutType.SWIZZLE_128B_BASE32B else 8
+        canonical_layout = cute.logical_divide(layout, (swizzle_atom_mn_size, swizzle_atom_k_size))
+        if not cute.is_congruent(canonical_layout, ((1, 1), (1, 1))):
+            raise ValueError("Not a canonical UMMA_MN Layout: Expected profile failure.")
+        stride_00 = canonical_layout.stride[0][0]
+        if layout_type is not LayoutType.SWIZZLE_NONE and stride_00 != 1:
+            raise ValueError("Not a canonical UMMA_MN Layout: Expected stride failure.")
+        stride_10 = canonical_layout.stride[1][0]
+        if stride_10 != swizzle_atom_mn_size:
+            raise ValueError("Not a canonical UMMA_MN Layout: Expected stride failure.")
+        stride_01, stride_11 = canonical_layout.stride[0][1], canonical_layout.stride[1][1]
+        if layout_type is LayoutType.SWIZZLE_NONE:
+            stride_byte_offset, leading_byte_offset = stride_01, stride_11
+        else:
+            stride_byte_offset, leading_byte_offset = stride_11, stride_01
+    else:
+        if layout_type == LayoutType.SWIZZLE_128B_BASE32B:
+            raise ValueError("SWIZZLE_128B_BASE32B is invalid for Major-K")
+        if not cute.size(layout.shape[0]) % 8 == 0:
+            raise ValueError("Not a canonical UMMA_K Layout: Expected MN-size multiple of 8.")
+        canonical_layout = cute.logical_divide(layout, (8, 2))
+        if not cute.is_congruent(canonical_layout, ((1, 1), (1, 1))):
+            raise ValueError("Not a canonical UMMA_K Layout: Expected profile failure.")
+        stride_00 = canonical_layout.stride[0][0]
+        if stride_00 != swizzle_atom_mn_size:
+            raise ValueError("Not a canonical UMMA_K Layout: Expected stride failure.")
+        stride_10 = canonical_layout.stride[1][0]
+        if layout_type is not LayoutType.SWIZZLE_NONE and stride_10 != 1:
+            raise ValueError("Not a canonical UMMA_K Layout: Expected stride failure.")
+        stride_01 = canonical_layout.stride[0][1]
+        stride_byte_offset, leading_byte_offset = stride_01, stride_10
+
+    # ------------------------------------------------------------------ pack
+    desc = 0
+    # leading_byte_offset_  [16:30)
+    desc |= (leading_byte_offset & 0x3FFF) << 16
+    # stride_byte_offset_   [32:46)
+    desc |= (stride_byte_offset & 0x3FFF) << 32
+    # version_             [46:48)
+    desc |= (VERSION & 0x3) << 46
+    # base_offset_         [49:52)
+    desc |= (BASE_OFFSET & 0x7) << 49
+    # lbo_mode_            [52:53)
+    desc |= (LBO_MODE & 0x1) << 52
+    # layout_type_         [61:64)
+    desc |= (int(layout_type) & 0x7) << 61
+
+    return desc & 0xFFFF_FFFF_FFFF_FFFF  # force 64-bit width
+
+
+def make_smem_desc_start_addr(start_addr: cute.Pointer) -> cutlass.Int32:
+    # 14 bits, remove 4 LSB (bits 0-13 in desc)
+    return (start_addr.toint() & 0x3FFFF) >> 4

build/torch-universal/named_barrier.py

@@ -0,0 +1,31 @@
+# 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 NamedBarrierBwd(enum.IntEnum):
+    Epilogue = enum.auto()
+    WarpSchedulerWG1 = enum.auto()
+    WarpSchedulerWG2 = enum.auto()
+    WarpSchedulerWG3 = enum.auto()
+    PdS = enum.auto()
+    dQFullWG0 = enum.auto()
+    dQFullWG1 = enum.auto()
+    dQEmptyWG0 = enum.auto()
+    dQEmptyWG1 = enum.auto()
+
+
+class NamedBarrierBwdSm100(enum.IntEnum):
+    EpilogueWG1 = enum.auto()
+    EpilogueWG2 = enum.auto()
+    Compute = enum.auto()
+    dQaccReduce = enum.auto()

build/torch-universal/pack_gqa.py

@@ -0,0 +1,166 @@
+# Copyright (c) 2025, Tri Dao.
+
+import math
+import operator
+
+import cutlass
+import cutlass.cute as cute
+
+from . import utils
+
+
+class PackGQA:
+    def __init__(
+        self,
+        m_block_size: cutlass.Constexpr[int],
+        head_dim_padded: cutlass.Constexpr[int],
+        check_hdim_oob: cutlass.Constexpr[bool],
+        qhead_per_kvhead: cutlass.Constexpr[bool],
+    ):
+        self.m_block_size = m_block_size
+        self.head_dim_padded = head_dim_padded
+        self.check_hdim_oob = check_hdim_oob
+        self.qhead_per_kvhead = qhead_per_kvhead
+
+    @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_fragment(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 = utils.make_acc_tensor_mn_view(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,
+                    )

build/torch-universal/pipeline.py

@@ -0,0 +1,365 @@
+# Copyright (c) 2025, Tri Dao.
+
+# import math
+from typing import Optional
+from dataclasses import dataclass
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Boolean, Int32, const_expr
+from cutlass.cutlass_dsl import if_generate
+from cutlass.pipeline import PipelineAsync, PipelineState, Agent, CooperativeGroup
+from cutlass.pipeline import PipelineUserType, PipelineOp
+from cutlass.pipeline import PipelineTmaAsync as PipelineTmaAsyncOg
+from cutlass.pipeline import PipelineTmaUmma as PipelineTmaUmmaOg
+
+
+# We deviate from cute-dsl implementation to use cute.arch.cluster_arrive_relaxed
+def pipeline_init_wait(cta_layout_vmnk: Optional[cute.Layout] = None):
+    """
+    Fences the mbarrier init and syncs the threadblock or cluster
+    """
+    cute.arch.mbarrier_init_fence()
+
+    if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1:
+        # If not using clusters, sync the threadblock
+        _sync(Agent.ThreadBlock)
+    else:
+        # If using clusters, sync the cluster
+        _sync(Agent.ThreadBlockCluster)
+
+
+def _sync(group: Agent):
+    """
+    Syncs all threads within an agent.
+    """
+    if group is Agent.Thread:
+        raise NotImplementedError("Error: Not supported.")
+    elif group is Agent.ThreadBlock:
+        cute.arch.sync_threads()
+    elif group is Agent.ThreadBlockCluster:
+        cute.arch.cluster_arrive_relaxed()
+        cute.arch.cluster_wait()
+    else:
+        assert False, (
+            "Error: No explicit sync instruction exists. Please use barriers (named / mbarrier) instead."
+        )
+
+
+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):
+        # assert stages < 2**16
+        # self._log_stages = int(math.log2(stages))
+        # assert 1 << self._log_stages == stages, "Number of stages must be a power of 2."
+        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 1 << self._log_stages
+        return self._stages
+
+    @property
+    def index(self) -> Int32:
+        # return self._phase_index & 0xFFFF
+        # return self._phase_index & ((1 << self._log_stages) - 1)
+        if const_expr(self._stages == 1):
+            return Int32(0)
+        else:
+            return self._phase_index % self._stages
+
+    @property
+    def phase(self) -> Int32:
+        # return self._phase_index >> 16
+        # 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.
+        # return (self._phase_index >> self._log_stages) % 2
+        # return self._phase_index >> self._log_stages
+        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 then_body(phase_index):
+        #     # XOR the phase bit and set the index to 0
+        #     return (phase_index & 0xFFFF0000) ^ (1 << 16)
+
+        # def else_body(phase_index):
+        #     return phase_index
+
+        # self._phase_index = if_generate(
+        #     (self._phase_index & 0xFFFF) == self.stages,
+        #     then_body,
+        #     else_body,
+        #     [self._phase_index],
+        #     [Int32],
+        # )
+
+    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(1 << 16))
+        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."
+
+
+@dataclass(frozen=True)
+class PipelineTmaAsync(PipelineTmaAsyncOg):
+    """
+    If size(ClusterShape) == 1, PipelineTmaAsync has all threads
+    signaling the barrier during consumer_release. This causes a perf regression in FA3
+    forward pass (especially hdim 128 causal). We instead implement a version of
+    PipelineTmaAsync where only 1 out of 128 threads signals the barrier.
+
+    Assumptions:
+    (1) num_consumers % NumThreadsPerWarpGroup == 0
+    (2) all 128 threads in the warp group are sync'ed right before calling consumer_release
+    """
+
+    @staticmethod
+    def create(
+        *,
+        num_stages: int,
+        producer_group: CooperativeGroup,
+        consumer_group: CooperativeGroup,
+        tx_count: int,
+        barrier_storage: cute.Pointer = None,
+        cta_layout_vmnk: Optional[cute.Layout] = None,
+        tidx: Optional[Int32] = None,
+        mcast_mode_mn: tuple[int, int] = (1, 1),
+        init_wait: cutlass.Constexpr[bool] = True,
+    ):
+        """
+        This helper function computes any necessary attributes and returns an instance of PipelineTmaAsync.
+        :param barrier_storage: Pointer to the smem address for this pipeline's mbarriers
+        :type barrier_storage: cute.Pointer
+        :param num_stages: Number of buffer stages for this pipeline
+        :type num_stages: Int32
+        :param producer_group: `CooperativeGroup` for the producer agent
+        :type producer_group: CooperativeGroup
+        :param consumer_group: `CooperativeGroup` for the consumer agent
+        :type consumer_group: CooperativeGroup
+        :param tx_count: Number of bytes expected to be written to the transaction barrier for one stage
+        :type tx_count: int
+        :param cta_layout_vmnk: Layout of the cluster shape
+        :type cta_layout_vmnk: cute.Layout | None
+        :param tidx: thread index to consumer async threads
+        :type tidx: Int32 | None
+        :param mcast_mode_mn: Tuple of two integers, specifying whether mcast is enabled for the m and n modes. At least one of the two integers must be 1.
+        :type mcast_mode_mn: tuple[int, int]
+        """
+        if not isinstance(barrier_storage, cute.Pointer):
+            raise ValueError(
+                f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}"
+            )
+
+        producer_type = PipelineOp.TmaLoad
+        consumer_type = PipelineOp.AsyncThread
+
+        producer = (producer_type, producer_group)
+        consumer = (consumer_type, consumer_group)
+
+        sync_object_full = PipelineAsync._make_sync_object(
+            barrier_storage.align(min_align=8), num_stages, producer, tx_count
+        )
+        sync_object_empty = PipelineAsync._make_sync_object(
+            barrier_storage.align(min_align=8) + num_stages, num_stages, consumer
+        )
+        if tidx is None:
+            tidx, _, _ = cute.arch.thread_idx()
+        if cta_layout_vmnk is None:
+            cta_layout_vmnk = cute.make_layout((1, 1, 1, 1))
+        if const_expr(cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1):
+            dst_rank = None
+            is_signalling_thread = tidx % 128 == 0
+        else:
+            (
+                dst_rank,
+                is_signalling_thread,
+            ) = PipelineTmaAsync.init_empty_barrier_arrive_signal(
+                cta_layout_vmnk, tidx, mcast_mode_mn
+            )
+
+        producer_mask = None
+
+        if const_expr(init_wait):
+            pipeline_init_wait()
+
+        return PipelineTmaAsync(
+            sync_object_full,
+            sync_object_empty,
+            num_stages,
+            producer_mask,
+            dst_rank,
+            is_signalling_thread,
+        )
+
+    def producer_acquire(
+        self,
+        state: PipelineState,
+        try_acquire_token: Optional[Boolean] = None,
+        extra_tx_count: int = 0,
+    ):
+        """
+        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),
+        )
+        if const_expr(extra_tx_count == 0):
+            self.sync_object_full.arrive(state.index, self.producer_mask)
+        else:
+            tx_count = self.sync_object_full.tx_count + extra_tx_count
+            self.sync_object_full.arrive_and_expect_tx(state.index, tx_count)
+
+    def consumer_release(self, state: PipelineState):
+        """
+        TMA consumer release conditionally signals the empty buffer to the producer.
+        """
+        # Only 1 thread per warp group signals the empty buffer.
+        if self.consumer_mask is None:  # No cluster, 1 thread per warp group to signal
+            if_generate(
+                cute.arch.thread_idx()[0] % 128 == 0,
+                lambda: self.sync_object_empty.arrive(state.index, self.consumer_mask),
+            )
+        else:
+            if_generate(
+                self.is_signalling_thread,
+                lambda: self.sync_object_empty.arrive(state.index, self.consumer_mask),
+            )
+
+
+@dataclass(frozen=True)
+class PipelineTmaUmma(PipelineTmaUmmaOg):
+    @staticmethod
+    def create(
+        *,
+        num_stages: int,
+        producer_group: CooperativeGroup,
+        consumer_group: CooperativeGroup,
+        tx_count: int,
+        barrier_storage: cute.Pointer = None,
+        cta_layout_vmnk: Optional[cute.Layout] = None,
+        mcast_mode_mn: tuple[int, int] = (1, 1),
+        init_wait: cutlass.Constexpr[bool] = True,
+    ):
+        """
+        This helper function computes any necessary attributes and returns an instance of PipelineTmaUmma.
+        :param barrier_storage: Pointer to the smem address for this pipeline's mbarriers
+        :type barrier_storage: cute.Pointer
+        :param num_stages: Number of buffer stages for this pipeline
+        :type num_stages: Int32
+        :param producer_group: `CooperativeGroup` for the producer agent
+        :type producer_group: CooperativeGroup
+        :param consumer_group: `CooperativeGroup` for the consumer agent
+        :type consumer_group: CooperativeGroup
+        :param tx_count: Number of bytes expected to be written to the transaction barrier for one stage
+        :type tx_count: int
+        :param cta_layout_vmnk: Layout of the cluster shape
+        :type cta_layout_vmnk: cute.Layout | None
+        :param mcast_mode_mn: Tuple of two integers, specifying whether mcast is enabled for the m and n modes. At least one of the two integers must be 1.
+        :type mcast_mode_mn: tuple[int, int]
+        """
+        if not isinstance(barrier_storage, cute.Pointer):
+            raise ValueError(
+                f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}"
+            )
+
+        producer_type = PipelineOp.TmaLoad
+        consumer_type = PipelineOp.TCGen05Mma
+
+        producer = (producer_type, producer_group)
+        consumer = (consumer_type, consumer_group)
+
+        sync_object_full = PipelineAsync._make_sync_object(
+            barrier_storage.align(min_align=8), num_stages, producer, tx_count
+        )
+        sync_object_empty = PipelineAsync._make_sync_object(
+            barrier_storage.align(min_align=8) + num_stages, num_stages, consumer
+        )
+
+        if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1:
+            # No mcast mask if not using clusters
+            producer_mask = None
+            # All threadblocks are leaders if not using clusters
+            is_leader_cta = True
+        else:
+            producer_mask = PipelineTmaUmma._compute_mcast_arrival_mask(
+                cta_layout_vmnk, mcast_mode_mn
+            )
+            is_leader_cta = PipelineTmaUmma._compute_is_leader_cta(cta_layout_vmnk)
+
+        cta_group = (
+            cute.nvgpu.tcgen05.CtaGroup.ONE
+            if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, mode=[0]) == 1
+            else cute.nvgpu.tcgen05.CtaGroup.TWO
+        )
+
+        consumer_mask = producer_mask
+
+        if const_expr(init_wait):
+            pipeline_init_wait(cta_layout_vmnk)
+
+        return PipelineTmaUmma(
+            sync_object_full,
+            sync_object_empty,
+            num_stages,
+            producer_mask,
+            consumer_mask,
+            is_leader_cta,
+            cta_group,
+        )
+
+    def producer_acquire(
+        self,
+        state: PipelineState,
+        try_acquire_token: Optional[Boolean] = None,
+        extra_tx_count: int = 0,
+    ):
+        """
+        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),
+        )
+        if const_expr(extra_tx_count == 0):
+            if_generate(
+                self.is_leader_cta,
+                lambda: self.sync_object_full.arrive(state.index, self.producer_mask),
+            )
+        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),
+            )

build/torch-universal/seqlen_info.py

@@ -0,0 +1,99 @@
+from typing import Optional
+from dataclasses import dataclass
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32, const_expr
+
+"""
+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: cutlass.Int32
+    seqlen: cutlass.Int32
+
+    @staticmethod
+    def create(
+        batch_idx: cutlass.Int32,
+        seqlen_static: cutlass.Int32,
+        cu_seqlens: Optional[cute.Tensor] = None,
+        seqused: Optional[cute.Tensor] = None,
+    ):
+        offset = 0 if const_expr(cu_seqlens is None) else cu_seqlens[batch_idx]
+        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, seqlen)
+
+
+@dataclass(frozen=True)
+class SeqlenInfoQK:
+    offset_q: cutlass.Int32
+    offset_k: cutlass.Int32
+    seqlen_q: cutlass.Int32
+    seqlen_k: cutlass.Int32
+    has_cu_seqlens_q: cutlass.Constexpr[bool]
+    has_cu_seqlens_k: cutlass.Constexpr[bool]
+
+    @staticmethod
+    def create(
+        batch_idx: cutlass.Int32,
+        seqlen_q_static: cutlass.Int32,
+        seqlen_k_static: cutlass.Int32,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        mSeqUsedQ: Optional[cute.Tensor] = None,
+        mSeqUsedK: Optional[cute.Tensor] = None,
+    ):
+        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]
+        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
+            )
+        has_cu_seqlens_q: int = mCuSeqlensQ is not None
+        has_cu_seqlens_k: int = mCuSeqlensK is not None
+        return SeqlenInfoQK(
+            offset_q, offset_k, seqlen_q, seqlen_k, has_cu_seqlens_q, has_cu_seqlens_k
+        )
+
+    def offset_batch_Q(self, mQ: cute.Tensor, batch_idx: Int32, dim: int) -> cute.Tensor:
+        """Seqlen must be the first dimension of mQ"""
+        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 = (
+                self.offset_q if const_expr(cute.rank(mQ.shape[0]) == 1) else (0, self.offset_q)
+            )
+            idx = (offset,) + (0,) * (cute.rank(mQ) - 1)
+            return cute.domain_offset(idx, mQ)
+
+    def offset_batch_K(self, mK: cute.Tensor, batch_idx: Int32, dim: int) -> cute.Tensor:
+        """Seqlen must be the first dimension of mK"""
+        if const_expr(not self.has_cu_seqlens_k):
+            idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mK) - 1 - dim)
+            return mK[idx]
+        else:
+            idx = (self.offset_k,) + (0,) * (cute.rank(mK) - 1)
+            return cute.domain_offset(idx, mK)

build/torch-universal/softmax.py

@@ -0,0 +1,440 @@
+# 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 . import utils
+from .cute_dsl_utils import ParamsBase
+
+
+@dataclass
+class Softmax(ParamsBase):
+    scale_log2: Float32
+    num_rows: cutlass.Constexpr[int]
+    row_max: cute.Tensor
+    row_sum: 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_fragment(num_rows, Float32)
+        row_sum = cute.make_fragment(num_rows, Float32)
+        return Softmax(scale_log2, num_rows, row_max, row_sum, arch, softmax_scale)
+
+    def reset(self) -> None:
+        self.row_max.fill(-Float32.inf)
+        self.row_sum.fill(0.0)
+
+    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 = utils.make_acc_tensor_mn_view(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 = utils.warp_reduce(row_max_cur, cute.arch.fmax, width=4)
+            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 = utils.exp2f(acc_S_row * scale_log2 - row_max_cur_scaled)
+
+                acc_S_row_sum = utils.fadd_reduce(acc_S_row_exp, init_val=None, arch=arch)
+                row_scale[r] = 1.0
+            else:
+                row_max_prev = row_max[r]
+                row_max_cur_scaled = row_max_cur * scale_log2
+                acc_S_row_exp = utils.exp2f(acc_S_row * scale_log2 - row_max_cur_scaled)
+                # row_scale[r] = utils.exp2f(row_max_prev * self.scale_log2 - row_max_cur_scaled)
+                row_scale[r] = utils.exp2f((row_max_prev - row_max_cur) * scale_log2)
+
+                acc_S_row_sum = utils.fadd_reduce(
+                    acc_S_row_exp, init_val=row_sum[r] * row_scale[r], arch=arch
+                )
+
+            row_max[r] = row_max_cur
+            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] += utils.exp2f(sink_val_cur * LOG2_E - row_max[r] * scale_log2)
+
+            # 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 + utils.log2f(row_sum_cur)) * 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 = utils.make_acc_tensor_mn_view(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
+
+    @staticmethod
+    def create(
+        scale_log2: Float32,
+        rescale_threshold: cutlass.Constexpr[float] = 0.0,
+        softmax_scale: Float32 | None = None,
+    ):
+        num_rows = 1
+        arch = 100
+        row_max = cute.make_fragment(num_rows, Float32)
+        row_sum = cute.make_fragment(num_rows, Float32)
+        return SoftmaxSm100(
+            scale_log2,
+            num_rows,
+            row_max,
+            row_sum,
+            arch,
+            softmax_scale,
+            rescale_threshold=rescale_threshold,
+        )
+
+    @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 = utils.exp2f(acc_scale_)
+            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] = utils.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,
+        e2e: cutlass.Constexpr[bool] = False,
+        e2e_freq: cutlass.Constexpr[int] = 16,
+        e2e_res: cutlass.Constexpr[int] = 4,
+        e2e_frg_limit: cutlass.Constexpr[int] = 1,
+    ):
+        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] = utils.exp2f(acc_S_row_frg[k, j])
+                # acc_S_row_frg[k + 1, j] = utils.exp2f(acc_S_row_frg[k + 1, j])
+                if cutlass.const_expr(not e2e):
+                    acc_S_row_frg[k, j] = cute.arch.exp2(acc_S_row_frg[k, j])
+                    acc_S_row_frg[k + 1, j] = cute.arch.exp2(acc_S_row_frg[k + 1, j])
+                else:
+                    if cutlass.const_expr(
+                        k % e2e_freq < e2e_freq - e2e_res or j >= frg_cnt - e2e_frg_limit
+                    ):
+                        acc_S_row_frg[k, j] = cute.arch.exp2(acc_S_row_frg[k, j])
+                        acc_S_row_frg[k + 1, j] = cute.arch.exp2(acc_S_row_frg[k + 1, j])
+                    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] = utils.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] = utils.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.arch.exp2(acc_S_row[i])
+        #     acc_S_row[i + 1] = cute.arch.exp2(acc_S_row[i + 1])
+
+        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] = (
+                #     utils.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] = utils.exp2f(acc_S_row_frg[k, j])
+                # acc_S_row_frg[k + 1, j] = utils.exp2f(acc_S_row_frg[k + 1, j])
+                acc_S_row_frg[k, j] = cute.arch.exp2(acc_S_row_frg[k, j])
+                acc_S_row_frg[k + 1, j] = cute.arch.exp2(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 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,
+    constant_q_idx: cutlass.Constexpr,
+    qhead_per_kvhead: cutlass.Constexpr[int] = 1,
+):
+    """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
+        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.
+    """
+    n_vals = cutlass.const_expr(cute.size(score_tensor.shape))
+    score_vec = cute.make_fragment(vec_size, qk_acc_dtype)
+    kv_idx_vec = cute.make_fragment(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_fragment(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_fragment(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][0]
+                # 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][0], qhead_per_kvhead)
+                    _, q_idx_wrapped = seqlen_q_divmod.divmod(q_idx_floored)
+                    q_idx_vec[j] = q_idx_wrapped
+                else:
+                    _, seqlen_k_divmod = fastdiv_mods
+
+                _, kv_idx_wrapped = seqlen_k_divmod.divmod(index_tensor[i + j][1])
+                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][0], qhead_per_kvhead)
+                kv_idx_vec[j] = index_tensor[i + j][1]
+
+        # 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,
+            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]

build/torch-universal/testing.py

@@ -0,0 +1,404 @@
+import math
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+
+
+class IndexFirstAxis(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, indices):
+        ctx.save_for_backward(indices)
+        assert input.ndim >= 2
+        ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]
+        second_dim = other_shape.numel()
+        return torch.gather(
+            rearrange(input, "b ... -> b (...)"),
+            0,
+            repeat(indices, "z -> z d", d=second_dim),
+        ).reshape(-1, *other_shape)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (indices,) = ctx.saved_tensors
+        assert grad_output.ndim >= 2
+        other_shape = grad_output.shape[1:]
+        grad_output = rearrange(grad_output, "b ... -> b (...)")
+        grad_input = torch.zeros(
+            [ctx.first_axis_dim, grad_output.shape[1]],
+            device=grad_output.device,
+            dtype=grad_output.dtype,
+        )
+        grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output)
+        return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
+
+
+index_first_axis = IndexFirstAxis.apply
+
+
+class IndexPutFirstAxis(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, values, indices, first_axis_dim):
+        ctx.save_for_backward(indices)
+        assert indices.ndim == 1
+        assert values.ndim >= 2
+        output = torch.zeros(
+            first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype
+        )
+        output[indices] = values
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (indices,) = ctx.saved_tensors
+        grad_values = grad_output[indices]
+        return grad_values, None, None
+
+
+index_put_first_axis = IndexPutFirstAxis.apply
+
+
+def unpad_input(hidden_states, attention_mask, unused_mask=None):
+    all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask
+    seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32)
+    used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
+    return (
+        index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices),
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+        used_seqlens_in_batch,
+    )
+
+
+def pad_input(hidden_states, indices, batch, seqlen):
+    output = index_put_first_axis(hidden_states, indices, batch * seqlen)
+    return rearrange(output, "(b s) ... -> b s ...", b=batch)
+
+
+def generate_random_padding_mask(max_seqlen, batch_size, device, mode="random", zero_lengths=False):
+    assert mode in ["full", "random", "third"]
+    if mode == "full":
+        lengths = torch.full((batch_size, 1), max_seqlen, device=device, dtype=torch.int32)
+    elif mode == "random":
+        lengths = torch.randint(
+            max(0 if zero_lengths else 1, max_seqlen - 20),
+            max_seqlen + 1,
+            (batch_size, 1),
+            device=device,
+        )
+    else:
+        lengths = torch.randint(max_seqlen // 3, max_seqlen + 1, (batch_size, 1), device=device)
+
+    if zero_lengths:
+        for i in range(batch_size):
+            if i % 5 == 0:
+                lengths[i] = 0
+        lengths[-1] = 0
+    padding_mask = repeat(torch.arange(max_seqlen, device=device), "s -> b s", b=batch_size) < lengths
+    return padding_mask
+
+
+def generate_qkv(
+    q,
+    k,
+    v,
+    query_padding_mask=None,
+    key_padding_mask=None,
+    qv=None,
+    kvpacked=False,
+    qkvpacked=False,
+    query_unused_mask=None,
+    key_unused_mask=None,
+):
+    assert not (kvpacked and qkvpacked)
+    batch_size, seqlen_q, nheads, d = q.shape
+    d_v = v.shape[-1]
+    _, seqlen_k, nheads_k, _ = k.shape
+    assert k.shape == (batch_size, seqlen_k, nheads_k, d)
+    assert v.shape == (batch_size, seqlen_k, nheads_k, d_v)
+    if query_unused_mask is not None or key_unused_mask is not None:
+        assert not kvpacked
+        assert not qkvpacked
+
+    if query_padding_mask is not None:
+        q_unpad, indices_q, cu_seqlens_q, max_seqlen_q, seqused_q = unpad_input(
+            q, query_padding_mask, query_unused_mask
+        )
+        output_pad_fn = lambda output_unpad: pad_input(output_unpad, indices_q, batch_size, seqlen_q)
+        qv_unpad = rearrange(qv, "b s ... -> (b s) ...")[indices_q] if qv is not None else None
+    else:
+        q_unpad = rearrange(q, "b s h d -> (b s) h d")
+        cu_seqlens_q = torch.arange(
+            0, (batch_size + 1) * seqlen_q, step=seqlen_q, dtype=torch.int32, device=q_unpad.device
+        )
+        seqused_q = None
+        max_seqlen_q = seqlen_q
+        output_pad_fn = lambda output_unpad: rearrange(output_unpad, "(b s) h d -> b s h d", b=batch_size)
+        qv_unpad = rearrange(qv, "b s ... -> (b s) ...") if qv is not None else None
+
+    if key_padding_mask is not None:
+        k_unpad, indices_k, cu_seqlens_k, max_seqlen_k, seqused_k = unpad_input(
+            k, key_padding_mask, key_unused_mask
+        )
+        v_unpad, *_ = unpad_input(v, key_padding_mask, key_unused_mask)
+    else:
+        k_unpad = rearrange(k, "b s h d -> (b s) h d")
+        v_unpad = rearrange(v, "b s h d -> (b s) h d")
+        cu_seqlens_k = torch.arange(
+            0, (batch_size + 1) * seqlen_k, step=seqlen_k, dtype=torch.int32, device=k_unpad.device
+        )
+        seqused_k = None
+        max_seqlen_k = seqlen_k
+
+    if qkvpacked:
+        assert (query_padding_mask == key_padding_mask).all()
+        assert nheads == nheads_k
+        qkv_unpad = torch.stack([q_unpad, k_unpad, v_unpad], dim=1)
+        qkv = torch.stack([q, k, v], dim=2)
+        if query_padding_mask is not None:
+            dqkv_pad_fn = lambda dqkv_unpad: pad_input(dqkv_unpad, indices_q, batch_size, seqlen_q)
+        else:
+            dqkv_pad_fn = lambda dqkv_unpad: rearrange(
+                dqkv_unpad, "(b s) t h d -> b s t h d", b=batch_size
+            )
+        return (
+            qkv_unpad.detach().requires_grad_(),
+            cu_seqlens_q,
+            max_seqlen_q,
+            qkv.detach().requires_grad_(),
+            output_pad_fn,
+            dqkv_pad_fn,
+        )
+    elif kvpacked:
+        kv_unpad = torch.stack([k_unpad, v_unpad], dim=1)
+        kv = torch.stack([k, v], dim=2)
+        dq_pad_fn = output_pad_fn
+        if key_padding_mask is not None:
+            dkv_pad_fn = lambda dkv_unpad: pad_input(dkv_unpad, indices_k, batch_size, seqlen_k)
+        else:
+            dkv_pad_fn = lambda dkv_unpad: rearrange(
+                dkv_unpad, "(b s) t h d -> b s t h d", b=batch_size
+            )
+        return (
+            q_unpad.detach().requires_grad_(),
+            kv_unpad.detach().requires_grad_(),
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            q.detach().requires_grad_(),
+            kv.detach().requires_grad_(),
+            output_pad_fn,
+            dq_pad_fn,
+            dkv_pad_fn,
+        )
+    else:
+        dq_pad_fn = output_pad_fn
+        if key_padding_mask is not None:
+            dk_pad_fn = lambda dk_unpad: pad_input(dk_unpad, indices_k, batch_size, seqlen_k)
+        else:
+            dk_pad_fn = lambda dk_unpad: rearrange(dk_unpad, "(b s) h d -> b s h d", b=batch_size)
+        return (
+            q_unpad.detach().requires_grad_(),
+            k_unpad.detach().requires_grad_(),
+            v_unpad.detach().requires_grad_(),
+            qv_unpad.detach() if qv is not None else None,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqused_q,
+            seqused_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            q.detach().requires_grad_(),
+            k.detach().requires_grad_(),
+            v.detach().requires_grad_(),
+            qv.detach() if qv is not None else None,
+            output_pad_fn,
+            dq_pad_fn,
+            dk_pad_fn,
+        )
+
+
+def construct_local_mask(
+    seqlen_q,
+    seqlen_k,
+    window_size=(None, None),
+    sink_token_length=0,
+    query_padding_mask=None,
+    key_padding_mask=None,
+    key_leftpad=None,
+    device=None,
+):
+    row_idx = rearrange(torch.arange(seqlen_q, device=device, dtype=torch.long), "s -> s 1")
+    col_idx = torch.arange(seqlen_k, device=device, dtype=torch.long)
+    if key_leftpad is not None:
+        key_leftpad = rearrange(key_leftpad, "b -> b 1 1 1")
+        col_idx = repeat(col_idx, "s -> b 1 1 s", b=key_leftpad.shape[0])
+        col_idx = torch.where(col_idx >= key_leftpad, col_idx - key_leftpad, 2**32)
+    sk = (
+        seqlen_k
+        if key_padding_mask is None
+        else rearrange(key_padding_mask.sum(-1), "b -> b 1 1 1")
+    )
+    sq = (
+        seqlen_q
+        if query_padding_mask is None
+        else rearrange(query_padding_mask.sum(-1), "b -> b 1 1 1")
+    )
+    if window_size[0] is None:
+        return col_idx > row_idx + sk - sq + window_size[1]
+    else:
+        sk = torch.full_like(col_idx, seqlen_k) if key_padding_mask is None else sk
+        return torch.logical_or(
+            col_idx > torch.minimum(row_idx + sk - sq + window_size[1], sk),
+            torch.logical_and(col_idx < row_idx + sk - sq - window_size[0], col_idx >= sink_token_length),
+        )
+
+
+def construct_chunk_mask(
+    seqlen_q,
+    seqlen_k,
+    attention_chunk,
+    query_padding_mask=None,
+    key_padding_mask=None,
+    key_leftpad=None,
+    device=None,
+):
+    row_idx = rearrange(torch.arange(seqlen_q, device=device, dtype=torch.long), "s -> s 1")
+    col_idx = torch.arange(seqlen_k, device=device, dtype=torch.long)
+    if key_leftpad is not None:
+        key_leftpad = rearrange(key_leftpad, "b -> b 1 1 1")
+        col_idx = repeat(col_idx, "s -> b 1 1 s", b=key_leftpad.shape[0])
+        col_idx = torch.where(col_idx >= key_leftpad, col_idx - key_leftpad, 2**32)
+    sk = (
+        seqlen_k
+        if key_padding_mask is None
+        else rearrange(key_padding_mask.sum(-1), "b -> b 1 1 1")
+    )
+    sq = (
+        seqlen_q
+        if query_padding_mask is None
+        else rearrange(query_padding_mask.sum(-1), "b -> b 1 1 1")
+    )
+    sk = torch.full_like(col_idx, seqlen_k) if key_padding_mask is None else sk
+    col_limit_left_chunk = row_idx + sk - sq - (row_idx + sk - sq) % attention_chunk
+    return torch.logical_or(
+        col_idx < col_limit_left_chunk, col_idx >= col_limit_left_chunk + attention_chunk
+    )
+
+
+def attention_ref(
+    q,
+    k,
+    v,
+    query_padding_mask=None,
+    key_padding_mask=None,
+    key_leftpad=None,
+    attn_bias=None,
+    dropout_p=0.0,
+    dropout_mask=None,
+    causal=False,
+    qv=None,
+    q_descale=None,
+    k_descale=None,
+    v_descale=None,
+    window_size=(None, None),
+    attention_chunk=0,
+    sink_token_length=0,
+    learnable_sink: Optional[torch.Tensor] = None,
+    softcap=0.0,
+    upcast=True,
+    reorder_ops=False,
+    intermediate_dtype=None,
+):
+    if causal:
+        window_size = (window_size[0], 0)
+    dtype_og = q.dtype
+    if upcast:
+        q, k, v = q.float(), k.float(), v.float()
+        qv = qv.float() if qv is not None else None
+    if q_descale is not None:
+        q_descale = repeat(q_descale, "b h -> b 1 (h g) 1", g=q.shape[2] // k.shape[2])
+        q = (q.float() * q_descale).to(q.dtype)
+        qv = (qv.float() * q_descale).to(qv.dtype) if qv is not None else None
+    if k_descale is not None:
+        k = (k.float() * rearrange(k_descale, "b h -> b 1 h 1")).to(dtype=k.dtype)
+    if v_descale is not None:
+        v = (v.float() * rearrange(v_descale, "b h -> b 1 h 1")).to(dtype=v.dtype)
+    seqlen_q, seqlen_k = q.shape[1], k.shape[1]
+    k = repeat(k, "b s h d -> b s (h g) d", g=q.shape[2] // k.shape[2])
+    v = repeat(v, "b s h d -> b s (h g) d", g=q.shape[2] // v.shape[2])
+    d = q.shape[-1]
+    dv = v.shape[-1]
+    softmax_scale = 1.0 / math.sqrt(d if qv is None else d + dv)
+    if not reorder_ops:
+        scores = torch.einsum("bthd,bshd->bhts", q * softmax_scale, k)
+    else:
+        scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
+    if qv is not None:
+        scores = scores + torch.einsum("bthd,bshd->bhts", qv * softmax_scale, v)
+    if softcap > 0:
+        scores = torch.tanh(scores / softcap) * softcap
+    if key_padding_mask is not None:
+        scores.masked_fill_(rearrange(~key_padding_mask, "b s -> b 1 1 s"), float("-inf"))
+    local_mask = None
+    if window_size[0] is not None or window_size[1] is not None:
+        local_mask = construct_local_mask(
+            seqlen_q,
+            seqlen_k,
+            window_size,
+            sink_token_length,
+            query_padding_mask,
+            key_padding_mask,
+            key_leftpad=key_leftpad,
+            device=q.device,
+        )
+    if attention_chunk > 0:
+        chunk_mask = construct_chunk_mask(
+            seqlen_q,
+            seqlen_k,
+            attention_chunk,
+            query_padding_mask,
+            key_padding_mask,
+            key_leftpad=key_leftpad,
+            device=q.device,
+        )
+        local_mask = torch.logical_or(local_mask, chunk_mask) if local_mask is not None else chunk_mask
+    if local_mask is not None:
+        scores.masked_fill_(local_mask, float("-inf"))
+    if attn_bias is not None:
+        scores = scores + attn_bias
+    if learnable_sink is None:
+        attention = torch.softmax(scores, dim=-1).to(v.dtype)
+    else:
+        scores_fp32 = scores.to(torch.float32)
+        logits_max = torch.amax(scores_fp32, dim=-1, keepdim=True)
+        learnable_sink = rearrange(learnable_sink, "h -> h 1 1")
+        logits_or_sinks_max = torch.maximum(learnable_sink, logits_max)
+        unnormalized_scores = torch.exp(scores_fp32 - logits_or_sinks_max)
+        normalizer = unnormalized_scores.sum(dim=-1, keepdim=True) + torch.exp(
+            learnable_sink - logits_or_sinks_max
+        )
+        attention = (unnormalized_scores / normalizer).to(v.dtype)
+    if query_padding_mask is not None:
+        attention = attention.masked_fill(rearrange(~query_padding_mask, "b s -> b 1 s 1"), 0.0)
+    if key_padding_mask is not None:
+        attention = attention.masked_fill(rearrange(~key_padding_mask, "b s -> b 1 1 s"), 0.0)
+    if local_mask is not None:
+        attention = attention.masked_fill(torch.all(local_mask, dim=-1, keepdim=True), 0.0)
+    dropout_scaling = 1.0 / (1 - dropout_p)
+    if dropout_mask is not None:
+        attention_drop = attention.masked_fill(~dropout_mask, 0.0)
+    else:
+        attention_drop = attention
+    if intermediate_dtype is not None:
+        attention_drop = attention_drop.to(intermediate_dtype).to(attention_drop.dtype)
+    output = torch.einsum("bhts,bshd->bthd", attention_drop, v * dropout_scaling)
+    if query_padding_mask is not None:
+        output.masked_fill_(rearrange(~query_padding_mask, "b s -> b s 1 1"), 0.0)
+    return output.to(dtype=dtype_og), attention.to(dtype=dtype_og)

build/torch-universal/tile_scheduler.py

@@ -0,0 +1,695 @@
+# Copyright (c) 2025, Tri Dao.
+
+from typing import Optional, Tuple
+from dataclasses import dataclass, fields
+
+try:
+    from typing import override
+except ImportError:  # Python < 3.12
+    from typing_extensions import override
+
+import cutlass
+from cutlass._mlir import ir
+import cutlass.cute as cute
+from cutlass import Int32, const_expr
+
+from . import utils
+from .fast_math import FastDivmod, clz
+
+
+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)
+
+
+@dataclass
+class ParamsBase:
+    def __extract_mlir_values__(self):
+        all_fields = [getattr(self, field.name) for field in fields(self)]
+        non_constexpr_fields = [f for f in all_fields if not isinstance(f, cutlass.Constexpr)]
+        values, self._values_pos = [], []
+        for obj in non_constexpr_fields:
+            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):
+        all_fields = {field.name: getattr(self, field.name) for field in fields(self)}
+        constexpr_fields = {n: f for n, f in all_fields.items() if isinstance(f, cutlass.Constexpr)}
+        non_constexpr_fields = {
+            n: f for n, f in all_fields.items() if not isinstance(f, cutlass.Constexpr)
+        }
+        for (name, field), n_items in zip(non_constexpr_fields.items(), self._values_pos):
+            non_constexpr_fields[name] = cutlass.new_from_mlir_values(field, values[:n_items])
+            values = values[n_items:]
+        return self.__class__(**non_constexpr_fields, **constexpr_fields)
+
+
+@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
+
+
+class SingleTileScheduler:
+    @dataclass
+    class Params(ParamsBase):
+        num_block: Int32
+        num_head: Int32
+        num_batch: Int32
+        num_splits: Int32
+        num_splits_divmod: FastDivmod
+        is_split_kv: cutlass.Constexpr[bool] = False
+        cluster_shape_mn: cutlass.Constexpr[Tuple[int, int]] = (1, 1)
+
+        @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,
+                FastDivmod.create(args.num_splits),
+                args.is_split_kv,
+                args.cluster_shape_mn,
+            )
+
+    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, *, loc=None, ip=None) -> Params:
+        return SingleTileScheduler.Params.create(args, loc=loc, ip=ip)
+
+    @staticmethod
+    def create(params: Params, *, loc=None, ip=None) -> "SingleTileScheduler":
+        blk_coord = cute.arch.block_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"
+        return (
+            cute.round_up(params.num_block, params.cluster_shape_mn[0]),
+            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 = self.params.num_splits_divmod.divmod(head_idx)
+        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
+
+    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_divmod: FastDivmod
+        num_head_divmod: FastDivmod
+        total_blocks: Int32
+
+        @staticmethod
+        def create(
+            args: TileSchedulerArguments, *, loc=None, ip=None
+        ) -> "StaticPersistentTileScheduler.Params":
+            total_blocks = args.num_block * args.num_head * args.num_batch
+            return StaticPersistentTileScheduler.Params(
+                FastDivmod.create(args.num_block), FastDivmod.create(args.num_head), total_blocks
+            )
+
+    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, *, loc=None, ip=None) -> Params:
+        return StaticPersistentTileScheduler.Params.create(args, loc=loc, ip=ip)
+
+    @staticmethod
+    def create(params: Params, *, loc=None, ip=None) -> "StaticPersistentTileScheduler":
+        tile_idx = cute.arch.block_idx()[0]
+        return StaticPersistentTileScheduler(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]:
+        hardware_info = cutlass.utils.HardwareInfo()
+        sm_count = hardware_info.get_device_multiprocessor_count()
+        return (cutlass.min(sm_count, params.total_blocks), Int32(1), Int32(1))
+
+    # @cute.jit
+    def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo:
+        hn_idx, block_idx = self.params.num_block_divmod.divmod(self._tile_idx)
+        batch_idx, head_idx = self.params.num_head_divmod.divmod(hn_idx)
+        is_valid = self._tile_idx < self.params.total_blocks
+        # if cute.arch.thread_idx()[0] == 0:
+        #     cute.printf("TileScheduler: tile_idx=%d, hn_idx=%d, block_idx=%d, batch_idx=%d, head_idx=%d, is_valid=%d", self._tile_idx, hn_idx, block_idx, batch_idx, head_idx, is_valid)
+        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):
+        self._tile_idx += cute.arch.grid_dim()[0]
+
+    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_divmod: FastDivmod
+        num_head_divmod: FastDivmod
+        l2_minor_divmod: FastDivmod
+        l2_major_divmod: FastDivmod
+        l2_minor_residual_divmod: FastDivmod
+        num_hb_quotient: Int32
+        is_split_kv: cutlass.Constexpr[bool] = False
+
+        @staticmethod
+        @cute.jit
+        def create(
+            args: TileSchedulerArguments, *, loc=None, ip=None
+        ) -> "SingleTileLPTScheduler.Params":
+            # cute.printf(args.num_block, args.num_head, args.num_batch, args.seqlen_k, args.headdim, args.headdim_v, args.total_q, args.tile_shape_mn, args.qhead_per_kvhead_packgqa, args.element_size)
+            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
+            # swizzle = 1 if size_l2 < size_one_head else (size_l2 // size_one_head)
+            # Seems faster if swizzle if 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))
+            # swizzle = 1 if size_l2 < size_one_head else (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_divmod=FastDivmod.create(args.num_block),
+                num_head_divmod=FastDivmod.create(args.num_head),
+                l2_minor_divmod=FastDivmod.create(swizzle),
+                l2_major_divmod=FastDivmod.create(swizzle * args.num_block),
+                l2_minor_residual_divmod=FastDivmod.create(
+                    max(num_hb_remainder, 1)
+                ),  # don't divide by 0
+                num_hb_quotient=Int32(num_hb_quotient),
+                num_splits=args.num_splits,
+                is_split_kv=args.is_split_kv,
+            )
+
+    def __init__(self, params: Params, tile_idx: Int32, split_idx: Int32, *, loc=None, ip=None):
+        self.params = params
+        self._tile_idx = tile_idx
+        self._split_idx = split_idx
+        self._loc = loc
+        self._ip = ip
+
+    @staticmethod
+    def to_underlying_arguments(args: TileSchedulerArguments, *, loc=None, ip=None) -> Params:
+        return SingleTileLPTScheduler.Params.create(args, loc=loc, ip=ip)
+
+    @staticmethod
+    @cute.jit
+    def create(params: Params, *, loc=None, ip=None) -> "SingleTileLPTScheduler":
+        tile_idx, split_idx, _ = cute.arch.block_idx()
+        return SingleTileLPTScheduler(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]:
+        return (params.total_blocks, params.num_splits, Int32(1))
+
+    @cute.jit
+    def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo:
+        params = self.params
+        # Implement LPT scheduling coordinate calculation
+        bidhb, l2_mod = params.l2_major_divmod.divmod(self._tile_idx)
+        # 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 = params.l2_minor_divmod.divmod(l2_mod)
+        else:
+            block, bidhb_residual = params.l2_minor_residual_divmod.divmod(l2_mod)
+        bidhb_actual = bidhb * params.l2_minor_divmod.divisor + bidhb_residual
+        batch_idx, head_idx = params.num_head_divmod.divmod(bidhb_actual)
+        # Longest-processing-time-first
+        block = params.num_block_divmod.divisor - 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
+        )
+
+    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
+
+    def __extract_mlir_values__(self):
+        values, self._values_pos = [], []
+        for obj in [self.params, self._tile_idx, self._split_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._split_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 SingleTileLPTBwdScheduler:
+    @dataclass
+    class Params(ParamsBase):
+        total_blocks: Int32
+        num_head_divmod: FastDivmod
+        l2_minor_divmod: FastDivmod
+        l2_major_divmod: FastDivmod
+        l2_minor_residual_divmod: FastDivmod
+        num_hb_quotient: Int32
+        cluster_shape_mn: cutlass.Constexpr[Tuple[int, int]] = (1, 1)
+
+        @staticmethod
+        @cute.jit
+        def create(
+            args: TileSchedulerArguments, *, loc=None, ip=None
+        ) -> "SingleTileLPTBwdScheduler.Params":
+            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_head_divmod=FastDivmod.create(args.num_head),
+                l2_minor_divmod=FastDivmod.create(swizzle),
+                l2_major_divmod=FastDivmod.create(swizzle * num_block),
+                l2_minor_residual_divmod=FastDivmod.create(
+                    max(num_hb_remainder, 1)
+                ),  # don't divide by 0
+                num_hb_quotient=Int32(num_hb_quotient),
+                cluster_shape_mn=args.cluster_shape_mn,
+            )
+
+    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, *, loc=None, ip=None) -> Params:
+        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 = params.l2_major_divmod.divmod(cluster_idx)
+        # 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 = params.l2_minor_divmod.divmod(l2_mod)
+        else:
+            block, bidhb_residual = params.l2_minor_residual_divmod.divmod(l2_mod)
+        bidhb_actual = bidhb * params.l2_minor_divmod.divisor + bidhb_residual
+        batch_idx, head_idx = params.num_head_divmod.divmod(bidhb_actual)
+        is_valid = self._tile_idx < params.total_blocks
+        bidx_in_cluster = cute.arch.block_in_cluster_idx()
+        block = block * params.cluster_shape_mn[0] + bidx_in_cluster[0]
+        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
+
+    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
+
+        @staticmethod
+        @cute.jit
+        def create(
+            args: TileSchedulerArguments, *, loc=None, ip=None
+        ) -> "SingleTileVarlenScheduler.Params":
+            size_l2 = 50 * 1024 * 1024  # 50 MB for K & V
+            max_kvblock_in_l2 = size_l2 // (
+                (args.headdim + args.headdim_v) * args.element_size * args.tile_shape_mn[1]
+            )
+            assert args.mCuSeqlensQ is not None or args.mSeqUsedQ is not None, (
+                "At least one of mCuSeqlensQ or mSeqUsedQ must be provided"
+            )
+            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,
+            )
+
+    def __init__(self, params: Params, tile_idx: Int32, split_idx: Int32, *, loc=None, ip=None):
+        self.params = params
+        self._tile_idx = tile_idx
+        self._split_idx = split_idx
+        self._is_first_block = True
+        self._loc = loc
+        self._ip = ip
+
+    @staticmethod
+    def to_underlying_arguments(args: TileSchedulerArguments, *, loc=None, ip=None) -> Params:
+        return SingleTileVarlenScheduler.Params.create(args, loc=loc, ip=ip)
+
+    @staticmethod
+    def create(params: Params, *, loc=None, ip=None) -> "SingleTileVarlenScheduler":
+        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.tile_shape_mn[0] - 1)
+        ) // params.tile_shape_mn[0]
+        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(seqlen, params.tile_shape_mn[0])
+            if batch_idx < params.num_batch and lane < cute.arch.WARP_SIZE - 1
+            else Int32(0)
+        )
+
+    @cute.jit
+    def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo:
+        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
+        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):
+                # 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.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
+                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 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)
+
+    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._is_first_block = False
+
+    def __extract_mlir_values__(self):
+        values, self._values_pos = [], []
+        for obj in [self.params, self._tile_idx, self._split_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._split_idx],
+            self._values_pos,
+        ):
+            obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items]))
+            values = values[n_items:]
+        return SingleTileVarlenScheduler(*(tuple(obj_list)), loc=self._loc)

build/torch-universal/utils.py

@@ -0,0 +1,788 @@
+# Copyright (c) 2025, Tri Dao.
+
+import math
+import hashlib
+import inspect
+import re
+from typing import Type, Callable, Optional, Tuple, overload
+from functools import partial
+
+import cutlass
+import cutlass.cute as cute
+
+from cutlass import Float32, Int32, const_expr
+from cutlass.cutlass_dsl import T, dsl_user_op
+from cutlass._mlir.dialects import nvvm, llvm
+from cutlass.cute.runtime import from_dlpack
+
+
+# cute.arch.{fma,mul,add}_packed_f32x2 uses RZ rounding mode by default
+fma_packed_f32x2 = partial(cute.arch.fma_packed_f32x2, rnd=nvvm.RoundingModeKind.RN)
+mul_packed_f32x2 = partial(cute.arch.mul_packed_f32x2, rnd=nvvm.RoundingModeKind.RN)
+add_packed_f32x2 = partial(cute.arch.add_packed_f32x2, rnd=nvvm.RoundingModeKind.RN)
+sub_packed_f32x2 = partial(
+    cute.arch.calc_packed_f32x2_op,
+    src_c=None,
+    calc_func=nvvm.sub_packed_f32x2,
+    rnd=nvvm.RoundingModeKind.RN
+)
+
+def hash_callable(func: Callable) -> str:
+    """Hash a callable based on the source code or bytecode and closure values."""
+    if hasattr(func, "__wrapped__"):
+        # cute.jit returns a wrapper whose repr/closure changes per compile; hash the undecorated function.
+        base_func = func.__wrapped__
+        func = base_func
+
+    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 idx, cell in enumerate(func.__closure__):
+            cell_value = cell.cell_contents
+            hasher.update(repr(cell_value).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, buffers):
+        scores = acc_S_SSA * inv_softcap
+        return scores * cute.math.tanh(scores, fastmath=True)
+
+    return scoremod_premask_fn
+
+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 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.core.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.core.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_fragment(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
+
+
+def convert_layout_acc_mn(acc_layout: cute.Layout, transpose: bool = False) -> cute.Layout:
+    """
+    For Sm80, convert ((2, 2), MMA_M, MMA_N, ...) to ((2, MMA_M), (2, MMA_N), ...).
+    For Sm90, convert ((2, 2, V), MMA_M, MMA_N, ...) to ((2, MMA_M), (2, V, MMA_N), ...).
+    """
+    acc_layout_col_major = cute.make_layout(acc_layout.shape)
+    shape = (
+        (acc_layout_col_major.shape[0][1], acc_layout_col_major.shape[1]),  # MMA_M
+        (
+            acc_layout_col_major.shape[0][0],
+            *acc_layout_col_major.shape[0][2:],
+            acc_layout_col_major.shape[2],
+        ),  # MMA_N
+        *acc_layout_col_major.shape[3:],
+    )
+    stride = (
+        (acc_layout_col_major.stride[0][1], acc_layout_col_major.stride[1]),  # MMA_M
+        (
+            acc_layout_col_major.stride[0][0],
+            *acc_layout_col_major.stride[0][2:],
+            acc_layout_col_major.stride[2],
+        ),  # MMA_N
+        *acc_layout_col_major.stride[3:],
+    )
+    if const_expr(transpose):
+        shape = (shape[1], shape[0], *shape[2:])
+        stride = (stride[1], stride[0], *stride[2:])
+    acc_layout_mn = cute.make_layout(shape, stride=stride)
+    return cute.composition(acc_layout, acc_layout_mn)
+
+
+def make_acc_tensor_mn_view(acc: cute.Tensor, transpose: bool = False) -> cute.Tensor:
+    return cute.make_tensor(acc.iterator, convert_layout_acc_mn(acc.layout, transpose=transpose))
+
+
+@cute.jit
+def convert_layout_acc_frgA(acc_layout: cute.Layout) -> cute.Layout:
+    # For back to back gemm, convert layout of acc0 to gemm 1 accept layout.
+    # For Sm80, as the mma instruction shape is 16x8x16, we need to convert from (4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2)
+    # For Sm90, FP16/BF16, convert acc_layout from ((2, 2, N / 8), MMA_M, MMA_N) to ((2, 2, 2), MMA_M, (N / 16, MMA_N))
+    # TODO: Sm90 FP8
+    if const_expr(cute.rank(acc_layout.shape[0]) == 3):  # Sm90
+        l = cute.logical_divide(
+            acc_layout, ((None, None, 2), None, None)
+        )  # ((2, 2, (2, N / 16)), MMA_M, MMA_N)
+        rA_mma_view = cute.make_layout(
+            (
+                (l.shape[0][0], l.shape[0][1], l.shape[0][2][0]),
+                l.shape[1],
+                (l.shape[0][2][1], l.shape[2]),
+            ),
+            stride=(
+                (l.stride[0][0], l.stride[0][1], l.stride[0][2][0]),
+                l.stride[1],
+                (l.stride[0][2][1], l.stride[2]),
+            ),
+        )
+    else:  # Sm80
+        # (4, MMA_M, MMA_N) -> (4, MMA_M, (2, MMA_N / 2))
+        l = cute.logical_divide(acc_layout, (None, None, 2))
+        rA_mma_view = cute.make_layout(
+            (
+                (l.shape[0], l.shape[2][0]),
+                l.shape[1],
+                l.shape[2][1],
+            ),
+            stride=(
+                (l.stride[0], l.stride[2][0]),
+                l.stride[1],
+                l.stride[2][1],
+            ),
+        )
+    return rA_mma_view
+
+
+def make_acc_tensor_frgA_view(acc: cute.Tensor) -> cute.Tensor:
+    return cute.make_tensor(acc.iterator, convert_layout_acc_frgA(acc.layout))
+
+
+def select(a: cute.Tensor, mode: list[int]) -> cute.Tensor:
+    return cute.make_tensor(a.iterator, cute.select(a.layout, mode))
+
+
+def transpose_view(a: cute.Tensor) -> cute.Tensor:
+    """Transpose the first two dimensions of a tensor on smem."""
+    shape = (a.shape[1], a.shape[0], *a.shape[2:])
+    order = (1, 0, *range(2, cute.rank(a)))
+    return cute.composition(a, cute.make_ordered_layout(shape, order=order))
+    # stride = (a.layout.stride[1], a.layout.stride[0], *a.layout.stride[2:])
+    # return cute.make_tensor(a.iterator, cute.make_layout(shape, stride=stride))
+
+
+def parse_swizzle_from_pointer(ptr: cute.Pointer) -> cute.Swizzle:
+    """Extract swizzle parameters from a pointer's swizzle_type.
+
+    The swizzle_type string has the form '!cute.swizzle<"S<b,m,s>">' where
+    b, m, s are the swizzle parameters (bits, base, shift).
+
+    Returns:
+        A cute.Swizzle object constructed from the extracted parameters
+
+    Raises:
+        ValueError: If the swizzle_type string cannot be parsed
+    """
+    # Ideally there should be a better API to get swizzle parameters, but we'll just parse
+    # the string here.
+    swizzle_str = str(ptr.type.swizzle_type)
+    # Extract the inner part "S<b,m,s>"
+    match = re.search(r'S<(\d+),(\d+),(\d+)>', swizzle_str)
+    if match:
+        b, m, s = int(match.group(1)), int(match.group(2)), int(match.group(3))
+        return cute.make_swizzle(b, m, s)
+    else:
+        raise ValueError(f"Could not parse swizzle_type: {swizzle_str}")
+
+
+@cute.jit
+def exp2f(x: cute.TensorSSA | Float32) -> cute.TensorSSA | Float32:
+    """exp2f calculation for both vector and scalar.
+    :param x: input value
+    :type x: cute.TensorSSA or Float32
+    :return: exp2 value
+    :rtype: cute.TensorSSA or Float32
+    """
+    if const_expr(isinstance(x, cute.TensorSSA)):
+        res = cute.make_fragment(x.shape, Float32)
+        res.store(x)
+        for i in cutlass.range_constexpr(cute.size(x.shape)):
+            res[i] = cute.arch.exp2(res[i])
+        return res.load()
+    else:
+        return cute.arch.exp2(x)
+
+
+@dsl_user_op
+def log2f(a: float | Float32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        llvm.inline_asm(
+            T.f32(),
+            [Float32(a).ir_value(loc=loc, ip=ip)],
+            "lg2.approx.ftz.f32 $0, $1;",
+            "=f,f",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+@dsl_user_op
+def logf(a: float | Float32, *, loc=None, ip=None) -> Float32:
+    return log2f(a, loc=loc, ip=ip) * math.log(2.0)
+
+
+@dsl_user_op
+def fmax(
+    a: float | Float32, b: float | Float32, c: float | Float32 | None = None, *, loc=None, ip=None
+) -> Float32:
+    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,
+        )
+    )
+
+
+@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_fragment(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_fragment(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_fragment(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_fragment(x.shape, Float32)
+        res.store(x)
+        local_sum_0 = (
+            add_packed_f32x2((init_val, 0.0), (res[0], res[1]))
+            # 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] = add_packed_f32x2(local_sum[0], (res[i + 0], res[i + 1]))
+            local_sum[1] = add_packed_f32x2(local_sum[1], (res[i + 2], res[i + 3]))
+            local_sum[2] = add_packed_f32x2(local_sum[2], (res[i + 4], res[i + 5]))
+            local_sum[3] = add_packed_f32x2(local_sum[3], (res[i + 6], res[i + 7]))
+        local_sum[0] = add_packed_f32x2(local_sum[0], local_sum[1])
+        local_sum[2] = add_packed_f32x2(local_sum[2], local_sum[3])
+        local_sum[0] = 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)
+
+
+@dsl_user_op
+def elem_pointer_i64(x: cute.Tensor, coord: cute.Coord, *, loc=None, ip=None) -> cute.Pointer:
+    flat_coord_i64 = tuple(cutlass.Int64(c) for c in cute.flatten(coord))
+    flat_stride = cute.flatten_to_tuple(x.stride)
+    assert len(flat_coord_i64) == len(flat_stride), "Coordinate and stride must have the same length"
+    offset = sum(c * s for c, s in zip(flat_coord_i64, flat_stride))
+    # HACK: we assume that applying the offset does not change the pointer alignment
+    byte_offset = offset * x.element_type.width // 8
+    return cute.make_ptr(
+        x.element_type,
+        x.iterator.toint() + byte_offset,
+        x.memspace,
+        assumed_align=x.iterator.alignment,
+    )
+
+
+@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_fragment(
+        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
+    val = cute.make_fragment(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 shr_u32(val: cutlass.Uint32, shift: cutlass.Uint32, *, loc=None, ip=None) -> cutlass.Uint32:
+    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.s32 $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_fragment(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)
+
+
+@cute.jit
+@dsl_user_op
+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
+
+
+@cute.jit
+@dsl_user_op
+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 = 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)],
+            f"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, *, loc=None, ip=None) -> Float32:
+    # We assume x <= 127.0
+    poly_ex2_deg3 = (1.0, 0.695146143436431884765625, 0.227564394474029541015625, 0.077119089663028717041015625)
+    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_deg3, 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, *, loc=None, ip=None) -> Tuple[Float32, Float32]:
+    # We assume x <= 127.0 and y <= 127.0
+    poly_ex2_deg3 = (1.0, 0.695146143436431884765625, 0.227564394474029541015625, 0.077119089663028717041015625)
+    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=nvvm.RoundingModeKind.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 = sub_packed_f32x2(xy_rounded, (fp32_round_int, fp32_round_int))
+    xy_frac = sub_packed_f32x2(xy_clamped, xy_rounded_back)
+    xy_frac_ex2 = evaluate_polynomial_2(*xy_frac, poly_ex2_deg3, 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)
+
+
+@dsl_user_op
+def domain_offset_i64(coord: cute.Coord, tensor: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor:
+    flat_coord_i64 = tuple(cutlass.Int64(c) for c in cute.flatten(coord))
+    flat_stride = cute.flatten_to_tuple(tensor.stride)
+    assert len(flat_coord_i64) == len(
+        flat_stride
+    ), "Coordinate and stride must have the same length"
+    offset = sum(c * s for c, s in zip(flat_coord_i64, flat_stride))
+    assert isinstance(tensor.iterator, cute.Pointer)
+    # HACK: we assume that applying the offset does not change the pointer alignment
+    new_ptr = cute.make_ptr(
+        tensor.element_type,
+        tensor.iterator.toint() + offset * tensor.element_type.width // 8,
+        tensor.memspace,
+        assumed_align=tensor.iterator.max_alignment,
+    )
+    return cute.make_tensor(new_ptr, tensor.layout)
+
+
+@dsl_user_op
+def coord_offset_i64(
+    tensor: cute.Tensor, idx: cute.typing.Int, dim: int, *, loc=None, ip=None
+) -> cute.Tensor:
+    offset = cutlass.Int64(idx) * cute.size(tensor.stride[dim])
+    assert isinstance(tensor.iterator, cute.Pointer)
+    # HACK: we assume that applying the offset does not change the pointer alignment
+    new_ptr = cute.make_ptr(
+        tensor.element_type,
+        tensor.iterator.toint() + offset * tensor.element_type.width // 8,
+        tensor.memspace,
+        assumed_align=tensor.iterator.max_alignment,
+    )
+    new_layout = cute.slice_(tensor.layout, (*[None] * dim, 0, *[None] * (cute.rank(tensor) - dim - 1)))
+    return cute.make_tensor(new_ptr, new_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_fragment(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]