Kernels:

kernels-community
/

flash-attn4

Trusted publisher

Kernel card Files Files and versions

xet

Community

kernels-bot commited on 4 days ago

Commit

abbad15

verified ·

1 Parent(s): 2667923

Uploaded using `kernel-builder`.

Browse files

This view is limited to 50 files because it contains too many changes. See raw diff

Files changed (50) hide show

README.md +8 -3
build/torch-cuda/__init__.py +5 -11
build/torch-cuda/_ops.py +39 -3
build/torch-cuda/_private_for_testing.py +16 -0
build/torch-cuda/benchmark.py +0 -268
build/torch-cuda/blackwell_helpers.py +48 -22
build/torch-cuda/block_info.py +69 -8
build/torch-cuda/block_sparse_utils.py +790 -312
build/torch-cuda/block_sparsity.py +359 -73
build/torch-cuda/cache_utils.py +9 -16
build/torch-cuda/compute_block_sparsity.py +330 -157
build/torch-cuda/cu_blocks_kernel.py +181 -0
build/torch-cuda/cute_dsl_utils.py +155 -67
build/torch-cuda/fa_logging.py +97 -0
build/torch-cuda/flash_attn4/__init__.py +0 -26
build/torch-cuda/flash_bwd.py +161 -79
build/torch-cuda/flash_bwd_mla_dk_sm100.py +1185 -0
build/torch-cuda/flash_bwd_mla_dq_dqv_sm100.py +1297 -0
build/torch-cuda/flash_bwd_mla_sm100.py +2135 -0
build/torch-cuda/flash_bwd_postprocess.py +133 -35
build/torch-cuda/flash_bwd_preprocess.py +274 -164
build/torch-cuda/flash_bwd_sm100.py +621 -259
build/torch-cuda/flash_bwd_sm120.py +55 -0
build/torch-cuda/flash_bwd_sm90.py +556 -192
build/torch-cuda/flash_fwd.py +0 -0
build/torch-cuda/flash_fwd_combine.py +356 -310
build/torch-cuda/flash_fwd_mla_sm100.py +0 -0
build/torch-cuda/flash_fwd_sm100.py +0 -0
build/torch-cuda/flash_fwd_sm120.py +61 -0
build/torch-cuda/flash_fwd_sm90.py +1550 -0
build/torch-cuda/interface.py +0 -0
build/torch-cuda/mask.py +1221 -163
build/torch-cuda/metadata.json +86 -2
build/torch-cuda/metadata.json.sigstore +1 -0
build/torch-cuda/mma_sm100_desc.py +2 -2
build/torch-cuda/named_barrier.py +31 -0
build/torch-cuda/pack_gqa.py +111 -13
build/torch-cuda/paged_kv.py +62 -28
build/torch-cuda/pipeline.py +198 -236
build/torch-cuda/prepare_scheduler.py +387 -0
build/torch-cuda/quack/__init__.py +9 -0
build/torch-cuda/quack/activation.py +0 -568
build/torch-cuda/quack/copy_utils.py +529 -73
build/torch-cuda/quack/cute_dsl_utils.py +55 -29
build/torch-cuda/quack/dsl/__init__.py +12 -0
build/torch-cuda/quack/dsl/cute_tensor_indexing.py +139 -0
build/torch-cuda/quack/layout_utils.py +92 -4
build/torch-cuda/quack/rounding.py +231 -0
build/torch-cuda/quack/sm90_utils.py +33 -1
build/torch-cuda/quack/utils.py +201 -0

README.md CHANGED Viewed

@@ -11,16 +11,21 @@ This is the repository card of kernels-community/flash-attn4 that has been pushe
 # make sure `kernels` is installed: `pip install -U kernels`
 from kernels import get_kernel
-kernel_module = get_kernel("kernels-community/flash-attn4", version=0)
-flash_attn_func = kernel_module.flash_attn_func
-flash_attn_func(...)
 ```
 ## Available functions
 - `flash_attn_func`
 - `flash_attn_varlen_func`
 ## Benchmarks
 No benchmark available yet.

 # make sure `kernels` is installed: `pip install -U kernels`
 from kernels import get_kernel
+kernel_module = get_kernel("kernels-community/flash-attn4", version=1)
+_private_for_testing = kernel_module._private_for_testing
+_private_for_testing(...)
 ```
 ## Available functions
+- `_private_for_testing`
 - `flash_attn_func`
 - `flash_attn_varlen_func`
 ## Benchmarks
 No benchmark available yet.
+## Source code
+Source code of this kernel originally comes from https://github.com/Dao-AILab/flash-attention.git and it was repurposed for compatibility with `kernels`.

build/torch-cuda/__init__.py CHANGED Viewed

@@ -1,24 +1,18 @@
 """Flash Attention CUTE (CUDA Template Engine) implementation."""
-from importlib.metadata import PackageNotFoundError, version
-# Update when syncing again.
-__version__ = "4.0.0.beta4"
-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",
 ]

 """Flash Attention CUTE (CUDA Template Engine) implementation."""
+__version__ = "4.0.0.beta30"
 from .interface import (
     flash_attn_func,
     flash_attn_varlen_func,
 )
+# Internal symbols the tests need to reach. Not public API -- see the module
+# docstring.
+from . import _private_for_testing  # noqa: F401
 __all__ = [
+    "_private_for_testing",
     "flash_attn_func",
     "flash_attn_varlen_func",
 ]

build/torch-cuda/_ops.py CHANGED Viewed

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

 import torch
+def get_backend() -> str:
+    """Detect the backend by inspecting torch."""
+    import torch
+    if hasattr(torch.backends, "tpu"):
+        # torch_tpu sets torch.backends.tpu when it is imported (via
+        # torch's device-backend autoload), regardless of whether TPU
+        # hardware is present — analogous to torch.version.cuda being
+        # set on CUDA builds without a GPU.
+        return "tpu"
+    elif hasattr(torch, "neuron"):
+        # Needs to be sorted before specific Torch builds, since Neuron
+        # extension can be loaded into e.g. CUDA Torch builds.
+        return "neuron"
+    elif torch.version.cuda is not None:
+        return "cuda"
+    elif torch.version.hip is not None:
+        return "rocm"
+    elif torch.backends.mps.is_available():
+        return "metal"
+    elif hasattr(torch.version, "xpu") and torch.version.xpu is not None:
+        return "xpu"
+    else:
+        return "cpu"
+def _find_ops_name() -> str:
+    kernel_name = "flash_attn4"
+    unique_id = "d8eaf3c"
+    backend = get_backend()
+    return f"_{kernel_name}_{backend}_{unique_id}"
+_OPS_NAME = _find_ops_name()
+ops = getattr(torch.ops, _OPS_NAME)
+def add_op_namespace_prefix(op_name: str) -> str:
     """
     Prefix op by namespace.
     """
+    return f"{_OPS_NAME}::{op_name}"

build/torch-cuda/_private_for_testing.py ADDED Viewed

	@@ -0,0 +1,16 @@

+"""Internal symbols reachable for the test suite.
+Modules re-exported here are **not** public API. They carry no stability or
+versioning guarantee and must not be used by downstream code -- they are
+exposed only so that the tests can reach them through ``get_kernel``, which
+loads the kernel under a generated module name.
+Anything genuinely public belongs in ``__init__.py``'s ``__all__`` instead,
+where changing it requires a kernel version bump.
+"""
+from . import compute_block_sparsity  # noqa: F401
+__all__ = [
+    "compute_block_sparsity",
+]

build/torch-cuda/benchmark.py DELETED Viewed

@@ -1,268 +0,0 @@
-# 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-cuda/blackwell_helpers.py CHANGED Viewed

@@ -10,6 +10,26 @@ from cutlass._mlir.dialects import llvm
 from . import mma_sm100_desc as sm100_desc
 @cute.jit
 def gemm_w_idx(
     tiled_mma: cute.TiledMma,
@@ -108,6 +128,7 @@ def gemm_ptx(
     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 = sA.iterator.type.swizzle_type
         smem_desc_base_a: int = const_expr(
@@ -177,7 +198,7 @@ def gemm_ptx(
                     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,
@@ -198,7 +219,7 @@ def gemm_ptx(
                     ".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,
@@ -223,6 +244,7 @@ def gemm_ptx_loop(
     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 = sA.iterator.type.swizzle_type
         smem_desc_base_a: int = const_expr(
@@ -310,14 +332,14 @@ def gemm_ptx_loop(
             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]))
             )
@@ -351,14 +373,14 @@ def gemm_ptx_loop(
             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]))
             )
@@ -394,6 +416,7 @@ def gemm_ptx_partial(
     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 = sA.iterator.type.swizzle_type
         smem_desc_base_a: int = const_expr(
@@ -477,7 +500,7 @@ def gemm_ptx_partial(
             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::{cta_group}.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"
@@ -486,7 +509,7 @@ def gemm_ptx_partial(
                     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::{cta_group}.kind::f16 [tmem_acc], smem_desc_a, smem_desc_b, idesc, 1;\n\t"
                 )
                 for k in range(1, cute.size(tCrA.shape[2]))
             )
@@ -554,7 +577,7 @@ def gemm_ptx_partial(
             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::{cta_group}.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"
@@ -562,7 +585,7 @@ def gemm_ptx_partial(
                     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::{cta_group}.kind::f16 [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t"
                 )
                 for k in range(
                     1,
@@ -575,7 +598,7 @@ def gemm_ptx_partial(
                     (
                         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::{cta_group}.kind::f16 [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t"
                     )
                     for k in range(split_arrive_idx, cute.size(tCrA.shape[2]))
                 )
@@ -613,6 +636,7 @@ def gemm_ptx_partial1(
         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(
@@ -706,14 +730,14 @@ def gemm_ptx_partial1(
             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]))
             )
@@ -751,13 +775,13 @@ def gemm_ptx_partial1(
             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]))
             )
@@ -783,6 +807,7 @@ def gemm_ptx_precomputed(
     mbar_phase: Optional[Int32] = None,
     zero_init: bool | Boolean = False,
     cta_group: int = 1,
 ) -> None:
     # acc_tmem_addr += acc_offset
     is_ts = const_expr(smem_desc_base_a is None)
@@ -842,7 +867,7 @@ def gemm_ptx_precomputed(
             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::{cta_group}.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"
@@ -851,7 +876,7 @@ def gemm_ptx_precomputed(
                     f"add.s32 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::{cta_group}.kind::f16 [tmem_acc], smem_desc_a, smem_desc_b, idesc, 1;\n\t"
                 )
                 for k in range(1, num_k_tile)
             )
@@ -911,7 +936,7 @@ def gemm_ptx_precomputed(
             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::{cta_group}.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"
@@ -919,7 +944,7 @@ def gemm_ptx_precomputed(
                     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::{cta_group}.kind::f16 [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t"
                 )
                 for k in range(
                     1,
@@ -933,7 +958,7 @@ def gemm_ptx_precomputed(
                         # 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::{cta_group}.kind::f16 [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t"
                     )
                     for k in range(num_k_tile // 4 * 3, num_k_tile)
                 )
@@ -1019,6 +1044,7 @@ def gemm_ptx_precomputed_varname(
     smem_offset: int,
     zero_init: bool | Boolean = False,
     cta_group: int = 1,
 ) -> None:
     is_ts = False
     num_k_tile = cute.size(tCrB_layout.shape[2])
@@ -1067,7 +1093,7 @@ def gemm_ptx_precomputed_varname(
             )
             + "setp.ne.b32 p, $1, 0;\n\t"
             # f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], {smem_var_name_prefix}_0, smem_desc_b, idesc, {pred_str};\n\t"
-            f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], {smem_var_name_prefix}_0, smem_desc_b_0, {idesc_var_name}, {pred_str};\n\t"
             + "".join(
                 (
                     # f"mov.b64 {{smem_desc_a_lo, smem_desc_a_hi}}, {smem_var_name_prefix}_{k};\n\t"
@@ -1077,7 +1103,7 @@ def gemm_ptx_precomputed_varname(
                     # f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t"
                     # f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], {smem_var_name_prefix}_{k}, smem_desc_b, idesc, 1;\n\t"
                     # f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], {smem_var_name_prefix}_{k}, smem_desc_b, {idesc_var_name}, 1;\n\t"
-                    f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], {smem_var_name_prefix}_{k}, smem_desc_b_{k}, {idesc_var_name}, 1;\n\t"
                 )
                 for k in range(1, num_k_tile)
             )

 from . import mma_sm100_desc as sm100_desc
+def _tcgen05_mma_kind(op: cute.nvgpu.tcgen05.mma.MmaOp) -> str:
+    if isinstance(op, tcgen05.mma.MmaF16BF16Op):
+        return "f16"
+    if isinstance(op, tcgen05.mma.MmaTF32Op):
+        return "tf32"
+    if isinstance(op, tcgen05.mma.MmaI8Op):
+        return "i8"
+    # cutlass-dsl >=4.5.2 builds plain FP8 MMAs as MmaF8F6F4Op (make_trivial_tiled_mma's
+    # _F8F6F4_TYPES branch); <4.4.x returned the now-legacy MmaFP8Op. Both map to kind::f8f6f4.
+    if isinstance(op, (tcgen05.mma.MmaFP8Op, tcgen05.mma.MmaF8F6F4Op)):
+        return "f8f6f4"
+    if isinstance(op, tcgen05.mma.MmaMXF8Op):
+        return "mxf8f6f4"
+    if isinstance(op, tcgen05.mma.MmaMXF4Op):
+        return "mxf4"
+    if isinstance(op, tcgen05.mma.MmaMXF4NVF4Op):
+        return "mxf4nvf4"
+    raise TypeError(f"Unsupported tcgen05 MMA op kind: {type(op).__name__}")
 @cute.jit
 def gemm_w_idx(
     tiled_mma: cute.TiledMma,
     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))
+    kind = _tcgen05_mma_kind(op)
     if const_expr(not is_ts):
         sA_swizzle = sA.iterator.type.swizzle_type
         smem_desc_base_a: int = const_expr(
                     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::{kind} [$0], smem_desc_a, smem_desc_b, idesc, p;\n\t"
                     "}\n",
                     "r,r,r,r",
                     has_side_effects=True,
                     ".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::{kind} [$0], [$1], smem_desc_b, {hex(idesc)}, p;\n\t"
                     "}\n",
                     "r,r,r,r",
                     has_side_effects=True,
     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))
+    kind = _tcgen05_mma_kind(op)
     if const_expr(not is_ts):
         sA_swizzle = sA.iterator.type.swizzle_type
         smem_desc_base_a: int = const_expr(
             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::{kind} [$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::{kind} [$0], smem_desc_a, smem_desc_b, idesc, 1;\n\t"
                 )
                 for k in cutlass.range_constexpr(1, cute.size(tCrA.shape[2]))
             )
             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::{kind} [$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::{kind} [$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]))
             )
     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))
+    kind = _tcgen05_mma_kind(op)
     if const_expr(not is_ts):
         sA_swizzle = sA.iterator.type.swizzle_type
         smem_desc_base_a: int = const_expr(
             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::{cta_group}.kind::{kind} [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_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::{cta_group}.kind::{kind} [tmem_acc], smem_desc_a, smem_desc_b, idesc, 1;\n\t"
                 )
                 for k in range(1, cute.size(tCrA.shape[2]))
             )
             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::{cta_group}.kind::{kind} [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_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::{cta_group}.kind::{kind} [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t"
                 )
                 for k in range(
                     1,
                     (
                         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::{cta_group}.kind::{kind} [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t"
                     )
                     for k in range(split_arrive_idx, cute.size(tCrA.shape[2]))
                 )
         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))
+    kind = _tcgen05_mma_kind(op)
     if const_expr(not is_ts):
         smem_desc_base_a: int = const_expr(
             sm100_desc.make_smem_desc_base(
             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::{kind} [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::{kind} [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]))
             )
             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::{kind} [$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::{kind} [$0], [tmem_a], smem_desc_b, idesc, {{$4, $5, $6, $7}}, 1;\n\t"
                 )
                 for k in range(1, cute.size(tCrA.shape[2]))
             )
     mbar_phase: Optional[Int32] = None,
     zero_init: bool | Boolean = False,
     cta_group: int = 1,
+    kind: str = "f16",
 ) -> None:
     # acc_tmem_addr += acc_offset
     is_ts = const_expr(smem_desc_base_a is None)
             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::{cta_group}.kind::{kind} [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.s32 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::{cta_group}.kind::{kind} [tmem_acc], smem_desc_a, smem_desc_b, idesc, 1;\n\t"
                 )
                 for k in range(1, num_k_tile)
             )
             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::{cta_group}.kind::{kind} [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_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::{cta_group}.kind::{kind} [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t"
                 )
                 for k in range(
                     1,
                         # 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::{cta_group}.kind::{kind} [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t"
                     )
                     for k in range(num_k_tile // 4 * 3, num_k_tile)
                 )
     smem_offset: int,
     zero_init: bool | Boolean = False,
     cta_group: int = 1,
+    kind: str = "f16",
 ) -> None:
     is_ts = False
     num_k_tile = cute.size(tCrB_layout.shape[2])
             )
             + "setp.ne.b32 p, $1, 0;\n\t"
             # f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], {smem_var_name_prefix}_0, smem_desc_b, idesc, {pred_str};\n\t"
+            f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], {smem_var_name_prefix}_0, smem_desc_b_0, {idesc_var_name}, {pred_str};\n\t"
             + "".join(
                 (
                     # f"mov.b64 {{smem_desc_a_lo, smem_desc_a_hi}}, {smem_var_name_prefix}_{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::{cta_group}.kind::f16 [tmem_acc], {smem_var_name_prefix}_{k}, smem_desc_b, idesc, 1;\n\t"
                     # f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], {smem_var_name_prefix}_{k}, smem_desc_b, {idesc_var_name}, 1;\n\t"
+                    f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], {smem_var_name_prefix}_{k}, smem_desc_b_{k}, {idesc_var_name}, 1;\n\t"
                 )
                 for k in range(1, num_k_tile)
             )

build/torch-cuda/block_info.py CHANGED Viewed

@@ -6,7 +6,7 @@ import cutlass
 import cutlass.cute as cute
 from cutlass import Int32, const_expr
-from .seqlen_info import SeqlenInfoQK
 @dataclass(frozen=True)
@@ -19,14 +19,18 @@ class BlockInfo:
     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)):
@@ -45,11 +49,20 @@ class BlockInfo:
             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
@@ -70,6 +83,37 @@ class BlockInfo:
             m_block_max = min(m_block_max, cute.ceil_div(m_idx_left, self.tile_m))
         return m_block_min, m_block_max
     @cute.jit
     def get_n_block_min_causal_local_mask(
         self,
@@ -106,3 +150,20 @@ class BlockInfo:
             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))

 import cutlass.cute as cute
 from cutlass import Int32, const_expr
+from .seqlen_info import SeqlenInfoQK, SeqlenInfoQKNewK
 @dataclass(frozen=True)
     window_size_left: Optional[Int32] = None
     window_size_right: Optional[Int32] = None
     qhead_per_kvhead_packgqa: cutlass.Constexpr[int] = 1
+    num_splits: Int32 = 1
+    # If True, the scheduler packs num_splits into the top 16 bits of split_idx
+    pack_split_idx: cutlass.Constexpr[bool] = False
+    num_n_blocks_per_split: Optional[cutlass.Constexpr[Int32]] = None
     @cute.jit
     def get_n_block_min_max(
         self,
         seqlen_info: SeqlenInfoQK,
         m_block: Int32,
+        split_idx: Int32 = 0,
+        num_splits: Int32 = 1,
     ) -> Tuple[Int32, Int32]:
         n_block_max = cute.ceil_div(seqlen_info.seqlen_k, self.tile_n)
         if const_expr(self.is_causal or (self.is_local and self.window_size_right is not None)):
             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):
+            if const_expr(self.pack_split_idx):
+                # Unpack num_splits from top 16 bits of split_idx (packed by scheduler)
+                num_splits = split_idx >> 16
+                split_idx = split_idx & 0xFFFF
+            else:
+                num_splits = self.num_splits
+            if const_expr(self.num_n_blocks_per_split is not None):
+                num_n_blocks_per_split = self.num_n_blocks_per_split
+            else:
+                num_n_blocks_per_split = (
+                    Int32(0)
+                    if n_block_max <= n_block_min
+                    else (n_block_max - n_block_min + num_splits - 1) // num_splits
+                )
             n_block_min = n_block_min + split_idx * num_n_blocks_per_split
             n_block_max = cutlass.min(n_block_min + num_n_blocks_per_split, n_block_max)
         return n_block_min, n_block_max
             m_block_max = min(m_block_max, cute.ceil_div(m_idx_left, self.tile_m))
         return m_block_min, m_block_max
+    @cute.jit
+    def get_n_block_k_new_min_max(
+        self,
+        seqlen_info: SeqlenInfoQKNewK,
+        m_block: Int32,
+        split_idx: Int32 = 0,
+        num_splits: Int32 = 1,
+    ) -> Tuple[Int32, Int32]:
+        """Get the block range for new K tokens (append KV).
+        First computes the full n_block range via get_n_block_min_max, then maps
+        those blocks into the new-K index space by subtracting seqlen_k_og.
+        """
+        n_block_min, n_block_max = self.get_n_block_min_max(
+            seqlen_info,
+            m_block,
+            split_idx,
+            num_splits,
+        )
+        idx_k_new_min = cutlass.max(n_block_min * self.tile_n - seqlen_info.seqlen_k_og, 0)
+        idx_k_new_max = cutlass.min(
+            n_block_max * self.tile_n - seqlen_info.seqlen_k_og, seqlen_info.seqlen_k_new
+        )
+        n_block_new_min = idx_k_new_min // self.tile_n
+        n_block_new_max = (
+            cute.ceil_div(idx_k_new_max, self.tile_n)
+            if idx_k_new_max > idx_k_new_min
+            else n_block_new_min
+        )
+        return n_block_new_min, n_block_new_max
     @cute.jit
     def get_n_block_min_causal_local_mask(
         self,
             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))
+    @cute.jit
+    def get_n_block_max_for_m_block(
+        self,
+        seqlen_info: SeqlenInfoQK,
+        m_block: Int32,
+    ) -> Int32:
+        n_block_max = cute.ceil_div(seqlen_info.seqlen_k, self.tile_n)
+        if const_expr(self.is_causal or 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_right = m_idx_max + seqlen_info.seqlen_k - seqlen_info.seqlen_q
+            if const_expr(self.window_size_right is not None):
+                n_idx_right += self.window_size_right
+            n_block_max = min(n_block_max, cute.ceil_div(n_idx_right, self.tile_n))
+        return n_block_max

build/torch-cuda/block_sparse_utils.py CHANGED Viewed

@@ -5,7 +5,7 @@ This module contains runtime execution functions for block-sparse attention kern
 These utilities are used by CUTE DSL kernels to produce and consume block-sparse loads.
 """
-from typing import Callable, Optional
 from functools import partial
 import math
 import cutlass
@@ -17,6 +17,80 @@ from .quack import copy_utils
 # Import data structures from block_sparsity
 from .block_sparsity import BlockSparseTensors
 from .named_barrier import NamedBarrierBwd
 # NOTE [SM100 block-sparse empty tiles: mbarrier contract]
@@ -61,35 +135,51 @@ from .named_barrier import NamedBarrierBwd
 #   to ack/advance, and arrives `mbar_P_full_O_rescaled` when MMA can proceed.
 #
 # Backward (SM100):
-# - Empty KV tile: for a given `n_block`, `total_m_block_cnt == 0` means no Q tiles contribute.
-# - Both the load and compute loops guard all pipeline work on `process_tile`, so empty tiles
-#   skip producer/consumer operations entirely (no per-tile mbarrier phase handshake like forward).
 # - In the `not dKV_postprocess` path, dK/dV for empty KV tiles are explicitly written as zeros
 #   even when `process_tile == False` (see `flash_bwd_sm100.py` `should_zero_dKV`).
 @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.
@@ -98,44 +188,40 @@ def load_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)
@@ -153,10 +239,11 @@ def finish_overlap_v_load(
     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()
@@ -185,20 +272,21 @@ def produce_block_sparse_loads(
     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,
     qhead_per_kvhead: cutlass.Constexpr[int] = 1,
     q_subtile_factor: cutlass.Constexpr[int] = 1,
 ):
     """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.
@@ -211,40 +299,42 @@ def produce_block_sparse_loads(
         qhead_per_kvhead: Pack-GQA factor. When > 1, m_block is in packed space and
             must be converted to unpacked for sparse tensor indexing.
     """
-    mask_block_cnt, mask_block_idx, full_block_cnt, full_block_idx = blocksparse_tensors
     m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor)
-    curr_mask_block_cnt = mask_block_cnt[batch_idx, head_idx, m_block_sparse]
-    curr_mask_block_idx = mask_block_idx[batch_idx, head_idx, m_block_sparse, None]
-    if const_expr(full_block_cnt is not None):
-        curr_full_block_cnt = full_block_cnt[batch_idx, head_idx, m_block_sparse]
-        curr_full_block_idx = full_block_idx[batch_idx, head_idx, m_block_sparse, None]
-    else:
-        curr_full_block_cnt = Int32(0)
-        curr_full_block_idx = 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:
@@ -254,24 +344,21 @@ def produce_block_sparse_loads(
                 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:
@@ -282,13 +369,18 @@ def produce_block_sparse_loads(
                     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)
@@ -299,17 +391,14 @@ def produce_block_sparse_loads(
                 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(
@@ -318,24 +407,21 @@ def produce_block_sparse_loads(
                     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
@@ -347,7 +433,7 @@ def consume_block_sparse_loads(
     batch_idx,
     head_idx,
     m_block,
-    seqlen,
     kv_consumer_state,
     mma_pv_fn,
     mma_one_n_block,
@@ -373,15 +459,20 @@ def consume_block_sparse_loads(
         qhead_per_kvhead: Pack-GQA factor. When > 1, m_block is in packed space and
             must be converted to unpacked for sparse tensor indexing.
     """
-    mask_block_cnt, mask_block_idx, full_block_cnt, full_block_idx = blocksparse_tensors
     m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor)
-    curr_mask_block_cnt = mask_block_cnt[batch_idx, head_idx, m_block_sparse]
-    curr_mask_block_idx = mask_block_idx[batch_idx, head_idx, m_block_sparse, None]
-    curr_full_block_cnt = full_block_cnt[batch_idx, head_idx, m_block_sparse]
-    curr_full_block_idx = full_block_idx[batch_idx, head_idx, m_block_sparse, None]
     processed_any = curr_mask_block_cnt + curr_full_block_cnt > 0
@@ -462,7 +553,7 @@ def consume_block_sparse_loads(
             mask_n_block = curr_mask_block_idx[curr_mask_block_cnt - 1]
             kv_consumer_state = process_first_half_block(
                 n_block=mask_n_block,
-                seqlen=seqlen,
                 kv_consumer_state=kv_consumer_state,
                 mask_fn=partial(
                     mask_fn,
@@ -478,7 +569,7 @@ def consume_block_sparse_loads(
                 kv_consumer_state = mma_one_n_block(
                     kv_consumer_state,
                     n_block=mask_n_block,
-                    seqlen=seqlen,
                     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),
                 )
@@ -489,7 +580,7 @@ def consume_block_sparse_loads(
             if curr_mask_block_cnt == 0:
                 kv_consumer_state = process_first_half_block(
                     n_block=full_n_block,
-                    seqlen=seqlen,
                     kv_consumer_state=kv_consumer_state,
                     mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=True),
                     score_mod_fn=score_mod_fn,
@@ -499,7 +590,7 @@ def consume_block_sparse_loads(
                 kv_consumer_state = mma_one_n_block(
                     kv_consumer_state,
                     n_block=full_n_block,
-                    seqlen=seqlen,
                     mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate),
                     mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=True),
                 )
@@ -509,7 +600,7 @@ def consume_block_sparse_loads(
                 kv_consumer_state = mma_one_n_block(
                     kv_consumer_state,
                     n_block=full_n_block,
-                    seqlen=seqlen,
                     mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate),
                     mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=False),
                 )
@@ -525,10 +616,20 @@ def consume_block_sparse_loads(
     return kv_consumer_state, O_should_accumulate, processed_any
 @cute.jit
 def load_block_list_sm100(
     block_indices: cute.Tensor,
-    block_count,
     load_q_with_first: cutlass.Constexpr,
     q_stage: cutlass.Constexpr,
     kv_producer_state,
@@ -536,11 +637,16 @@ def load_block_list_sm100(
     load_K,
     load_V,
     pipeline_kv,
 ):
-    """SM100 version of load_block_list (no intra_wg_overlap, no extra_tx_count)."""
-    if block_count > 0:
         # First iteration: load Q alongside K if requested
-        n_block_first = block_indices[block_count - 1]
         if const_expr(load_q_with_first):
             # SM100 loads Q0 and optionally Q1
@@ -556,8 +662,8 @@ def load_block_list_sm100(
         kv_producer_state.advance()
         # Remaining blocks
-        for offset in cutlass.range(1, block_count):
-            n_block = block_indices[block_count - 1 - offset]
             load_K(block=n_block, producer_state=kv_producer_state, page_idx=None)
             kv_producer_state.advance()
             load_V(block=n_block, producer_state=kv_producer_state, page_idx=None)
@@ -573,6 +679,9 @@ def produce_block_sparse_loads_sm100(
     batch_idx,
     head_idx,
     m_block,
     kv_producer_state,
     load_Q,
     load_K,
@@ -582,6 +691,7 @@ def produce_block_sparse_loads_sm100(
     q_producer_phase: Int32,
     qhead_per_kvhead: cutlass.Constexpr,
     q_subtile_factor: cutlass.Constexpr,
 ):
     """SM100 entry point for sparse block iteration.
@@ -594,20 +704,32 @@ def produce_block_sparse_loads_sm100(
     """
     m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor)
-    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_sparse]
-    curr_mask_block_idx = mask_block_idx[batch_idx, head_idx, m_block_sparse, None]
-    if const_expr(full_block_cnt is not None):
-        curr_full_block_cnt = full_block_cnt[batch_idx, head_idx, m_block_sparse]
-        curr_full_block_idx = full_block_idx[batch_idx, head_idx, m_block_sparse, None]
-    else:
-        curr_full_block_cnt = Int32(0)
-        curr_full_block_idx = None
-    mask_empty = curr_mask_block_cnt == 0
-    full_empty = curr_full_block_cnt == 0
     q_phase_flipped = False
@@ -615,7 +737,8 @@ def produce_block_sparse_loads_sm100(
         # No masked blocks: process full list with Q loading
         kv_producer_state = load_block_list_sm100(
             curr_full_block_idx,
-            curr_full_block_cnt,
             load_q_with_first=True,
             q_stage=q_stage,
             kv_producer_state=kv_producer_state,
@@ -623,13 +746,15 @@ def produce_block_sparse_loads_sm100(
             load_K=load_K,
             load_V=load_V,
             pipeline_kv=pipeline_kv,
         )
         q_phase_flipped = not full_empty
     else:
         # Process masked blocks with Q loading
         kv_producer_state = load_block_list_sm100(
             curr_mask_block_idx,
-            curr_mask_block_cnt,
             load_q_with_first=True,
             q_stage=q_stage,
             kv_producer_state=kv_producer_state,
@@ -637,6 +762,7 @@ def produce_block_sparse_loads_sm100(
             load_K=load_K,
             load_V=load_V,
             pipeline_kv=pipeline_kv,
         )
         q_phase_flipped = True
@@ -644,7 +770,8 @@ def produce_block_sparse_loads_sm100(
             # Process full blocks without Q loading
             kv_producer_state = load_block_list_sm100(
                 curr_full_block_idx,
-                curr_full_block_cnt,
                 load_q_with_first=False,
                 q_stage=q_stage,
                 kv_producer_state=kv_producer_state,
@@ -652,6 +779,7 @@ def produce_block_sparse_loads_sm100(
                 load_K=load_K,
                 load_V=load_V,
                 pipeline_kv=pipeline_kv,
             )
     if q_phase_flipped:
@@ -666,19 +794,39 @@ def get_total_block_count(
     batch_idx,
     head_idx,
     m_block,
     qhead_per_kvhead: cutlass.Constexpr,
     q_subtile_factor: cutlass.Constexpr,
 ):
     m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor)
-    mask_block_cnt, mask_block_idx, full_block_cnt, full_block_idx = blocksparse_tensors
-    if const_expr(full_block_cnt is not None):
-        return (
-            mask_block_cnt[batch_idx, head_idx, m_block_sparse]
-            + full_block_cnt[batch_idx, head_idx, m_block_sparse]
-        )
-    else:
-        return mask_block_cnt[batch_idx, head_idx, m_block_sparse]
 @cute.jit
@@ -691,7 +839,7 @@ def handle_block_sparse_empty_tile_correction_sm100(
     is_split_kv: cutlass.Constexpr,
     learnable_sink,
     mLSE,
-    seqlen,
     m_block: Int32,
     head_idx: Int32,
     batch_idx: Int32,
@@ -699,7 +847,7 @@ def handle_block_sparse_empty_tile_correction_sm100(
     sScale: cute.Tensor,
     stats: list,
     correction_epilogue: Callable,
-    thr_mma_pv: cute.core.ThrMma,
     tOtO: cute.Tensor,
     sO: cute.Tensor,
     pipeline_sm_stats: cutlass.pipeline.PipelineAsync,
@@ -709,6 +857,8 @@ def handle_block_sparse_empty_tile_correction_sm100(
     o_corr_consumer_phase: Int32,
     corr_epi_producer_phase: Int32,
     softmax_scale_log2: Float32,
     mO_cur: Optional[cute.Tensor] = None,
     gO: Optional[cute.Tensor] = None,
     gmem_tiled_copy_O: Optional[cute.TiledCopy] = None,
@@ -748,10 +898,11 @@ def handle_block_sparse_empty_tile_correction_sm100(
             if sink_val != -Float32.inf and (const_expr(not is_split_kv) or split_idx == 0):
                 if row_max_value == -Float32.inf:
                     row_max_value = sink_val * (LOG2_E / softmax_scale_log2)
-                    row_sum_value = Float32(1.0)
                 else:
                     row_sum_value = row_sum_value + cute.math.exp2(
-                        sink_val * LOG2_E - row_max_value * softmax_scale_log2, fastmath=True
                     )
         if tidx < m_block_size:
             scale_row_idx = tidx + stage * m_block_size
@@ -768,17 +919,19 @@ def handle_block_sparse_empty_tile_correction_sm100(
         if const_expr(gmem_tiled_copy_O is None):
             pipeline_o_epi.producer_acquire_w_index_phase(stage, corr_epi_producer_phase)
         correction_epilogue(
             thr_mma_pv,
             tOtO[None, None, None, stage],
             tidx,
             stage,
             m_block,
-            seqlen.seqlen_q,
             Float32(0.0),  # zero scale ensures empty tile writes zeros into staged outputs
             sO[None, None, stage],
             mO_cur,
-            gO[None, None, stage],
             gmem_tiled_copy_O,
         )
         if const_expr(gmem_tiled_copy_O is None):
@@ -794,12 +947,103 @@ def handle_block_sparse_empty_tile_correction_sm100(
     )
 @cute.jit
 def softmax_block_sparse_sm100(
     blocksparse_tensors: BlockSparseTensors,
     batch_idx,
     head_idx,
     m_block,
     softmax_step: Callable,
     mask_fn: Callable,
     mask_fn_none: Callable,
@@ -813,104 +1057,81 @@ def softmax_block_sparse_sm100(
     check_m_boundary: bool,
     qhead_per_kvhead: cutlass.Constexpr,
     q_subtile_factor: cutlass.Constexpr[int] = 1,
 ):
     warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4
     m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor)
-    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_sparse]
-    curr_mask_block_idx = mask_block_idx[batch_idx, head_idx, m_block_sparse, None]
-    if const_expr(full_block_cnt is not None):
-        curr_full_block_cnt = full_block_cnt[batch_idx, head_idx, m_block_sparse]
-        curr_full_block_idx = full_block_idx[batch_idx, head_idx, m_block_sparse, None]
-    else:
-        curr_full_block_cnt = Int32(0)
-        curr_full_block_idx = None
-    total_block_cnt = curr_mask_block_cnt + curr_full_block_cnt
     if total_block_cnt == 0:
-        # See NOTE [SM100 block-sparse empty tiles: mbarrier contract].
-        # pipeline_sm_stats.producer_commit_w_index(stage_idx)
         sm_stats_barrier.arrive_w_index(index=stage_idx * 4 + warp_idx)
     else:
-        if curr_mask_block_cnt > 0:
-            mask_n_block = curr_mask_block_idx[curr_mask_block_cnt - 1]
             (
                 mma_si_consumer_phase,
                 si_corr_producer_phase,
                 s0_s1_sequence_phase,
-            ) = softmax_step(
                 mma_si_consumer_phase,
                 si_corr_producer_phase,
                 s0_s1_sequence_phase,
-                mask_n_block,
-                is_first=True,
-                mask_fn=partial(mask_fn, mask_seqlen=True, check_q_boundary=check_m_boundary),
             )
-            for i in cutlass.range(1, curr_mask_block_cnt):
-                mask_n_block = curr_mask_block_idx[curr_mask_block_cnt - 1 - i]
-                (
-                    mma_si_consumer_phase,
-                    si_corr_producer_phase,
-                    s0_s1_sequence_phase,
-                ) = softmax_step(
-                    mma_si_consumer_phase,
-                    si_corr_producer_phase,
-                    s0_s1_sequence_phase,
-                    mask_n_block,
-                    mask_fn=partial(mask_fn, mask_seqlen=False, check_q_boundary=check_m_boundary),
-                )
-        if curr_full_block_cnt > 0:
-            full_n_block = curr_full_block_idx[curr_full_block_cnt - 1]
-            if curr_mask_block_cnt == 0:
-                (
-                    mma_si_consumer_phase,
-                    si_corr_producer_phase,
-                    s0_s1_sequence_phase,
-                ) = softmax_step(
-                    mma_si_consumer_phase,
-                    si_corr_producer_phase,
-                    s0_s1_sequence_phase,
-                    full_n_block,
-                    is_first=True,
-                    mask_fn=partial(
-                        mask_fn_none, mask_seqlen=True, check_q_boundary=check_m_boundary
-                    ),
-                )
-            else:
-                (
-                    mma_si_consumer_phase,
-                    si_corr_producer_phase,
-                    s0_s1_sequence_phase,
-                ) = softmax_step(
-                    mma_si_consumer_phase,
-                    si_corr_producer_phase,
-                    s0_s1_sequence_phase,
-                    full_n_block,
-                    is_first=False,
-                    mask_fn=partial(
-                        mask_fn_none, mask_seqlen=False, check_q_boundary=check_m_boundary
-                    ),
-                )
-            for i in cutlass.range(1, curr_full_block_cnt):
-                full_n_block = curr_full_block_idx[curr_full_block_cnt - 1 - i]
-                (
-                    mma_si_consumer_phase,
-                    si_corr_producer_phase,
-                    s0_s1_sequence_phase,
-                ) = softmax_step(
-                    mma_si_consumer_phase,
-                    si_corr_producer_phase,
-                    s0_s1_sequence_phase,
-                    full_n_block,
-                    mask_fn=partial(
-                        mask_fn_none, mask_seqlen=False, check_q_boundary=check_m_boundary
-                    ),
-                )
     return (
         mma_si_consumer_phase,
@@ -940,24 +1161,24 @@ def get_total_q_block_count_bwd(
     batch_idx,
     head_idx,
     n_block,
-    subtile_factor: cutlass.Constexpr = 1,
     m_block_max: int = 0,
 ):
     """Count total tile iterations for given n_block (KV tile) in backward."""
-    q_block_cnt, _, full_block_cnt, _ = blocksparse_tensors
     total = q_block_cnt[batch_idx, head_idx, n_block]
     if const_expr(full_block_cnt is not None):
         total = total + full_block_cnt[batch_idx, head_idx, n_block]
-    return total * subtile_factor
 @cute.jit
-def produce_block_sparse_q_loads_bwd_sm100(
     blocksparse_tensors: BlockSparseTensors,
     batch_idx,
     head_idx,
     n_block,
-    # Pipeline states (will be returned after advancing)
     producer_state_Q_LSE,
     producer_state_dO_dPsum,
     # Pipelines
@@ -983,14 +1204,20 @@ def produce_block_sparse_q_loads_bwd_sm100(
     should_load_Q: cutlass.Constexpr,
     should_load_dO: cutlass.Constexpr,
     # Subtiling factor and bounds
-    subtile_factor: cutlass.Constexpr = 1,
     m_block_max: int = 0,
 ):
-    """SM100 backward block sparse loading with subtiling.
-    Returns updated (producer_state_Q_LSE, producer_state_dO_dPsum).
-    First iteration loads K/V alongside Q/dO; subsequent iterations load only Q/dO.
-    """
     (
         curr_q_cnt,
         curr_q_idx,
@@ -998,81 +1225,325 @@ def produce_block_sparse_q_loads_bwd_sm100(
         curr_full_idx,
         loop_count,
     ) = get_block_sparse_iteration_info_bwd(
-        blocksparse_tensors, batch_idx, head_idx, n_block, subtile_factor, m_block_max
     )
-    for iter_idx in cutlass.range(loop_count, unroll=1):
-        m_block, _ = get_m_block_from_iter_bwd(
-            iter_idx,
             curr_q_cnt,
             curr_q_idx,
             curr_full_cnt,
             curr_full_idx,
-            subtile_factor,
-            m_block_max,
         )
-        m_block_safe = m_block
         if m_block_max > 0:
-            m_block_safe = cutlass.min(m_block, m_block_max - 1)
-        if iter_idx == 0:
-            # First block: load K/V alongside Q/dO
-            if const_expr(should_load_Q):
-                pipeline_Q.producer_acquire(producer_state_Q_LSE, extra_tx_count=tma_copy_bytes_K)
-                load_K(tma_bar_ptr=pipeline_Q.producer_get_barrier(producer_state_Q_LSE))
-                load_Q(m_block_safe, producer_state=producer_state_Q_LSE)
-                pipeline_Q.producer_commit(producer_state_Q_LSE)
-                pipeline_LSE.producer_acquire(producer_state_Q_LSE)
-                with cute.arch.elect_one():
-                    copy_stats(
-                        gLSE[None, m_block_safe],
-                        sLSE[None, producer_state_Q_LSE.index],
-                        mbar_ptr=pipeline_LSE.producer_get_barrier(producer_state_Q_LSE),
-                    )
-                producer_state_Q_LSE.advance()
-            if const_expr(should_load_dO):
-                pipeline_dO.producer_acquire(
-                    producer_state_dO_dPsum, extra_tx_count=tma_copy_bytes_V
                 )
-                load_V(tma_bar_ptr=pipeline_dO.producer_get_barrier(producer_state_dO_dPsum))
-                load_dO(m_block_safe, producer_state=producer_state_dO_dPsum)
-                pipeline_dO.producer_commit(producer_state_dO_dPsum)
-                pipeline_dPsum.producer_acquire(producer_state_dO_dPsum)
-                with cute.arch.elect_one():
-                    copy_stats(
-                        gdPsum[None, m_block_safe],
-                        sdPsum[None, producer_state_dO_dPsum.index],
-                        mbar_ptr=pipeline_dPsum.producer_get_barrier(producer_state_dO_dPsum),
-                    )
-                producer_state_dO_dPsum.advance()
-        else:
-            # Subsequent blocks: just load Q/dO (K/V already loaded)
             if const_expr(should_load_Q):
                 pipeline_Q.producer_acquire(producer_state_Q_LSE)
-                load_Q(m_block_safe, producer_state=producer_state_Q_LSE)
                 pipeline_Q.producer_commit(producer_state_Q_LSE)
                 pipeline_LSE.producer_acquire(producer_state_Q_LSE)
                 with cute.arch.elect_one():
                     copy_stats(
-                        gLSE[None, m_block_safe],
                         sLSE[None, producer_state_Q_LSE.index],
                         mbar_ptr=pipeline_LSE.producer_get_barrier(producer_state_Q_LSE),
                     )
                 producer_state_Q_LSE.advance()
             if const_expr(should_load_dO):
-                pipeline_dO.producer_acquire(producer_state_dO_dPsum)
-                load_dO(m_block_safe, producer_state=producer_state_dO_dPsum)
                 pipeline_dO.producer_commit(producer_state_dO_dPsum)
                 pipeline_dPsum.producer_acquire(producer_state_dO_dPsum)
                 with cute.arch.elect_one():
                     copy_stats(
-                        gdPsum[None, m_block_safe],
                         sdPsum[None, producer_state_dO_dPsum.index],
                         mbar_ptr=pipeline_dPsum.producer_get_barrier(producer_state_dO_dPsum),
                     )
                 producer_state_dO_dPsum.advance()
-    return producer_state_Q_LSE, producer_state_dO_dPsum
 @cute.jit
@@ -1081,14 +1552,14 @@ def get_block_sparse_iteration_info_bwd(
     batch_idx,
     head_idx,
     n_block,
-    subtile_factor: cutlass.Constexpr = 1,
     m_block_max: int = 0,
 ):
     """Extract block-sparse iteration info for backward pass.
     Returns (curr_q_cnt, curr_q_idx, curr_full_cnt, curr_full_idx, total_count).
     """
-    q_cnt, q_idx, full_cnt, full_idx = blocksparse_tensors
     curr_q_cnt = q_cnt[batch_idx, head_idx, n_block]
     curr_q_idx = q_idx[batch_idx, head_idx, n_block, None]
@@ -1102,7 +1573,7 @@ def get_block_sparse_iteration_info_bwd(
     sparse_block_count = curr_q_cnt
     if const_expr(full_cnt is not None):
         sparse_block_count = sparse_block_count + curr_full_cnt
-    total_count = sparse_block_count * subtile_factor
     return curr_q_cnt, curr_q_idx, curr_full_cnt, curr_full_idx, total_count
@@ -1114,7 +1585,7 @@ def get_m_block_from_iter_bwd(
     curr_q_idx: cute.Tensor,
     curr_full_cnt,
     curr_full_idx: Optional[cute.Tensor],
-    subtile_factor: cutlass.Constexpr = 1,
     m_block_max: int = 0,
 ):
     """Derive m_block index and is_full_block flag from iteration index.
@@ -1123,8 +1594,8 @@ def get_m_block_from_iter_bwd(
         - m_block: The actual Q-tile block index
         - is_full_block: True if this is a full block (no mask_mod needed)
     """
-    sparse_iter_idx = iter_idx // subtile_factor
-    subtile_offset = iter_idx % subtile_factor
     sparse_m_block = Int32(0)
     is_full_block = False
@@ -1137,7 +1608,7 @@ def get_m_block_from_iter_bwd(
     else:
         sparse_m_block = curr_q_idx[sparse_iter_idx]
-    return sparse_m_block * subtile_factor + subtile_offset, is_full_block
 @cute.jit
@@ -1165,6 +1636,13 @@ def _load_q_do_block_sm90(
     else:
         pipeline_Q.producer_acquire(producer_state_Q)
     load_Q(m_block, producer_state=producer_state_Q)
     load_LSE(m_block, producer_state=producer_state_Q)
     producer_state_dO_cur = (
@@ -1202,7 +1680,7 @@ def produce_block_sparse_q_loads_bwd_sm90(
     tma_copy_bytes_K,
     tma_copy_bytes_V,
     Q_stage_eq_dO_stage: cutlass.Constexpr,
-    subtile_factor: cutlass.Constexpr,
     m_block_max: int,
 ):
     """SM90 backward block sparse loading with separate partial/full loops.
@@ -1212,7 +1690,7 @@ def produce_block_sparse_q_loads_bwd_sm90(
     Returns updated (producer_state_Q, producer_state_dO).
     """
-    q_cnt, q_idx, full_cnt, full_idx = blocksparse_tensors
     curr_q_cnt = q_cnt[batch_idx, head_idx, n_block]
     curr_q_idx = q_idx[batch_idx, head_idx, n_block, None]
@@ -1225,10 +1703,10 @@ def produce_block_sparse_q_loads_bwd_sm90(
     kv_loaded = False
-    for iter_idx in cutlass.range(curr_q_cnt * subtile_factor, unroll=1):
-        sparse_idx = iter_idx // subtile_factor
-        subtile_offset = iter_idx % subtile_factor
-        m_block = curr_q_idx[sparse_idx] * subtile_factor + subtile_offset
         if m_block < m_block_max:
             producer_state_Q, producer_state_dO = _load_q_do_block_sm90(
@@ -1251,10 +1729,10 @@ def produce_block_sparse_q_loads_bwd_sm90(
             kv_loaded = True
     if const_expr(full_cnt is not None):
-        for iter_idx in cutlass.range(curr_full_cnt * subtile_factor, unroll=1):
-            sparse_idx = iter_idx // subtile_factor
-            subtile_offset = iter_idx % subtile_factor
-            m_block = curr_full_idx[sparse_idx] * subtile_factor + subtile_offset
             if m_block < m_block_max:
                 producer_state_Q, producer_state_dO = _load_q_do_block_sm90(
@@ -1295,9 +1773,9 @@ def consume_block_sparse_mma_bwd_sm90(
     thr_mma_SdP,
     score_mod_fn=None,
     score_mod_bwd_fn=None,
-    subtile_factor: cutlass.Constexpr = 1,
     m_block_max: int = 0,
-    aux_tensors=None,
     fastdiv_mods=(None, None),
 ):
     """SM90 backward block sparse MMA consumption with separate partial/full loops.
@@ -1307,7 +1785,7 @@ def consume_block_sparse_mma_bwd_sm90(
     Returns updated (consumer_state_Q, consumer_state_dO).
     """
-    q_cnt, q_idx, full_cnt, full_idx = blocksparse_tensors
     curr_q_cnt = q_cnt[batch_idx, head_idx, n_block]
     curr_q_idx = q_idx[batch_idx, head_idx, n_block, None]
@@ -1330,7 +1808,7 @@ def consume_block_sparse_mma_bwd_sm90(
         mask_causal=is_causal,
         mask_local=is_local,
         mask_mod=mask_mod,
-        aux_tensors=aux_tensors,
         fastdiv_mods=fastdiv_mods,
     )
@@ -1343,14 +1821,14 @@ def consume_block_sparse_mma_bwd_sm90(
         mask_seqlen=True,
         mask_causal=is_causal,
         mask_local=is_local,
-        aux_tensors=aux_tensors,
         fastdiv_mods=fastdiv_mods,
     )
-    for iter_idx in cutlass.range(curr_q_cnt * subtile_factor, unroll=1):
-        sparse_idx = iter_idx // subtile_factor
-        subtile_offset = iter_idx % subtile_factor
-        m_block = curr_q_idx[sparse_idx] * subtile_factor + subtile_offset
         if m_block < m_block_max:
             consumer_state_Q, consumer_state_dO = mma_one_m_block_fn(
@@ -1365,10 +1843,10 @@ def consume_block_sparse_mma_bwd_sm90(
             dKV_accumulate = True
     if const_expr(full_cnt is not None):
-        for iter_idx in cutlass.range(curr_full_cnt * subtile_factor, unroll=1):
-            sparse_idx = iter_idx // subtile_factor
-            subtile_offset = iter_idx % subtile_factor
-            m_block = curr_full_idx[sparse_idx] * subtile_factor + subtile_offset
             if m_block < m_block_max:
                 consumer_state_Q, consumer_state_dO = mma_one_m_block_fn(
@@ -1390,18 +1868,18 @@ def _store_one_dQaccum_sm90(
     m_block,
     sdQaccum: cute.Tensor,
     gdQaccum: cute.Tensor,
-    num_mma_warp_groups: cutlass.Constexpr,
     num_threads_per_warp_group: cutlass.Constexpr,
     tma_copy_bytes_dQ,
 ):
     """Store dQaccum for a single m_block."""
-    for warp_group_idx in cutlass.range_constexpr(num_mma_warp_groups):
-        cute.arch.cp_async_bulk_wait_group(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=num_threads_per_warp_group + cute.arch.WARP_SIZE,
         )
-    for warp_group_idx in cutlass.range_constexpr(num_mma_warp_groups):
         cute.arch.barrier(
             barrier_id=int(NamedBarrierBwd.dQFullWG0) + warp_group_idx,
             number_of_threads=num_threads_per_warp_group + cute.arch.WARP_SIZE,
@@ -1409,7 +1887,7 @@ def _store_one_dQaccum_sm90(
         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,
                 tma_copy_bytes_dQ,
             )
         cute.arch.cp_async_bulk_commit_group()
@@ -1423,9 +1901,9 @@ def dQaccum_store_block_sparse_bwd_sm90(
     n_block,
     sdQaccum: cute.Tensor,
     gdQaccum: cute.Tensor,
-    subtile_factor: cutlass.Constexpr,
     m_block_max: int,
-    num_mma_warp_groups: cutlass.Constexpr,
     num_threads_per_warp_group: cutlass.Constexpr,
     tma_copy_bytes_dQ,
 ):
@@ -1433,7 +1911,7 @@ def dQaccum_store_block_sparse_bwd_sm90(
     Iterates partial blocks first, then full blocks, matching producer/consumer order.
     """
-    q_cnt, q_idx, full_cnt, full_idx = blocksparse_tensors
     curr_q_cnt = q_cnt[batch_idx, head_idx, n_block]
     curr_q_idx = q_idx[batch_idx, head_idx, n_block, None]
@@ -1444,33 +1922,33 @@ def dQaccum_store_block_sparse_bwd_sm90(
         curr_full_cnt = Int32(0)
         curr_full_idx = None
-    for iter_idx in cutlass.range(curr_q_cnt * subtile_factor, unroll=1):
-        sparse_idx = iter_idx // subtile_factor
-        subtile_offset = iter_idx % subtile_factor
-        m_block = curr_q_idx[sparse_idx] * subtile_factor + subtile_offset
         if m_block < m_block_max:
             _store_one_dQaccum_sm90(
                 m_block,
                 sdQaccum,
                 gdQaccum,
-                num_mma_warp_groups,
                 num_threads_per_warp_group,
                 tma_copy_bytes_dQ,
             )
     if const_expr(full_cnt is not None):
-        for iter_idx in cutlass.range(curr_full_cnt * subtile_factor, unroll=1):
-            sparse_idx = iter_idx // subtile_factor
-            subtile_offset = iter_idx % subtile_factor
-            m_block = curr_full_idx[sparse_idx] * subtile_factor + subtile_offset
             if m_block < m_block_max:
                 _store_one_dQaccum_sm90(
                     m_block,
                     sdQaccum,
                     gdQaccum,
-                    num_mma_warp_groups,
                     num_threads_per_warp_group,
                     tma_copy_bytes_dQ,
                 )

 These utilities are used by CUTE DSL kernels to produce and consume block-sparse loads.
 """
+from typing import Callable, Optional, Tuple
 from functools import partial
 import math
 import cutlass
 # Import data structures from block_sparsity
 from .block_sparsity import BlockSparseTensors
 from .named_barrier import NamedBarrierBwd
+from .seqlen_info import SeqlenInfoQK
+from .utils import AuxData
+@cute.jit
+def _get_curr_blocksparse_tensors_varlen(
+    head_idx: cutlass.Int32,
+    m_block: cutlass.Int32,
+    blocksparse_tensors: BlockSparseTensors,
+    seqlen_info: SeqlenInfoQK,
+    kv_subtile_factor: cutlass.Constexpr[int] = 1,
+) -> Tuple[cutlass.Int32, cute.Tensor, cutlass.Int32, Optional[cute.Tensor]]:
+    """Varlen path: tensors are 2D [nheads, total_m_blocks] / [nheads, total_n_blocks]."""
+    mask_block_cnt, mask_block_idx, full_block_cnt, full_block_idx, *_ = blocksparse_tensors
+    curr_m_block = seqlen_info.m_block_offset + m_block
+    sparse_num_n_blocks = (
+        seqlen_info.num_n_blocks
+        if const_expr(kv_subtile_factor == 1)
+        else (seqlen_info.num_n_blocks + kv_subtile_factor - 1) // kv_subtile_factor
+    )
+    batch_block_idx_offset = (
+        seqlen_info.block_idx_offset
+        if const_expr(seqlen_info.has_cu_block_idx_offsets)
+        else seqlen_info.m_block_offset * sparse_num_n_blocks
+    )
+    curr_block_idx_offset = batch_block_idx_offset + m_block * sparse_num_n_blocks
+    curr_mask_block_cnt = mask_block_cnt[head_idx, curr_m_block]
+    curr_mask_block_idx = cute.domain_offset(curr_block_idx_offset, mask_block_idx[head_idx, None])
+    if const_expr(full_block_cnt is not None):
+        curr_full_block_cnt = full_block_cnt[head_idx, curr_m_block]
+        curr_full_block_idx = cute.domain_offset(
+            curr_block_idx_offset, full_block_idx[head_idx, None]
+        )
+    else:
+        curr_full_block_cnt = Int32(0)
+        curr_full_block_idx = None
+    return (curr_mask_block_cnt, curr_mask_block_idx, curr_full_block_cnt, curr_full_block_idx)
+@cute.jit
+def _get_curr_blocksparse_tensors(
+    batch_idx: cutlass.Int32,
+    head_idx: cutlass.Int32,
+    m_block: cutlass.Int32,
+    blocksparse_tensors: BlockSparseTensors,
+) -> Tuple[cutlass.Int32, cute.Tensor, cutlass.Int32, Optional[cute.Tensor]]:
+    """Fixed-length path: tensors are 4D [batch, nheads, m_block, n_block]."""
+    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]
+    if const_expr(full_block_cnt is not 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]
+    else:
+        curr_full_block_cnt = Int32(0)
+        curr_full_block_idx = None
+    return (curr_mask_block_cnt, curr_mask_block_idx, curr_full_block_cnt, curr_full_block_idx)
+@cute.jit
+def get_curr_blocksparse_tensors(
+    batch_idx: cutlass.Int32,
+    head_idx: cutlass.Int32,
+    m_block: cutlass.Int32,
+    blocksparse_tensors: BlockSparseTensors,
+    seqlen_info: SeqlenInfoQK,
+    kv_subtile_factor: cutlass.Constexpr[int] = 1,
+) -> Tuple[cutlass.Int32, cute.Tensor, cutlass.Int32, Optional[cute.Tensor]]:
+    """Extract head, m_block, and batch-local blocksparsity data from blocksparse_tensors"""
+    if const_expr(len(blocksparse_tensors.mask_block_cnt.shape) == 2):
+        return _get_curr_blocksparse_tensors_varlen(
+            head_idx, m_block, blocksparse_tensors, seqlen_info, kv_subtile_factor
+        )
+    return _get_curr_blocksparse_tensors(batch_idx, head_idx, m_block, blocksparse_tensors)
 # NOTE [SM100 block-sparse empty tiles: mbarrier contract]
 #   to ack/advance, and arrives `mbar_P_full_O_rescaled` when MMA can proceed.
 #
 # Backward (SM100):
+# - Empty KV tile: for a given sparse `n_block`, `loop_count == 0` means no Q tiles contribute.
+# - Load helpers guard their prologue/mainloop/tail on `loop_count > 0`; MMA/softmax/relay use
+#   the same block-sparse count as `process_tile`. Empty tiles therefore skip both producer and
+#   consumer operations entirely, with no synthetic per-tile mbarrier handshake like forward.
+# - For non-empty tiles, producer tails are part of the non-empty load contract: the load warp emits
+#   exactly the Q/dO/LSE/dPsum/Qt/Kt items that the MMA and softmax consumers wait on, then tails the
+#   advanced producer states. Empty tiles leave producer and consumer states unchanged for the next tile.
 # - In the `not dKV_postprocess` path, dK/dV for empty KV tiles are explicitly written as zeros
 #   even when `process_tile == False` (see `flash_bwd_sm100.py` `should_zero_dKV`).
+@cute.jit
+def sparse_physical_n_block_forward(
+    block_indices: cute.Tensor,
+    offset,
+    kv_subtile_factor: cutlass.Constexpr[int],
+):
+    """Map forward physical N offsets to physical N tile indices."""
+    sparse_offset = offset // kv_subtile_factor
+    subtile_offset = offset - sparse_offset * kv_subtile_factor
+    coarse_block = block_indices[sparse_offset]
+    return coarse_block * kv_subtile_factor + subtile_offset
 @cute.jit
 def load_block_list(
     block_indices: cute.Tensor,
     block_count,
     first_block_preloaded: cutlass.Constexpr,
     kv_producer_state,
     load_K,
     load_V,
     pipeline_k,
     pipeline_v,
     intra_wg_overlap: cutlass.Constexpr,
+    kv_subtile_factor: cutlass.Constexpr[int] = 1,
 ):
+    """Iterate over the sparse blocks and load K, V 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.
+    Q is loaded separately on its own mbarrier before this function is called.
     Note:
         we iterate along the block_n indices in reverse.
     """
     if block_count > 0:
+        total_blocks = block_count * kv_subtile_factor
         if const_expr(not intra_wg_overlap):
+            for offset in cutlass.range(total_blocks):
+                n_block = sparse_physical_n_block_forward(
+                    block_indices,
+                    total_blocks - 1 - offset,
+                    kv_subtile_factor,
+                )
                 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 = sparse_physical_n_block_forward(
+                block_indices,
+                total_blocks - 1,
+                kv_subtile_factor,
+            )
             if const_expr(not first_block_preloaded):
+                pipeline_k.producer_acquire(kv_producer_state)
                 load_K(src_idx=n_block_first, producer_state=kv_producer_state)
+            for idx in cutlass.range(total_blocks - 1, unroll=1):
+                n_block_prev = sparse_physical_n_block_forward(
+                    block_indices,
+                    total_blocks - 1 - idx,
+                    kv_subtile_factor,
+                )
+                n_block = sparse_physical_n_block_forward(
+                    block_indices,
+                    total_blocks - 1 - (idx + 1),
+                    kv_subtile_factor,
+                )
                 kv_producer_state_prev = kv_producer_state.clone()
                 kv_producer_state.advance()
                 pipeline_k.producer_acquire(kv_producer_state)
     load_V,
     pipeline_v,
     kv_producer_state,
+    kv_subtile_factor: cutlass.Constexpr[int] = 1,
 ):
     """Load the final V block after overlapped K/V loads."""
     if block_count > 0:
+        n_block_last = block_indices[0] * kv_subtile_factor
         pipeline_v.producer_acquire(kv_producer_state)
         load_V(src_idx=n_block_last, producer_state=kv_producer_state)
         kv_producer_state.advance()
     batch_idx,
     head_idx,
     m_block,
+    seqlen_info: SeqlenInfoQK,
     kv_producer_state,
     load_K,
     load_V,
     pipeline_k,
     pipeline_v,
     intra_wg_overlap: cutlass.Constexpr,
     qhead_per_kvhead: cutlass.Constexpr[int] = 1,
     q_subtile_factor: cutlass.Constexpr[int] = 1,
+    kv_subtile_factor: cutlass.Constexpr[int] = 1,
 ):
     """Iterate over the mask and full block lists for a single tile.
+    Q is loaded separately on its own mbarrier before this function is called.
     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.
         qhead_per_kvhead: Pack-GQA factor. When > 1, m_block is in packed space and
             must be converted to unpacked for sparse tensor indexing.
     """
+    # The SM90 consumer (consume_block_sparse_loads) has no kv_subtile support yet.
+    assert kv_subtile_factor == 1, (
+        "Coarse KV blocks (kv_subtile_factor > 1) are not supported on the SM90 forward path yet."
+    )
     m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor)
+    (
+        curr_mask_block_cnt,
+        curr_mask_block_idx,
+        curr_full_block_cnt,
+        curr_full_block_idx,
+    ) = get_curr_blocksparse_tensors(
+        batch_idx,
+        head_idx,
+        m_block_sparse,
+        blocksparse_tensors,
+        seqlen_info,
+        kv_subtile_factor,
+    )
     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 K load.
         kv_producer_state = load_block_list(
             curr_full_block_idx,
             curr_full_block_cnt,
             first_block_preloaded=False,
             kv_producer_state=kv_producer_state,
             load_K=load_K,
             load_V=load_V,
             pipeline_k=pipeline_k,
             pipeline_v=pipeline_v,
             intra_wg_overlap=intra_wg_overlap,
+            kv_subtile_factor=kv_subtile_factor,
         )
         if const_expr(intra_wg_overlap) and curr_full_block_cnt > 0:
                 load_V,
                 pipeline_v,
                 kv_producer_state,
+                kv_subtile_factor=kv_subtile_factor,
             )
     else:
+        # Masked blocks present. When overlap is disabled this fully drains the list.
         kv_producer_state = load_block_list(
             curr_mask_block_idx,
             curr_mask_block_cnt,
             first_block_preloaded=False,
             kv_producer_state=kv_producer_state,
             load_K=load_K,
             load_V=load_V,
             pipeline_k=pipeline_k,
             pipeline_v=pipeline_v,
             intra_wg_overlap=intra_wg_overlap,
+            kv_subtile_factor=kv_subtile_factor,
         )
         if full_empty:
                     load_V,
                     pipeline_v,
                     kv_producer_state,
+                    kv_subtile_factor=kv_subtile_factor,
                 )
         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] * kv_subtile_factor
+                n_block_full_first = sparse_physical_n_block_forward(
+                    curr_full_block_idx,
+                    curr_full_block_cnt * kv_subtile_factor - 1,
+                    kv_subtile_factor,
+                )
                 kv_producer_state_prev = kv_producer_state.clone()
                 kv_producer_state.advance()
                 pipeline_k.producer_acquire(kv_producer_state)
                 kv_producer_state = load_block_list(
                     curr_full_block_idx,
                     curr_full_block_cnt,
                     first_block_preloaded=True,
                     kv_producer_state=kv_producer_state,
                     load_K=load_K,
                     load_V=load_V,
                     pipeline_k=pipeline_k,
                     pipeline_v=pipeline_v,
                     intra_wg_overlap=intra_wg_overlap,
+                    kv_subtile_factor=kv_subtile_factor,
                 )
                 kv_producer_state = finish_overlap_v_load(
                     load_V,
                     pipeline_v,
                     kv_producer_state,
+                    kv_subtile_factor=kv_subtile_factor,
                 )
             else:
+                # Non-overlap path with both lists: run the full list normally.
                 kv_producer_state = load_block_list(
                     curr_full_block_idx,
                     curr_full_block_cnt,
                     first_block_preloaded=False,
                     kv_producer_state=kv_producer_state,
                     load_K=load_K,
                     load_V=load_V,
                     pipeline_k=pipeline_k,
                     pipeline_v=pipeline_v,
                     intra_wg_overlap=intra_wg_overlap,
+                    kv_subtile_factor=kv_subtile_factor,
                 )
     return kv_producer_state
     batch_idx,
     head_idx,
     m_block,
+    seqlen_info,
     kv_consumer_state,
     mma_pv_fn,
     mma_one_n_block,
         qhead_per_kvhead: Pack-GQA factor. When > 1, m_block is in packed space and
             must be converted to unpacked for sparse tensor indexing.
     """
     m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor)
+    (
+        curr_mask_block_cnt,
+        curr_mask_block_idx,
+        curr_full_block_cnt,
+        curr_full_block_idx,
+    ) = get_curr_blocksparse_tensors(
+        batch_idx,
+        head_idx,
+        m_block_sparse,
+        blocksparse_tensors,
+        seqlen_info,
+    )
     processed_any = curr_mask_block_cnt + curr_full_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,
+                seqlen=seqlen_info,
                 kv_consumer_state=kv_consumer_state,
                 mask_fn=partial(
                     mask_fn,
                 kv_consumer_state = mma_one_n_block(
                     kv_consumer_state,
                     n_block=mask_n_block,
+                    seqlen=seqlen_info,
                     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),
                 )
             if curr_mask_block_cnt == 0:
                 kv_consumer_state = process_first_half_block(
                     n_block=full_n_block,
+                    seqlen=seqlen_info,
                     kv_consumer_state=kv_consumer_state,
                     mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=True),
                     score_mod_fn=score_mod_fn,
                 kv_consumer_state = mma_one_n_block(
                     kv_consumer_state,
                     n_block=full_n_block,
+                    seqlen=seqlen_info,
                     mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate),
                     mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=True),
                 )
                 kv_consumer_state = mma_one_n_block(
                     kv_consumer_state,
                     n_block=full_n_block,
+                    seqlen=seqlen_info,
                     mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate),
                     mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=False),
                 )
     return kv_consumer_state, O_should_accumulate, processed_any
+@cute.jit
+def split_block_range(block_count, split_idx: Int32, num_splits: Int32):
+    """Return the half-open block-list range assigned to one SplitKV partition."""
+    blocks_per_split = cute.ceil_div(block_count, num_splits)
+    block_begin = cutlass.min(split_idx * blocks_per_split, block_count)
+    block_end = cutlass.min(block_begin + blocks_per_split, block_count)
+    return block_begin, block_end
 @cute.jit
 def load_block_list_sm100(
     block_indices: cute.Tensor,
+    block_begin,
+    block_end,
     load_q_with_first: cutlass.Constexpr,
     q_stage: cutlass.Constexpr,
     kv_producer_state,
     load_K,
     load_V,
     pipeline_kv,
+    kv_subtile_factor: cutlass.Constexpr[int] = 1,
 ):
+    """SM100 sparse load loop over physical N tiles from coarse metadata."""
+    physical_n_block = partial(
+        sparse_physical_n_block_forward, block_indices, kv_subtile_factor=kv_subtile_factor
+    )
+    physical_block_count = block_end - block_begin
+    if physical_block_count > 0:
         # First iteration: load Q alongside K if requested
+        n_block_first = physical_n_block(block_end - 1)
         if const_expr(load_q_with_first):
             # SM100 loads Q0 and optionally Q1
         kv_producer_state.advance()
         # Remaining blocks
+        for offset in cutlass.range(1, physical_block_count):
+            n_block = physical_n_block(block_end - 1 - offset)
             load_K(block=n_block, producer_state=kv_producer_state, page_idx=None)
             kv_producer_state.advance()
             load_V(block=n_block, producer_state=kv_producer_state, page_idx=None)
     batch_idx,
     head_idx,
     m_block,
+    seqlen_info: SeqlenInfoQK,
+    split_idx: Int32,
+    num_splits: Int32,
     kv_producer_state,
     load_Q,
     load_K,
     q_producer_phase: Int32,
     qhead_per_kvhead: cutlass.Constexpr,
     q_subtile_factor: cutlass.Constexpr,
+    kv_subtile_factor: cutlass.Constexpr[int] = 1,
 ):
     """SM100 entry point for sparse block iteration.
     """
     m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor)
+    (
+        curr_mask_block_cnt,
+        curr_mask_block_idx,
+        curr_full_block_cnt,
+        curr_full_block_idx,
+    ) = get_curr_blocksparse_tensors(
+        batch_idx,
+        head_idx,
+        m_block_sparse,
+        blocksparse_tensors,
+        seqlen_info,
+        kv_subtile_factor,
+    )
+    mask_begin, mask_end = split_block_range(
+        curr_mask_block_cnt * kv_subtile_factor,
+        split_idx,
+        num_splits,
+    )
+    full_begin, full_end = split_block_range(
+        curr_full_block_cnt * kv_subtile_factor,
+        split_idx,
+        num_splits,
+    )
+    mask_empty = mask_begin == mask_end
+    full_empty = full_begin == full_end
     q_phase_flipped = False
         # No masked blocks: process full list with Q loading
         kv_producer_state = load_block_list_sm100(
             curr_full_block_idx,
+            full_begin,
+            full_end,
             load_q_with_first=True,
             q_stage=q_stage,
             kv_producer_state=kv_producer_state,
             load_K=load_K,
             load_V=load_V,
             pipeline_kv=pipeline_kv,
+            kv_subtile_factor=kv_subtile_factor,
         )
         q_phase_flipped = not full_empty
     else:
         # Process masked blocks with Q loading
         kv_producer_state = load_block_list_sm100(
             curr_mask_block_idx,
+            mask_begin,
+            mask_end,
             load_q_with_first=True,
             q_stage=q_stage,
             kv_producer_state=kv_producer_state,
             load_K=load_K,
             load_V=load_V,
             pipeline_kv=pipeline_kv,
+            kv_subtile_factor=kv_subtile_factor,
         )
         q_phase_flipped = True
             # Process full blocks without Q loading
             kv_producer_state = load_block_list_sm100(
                 curr_full_block_idx,
+                full_begin,
+                full_end,
                 load_q_with_first=False,
                 q_stage=q_stage,
                 kv_producer_state=kv_producer_state,
                 load_K=load_K,
                 load_V=load_V,
                 pipeline_kv=pipeline_kv,
+                kv_subtile_factor=kv_subtile_factor,
             )
     if q_phase_flipped:
     batch_idx,
     head_idx,
     m_block,
+    split_idx: Int32,
+    num_splits: Int32,
     qhead_per_kvhead: cutlass.Constexpr,
     q_subtile_factor: cutlass.Constexpr,
+    seqlen_info: SeqlenInfoQK,
+    kv_subtile_factor: cutlass.Constexpr[int] = 1,
 ):
     m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor)
+    (
+        curr_mask_block_cnt,
+        _,
+        curr_full_block_cnt,
+        _,
+    ) = get_curr_blocksparse_tensors(
+        batch_idx,
+        head_idx,
+        m_block_sparse,
+        blocksparse_tensors,
+        seqlen_info,
+        kv_subtile_factor,
+    )
+    mask_begin, mask_end = split_block_range(
+        curr_mask_block_cnt * kv_subtile_factor,
+        split_idx,
+        num_splits,
+    )
+    full_begin, full_end = split_block_range(
+        curr_full_block_cnt * kv_subtile_factor,
+        split_idx,
+        num_splits,
+    )
+    return mask_end - mask_begin + full_end - full_begin
 @cute.jit
     is_split_kv: cutlass.Constexpr,
     learnable_sink,
     mLSE,
+    seqlen_info,
     m_block: Int32,
     head_idx: Int32,
     batch_idx: Int32,
     sScale: cute.Tensor,
     stats: list,
     correction_epilogue: Callable,
+    thr_mma_pv: cute.ThrMma,
     tOtO: cute.Tensor,
     sO: cute.Tensor,
     pipeline_sm_stats: cutlass.pipeline.PipelineAsync,
     o_corr_consumer_phase: Int32,
     corr_epi_producer_phase: Int32,
     softmax_scale_log2: Float32,
+    max_offset: Float32,
+    max_offset_scale: Float32,
     mO_cur: Optional[cute.Tensor] = None,
     gO: Optional[cute.Tensor] = None,
     gmem_tiled_copy_O: Optional[cute.TiledCopy] = None,
             if sink_val != -Float32.inf and (const_expr(not is_split_kv) or split_idx == 0):
                 if row_max_value == -Float32.inf:
                     row_max_value = sink_val * (LOG2_E / softmax_scale_log2)
+                    row_sum_value = max_offset_scale
                 else:
                     row_sum_value = row_sum_value + cute.math.exp2(
+                        sink_val * LOG2_E - row_max_value * softmax_scale_log2 + max_offset,
+                        fastmath=True,
                     )
         if tidx < m_block_size:
             scale_row_idx = tidx + stage * m_block_size
         if const_expr(gmem_tiled_copy_O is None):
             pipeline_o_epi.producer_acquire_w_index_phase(stage, corr_epi_producer_phase)
+        gO_stage = gO[None, None, stage] if const_expr(gO is not None) else None
         correction_epilogue(
             thr_mma_pv,
             tOtO[None, None, None, stage],
             tidx,
             stage,
             m_block,
+            seqlen_info.seqlen_q,
             Float32(0.0),  # zero scale ensures empty tile writes zeros into staged outputs
             sO[None, None, stage],
             mO_cur,
+            gO_stage,
             gmem_tiled_copy_O,
         )
         if const_expr(gmem_tiled_copy_O is None):
     )
+@cute.jit
+def softmax_block_sparse_sm100_list(
+    block_indices: cute.Tensor,
+    block_end,
+    split_block_cnt,
+    kv_subtile_factor: cutlass.Constexpr[int],
+    softmax_step: Callable,
+    mask_fn_base: Callable,
+    mma_si_consumer_phase: Int32,
+    si_corr_producer_phase: Int32,
+    s0_s1_sequence_phase: Int32,
+    check_m_boundary: bool,
+    is_first_block: bool,
+    allow_unmasked_inner_blocks: cutlass.Constexpr[bool],
+):
+    """Run one reverse sparse list while masking only the first coarse KV fragment."""
+    mask_fn_seqlen = partial(mask_fn_base, mask_seqlen=True, check_q_boundary=check_m_boundary)
+    physical_n_block = partial(
+        sparse_physical_n_block_forward, block_indices, kv_subtile_factor=kv_subtile_factor
+    )
+    n_block = physical_n_block(block_end - 1)
+    if is_first_block:
+        (
+            mma_si_consumer_phase,
+            si_corr_producer_phase,
+            s0_s1_sequence_phase,
+        ) = softmax_step(
+            mma_si_consumer_phase,
+            si_corr_producer_phase,
+            s0_s1_sequence_phase,
+            n_block,
+            is_first=True,
+            mask_fn=mask_fn_seqlen,
+        )
+    else:
+        (
+            mma_si_consumer_phase,
+            si_corr_producer_phase,
+            s0_s1_sequence_phase,
+        ) = softmax_step(
+            mma_si_consumer_phase,
+            si_corr_producer_phase,
+            s0_s1_sequence_phase,
+            n_block,
+            is_first=False,
+            mask_fn=mask_fn_seqlen,
+        )
+    first_fragment_count = cutlass.min(
+        split_block_cnt,
+        (block_end - 1) % kv_subtile_factor + 1,
+    )
+    for j in cutlass.range_constexpr(1, kv_subtile_factor):
+        if j < first_fragment_count:
+            (
+                mma_si_consumer_phase,
+                si_corr_producer_phase,
+                s0_s1_sequence_phase,
+            ) = softmax_step(
+                mma_si_consumer_phase,
+                si_corr_producer_phase,
+                s0_s1_sequence_phase,
+                physical_n_block(block_end - 1 - j),
+                mask_fn=mask_fn_seqlen,
+            )
+    for i in cutlass.range(first_fragment_count, split_block_cnt):
+        (
+            mma_si_consumer_phase,
+            si_corr_producer_phase,
+            s0_s1_sequence_phase,
+        ) = softmax_step(
+            mma_si_consumer_phase,
+            si_corr_producer_phase,
+            s0_s1_sequence_phase,
+            physical_n_block(block_end - 1 - i),
+            mask_fn=None
+            if const_expr(allow_unmasked_inner_blocks and check_m_boundary is False)
+            else partial(
+                mask_fn_base,
+                mask_seqlen=False,
+                check_q_boundary=check_m_boundary,
+            ),
+        )
+    return mma_si_consumer_phase, si_corr_producer_phase, s0_s1_sequence_phase
 @cute.jit
 def softmax_block_sparse_sm100(
     blocksparse_tensors: BlockSparseTensors,
     batch_idx,
     head_idx,
     m_block,
+    seqlen_info: SeqlenInfoQK,
+    split_idx: Int32,
+    num_splits: Int32,
     softmax_step: Callable,
     mask_fn: Callable,
     mask_fn_none: Callable,
     check_m_boundary: bool,
     qhead_per_kvhead: cutlass.Constexpr,
     q_subtile_factor: cutlass.Constexpr[int] = 1,
+    kv_subtile_factor: cutlass.Constexpr[int] = 1,
 ):
     warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4
     m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor)
+    (
+        curr_mask_block_cnt,
+        curr_mask_block_idx,
+        curr_full_block_cnt,
+        curr_full_block_idx,
+    ) = get_curr_blocksparse_tensors(
+        batch_idx,
+        head_idx,
+        m_block_sparse,
+        blocksparse_tensors,
+        seqlen_info,
+        kv_subtile_factor,
+    )
+    mask_begin, mask_end = split_block_range(
+        curr_mask_block_cnt * kv_subtile_factor,
+        split_idx,
+        num_splits,
+    )
+    full_begin, full_end = split_block_range(
+        curr_full_block_cnt * kv_subtile_factor,
+        split_idx,
+        num_splits,
+    )
+    split_mask_block_cnt = mask_end - mask_begin
+    split_full_block_cnt = full_end - full_begin
+    total_block_cnt = split_mask_block_cnt + split_full_block_cnt
     if total_block_cnt == 0:
         sm_stats_barrier.arrive_w_index(index=stage_idx * 4 + warp_idx)
     else:
+        if split_mask_block_cnt > 0:
             (
                 mma_si_consumer_phase,
                 si_corr_producer_phase,
                 s0_s1_sequence_phase,
+            ) = softmax_block_sparse_sm100_list(
+                curr_mask_block_idx,
+                mask_end,
+                split_mask_block_cnt,
+                kv_subtile_factor,
+                softmax_step,
+                mask_fn,
                 mma_si_consumer_phase,
                 si_corr_producer_phase,
                 s0_s1_sequence_phase,
+                check_m_boundary,
+                True,
+                False,
             )
+        if split_full_block_cnt > 0:
+            (
+                mma_si_consumer_phase,
+                si_corr_producer_phase,
+                s0_s1_sequence_phase,
+            ) = softmax_block_sparse_sm100_list(
+                curr_full_block_idx,
+                full_end,
+                split_full_block_cnt,
+                kv_subtile_factor,
+                softmax_step,
+                mask_fn_none,
+                mma_si_consumer_phase,
+                si_corr_producer_phase,
+                s0_s1_sequence_phase,
+                check_m_boundary,
+                split_mask_block_cnt == 0,
+                True,
+            )
     return (
         mma_si_consumer_phase,
     batch_idx,
     head_idx,
     n_block,
+    q_subtile_factor: cutlass.Constexpr = 1,
     m_block_max: int = 0,
 ):
     """Count total tile iterations for given n_block (KV tile) in backward."""
+    q_block_cnt, _, full_block_cnt, _, *_ = blocksparse_tensors
     total = q_block_cnt[batch_idx, head_idx, n_block]
     if const_expr(full_block_cnt is not None):
         total = total + full_block_cnt[batch_idx, head_idx, n_block]
+    return total * q_subtile_factor
 @cute.jit
+def produce_block_sparse_q_loads_bwd_sm100_default(
     blocksparse_tensors: BlockSparseTensors,
     batch_idx,
     head_idx,
     n_block,
+    # Pipeline states returned after advancing
     producer_state_Q_LSE,
     producer_state_dO_dPsum,
     # Pipelines
     should_load_Q: cutlass.Constexpr,
     should_load_dO: cutlass.Constexpr,
     # Subtiling factor and bounds
+    q_subtile_factor: cutlass.Constexpr = 1,
     m_block_max: int = 0,
+    # Optional 2CTA state for hdim <= 128
+    use_2cta_instrs: cutlass.Constexpr = False,
+    producer_state_Qt=None,
+    producer_state_Kt=None,
+    pipeline_Qt=None,
+    pipeline_Kt=None,
+    load_Qt=None,
+    load_Kt=None,
+    load_dOt=None,
+    tma_copy_bytes_dO=0,
 ):
+    """Produce SM100 backward block-sparse Q/dO loads for 1CTA and non-hdim192 2CTA."""
     (
         curr_q_cnt,
         curr_q_idx,
         curr_full_idx,
         loop_count,
     ) = get_block_sparse_iteration_info_bwd(
+        blocksparse_tensors, batch_idx, head_idx, n_block, q_subtile_factor, m_block_max
     )
+    # 2 cta peels the first loop for Qt path; so we guard the whole block if loopcount == 0
+    if loop_count > Int32(0):
+        first_m_block, _ = get_m_block_from_iter_bwd(
+            Int32(0),
             curr_q_cnt,
             curr_q_idx,
             curr_full_cnt,
             curr_full_idx,
+            q_subtile_factor=q_subtile_factor,
+            m_block_max=m_block_max,
         )
+        # with q_subtile > 1 we need to guard against fully OOB regions
         if m_block_max > 0:
+            first_m_block = cutlass.min(first_m_block, m_block_max - 1)
+        if const_expr(should_load_Q):
+            pipeline_Q.producer_acquire(producer_state_Q_LSE, extra_tx_count=tma_copy_bytes_K)
+            load_K(tma_bar_ptr=pipeline_Q.producer_get_barrier(producer_state_Q_LSE))
+            load_Q(first_m_block, producer_state=producer_state_Q_LSE)
+            pipeline_Q.producer_commit(producer_state_Q_LSE)
+            pipeline_LSE.producer_acquire(producer_state_Q_LSE)
+            with cute.arch.elect_one():
+                copy_stats(
+                    gLSE[None, first_m_block],
+                    sLSE[None, producer_state_Q_LSE.index],
+                    mbar_ptr=pipeline_LSE.producer_get_barrier(producer_state_Q_LSE),
                 )
+            producer_state_Q_LSE.advance()
+        if const_expr(should_load_dO):
+            pipeline_dO.producer_acquire(
+                producer_state_dO_dPsum,
+                extra_tx_count=tma_copy_bytes_V + tma_copy_bytes_dO
+                if const_expr(load_dOt is not None)
+                else tma_copy_bytes_V,
+            )
+            load_V(tma_bar_ptr=pipeline_dO.producer_get_barrier(producer_state_dO_dPsum))
+            load_dO(first_m_block, producer_state=producer_state_dO_dPsum)
+            if const_expr(load_dOt is not None):
+                load_dOt(first_m_block, producer_state=producer_state_dO_dPsum)
+            pipeline_dO.producer_commit(producer_state_dO_dPsum)
+            pipeline_dPsum.producer_acquire(producer_state_dO_dPsum)
+            with cute.arch.elect_one():
+                copy_stats(
+                    gdPsum[None, first_m_block],
+                    sdPsum[None, producer_state_dO_dPsum.index],
+                    mbar_ptr=pipeline_dPsum.producer_get_barrier(producer_state_dO_dPsum),
+                )
+            producer_state_dO_dPsum.advance()
+        if const_expr(use_2cta_instrs):
+            assert load_Kt is not None and pipeline_Kt is not None
+            assert producer_state_Kt is not None
+            pipeline_Kt.producer_acquire(producer_state_Kt)
+            load_Kt(tma_bar_ptr=pipeline_Kt.producer_get_barrier(producer_state_Kt))
+            pipeline_Kt.producer_commit(producer_state_Kt)
+            producer_state_Kt.advance()
+        prev_m_block = first_m_block
+        for iter_idx in cutlass.range(Int32(1), loop_count, unroll=1):
+            m_block, _ = get_m_block_from_iter_bwd(
+                iter_idx,
+                curr_q_cnt,
+                curr_q_idx,
+                curr_full_cnt,
+                curr_full_idx,
+                q_subtile_factor=q_subtile_factor,
+                m_block_max=m_block_max,
+            )
+            if m_block_max > 0:
+                m_block = cutlass.min(m_block, m_block_max - 1)
             if const_expr(should_load_Q):
+                if const_expr(load_Qt is not None):
+                    assert pipeline_Qt is not None and producer_state_Qt is not None
+                    pipeline_Qt.producer_acquire(producer_state_Qt)
+                    load_Qt(prev_m_block, producer_state=producer_state_Qt)
+                    pipeline_Qt.producer_commit(producer_state_Qt)
+                    producer_state_Qt.advance()
                 pipeline_Q.producer_acquire(producer_state_Q_LSE)
+                load_Q(m_block, producer_state=producer_state_Q_LSE)
                 pipeline_Q.producer_commit(producer_state_Q_LSE)
                 pipeline_LSE.producer_acquire(producer_state_Q_LSE)
                 with cute.arch.elect_one():
                     copy_stats(
+                        gLSE[None, m_block],
                         sLSE[None, producer_state_Q_LSE.index],
                         mbar_ptr=pipeline_LSE.producer_get_barrier(producer_state_Q_LSE),
                     )
                 producer_state_Q_LSE.advance()
             if const_expr(should_load_dO):
+                pipeline_dO.producer_acquire(
+                    producer_state_dO_dPsum,
+                    extra_tx_count=tma_copy_bytes_dO if const_expr(load_dOt is not None) else 0,
+                )
+                load_dO(m_block, producer_state=producer_state_dO_dPsum)
+                if const_expr(load_dOt is not None):
+                    load_dOt(m_block, producer_state=producer_state_dO_dPsum)
                 pipeline_dO.producer_commit(producer_state_dO_dPsum)
                 pipeline_dPsum.producer_acquire(producer_state_dO_dPsum)
                 with cute.arch.elect_one():
                     copy_stats(
+                        gdPsum[None, m_block],
                         sdPsum[None, producer_state_dO_dPsum.index],
                         mbar_ptr=pipeline_dPsum.producer_get_barrier(producer_state_dO_dPsum),
                     )
                 producer_state_dO_dPsum.advance()
+            prev_m_block = m_block
+        if const_expr(should_load_Q):
+            if const_expr(load_Qt is not None):
+                assert pipeline_Qt is not None and producer_state_Qt is not None
+                pipeline_Qt.producer_acquire(producer_state_Qt)
+                load_Qt(prev_m_block, producer_state=producer_state_Qt)
+                pipeline_Qt.producer_commit(producer_state_Qt)
+                producer_state_Qt.advance()
+            pipeline_Q.producer_tail(producer_state_Q_LSE.clone())
+            pipeline_LSE.producer_tail(producer_state_Q_LSE)
+            if const_expr(load_Qt is not None):
+                pipeline_Qt.producer_tail(producer_state_Qt)
+        if const_expr(should_load_dO):
+            pipeline_dO.producer_tail(producer_state_dO_dPsum.clone())
+            pipeline_dPsum.producer_tail(producer_state_dO_dPsum)
+    return producer_state_Q_LSE, producer_state_dO_dPsum, producer_state_Qt, producer_state_Kt
+@cute.jit
+def produce_block_sparse_q_loads_bwd_sm100_2cta_hdim192(
+    blocksparse_tensors: BlockSparseTensors,
+    batch_idx,
+    head_idx,
+    n_block,
+    # Pipeline states returned after advancing
+    producer_state_Q_Qt,
+    producer_state_O_Ot,
+    producer_state_LSE,
+    producer_state_dPsum,
+    # Pipelines
+    pipeline_Q,
+    pipeline_LSE,
+    pipeline_dO,
+    pipeline_dPsum,
+    pipeline_Qt,
+    # Load functions
+    load_K,
+    load_V,
+    load_Q,
+    load_dO,
+    load_Qt,
+    load_Kt,
+    load_dOt,
+    copy_stats,
+    # Global tensors for LSE/dPsum
+    gLSE,
+    sLSE,
+    gdPsum,
+    sdPsum,
+    # TMA copy bytes for extra_tx_count
+    tma_copy_bytes_K,
+    tma_copy_bytes_V,
+    # Subtiling factor and bounds
+    q_subtile_factor: cutlass.Constexpr = 1,
+    m_block_max: int = 0,
+):
+    """Produce SM100 backward block-sparse Q/dO loads for the hdim192 2CTA schedule."""
+    (
+        curr_q_cnt,
+        curr_q_idx,
+        curr_full_cnt,
+        curr_full_idx,
+        loop_count,
+    ) = get_block_sparse_iteration_info_bwd(
+        blocksparse_tensors, batch_idx, head_idx, n_block, q_subtile_factor, m_block_max
+    )
+    if loop_count > Int32(0):
+        first_m_block, _ = get_m_block_from_iter_bwd(
+            Int32(0),
+            curr_q_cnt,
+            curr_q_idx,
+            curr_full_cnt,
+            curr_full_idx,
+            q_subtile_factor=q_subtile_factor,
+            m_block_max=m_block_max,
+        )
+        # with q_subtile > 1 we need to guard against fully OOB regions
+        if m_block_max > 0:
+            first_m_block = cutlass.min(first_m_block, m_block_max - 1)
+        # K & Q (for S)
+        pipeline_Q.producer_acquire(
+            producer_state_Q_Qt,
+            extra_tx_count=tma_copy_bytes_K,
+        )
+        load_K(tma_bar_ptr=pipeline_Q.producer_get_barrier(producer_state_Q_Qt))
+        load_Q(first_m_block, producer_state=producer_state_Q_Qt)
+        pipeline_Q.producer_commit(producer_state_Q_Qt)
+        producer_state_Q_Qt.advance()
+        # LSE
+        pipeline_LSE.producer_acquire(producer_state_LSE)
+        with cute.arch.elect_one():
+            copy_stats(
+                gLSE[None, first_m_block],
+                sLSE[None, producer_state_LSE.index],
+                mbar_ptr=pipeline_LSE.producer_get_barrier(producer_state_LSE),
+            )
+        producer_state_LSE.advance()
+        # dOt + V, for dP.T = V @ dO.T
+        pipeline_dO.producer_acquire(
+            producer_state_O_Ot,
+            extra_tx_count=tma_copy_bytes_V,
+        )
+        load_V(tma_bar_ptr=pipeline_dO.producer_get_barrier(producer_state_O_Ot))
+        load_dOt(first_m_block, producer_state=producer_state_O_Ot)
+        pipeline_dO.producer_commit(producer_state_O_Ot)
+        producer_state_O_Ot.advance()
+        # dPsum
+        pipeline_dPsum.producer_acquire(producer_state_dPsum)
+        with cute.arch.elect_one():
+            copy_stats(
+                gdPsum[None, first_m_block],
+                sdPsum[None, producer_state_dPsum.index],
+                mbar_ptr=pipeline_dPsum.producer_get_barrier(producer_state_dPsum),
+            )
+        producer_state_dPsum.advance()
+        # Qt, for dK = dS.T @ Q
+        pipeline_Qt.producer_acquire(
+            producer_state_Q_Qt,
+            extra_tx_count=tma_copy_bytes_K,
+        )
+        load_Qt(first_m_block, producer_state=producer_state_Q_Qt)
+        load_Kt(tma_bar_ptr=pipeline_Qt.producer_get_barrier(producer_state_Q_Qt))
+        pipeline_Qt.producer_commit(producer_state_Q_Qt)
+        producer_state_Q_Qt.advance()
+        # dO, for dV = P.T @ dO
+        pipeline_dO.producer_acquire(producer_state_O_Ot)
+        load_dO(first_m_block, producer_state=producer_state_O_Ot)
+        pipeline_dO.producer_commit(producer_state_O_Ot)
+        producer_state_O_Ot.advance()
+        # 2CTA: [lse | Q | dOt | dPsum | Qt | dO]
+        for iter_idx in cutlass.range(Int32(1), loop_count, unroll=1):
+            m_block, _ = get_m_block_from_iter_bwd(
+                iter_idx,
+                curr_q_cnt,
+                curr_q_idx,
+                curr_full_cnt,
+                curr_full_idx,
+                q_subtile_factor=q_subtile_factor,
+                m_block_max=m_block_max,
+            )
+            if m_block_max > 0:
+                m_block = cutlass.min(m_block, m_block_max - 1)
+            # LSE
+            pipeline_LSE.producer_acquire(producer_state_LSE)
+            with cute.arch.elect_one():
+                copy_stats(
+                    gLSE[None, m_block],
+                    sLSE[None, producer_state_LSE.index],
+                    mbar_ptr=pipeline_LSE.producer_get_barrier(producer_state_LSE),
+                )
+            producer_state_LSE.advance()
+            # Q
+            pipeline_Q.producer_acquire(producer_state_Q_Qt)
+            load_Q(m_block, producer_state=producer_state_Q_Qt)
+            pipeline_Q.producer_commit(producer_state_Q_Qt)
+            producer_state_Q_Qt.advance()
+            # dPsum
+            pipeline_dPsum.producer_acquire(producer_state_dPsum)
+            with cute.arch.elect_one():
+                copy_stats(
+                    gdPsum[None, m_block],
+                    sdPsum[None, producer_state_dPsum.index],
+                    mbar_ptr=pipeline_dPsum.producer_get_barrier(producer_state_dPsum),
+                )
+            producer_state_dPsum.advance()
+            # dOt, for dP.T = V @ dO.T
+            pipeline_dO.producer_acquire(producer_state_O_Ot)
+            load_dOt(m_block, producer_state=producer_state_O_Ot)
+            pipeline_dO.producer_commit(producer_state_O_Ot)
+            producer_state_O_Ot.advance()
+            # Qt, for dK = dS.T @ Q
+            pipeline_Qt.producer_acquire(producer_state_Q_Qt)
+            load_Qt(m_block, producer_state=producer_state_Q_Qt)
+            pipeline_Qt.producer_commit(producer_state_Q_Qt)
+            producer_state_Q_Qt.advance()
+            # dO, for dV = P.T @ dO
+            pipeline_dO.producer_acquire(producer_state_O_Ot)
+            load_dO(m_block, producer_state=producer_state_O_Ot)
+            pipeline_dO.producer_commit(producer_state_O_Ot)
+            producer_state_O_Ot.advance()
+        pipeline_Q.producer_tail(producer_state_Q_Qt)
+        pipeline_LSE.producer_tail(producer_state_LSE)
+        pipeline_dO.producer_tail(producer_state_O_Ot)
+        pipeline_dPsum.producer_tail(producer_state_dPsum)
+    return producer_state_Q_Qt, producer_state_O_Ot, producer_state_LSE, producer_state_dPsum
 @cute.jit
     batch_idx,
     head_idx,
     n_block,
+    q_subtile_factor: cutlass.Constexpr = 1,
     m_block_max: int = 0,
 ):
     """Extract block-sparse iteration info for backward pass.
     Returns (curr_q_cnt, curr_q_idx, curr_full_cnt, curr_full_idx, total_count).
     """
+    q_cnt, q_idx, full_cnt, full_idx, *_ = blocksparse_tensors
     curr_q_cnt = q_cnt[batch_idx, head_idx, n_block]
     curr_q_idx = q_idx[batch_idx, head_idx, n_block, None]
     sparse_block_count = curr_q_cnt
     if const_expr(full_cnt is not None):
         sparse_block_count = sparse_block_count + curr_full_cnt
+    total_count = sparse_block_count * q_subtile_factor
     return curr_q_cnt, curr_q_idx, curr_full_cnt, curr_full_idx, total_count
     curr_q_idx: cute.Tensor,
     curr_full_cnt,
     curr_full_idx: Optional[cute.Tensor],
+    q_subtile_factor: cutlass.Constexpr = 1,
     m_block_max: int = 0,
 ):
     """Derive m_block index and is_full_block flag from iteration index.
         - m_block: The actual Q-tile block index
         - is_full_block: True if this is a full block (no mask_mod needed)
     """
+    sparse_iter_idx = iter_idx // q_subtile_factor
+    subtile_offset = iter_idx % q_subtile_factor
     sparse_m_block = Int32(0)
     is_full_block = False
     else:
         sparse_m_block = curr_q_idx[sparse_iter_idx]
+    return sparse_m_block * q_subtile_factor + subtile_offset, is_full_block
 @cute.jit
     else:
         pipeline_Q.producer_acquire(producer_state_Q)
     load_Q(m_block, producer_state=producer_state_Q)
+    if load_kv:
+        # Wait for bwd preprocess to finish writing LSE and dPsum, and zeroing dQaccum. Same
+        # point in the load order as the dense arm's wait in flash_bwd_sm90.load(): after this
+        # block's K/Q loads, before the first load that reads preprocess output.
+        # load_kv is the tile's first-block marker -- the caller flips kv_loaded in the same
+        # branch that loads a block -- so this runs once per tile, like the dense arm's.
+        cute.arch.griddepcontrol_wait()
     load_LSE(m_block, producer_state=producer_state_Q)
     producer_state_dO_cur = (
     tma_copy_bytes_K,
     tma_copy_bytes_V,
     Q_stage_eq_dO_stage: cutlass.Constexpr,
+    q_subtile_factor: cutlass.Constexpr,
     m_block_max: int,
 ):
     """SM90 backward block sparse loading with separate partial/full loops.
     Returns updated (producer_state_Q, producer_state_dO).
     """
+    q_cnt, q_idx, full_cnt, full_idx, *_ = blocksparse_tensors
     curr_q_cnt = q_cnt[batch_idx, head_idx, n_block]
     curr_q_idx = q_idx[batch_idx, head_idx, n_block, None]
     kv_loaded = False
+    for iter_idx in cutlass.range(curr_q_cnt * q_subtile_factor, unroll=1):
+        sparse_idx = iter_idx // q_subtile_factor
+        subtile_offset = iter_idx % q_subtile_factor
+        m_block = curr_q_idx[sparse_idx] * q_subtile_factor + subtile_offset
         if m_block < m_block_max:
             producer_state_Q, producer_state_dO = _load_q_do_block_sm90(
             kv_loaded = True
     if const_expr(full_cnt is not None):
+        for iter_idx in cutlass.range(curr_full_cnt * q_subtile_factor, unroll=1):
+            sparse_idx = iter_idx // q_subtile_factor
+            subtile_offset = iter_idx % q_subtile_factor
+            m_block = curr_full_idx[sparse_idx] * q_subtile_factor + subtile_offset
             if m_block < m_block_max:
                 producer_state_Q, producer_state_dO = _load_q_do_block_sm90(
     thr_mma_SdP,
     score_mod_fn=None,
     score_mod_bwd_fn=None,
+    q_subtile_factor: cutlass.Constexpr = 1,
     m_block_max: int = 0,
+    aux_data: AuxData = AuxData(),
     fastdiv_mods=(None, None),
 ):
     """SM90 backward block sparse MMA consumption with separate partial/full loops.
     Returns updated (consumer_state_Q, consumer_state_dO).
     """
+    q_cnt, q_idx, full_cnt, full_idx, *_ = blocksparse_tensors
     curr_q_cnt = q_cnt[batch_idx, head_idx, n_block]
     curr_q_idx = q_idx[batch_idx, head_idx, n_block, None]
         mask_causal=is_causal,
         mask_local=is_local,
         mask_mod=mask_mod,
+        aux_data=aux_data,
         fastdiv_mods=fastdiv_mods,
     )
         mask_seqlen=True,
         mask_causal=is_causal,
         mask_local=is_local,
+        aux_data=aux_data,
         fastdiv_mods=fastdiv_mods,
     )
+    for iter_idx in cutlass.range(curr_q_cnt * q_subtile_factor, unroll=1):
+        sparse_idx = iter_idx // q_subtile_factor
+        subtile_offset = iter_idx % q_subtile_factor
+        m_block = curr_q_idx[sparse_idx] * q_subtile_factor + subtile_offset
         if m_block < m_block_max:
             consumer_state_Q, consumer_state_dO = mma_one_m_block_fn(
             dKV_accumulate = True
     if const_expr(full_cnt is not None):
+        for iter_idx in cutlass.range(curr_full_cnt * q_subtile_factor, unroll=1):
+            sparse_idx = iter_idx // q_subtile_factor
+            subtile_offset = iter_idx % q_subtile_factor
+            m_block = curr_full_idx[sparse_idx] * q_subtile_factor + subtile_offset
             if m_block < m_block_max:
                 consumer_state_Q, consumer_state_dO = mma_one_m_block_fn(
     m_block,
     sdQaccum: cute.Tensor,
     gdQaccum: cute.Tensor,
+    num_dQ_warp_groups: cutlass.Constexpr,
     num_threads_per_warp_group: cutlass.Constexpr,
     tma_copy_bytes_dQ,
 ):
     """Store dQaccum for a single m_block."""
+    for warp_group_idx in cutlass.range_constexpr(num_dQ_warp_groups):
+        cute.arch.cp_async_bulk_wait_group(num_dQ_warp_groups - 1 - warp_group_idx, read=True)
         cute.arch.barrier_arrive(
             barrier_id=int(NamedBarrierBwd.dQEmptyWG0) + warp_group_idx,
             number_of_threads=num_threads_per_warp_group + cute.arch.WARP_SIZE,
         )
+    for warp_group_idx in cutlass.range_constexpr(num_dQ_warp_groups):
         cute.arch.barrier(
             barrier_id=int(NamedBarrierBwd.dQFullWG0) + warp_group_idx,
             number_of_threads=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,
                 tma_copy_bytes_dQ,
             )
         cute.arch.cp_async_bulk_commit_group()
     n_block,
     sdQaccum: cute.Tensor,
     gdQaccum: cute.Tensor,
+    q_subtile_factor: cutlass.Constexpr,
     m_block_max: int,
+    num_dQ_warp_groups: cutlass.Constexpr,
     num_threads_per_warp_group: cutlass.Constexpr,
     tma_copy_bytes_dQ,
 ):
     Iterates partial blocks first, then full blocks, matching producer/consumer order.
     """
+    q_cnt, q_idx, full_cnt, full_idx, *_ = blocksparse_tensors
     curr_q_cnt = q_cnt[batch_idx, head_idx, n_block]
     curr_q_idx = q_idx[batch_idx, head_idx, n_block, None]
         curr_full_cnt = Int32(0)
         curr_full_idx = None
+    for iter_idx in cutlass.range(curr_q_cnt * q_subtile_factor, unroll=1):
+        sparse_idx = iter_idx // q_subtile_factor
+        subtile_offset = iter_idx % q_subtile_factor
+        m_block = curr_q_idx[sparse_idx] * q_subtile_factor + subtile_offset
         if m_block < m_block_max:
             _store_one_dQaccum_sm90(
                 m_block,
                 sdQaccum,
                 gdQaccum,
+                num_dQ_warp_groups,
                 num_threads_per_warp_group,
                 tma_copy_bytes_dQ,
             )
     if const_expr(full_cnt is not None):
+        for iter_idx in cutlass.range(curr_full_cnt * q_subtile_factor, unroll=1):
+            sparse_idx = iter_idx // q_subtile_factor
+            subtile_offset = iter_idx % q_subtile_factor
+            m_block = curr_full_idx[sparse_idx] * q_subtile_factor + subtile_offset
             if m_block < m_block_max:
                 _store_one_dQaccum_sm90(
                     m_block,
                     sdQaccum,
                     gdQaccum,
+                    num_dQ_warp_groups,
                     num_threads_per_warp_group,
                     tma_copy_bytes_dQ,
                 )

build/torch-cuda/block_sparsity.py CHANGED Viewed

@@ -17,13 +17,23 @@ def ceildiv(a: int, b: int) -> int:
 class BlockSparseTensors(NamedTuple):
     mask_block_cnt: cute.Tensor
     mask_block_idx: cute.Tensor
-    full_block_cnt: cute.Tensor | None
-    full_block_idx: cute.Tensor | None
     def __new_from_mlir_values__(self, values):
-        if len(values) == 2:
-            values = (*values, None, None)
-        return BlockSparseTensors(*values)
 class BlockSparseTensorsTorch(NamedTuple):
@@ -31,7 +41,184 @@ class BlockSparseTensorsTorch(NamedTuple):
     mask_block_idx: torch.Tensor
     full_block_cnt: torch.Tensor | None = None
     full_block_idx: torch.Tensor | None = None
     block_size: tuple[int, int] | None = None
 def _expand_sparsity_tensor(
@@ -61,8 +248,8 @@ def _check_and_expand_block(
     name: str,
     cnt: torch.Tensor | None,
     idx: torch.Tensor | None,
-    expected_count_shape: Tuple[int, int, int],
-    expected_index_shape: Tuple[int, int, int, int],
     context: str | None,
     hint: str | Callable[[], str] | None,
 ) -> Tuple[torch.Tensor | None, torch.Tensor | None]:
@@ -81,12 +268,37 @@ def _check_and_expand_block(
     expanded_cnt = _expand_sparsity_tensor(
         cnt, expected_count_shape, f"{name}_block_cnt", context, hint
     )
     expanded_idx = _expand_sparsity_tensor(
         idx, expected_index_shape, f"{name}_block_idx", context, hint
     )
     return expanded_cnt, expanded_idx
 def get_block_sparse_expected_shapes(
     batch_size: int,
     num_head: int,
@@ -124,7 +336,7 @@ def infer_block_sparse_expected_shapes(
     Expectations:
     - mask_block_cnt is (B, H, M) and mask_block_idx is (B, H, M, N).
     - Batch/head dims may be 1 for broadcast, or match the requested sizes.
-    - sparse_block_size_kv must match tile_n.
     - sparse_block_size_q must be a multiple of q_stage * tile_m.
     - If sparse_block_size_q is omitted and seqlen_q/num_m_blocks is ambiguous,
       the caller must provide block_size to disambiguate. TODO will make this required in a future PR.
@@ -133,24 +345,23 @@ def infer_block_sparse_expected_shapes(
     base_n_block = n_block_size
     if sparse_block_size_kv is None:
         sparse_block_size_kv = base_n_block
-    if sparse_block_size_kv != base_n_block:
-        raise ValueError(f"Block sparse tensors{context} require BLOCK_SIZE_KV={base_n_block}.")
     if tensors.mask_block_idx is None:
         raise ValueError("mask_block_cnt and mask_block_idx must be provided for block sparsity.")
     num_m_blocks = tensors.mask_block_idx.shape[2]
     if sparse_block_size_q is None:
-        min_block_size = ceildiv(seqlen_q, num_m_blocks)
-        if num_m_blocks == 1:
-            max_block_size = seqlen_q
-        else:
-            max_block_size = (seqlen_q - 1) // (num_m_blocks - 1)
-        if max_block_size != min_block_size and base_m_block != 1:
             raise ValueError(
                 f"Block sparse tensors{context} require explicit sparse_block_size[0] "
                 f"to disambiguate block size for seqlen_q={seqlen_q} and num_m_blocks={num_m_blocks}."
             )
-        sparse_block_size_q = min_block_size
     if sparse_block_size_q % base_m_block != 0:
         raise ValueError(
@@ -186,9 +397,11 @@ def infer_block_sparse_expected_shapes(
             raise ValueError(f"Block sparse tensors{context} {dim_name} dim must be {tgt} or 1.")
     if mask_block_cnt.shape[2] != mask_block_idx.shape[2]:
         raise ValueError(f"Block sparse tensors{context} must share the same m-block dimension.")
-    if mask_block_idx.shape[3] != expected_n_blocks:
         raise ValueError(
-            f"Block sparse tensors{context} n-block dimension must be {expected_n_blocks}."
         )
     if expected_m_blocks != num_m_blocks:
         raise ValueError(
@@ -206,17 +419,19 @@ def get_block_sparse_expected_shapes_bwd(
     seqlen_k: int,
     m_block_size: int,
     n_block_size: int,
-    subtile_factor: int,
 ) -> Tuple[Tuple[int, int, int], Tuple[int, int, int, int]]:
     """Return (expected_count_shape, expected_index_shape) for backward block sparse normalization.
     Backward uses Q-direction indexing (transposed from forward), where shapes are
     indexed by N-blocks first, then M-blocks. The sparse_block_size_q is determined
-    by subtile_factor * m_block_size.
     """
-    sparse_block_size_q = subtile_factor * m_block_size
     expected_m_blocks = ceildiv(seqlen_q, sparse_block_size_q)
-    expected_n_blocks = ceildiv(seqlen_k, n_block_size)
     expected_count_shape = (batch_size, num_head, expected_n_blocks)
     expected_index_shape = (batch_size, num_head, expected_n_blocks, expected_m_blocks)
     return expected_count_shape, expected_index_shape
@@ -225,8 +440,8 @@ def get_block_sparse_expected_shapes_bwd(
 def normalize_block_sparse_tensors(
     tensors: BlockSparseTensorsTorch,
     *,
-    expected_count_shape: Tuple[int, int, int],
-    expected_index_shape: Tuple[int, int, int, int],
     context: str | None = None,
     hint: str | Callable[[], str] | None = None,
 ) -> BlockSparseTensorsTorch:
@@ -257,12 +472,39 @@ def normalize_block_sparse_tensors(
     if full_cnt is not None and mask_cnt.device != full_cnt.device:
         raise ValueError("All block sparse tensors must be on the same device")
     return BlockSparseTensorsTorch(
         mask_block_cnt=mask_cnt,
         mask_block_idx=mask_idx,
         full_block_cnt=full_cnt,
         full_block_idx=full_idx,
         block_size=tensors.block_size,
     )
@@ -294,6 +536,8 @@ def get_block_sparse_broadcast_pattern(
         tensors.mask_block_idx,
         tensors.full_block_cnt,
         tensors.full_block_idx,
     ):
         if tensor is not None:
             patterns.append(get_broadcast_dims(tensor))
@@ -302,6 +546,15 @@ def get_block_sparse_broadcast_pattern(
     return tuple(patterns)
 def normalize_block_sparse_config(
     tensors: BlockSparseTensorsTorch,
     *,
@@ -311,40 +564,67 @@ def normalize_block_sparse_config(
     seqlen_k: int,
     block_size: tuple[int, int],
     q_stage: int,
-) -> tuple[BlockSparseTensorsTorch, Tuple[Tuple[bool, ...], ...] | None, int]:
     m_block_size, n_block_size = block_size
     if tensors.block_size is None:
-        sparse_block_size_q, sparse_block_size_kv = q_stage * m_block_size, n_block_size
     else:
         sparse_block_size_q, sparse_block_size_kv = tensors.block_size
-    if sparse_block_size_kv != n_block_size:
         raise ValueError(
             f"Block sparsity requires sparse_block_size[1]={n_block_size} to match tile_n."
         )
-    expected_count_shape, expected_index_shape, q_subtile_factor = (
-        infer_block_sparse_expected_shapes(
-            tensors,
-            batch_size=batch_size,
-            num_head=num_head,
-            seqlen_q=seqlen_q,
-            seqlen_k=seqlen_k,
-            m_block_size=m_block_size,
-            n_block_size=n_block_size,
-            q_stage=q_stage,
-            context="forward",
-            sparse_block_size_q=sparse_block_size_q,
-            sparse_block_size_kv=sparse_block_size_kv,
         )
-    )
     normalized_tensors = normalize_block_sparse_tensors(
         tensors,
         expected_count_shape=expected_count_shape,
         expected_index_shape=expected_index_shape,
     )
-    return (
-        normalized_tensors,
-        get_block_sparse_broadcast_pattern(normalized_tensors),
-        q_subtile_factor,
     )
@@ -356,21 +636,24 @@ def normalize_block_sparse_config_bwd(
     seqlen_q: int,
     seqlen_k: int,
     block_size: tuple[int, int],
-    subtile_factor: int,
 ) -> tuple[BlockSparseTensorsTorch, Tuple[Tuple[bool, ...], ...] | None]:
     m_block_size, n_block_size = block_size
     if tensors.block_size is None:
-        sparse_block_size_q, sparse_block_size_kv = subtile_factor * m_block_size, n_block_size
     else:
         sparse_block_size_q, sparse_block_size_kv = tensors.block_size
-    if sparse_block_size_q != subtile_factor * m_block_size:
         raise ValueError(
-            f"Block sparsity expects sparse_block_size_q={subtile_factor * m_block_size} "
-            f"for subtile_factor={subtile_factor}."
         )
-    if sparse_block_size_kv != n_block_size:
         raise ValueError(
-            f"Block sparsity expects sparse_block_size[1]={n_block_size} to match tile_n."
         )
     expected_count_shape, expected_index_shape = get_block_sparse_expected_shapes_bwd(
         batch_size,
@@ -379,7 +662,8 @@ def normalize_block_sparse_config_bwd(
         seqlen_k,
         m_block_size,
         n_block_size,
-        subtile_factor,
     )
     normalized_tensors = normalize_block_sparse_tensors(
         tensors,
@@ -389,7 +673,7 @@ def normalize_block_sparse_config_bwd(
         hint=lambda: (
             f"Backward expects Q-direction block-sparse tensors (q_mask_cnt/q_mask_idx, "
             f"and optionally full_q_cnt/full_q_idx). Regenerate the backward BlockMask with "
-            f"BLOCK_SIZE=({subtile_factor * m_block_size}, {n_block_size})."
         ),
     )
     return normalized_tensors, get_block_sparse_broadcast_pattern(normalized_tensors)
@@ -401,29 +685,27 @@ def to_cute_block_sparse_tensors(
     """Convert torch block sparsity tensors to CuTe tensors, optionally for tvm ffi"""
     if not is_block_sparsity_enabled(tensors):
         return None
-    (
-        mask_block_cnt,
-        mask_block_idx,
-        full_block_cnt,
-        full_block_idx,
-        *_,
-    ) = tensors
-    (
-        mask_block_cnt_tensor,
-        mask_block_idx_tensor,
-    ) = [
         to_cute_tensor(t, assumed_align=4, leading_dim=-1, enable_tvm_ffi=enable_tvm_ffi)
-        for t in (mask_block_cnt, mask_block_idx)
     ]
-    (
-        full_block_cnt_tensor,
-        full_block_idx_tensor,
-    ) = [
         to_cute_tensor(t, assumed_align=4, leading_dim=-1, enable_tvm_ffi=enable_tvm_ffi)
         if t is not None
         else None
-        for t in (full_block_cnt, full_block_idx)
     ]
     return BlockSparseTensors(
@@ -431,6 +713,10 @@ def to_cute_block_sparse_tensors(
         mask_block_idx_tensor,
         full_block_cnt_tensor,
         full_block_idx_tensor,
     )

 class BlockSparseTensors(NamedTuple):
     mask_block_cnt: cute.Tensor
     mask_block_idx: cute.Tensor
+    full_block_cnt: cute.Tensor | None = None
+    full_block_idx: cute.Tensor | None = None
+    cu_total_m_blocks: cute.Tensor | None = None
+    cu_block_idx_offsets: cute.Tensor | None = None
+    dq_write_order: cute.Tensor | None = None
+    dq_write_order_full: cute.Tensor | None = None
     def __new_from_mlir_values__(self, values):
+        new_fields = []
+        idx = 0
+        for original in self:
+            if original is None:
+                new_fields.append(None)
+            else:
+                new_fields.append(values[idx])
+                idx += 1
+        return BlockSparseTensors(*new_fields)
 class BlockSparseTensorsTorch(NamedTuple):
     mask_block_idx: torch.Tensor
     full_block_cnt: torch.Tensor | None = None
     full_block_idx: torch.Tensor | None = None
+    cu_total_m_blocks: torch.Tensor | None = None
+    cu_block_idx_offsets: torch.Tensor | None = None
     block_size: tuple[int, int] | None = None
+    dq_write_order: torch.Tensor | None = None
+    dq_write_order_full: torch.Tensor | None = None
+    spt: bool | None = None
+def _ordered_to_dense_simple(
+    num_blocks: torch.Tensor,
+    indices: torch.Tensor,
+    num_cols: int,
+) -> torch.Tensor:
+    """Convert ordered sparse representation to dense binary matrix.
+    Args:
+        num_blocks: [B, H, num_rows] count of valid entries per row
+        indices: [B, H, num_rows, max_entries] column indices (valid entries packed left)
+        num_cols: total number of columns
+    Returns:
+        dense: [B, H, num_rows, num_cols] binary int32 matrix
+    """
+    B, H, num_rows, max_entries = indices.shape
+    device = indices.device
+    dense = torch.zeros(B, H, num_rows, num_cols + 1, dtype=torch.int32, device=device)
+    col_range = torch.arange(max_entries, device=device)
+    valid = col_range[None, None, None, :] < num_blocks[:, :, :, None]
+    safe_indices = torch.where(valid, indices.long(), num_cols)
+    row_idx = torch.arange(num_rows, device=device)[None, None, :, None].expand_as(indices)
+    b_idx = torch.arange(B, device=device)[:, None, None, None].expand_as(indices)
+    h_idx = torch.arange(H, device=device)[None, :, None, None].expand_as(indices)
+    dense[b_idx, h_idx, row_idx, safe_indices] = 1
+    return dense[:, :, :, :num_cols]
+def compute_dq_write_order(
+    fwd_mask_cnt: torch.Tensor,
+    fwd_mask_idx: torch.Tensor,
+    fwd_full_cnt: torch.Tensor | None,
+    fwd_full_idx: torch.Tensor | None,
+    bwd_mask_cnt: torch.Tensor,
+    bwd_mask_idx: torch.Tensor,
+    bwd_full_cnt: torch.Tensor | None,
+    bwd_full_idx: torch.Tensor | None,
+    spt: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Compute dQ write-order metadata for deterministic block-sparse backward.
+    For each (n_block, i) in the backward iteration, computes the semaphore
+    lock value: the rank of n_block in the combined (partial + full) sorted
+    contributor list for the target m_block.
+    Lock values are assigned in ascending n_block order (or descending if spt=True)
+    to guarantee deadlock-freedom with the CTA scheduling order.
+    Args:
+        fwd_mask_cnt: [B, H, num_m_blocks] partial contributor counts per m_block
+        fwd_mask_idx: [B, H, num_m_blocks, max_kv] partial contributor n_block indices (ascending)
+        fwd_full_cnt: [B, H, num_m_blocks] full contributor counts per m_block (optional)
+        fwd_full_idx: [B, H, num_m_blocks, max_kv] full contributor n_block indices (optional)
+        bwd_mask_cnt: [B, H, num_n_blocks] partial iteration counts per n_block
+        bwd_mask_idx: [B, H, num_n_blocks, max_q] partial iteration m_block indices
+        bwd_full_cnt: [B, H, num_n_blocks] full iteration counts per n_block (optional)
+        bwd_full_idx: [B, H, num_n_blocks, max_q] full iteration m_block indices (optional)
+        spt: if True, reverse ordering (highest n_block gets lock_value=0)
+    Returns:
+        (dq_write_order, dq_write_order_full): tensors parallel to bwd_mask_idx
+        and bwd_full_idx respectively, containing lock values.
+    """
+    device = fwd_mask_idx.device
+    B, H, num_m, max_kv_partial = fwd_mask_idx.shape
+    _, _, num_n, max_q_partial = bwd_mask_idx.shape
+    has_full = fwd_full_cnt is not None and fwd_full_idx is not None
+    dense_partial = _ordered_to_dense_simple(fwd_mask_cnt, fwd_mask_idx, num_n)
+    if has_full:
+        dense_full = _ordered_to_dense_simple(fwd_full_cnt, fwd_full_idx, num_n)
+        dense = (dense_partial + dense_full).clamp(max=1)
+    else:
+        dense = dense_partial
+    cumsum = dense.cumsum(dim=-1)
+    rank_table = (cumsum - dense).to(torch.int32)
+    if spt:
+        total_per_m = cumsum[:, :, :, -1:]
+        rank_table = (total_per_m - 1 - rank_table).to(torch.int32)
+    def _gather_write_order(bwd_idx, bwd_cnt):
+        b_i = torch.arange(B, device=device)[:, None, None, None].expand_as(bwd_idx)
+        h_i = torch.arange(H, device=device)[None, :, None, None].expand_as(bwd_idx)
+        n_i = torch.arange(bwd_idx.shape[2], device=device)[None, None, :, None].expand_as(bwd_idx)
+        m_vals = bwd_idx.long().clamp(0, num_m - 1)
+        return rank_table[b_i, h_i, m_vals, n_i].to(torch.int32)
+    dq_write_order = _gather_write_order(bwd_mask_idx, bwd_mask_cnt)
+    dq_write_order_full = None
+    if has_full and bwd_full_cnt is not None and bwd_full_idx is not None:
+        dq_write_order_full = _gather_write_order(bwd_full_idx, bwd_full_cnt)
+    return dq_write_order, dq_write_order_full
+def compute_dq_write_order_from_block_mask(
+    block_mask,
+    spt: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    (
+        _seq_q,
+        _seq_k,
+        kv_mask_cnt,
+        kv_mask_idx,
+        full_kv_cnt,
+        full_kv_idx,
+        q_mask_cnt,
+        q_mask_idx,
+        full_q_cnt,
+        full_q_idx,
+        *_,
+    ) = block_mask.as_tuple()
+    return compute_dq_write_order(
+        kv_mask_cnt,
+        kv_mask_idx,
+        full_kv_cnt,
+        full_kv_idx,
+        q_mask_cnt,
+        q_mask_idx,
+        full_q_cnt,
+        full_q_idx,
+        spt=spt,
+    )
+def get_sparse_q_block_size(
+    tensors: BlockSparseTensorsTorch | None,
+    seqlen_q: int,
+) -> int | None:
+    """Return the Q sparse block size, or None when sparsity is unset or ambiguous."""
+    if tensors is None:
+        return None
+    if tensors.block_size is not None:
+        return tensors.block_size[0]
+    num_m_blocks = tensors.mask_block_idx.shape[2]
+    min_block_size = ceildiv(seqlen_q, num_m_blocks)
+    max_block_size = seqlen_q if num_m_blocks == 1 else (seqlen_q - 1) // (num_m_blocks - 1)
+    if min_block_size != max_block_size:
+        return None
+    return min_block_size
+def get_kv_subtile_factor(
+    block_sparse_tensors: BlockSparseTensorsTorch | None,
+    n_block_size: int,
+) -> int:
+    """Return the number of physical KV tiles covered by one sparse KV block."""
+    if block_sparse_tensors is None or block_sparse_tensors.block_size is None:
+        return 1
+    sparse_block_size_kv = block_sparse_tensors.block_size[1]
+    if sparse_block_size_kv % n_block_size != 0:
+        raise ValueError(
+            "Block sparsity expects sparse_block_size[1] "
+            f"to be a multiple of tile_n={n_block_size}; got {sparse_block_size_kv}."
+        )
+    return sparse_block_size_kv // n_block_size
+def block_sparse_bwd_supports_2cta(
+    block_sparse_tensors: BlockSparseTensorsTorch | None,
+    n_block_size: int,
+) -> bool:
+    """Return whether sparse KV metadata constrains backward away from 2CTA."""
+    if block_sparse_tensors is None:
+        return True
+    return get_kv_subtile_factor(block_sparse_tensors, n_block_size) % 2 == 0
 def _expand_sparsity_tensor(
     name: str,
     cnt: torch.Tensor | None,
     idx: torch.Tensor | None,
+    expected_count_shape: Tuple[int, ...],
+    expected_index_shape: Tuple[int, ...],
     context: str | None,
     hint: str | Callable[[], str] | None,
 ) -> Tuple[torch.Tensor | None, torch.Tensor | None]:
     expanded_cnt = _expand_sparsity_tensor(
         cnt, expected_count_shape, f"{name}_block_cnt", context, hint
     )
+    # [Note] Allow Compact block sparse indices
+    # Allow the last dimension (n_blocks) of idx to be <= expected, since
+    # FA4 only accesses indices 0..cnt-1 per query tile. This enables compact
+    # index tensors that avoid O(N^2) memory at long sequence lengths.
+    if idx.ndim == 4 and idx.shape[3] <= expected_index_shape[3]:
+        expected_index_shape = (*expected_index_shape[:3], idx.shape[3])
     expanded_idx = _expand_sparsity_tensor(
         idx, expected_index_shape, f"{name}_block_idx", context, hint
     )
     return expanded_cnt, expanded_idx
+def _check_and_expand_metadata_tensor(
+    name: str,
+    tensor: torch.Tensor | None,
+    expected_shape: Tuple[int, ...],
+    context: str | None,
+    hint: str | Callable[[], str] | None,
+    device: torch.device,
+) -> torch.Tensor | None:
+    if tensor is None:
+        return None
+    if tensor.dtype != torch.int32:
+        raise ValueError(f"{name} must have dtype torch.int32")
+    if tensor.device != device:
+        raise ValueError(f"{name} must be on the same device as block sparse tensors")
+    if not tensor.is_cuda:
+        raise ValueError(f"{name} must live on CUDA")
+    return _expand_sparsity_tensor(tensor, expected_shape, name, context, hint)
 def get_block_sparse_expected_shapes(
     batch_size: int,
     num_head: int,
     Expectations:
     - mask_block_cnt is (B, H, M) and mask_block_idx is (B, H, M, N).
     - Batch/head dims may be 1 for broadcast, or match the requested sizes.
+    - sparse_block_size_kv must be a multiple of tile_n.
     - sparse_block_size_q must be a multiple of q_stage * tile_m.
     - If sparse_block_size_q is omitted and seqlen_q/num_m_blocks is ambiguous,
       the caller must provide block_size to disambiguate. TODO will make this required in a future PR.
     base_n_block = n_block_size
     if sparse_block_size_kv is None:
         sparse_block_size_kv = base_n_block
+    if sparse_block_size_kv % base_n_block != 0:
+        raise ValueError(
+            f"Block sparse tensors{context} require BLOCK_SIZE_KV to be a multiple of {base_n_block}."
+        )
     if tensors.mask_block_idx is None:
         raise ValueError("mask_block_cnt and mask_block_idx must be provided for block sparsity.")
     num_m_blocks = tensors.mask_block_idx.shape[2]
     if sparse_block_size_q is None:
+        sparse_block_size_q = get_sparse_q_block_size(tensors, seqlen_q)
+        if sparse_block_size_q is None and base_m_block != 1:
             raise ValueError(
                 f"Block sparse tensors{context} require explicit sparse_block_size[0] "
                 f"to disambiguate block size for seqlen_q={seqlen_q} and num_m_blocks={num_m_blocks}."
             )
+        if sparse_block_size_q is None:
+            sparse_block_size_q = ceildiv(seqlen_q, num_m_blocks)
     if sparse_block_size_q % base_m_block != 0:
         raise ValueError(
             raise ValueError(f"Block sparse tensors{context} {dim_name} dim must be {tgt} or 1.")
     if mask_block_cnt.shape[2] != mask_block_idx.shape[2]:
         raise ValueError(f"Block sparse tensors{context} must share the same m-block dimension.")
+    # [Note] Allow Compact block sparse indices: FA4 only accesses indices 0..cnt-1
+    # per query tile, so idx.shape[3] can be <= expected_n_blocks.
+    if mask_block_idx.shape[3] > expected_n_blocks:
         raise ValueError(
+            f"Block sparse tensors{context} n-block dimension must be <= {expected_n_blocks}."
         )
     if expected_m_blocks != num_m_blocks:
         raise ValueError(
     seqlen_k: int,
     m_block_size: int,
     n_block_size: int,
+    q_subtile_factor: int,
+    kv_subtile_factor: int = 1,
 ) -> Tuple[Tuple[int, int, int], Tuple[int, int, int, int]]:
     """Return (expected_count_shape, expected_index_shape) for backward block sparse normalization.
     Backward uses Q-direction indexing (transposed from forward), where shapes are
     indexed by N-blocks first, then M-blocks. The sparse_block_size_q is determined
+    by q_subtile_factor * m_block_size.
     """
+    sparse_block_size_q = q_subtile_factor * m_block_size
+    sparse_block_size_kv = kv_subtile_factor * n_block_size
     expected_m_blocks = ceildiv(seqlen_q, sparse_block_size_q)
+    expected_n_blocks = ceildiv(seqlen_k, sparse_block_size_kv)
     expected_count_shape = (batch_size, num_head, expected_n_blocks)
     expected_index_shape = (batch_size, num_head, expected_n_blocks, expected_m_blocks)
     return expected_count_shape, expected_index_shape
 def normalize_block_sparse_tensors(
     tensors: BlockSparseTensorsTorch,
     *,
+    expected_count_shape: Tuple[int, ...],
+    expected_index_shape: Tuple[int, ...],
     context: str | None = None,
     hint: str | Callable[[], str] | None = None,
 ) -> BlockSparseTensorsTorch:
     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")
+    dq_write_order = _check_and_expand_metadata_tensor(
+        "dq_write_order",
+        tensors.dq_write_order,
+        tuple(mask_idx.shape),
+        context,
+        hint,
+        mask_cnt.device,
+    )
+    dq_write_order_full = _check_and_expand_metadata_tensor(
+        "dq_write_order_full",
+        tensors.dq_write_order_full,
+        tuple(full_idx.shape) if full_idx is not None else expected_index_shape,
+        context,
+        hint,
+        mask_cnt.device,
+    )
+    spt = tensors.spt
+    if spt is not None and not isinstance(spt, bool):
+        raise ValueError("spt must be a bool when provided")
+    if spt is not None and dq_write_order is None:
+        raise ValueError("spt requires dq_write_order to be provided")
     return BlockSparseTensorsTorch(
         mask_block_cnt=mask_cnt,
         mask_block_idx=mask_idx,
         full_block_cnt=full_cnt,
         full_block_idx=full_idx,
+        cu_total_m_blocks=tensors.cu_total_m_blocks,
+        cu_block_idx_offsets=tensors.cu_block_idx_offsets,
         block_size=tensors.block_size,
+        dq_write_order=dq_write_order,
+        dq_write_order_full=dq_write_order_full,
+        spt=spt,
     )
         tensors.mask_block_idx,
         tensors.full_block_cnt,
         tensors.full_block_idx,
+        tensors.dq_write_order,
+        tensors.dq_write_order_full,
     ):
         if tensor is not None:
             patterns.append(get_broadcast_dims(tensor))
     return tuple(patterns)
+class NormalizedBlockSparseConfig(NamedTuple):
+    """Result of validating and normalizing a user block-sparse config."""
+    tensors: BlockSparseTensorsTorch
+    broadcast_pattern: Tuple[Tuple[bool, ...], ...] | None
+    q_subtile_factor: int
+    kv_subtile_factor: int
 def normalize_block_sparse_config(
     tensors: BlockSparseTensorsTorch,
     *,
     seqlen_k: int,
     block_size: tuple[int, int],
     q_stage: int,
+    allow_kv_subtile: bool = False,
+) -> NormalizedBlockSparseConfig:
+    """Validate the block-sparse config, infer expected shapes, and normalize.
+    Handles both fixed-length (3D `[B, H, M]` / 4D `[B, H, M, N]`) and varlen
+    (2D `[H, total_m_blocks]` / `[H, total_n_blocks]`) layouts. Varlen is
+    detected by `tensors.cu_total_m_blocks is not None` and forces
+    `q_subtile_factor == 1` (TODO: potentially remove this restriction).
+    """
     m_block_size, n_block_size = block_size
     if tensors.block_size is None:
+        sparse_block_size_q, sparse_block_size_kv = None, n_block_size
     else:
         sparse_block_size_q, sparse_block_size_kv = tensors.block_size
+    if sparse_block_size_kv % n_block_size != 0:
+        raise ValueError(
+            f"Block sparsity requires sparse_block_size[1] to be a multiple of tile_n={n_block_size}."
+        )
+    kv_subtile_factor = sparse_block_size_kv // n_block_size
+    if kv_subtile_factor != 1 and not allow_kv_subtile:
         raise ValueError(
             f"Block sparsity requires sparse_block_size[1]={n_block_size} to match tile_n."
         )
+    if tensors.cu_total_m_blocks is not None:
+        base_m_block = q_stage * m_block_size
+        if sparse_block_size_q is not None and sparse_block_size_q != base_m_block:
+            raise ValueError(
+                f"Varlen block sparsity requires sparse_block_size[0]={base_m_block} "
+                f"(= q_stage * tile_m); got {sparse_block_size_q}."
+            )
+        total_m_blocks = tensors.mask_block_cnt.shape[-1]
+        total_n_blocks = tensors.mask_block_idx.shape[-1]
+        expected_count_shape = (num_head, total_m_blocks)
+        expected_index_shape = (num_head, total_n_blocks)
+        q_subtile_factor = 1
+    else:
+        expected_count_shape, expected_index_shape, q_subtile_factor = (
+            infer_block_sparse_expected_shapes(
+                tensors,
+                batch_size=batch_size,
+                num_head=num_head,
+                seqlen_q=seqlen_q,
+                seqlen_k=seqlen_k,
+                m_block_size=m_block_size,
+                n_block_size=n_block_size,
+                q_stage=q_stage,
+                context="forward",
+                sparse_block_size_q=sparse_block_size_q,
+                sparse_block_size_kv=sparse_block_size_kv,
+            )
         )
     normalized_tensors = normalize_block_sparse_tensors(
         tensors,
         expected_count_shape=expected_count_shape,
         expected_index_shape=expected_index_shape,
     )
+    return NormalizedBlockSparseConfig(
+        tensors=normalized_tensors,
+        broadcast_pattern=get_block_sparse_broadcast_pattern(normalized_tensors),
+        q_subtile_factor=q_subtile_factor,
+        kv_subtile_factor=kv_subtile_factor,
     )
     seqlen_q: int,
     seqlen_k: int,
     block_size: tuple[int, int],
+    q_subtile_factor: int,
+    kv_subtile_factor: int = 1,
 ) -> tuple[BlockSparseTensorsTorch, Tuple[Tuple[bool, ...], ...] | None]:
     m_block_size, n_block_size = block_size
     if tensors.block_size is None:
+        sparse_block_size_q, sparse_block_size_kv = q_subtile_factor * m_block_size, n_block_size
     else:
         sparse_block_size_q, sparse_block_size_kv = tensors.block_size
+    if sparse_block_size_q != q_subtile_factor * m_block_size:
         raise ValueError(
+            f"Block sparsity expects sparse_block_size_q={q_subtile_factor * m_block_size} "
+            f"for q_subtile_factor={q_subtile_factor}."
         )
+    expected_sparse_block_size_kv = kv_subtile_factor * n_block_size
+    if sparse_block_size_kv != expected_sparse_block_size_kv:
         raise ValueError(
+            f"Block sparsity expects sparse_block_size[1]={expected_sparse_block_size_kv} "
+            f"for kv_subtile_factor={kv_subtile_factor}."
         )
     expected_count_shape, expected_index_shape = get_block_sparse_expected_shapes_bwd(
         batch_size,
         seqlen_k,
         m_block_size,
         n_block_size,
+        q_subtile_factor,
+        kv_subtile_factor,
     )
     normalized_tensors = normalize_block_sparse_tensors(
         tensors,
         hint=lambda: (
             f"Backward expects Q-direction block-sparse tensors (q_mask_cnt/q_mask_idx, "
             f"and optionally full_q_cnt/full_q_idx). Regenerate the backward BlockMask with "
+            f"BLOCK_SIZE=({q_subtile_factor * m_block_size}, {expected_sparse_block_size_kv})."
         ),
     )
     return normalized_tensors, get_block_sparse_broadcast_pattern(normalized_tensors)
     """Convert torch block sparsity tensors to CuTe tensors, optionally for tvm ffi"""
     if not is_block_sparsity_enabled(tensors):
         return None
+    mask_block_cnt_tensor, mask_block_idx_tensor = [
         to_cute_tensor(t, assumed_align=4, leading_dim=-1, enable_tvm_ffi=enable_tvm_ffi)
+        for t in (tensors.mask_block_cnt, tensors.mask_block_idx)
     ]
+    full_block_cnt_tensor, full_block_idx_tensor = [
+        to_cute_tensor(t, assumed_align=4, leading_dim=-1, enable_tvm_ffi=enable_tvm_ffi)
+        if t is not None
+        else None
+        for t in (tensors.full_block_cnt, tensors.full_block_idx)
+    ]
+    cu_total_m_blocks_tensor, cu_block_idx_offsets_tensor = [
+        to_cute_tensor(t, assumed_align=4, leading_dim=0, enable_tvm_ffi=enable_tvm_ffi)
+        if t is not None
+        else None
+        for t in (tensors.cu_total_m_blocks, tensors.cu_block_idx_offsets)
+    ]
+    dq_write_order_tensor, dq_write_order_full_tensor = [
         to_cute_tensor(t, assumed_align=4, leading_dim=-1, enable_tvm_ffi=enable_tvm_ffi)
         if t is not None
         else None
+        for t in (tensors.dq_write_order, tensors.dq_write_order_full)
     ]
     return BlockSparseTensors(
         mask_block_idx_tensor,
         full_block_cnt_tensor,
         full_block_idx_tensor,
+        cu_total_m_blocks_tensor,
+        cu_block_idx_offsets_tensor,
+        dq_write_order_tensor,
+        dq_write_order_full_tensor,
     )

build/torch-cuda/cache_utils.py CHANGED Viewed

@@ -1,7 +1,6 @@
 # Manage Ahead-of-Time (AOT) compiled kernels
 import fcntl
 import hashlib
-import logging
 import os
 import pickle
 import sys
@@ -18,6 +17,7 @@ import cutlass
 import cutlass.cute as cute
 import tvm_ffi
 from cutlass.cutlass_dsl import JitCompiledFunction
 # Pre-load cute DSL runtime libraries with RTLD_GLOBAL so that their symbols
 # (e.g. _cudaLibraryLoadData) are visible to .so modules loaded later via dlopen.
@@ -30,13 +30,6 @@ for _lib_path in cute.runtime.find_runtime_libraries(enable_tvm_ffi=False):
 CompileKeyType: TypeAlias = tuple[Hashable, ...]
 CallableFunction: TypeAlias = JitCompiledFunction | tvm_ffi.Function
-logger = logging.getLogger(__name__)
-_handler = logging.StreamHandler()
-_handler.setFormatter(logging.Formatter("%(asctime)s.%(msecs)03d %(levelname)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S"))
-logger.addHandler(_handler)
-logger.setLevel(logging.DEBUG)
 # Enable cache via `FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1`
 CUTE_DSL_CACHE_ENABLED: bool = os.getenv("FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED", "0") == "1"
@@ -222,13 +215,13 @@ class JITPersistentCache(JITCache):
             label=sha256_hex,
         ):
             if obj_path.exists():
-                logger.debug("Loading compiled function from disk: %s", obj_path)
                 m = cute.runtime.load_module(str(obj_path), enable_tvm_ffi=True)
                 fn = getattr(m, self.EXPORT_FUNCTION_PREFIX)
                 JITCache.__setitem__(self, key, fn)
                 return True
             else:
-                logger.debug("Cache miss on disk for key hash %s", sha256_hex)
         return False
     def _try_export_to_storage(self, key: CompileKeyType, fn: JitCompiledFunction) -> None:
@@ -243,14 +236,14 @@ class JITPersistentCache(JITCache):
             obj_path = self.cache_path / f"{sha256_hex}.o"
             if obj_path.exists():
                 # Another process already exported.
-                logger.debug("Skipping export, already on disk: %s", obj_path)
                 return
-            logger.debug("Exporting compiled function to disk: %s", obj_path)
             fn.export_to_c(
                 object_file_path=str(obj_path),
                 function_name=self.EXPORT_FUNCTION_PREFIX,
             )
-            logger.debug("Successfully exported compiled function to disk: %s", obj_path)
     def _key_to_hash(self, key: CompileKeyType) -> str:
         return hashlib.sha256(pickle.dumps(key)).hexdigest()
@@ -262,7 +255,7 @@ class JITPersistentCache(JITCache):
         """
         Not only clear the in-memory cache. Also purge persistent compilation cache.
         """
-        logger.debug("Clearing persistent cache at %s", self.cache_path)
         super().clear()
         for child in self.cache_path.iterdir():
             child.unlink()
@@ -281,8 +274,8 @@ def get_jit_cache(name: str | None = None) -> JITCache:
         path = get_cache_path() / _compute_source_fingerprint()
         if name:
             path = path / name
-        logger.debug("Creating persistent JIT cache at %s", path)
         return JITPersistentCache(path)
     else:
-        logger.debug("Persistent cache disabled, using in-memory JIT cache")
         return JITCache()

 # Manage Ahead-of-Time (AOT) compiled kernels
 import fcntl
 import hashlib
 import os
 import pickle
 import sys
 import cutlass.cute as cute
 import tvm_ffi
 from cutlass.cutlass_dsl import JitCompiledFunction
+from .fa_logging import fa_log
 # Pre-load cute DSL runtime libraries with RTLD_GLOBAL so that their symbols
 # (e.g. _cudaLibraryLoadData) are visible to .so modules loaded later via dlopen.
 CompileKeyType: TypeAlias = tuple[Hashable, ...]
 CallableFunction: TypeAlias = JitCompiledFunction | tvm_ffi.Function
 # Enable cache via `FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1`
 CUTE_DSL_CACHE_ENABLED: bool = os.getenv("FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED", "0") == "1"
             label=sha256_hex,
         ):
             if obj_path.exists():
+                fa_log(1, f"Loading compiled function from disk: {obj_path}")
                 m = cute.runtime.load_module(str(obj_path), enable_tvm_ffi=True)
                 fn = getattr(m, self.EXPORT_FUNCTION_PREFIX)
                 JITCache.__setitem__(self, key, fn)
                 return True
             else:
+                fa_log(1, f"Cache miss on disk for key hash {sha256_hex}")
         return False
     def _try_export_to_storage(self, key: CompileKeyType, fn: JitCompiledFunction) -> None:
             obj_path = self.cache_path / f"{sha256_hex}.o"
             if obj_path.exists():
                 # Another process already exported.
+                fa_log(1, f"Skipping export, already on disk: {obj_path}")
                 return
+            fa_log(1, f"Exporting compiled function to disk: {obj_path}")
             fn.export_to_c(
                 object_file_path=str(obj_path),
                 function_name=self.EXPORT_FUNCTION_PREFIX,
             )
+            fa_log(1, f"Successfully exported compiled function to disk: {obj_path}")
     def _key_to_hash(self, key: CompileKeyType) -> str:
         return hashlib.sha256(pickle.dumps(key)).hexdigest()
         """
         Not only clear the in-memory cache. Also purge persistent compilation cache.
         """
+        fa_log(1, f"Clearing persistent cache at {self.cache_path}")
         super().clear()
         for child in self.cache_path.iterdir():
             child.unlink()
         path = get_cache_path() / _compute_source_fingerprint()
         if name:
             path = path / name
+        fa_log(1, f"Creating persistent JIT cache at {path}")
         return JITPersistentCache(path)
     else:
+        fa_log(1, "Persistent cache disabled, using in-memory JIT cache")
         return JITCache()

build/torch-cuda/compute_block_sparsity.py CHANGED Viewed

@@ -11,8 +11,22 @@ from .block_sparsity import (
     BlockSparseTensorsTorch,
     to_cute_block_sparse_tensors,
 )
-from .utils import hash_callable, scalar_to_ssa, ssa_to_scalar
 from .seqlen_info import SeqlenInfoQK
 class BlockSparsityKernel:
@@ -28,7 +42,6 @@ class BlockSparsityKernel:
     TODO:
         - optimize mask_mod evaluation
-        - varlen support
         - transposed tensors for bwd pass
     """
@@ -52,18 +65,31 @@ class BlockSparsityKernel:
         blocksparse_tensors: BlockSparseTensors,
         seqlen_q: Int32,
         seqlen_k: Int32,
-        aux_tensors: Optional[list] = None,
     ):
-        self.mask_cnt, self.mask_idx, self.full_cnt, self.full_idx = blocksparse_tensors
         if const_expr(self.compute_full_blocks):
-            assert self.full_cnt is not None and self.full_idx is not None, (
                 "full block tensors must be provided when computing full blocks"
             )
-        batch_size, num_heads, num_m_blocks, num_n_blocks = self.mask_idx.shape
-        # launch 1 CTA per m block
-        grid = [num_m_blocks, num_heads, batch_size]
         if const_expr(self.use_fast_sampling):
             num_threads = 5
@@ -72,46 +98,46 @@ class BlockSparsityKernel:
             num_threads = self.tile_mn[0]
             self.num_warps = (num_threads + 32 - 1) // 32
         self.kernel(
-            self.mask_cnt,
-            self.mask_idx,
-            self.full_cnt,
-            self.full_idx,
-            num_n_blocks,
             seqlen_q,
             seqlen_k,
-            aux_tensors,
         ).launch(grid=grid, block=[num_threads, 1, 1])
     @cute.kernel
     def kernel(
         self,
-        mask_cnt: cute.Tensor,
-        mask_idx: cute.Tensor,
-        full_cnt: cute.Tensor,
-        full_idx: cute.Tensor,
-        num_n_blocks: Int32,
         seqlen_q: Int32,
         seqlen_k: Int32,
-        aux_tensors: Optional[list] = None,
     ):
         tidx, _, _ = cute.arch.thread_idx()
         warp_idx = cute.arch.warp_idx()
         lane_id = cute.arch.lane_idx()
-        m_block, head_idx, batch_idx = cute.arch.block_idx()
         ssa = partial(scalar_to_ssa, dtype=Int32)
-        seqlen = SeqlenInfoQK.create(
-            batch_idx,
-            seqlen_q,
-            seqlen_k,
-            mCuSeqlensQ=None,
-            mCuSeqlensK=None,
-            mSeqUsedQ=None,
-            mSeqUsedK=None,
-        )
         @cute.struct
         class SharedStorage:
             reduction_buffer_smem: cute.struct.Align[
@@ -124,132 +150,157 @@ class BlockSparsityKernel:
         reduction_buffer = storage.reduction_buffer_smem.get_tensor(
             cute.make_layout((self.num_warps, 2))
         )
         num_mask_blocks = Int32(0)
         num_full_blocks = Int32(0)
-        for n_block in cutlass.range(num_n_blocks, unroll_full=True):
-            m_base = m_block * self.tile_mn[0]
             n_base = n_block * self.tile_mn[1]
             if const_expr(self.use_fast_sampling):
-                # Fast path: 5-point sampling (4 corners + center)
-                # Clamps OOB indices to nearest in bounds.
-                thread_result = Boolean(False)
-                thread_is_valid = Boolean(False)
-                q_idx = Int32(0)
-                kv_idx = Int32(0)
-                if tidx == 0:
-                    # Top-left corner (0, 0); always in bounds
-                    q_idx = m_base
-                    kv_idx = n_base
-                elif tidx == 1:
-                    # Top-right corner
-                    q_idx = m_base
-                    kv_idx = cutlass.min(n_base + self.tile_mn[1] - 1, seqlen_k - 1)
-                elif tidx == 2:
-                    # Bottom-left corner
-                    q_idx = cutlass.min(m_base + self.tile_mn[0] - 1, seqlen_q - 1)
-                    kv_idx = n_base
-                elif tidx == 3:
-                    # Bottom-right corner
-                    q_idx = cutlass.min(m_base + self.tile_mn[0] - 1, seqlen_q - 1)
-                    kv_idx = cutlass.min(n_base + self.tile_mn[1] - 1, seqlen_k - 1)
                 elif tidx == 4:
-                    # Center point
-                    q_idx = m_base + (cutlass.min(seqlen_q - m_base, self.tile_mn[0])) // 2
-                    kv_idx = n_base + (cutlass.min(seqlen_k - n_base, self.tile_mn[1])) // 2
-                else:
-                    thread_is_valid = Boolean(False)
-                # Check bounds and determine if this thread has a valid index pair
-                if tidx < 5 and q_idx < seqlen_q and kv_idx < seqlen_k:
                     thread_is_valid = Boolean(True)
-                    q_idx_ssa = ssa(q_idx)
-                    kv_idx_ssa = ssa(kv_idx)
                     thread_result = ssa_to_scalar(
-                        self.mask_mod(
                             ssa(batch_idx),
                             ssa(head_idx),
-                            q_idx_ssa,
-                            kv_idx_ssa,
                             seqlen,
-                            aux_tensors,
                         )
                     )
-                else:
-                    thread_is_valid = Boolean(False)
-                # Use vote_any_sync to see if any valid thread found unmasked or masked
-                # Only count results from threads that checked valid indices
                 has_unmasked = cute.arch.vote_any_sync(thread_result & thread_is_valid)
-                has_masked = cute.arch.vote_any_sync((Boolean(not thread_result)) & thread_is_valid)
             else:
-                # Full path: check all elements in the block
-                # Track if this thread's row has any masked or unmasked elements
                 thread_has_unmasked = Boolean(False)
                 thread_has_masked = Boolean(False)
-                thread_is_valid = Boolean(False)
-                # Each thread handles 1 row
-                q_idx = m_base + tidx
                 kv_idx = Int32(0)
-                if tidx < self.tile_mn[0] and q_idx < seqlen_q:
-                    thread_is_valid = Boolean(True)
-                    q_idx_ssa = ssa(q_idx)
-                    # Loop over all columns in this row
-                    for c in cutlass.range(self.tile_mn[1], unroll_full=True):
-                        kv_idx = n_base + c
-                        kv_idx_ssa = ssa(kv_idx)
-                        # Only check elements within valid sequence bounds
-                        if kv_idx < seqlen_k:
-                            # Direct scalar call
                             mask_val = ssa_to_scalar(
-                                self.mask_mod(
                                     ssa(batch_idx),
                                     ssa(head_idx),
-                                    q_idx_ssa,
-                                    kv_idx_ssa,
                                     seqlen,
-                                    aux_tensors,
                                 )
                             )
-                            # Update tracking flags
-                            if mask_val:
-                                thread_has_unmasked = Boolean(True)
-                            else:
-                                thread_has_masked = Boolean(True)
-                # Block-level reduction to combine results across all threads
-                # Only count votes from threads that checked valid indices
-                warp_has_unmasked_mask = cute.arch.vote_any_sync(
-                    thread_has_unmasked & thread_is_valid
-                )
-                warp_has_masked_mask = cute.arch.vote_any_sync(thread_has_masked & thread_is_valid)
-                # lane 0 writes the ballot mask to shared memory
-                lane_id = tidx % 32
                 if lane_id == 0:
-                    # Store as Int8
-                    reduction_buffer[warp_idx, 0] = Int8(1) if warp_has_unmasked_mask else Int8(0)
-                    reduction_buffer[warp_idx, 1] = Int8(1) if warp_has_masked_mask else Int8(0)
                 cute.arch.sync_threads()
-                # Thread 0 ORs all warp results together
                 has_unmasked = Boolean(False)
                 has_masked = Boolean(False)
-                if tidx == 0:
-                    for w in cutlass.range(self.num_warps):
-                        if reduction_buffer[w, 0]:
-                            has_unmasked = Boolean(True)
-                        if reduction_buffer[w, 1]:
-                            has_masked = Boolean(True)
             # Only thread 0 updates the output arrays (common to both paths)
             if tidx == 0:
@@ -261,17 +312,23 @@ class BlockSparsityKernel:
                 is_full = Boolean(has_unmasked and (not has_masked))
                 if is_partial:
-                    mask_idx[batch_idx, head_idx, m_block, num_mask_blocks] = n_block
                     num_mask_blocks += 1
                 elif is_full and const_expr(self.compute_full_blocks):
-                    full_idx[batch_idx, head_idx, m_block, num_full_blocks] = n_block
                     num_full_blocks += 1
         # Only thread 0 writes back the counts
         if tidx == 0:
-            mask_cnt[batch_idx, head_idx, m_block] = num_mask_blocks
-            if const_expr(self.compute_full_blocks):
-                full_cnt[batch_idx, head_idx, m_block] = num_full_blocks
 def compute_block_sparsity(
@@ -282,11 +339,18 @@ def compute_block_sparsity(
     seqlen_q,
     seqlen_k,
     mask_mod: Callable,
-    aux_tensors: Optional[list],  # list[cute.Tensor]
     device,
     compute_full_blocks: bool = True,
     use_fast_sampling: bool = False,
-) -> Tuple[BlockSparseTensors, BlockSparseTensorsTorch]:
     """
     Computes block sparsity for a given `mask_mod`.
@@ -300,59 +364,145 @@ def compute_block_sparsity(
         mask_mod: The `mask_mod` callable to use.
         aux_tensors: A list of auxiliary tensors.
         device: The device to use.
         compute_full_blocks: Whether to compute full blocks. If False, only partially-masked blocks are computed.
         use_fast_sampling: Whether to use 5-point sampling (4 corners + center). This is much faster, but only suitable for masks where this check is sufficient.
     Returns:
-        A tuple of `BlockSparseTensors` and `BlockSparseTensorsTorch`.
     """
-    # Check if mask_mod is marked as suitable for 5-point fast sampling
     use_fast_sampling = getattr(mask_mod, "use_fast_sampling", use_fast_sampling)
     num_m_blocks = (seqlen_q + tile_m - 1) // tile_m
     num_n_blocks = (seqlen_k + tile_n - 1) // tile_n
-    mask_block_cnt = torch.zeros(
-        (batch_size, num_heads, num_m_blocks), device=device, dtype=torch.int32
-    )
-    mask_block_idx = torch.zeros(
-        (batch_size, num_heads, num_m_blocks, num_n_blocks), device=device, dtype=torch.int32
-    )
-    full_block_cnt = (
-        torch.zeros((batch_size, num_heads, num_m_blocks), device=device, dtype=torch.int32)
-        if compute_full_blocks
-        else None
-    )
-    full_block_idx = (
-        torch.zeros(
             (batch_size, num_heads, num_m_blocks, num_n_blocks), device=device, dtype=torch.int32
         )
-        if compute_full_blocks
-        else None
-    )
     blocksparse_tensors_torch = BlockSparseTensorsTorch(
         mask_block_cnt=mask_block_cnt,
         mask_block_idx=mask_block_idx,
         full_block_cnt=full_block_cnt,
         full_block_idx=full_block_idx,
         block_size=(tile_m, tile_n),
     )
     mask_mod_hash = hash_callable(mask_mod)
-    blocksparse_tensors = to_cute_block_sparse_tensors(
-        blocksparse_tensors_torch, enable_tvm_ffi=True
-    )
     compile_key = (
         tile_m,
         tile_n,
         mask_mod_hash,
         compute_full_blocks,
         aux_tensors is not None,
         use_fast_sampling,
     )
     if compile_key not in compute_block_sparsity.compile_cache:
         kernel = BlockSparsityKernel(
             mask_mod,
             tile_mn=(tile_m, tile_n),
@@ -362,17 +512,40 @@ def compute_block_sparsity(
         )
         compute_block_sparsity.compile_cache[compile_key] = cute.compile(
-            kernel, blocksparse_tensors, seqlen_q, seqlen_k, aux_tensors, options="--enable-tvm-ffi"
         )
-    compute_block_sparsity.compile_cache[compile_key](
-        blocksparse_tensors_torch[:4],
-        seqlen_q,
-        seqlen_k,
-        aux_tensors,
-    )
-    return blocksparse_tensors, blocksparse_tensors_torch
 compute_block_sparsity.compile_cache = {}

     BlockSparseTensorsTorch,
     to_cute_block_sparse_tensors,
 )
+from .block_sparse_utils import get_curr_blocksparse_tensors
+from .testing import is_fake_mode
+from .cute_dsl_utils import (
+    get_aux_tensor_metadata,
+    to_cute_aux_tensor,
+    to_cute_tensor,
+)
+from .utils import (
+    get_batch_from_cu_tensor,
+    hash_callable,
+    scalar_to_ssa,
+    ssa_to_scalar,
+)
+from .mask import call_mask_mod
 from .seqlen_info import SeqlenInfoQK
+from .utils import AuxData
 class BlockSparsityKernel:
     TODO:
         - optimize mask_mod evaluation
         - transposed tensors for bwd pass
     """
         blocksparse_tensors: BlockSparseTensors,
         seqlen_q: Int32,
         seqlen_k: Int32,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        mSeqUsedQ: Optional[cute.Tensor] = None,
+        mSeqUsedK: Optional[cute.Tensor] = None,
+        aux_data: AuxData = AuxData(),
     ):
+        mask_cnt, mask_idx, full_cnt, full_idx, mCuTotalMBlocks, mCuBlockIdxOffsets, *_ = (
+            blocksparse_tensors
+        )
+        self.is_varlen_q = const_expr(mCuSeqlensQ is not None)
         if const_expr(self.compute_full_blocks):
+            assert full_cnt is not None and full_idx is not None, (
                 "full block tensors must be provided when computing full blocks"
             )
+        if const_expr(not self.is_varlen_q):
+            batch_size, num_heads, num_m_blocks, _ = mask_idx.shape
+            total_m_blocks = batch_size * num_m_blocks
+        else:
+            assert const_expr(mCuTotalMBlocks is not None), (
+                "mCuTotalMBlocks must be provided when varlen q"
+            )
+            num_heads, total_m_blocks = mask_cnt.shape  # num_m_blocks is total_m_blocks
+            batch_size = mCuSeqlensQ.shape[0] - 1
         if const_expr(self.use_fast_sampling):
             num_threads = 5
             num_threads = self.tile_mn[0]
             self.num_warps = (num_threads + 32 - 1) // 32
+        if const_expr(not self.is_varlen_q):
+            grid = [num_m_blocks, num_heads, batch_size]
+        else:
+            grid = [total_m_blocks, num_heads, 1]
         self.kernel(
+            blocksparse_tensors,
             seqlen_q,
             seqlen_k,
+            batch_size,
+            mCuSeqlensQ,
+            mCuSeqlensK,
+            mSeqUsedQ,
+            mSeqUsedK,
+            mCuTotalMBlocks,
+            mCuBlockIdxOffsets,
+            aux_data,
         ).launch(grid=grid, block=[num_threads, 1, 1])
     @cute.kernel
     def kernel(
         self,
+        blocksparse_tensors: BlockSparseTensors,
         seqlen_q: Int32,
         seqlen_k: Int32,
+        batch_size: Int32,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        mSeqUsedQ: Optional[cute.Tensor] = None,
+        mSeqUsedK: Optional[cute.Tensor] = None,
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
+        mCuBlockIdxOffsets: Optional[cute.Tensor] = None,
+        aux_data: AuxData = AuxData(),
     ):
         tidx, _, _ = cute.arch.thread_idx()
         warp_idx = cute.arch.warp_idx()
         lane_id = cute.arch.lane_idx()
         ssa = partial(scalar_to_ssa, dtype=Int32)
         @cute.struct
         class SharedStorage:
             reduction_buffer_smem: cute.struct.Align[
         reduction_buffer = storage.reduction_buffer_smem.get_tensor(
             cute.make_layout((self.num_warps, 2))
         )
+        SeqlenInfoCls = partial(
+            SeqlenInfoQK.create,
+            seqlen_q_static=seqlen_q,
+            seqlen_k_static=seqlen_k,
+            mCuSeqlensQ=mCuSeqlensQ,
+            mCuSeqlensK=mCuSeqlensK,
+            mSeqUsedQ=mSeqUsedQ,
+            mSeqUsedK=mSeqUsedK,
+            mCuTotalMBlocks=mCuTotalMBlocks,
+            mCuBlockIdxOffsets=mCuBlockIdxOffsets,
+            tile_m=self.tile_mn[0],
+            tile_n=self.tile_mn[1],
+        )
+        if const_expr(not self.is_varlen_q):
+            m_block, head_idx, batch_idx = cute.arch.block_idx()
+        else:
+            global_m_block, head_idx, _ = cute.arch.block_idx()
+            batch_idx = get_batch_from_cu_tensor(global_m_block, mCuTotalMBlocks)
+            m_block = global_m_block - mCuTotalMBlocks[batch_idx]
+        seqlen = SeqlenInfoCls(batch_idx)
+        seqlen_q = seqlen.seqlen_q
+        seqlen_k = seqlen.seqlen_k
+        global_m_block = seqlen.m_block_offset + m_block
+        num_n_blocks = (seqlen_k + self.tile_mn[1] - 1) // self.tile_mn[1]
+        _, curr_mask_idx, _, curr_full_idx = get_curr_blocksparse_tensors(
+            batch_idx, head_idx, m_block, blocksparse_tensors, seqlen
+        )
         num_mask_blocks = Int32(0)
         num_full_blocks = Int32(0)
+        m_base = m_block * self.tile_mn[0]
+        if const_expr(self.use_fast_sampling):
+            # Loop-invariant per-thread q_idx for the 5 sample points
+            # (tidx 0, 1: top corners; 2, 3: bottom corners; 4: center).
+            q_idx_sample = m_base
+            if tidx == 2 or tidx == 3:
+                q_idx_sample = cutlass.min(m_base + self.tile_mn[0] - 1, seqlen_q - 1)
+            elif tidx == 4:
+                q_idx_sample = m_base + cutlass.min(seqlen_q - m_base, self.tile_mn[0]) // 2
+        else:
+            q_idx_thread = m_base + tidx
+            thread_in_bounds = Boolean(tidx < self.tile_mn[0] and q_idx_thread < seqlen_q)
+        for n_block in cutlass.range(num_n_blocks):
             n_base = n_block * self.tile_mn[1]
             if const_expr(self.use_fast_sampling):
+                # 5-point sampling (4 corners + center). Interior n_blocks
+                # (n_base + tile_n <= seqlen_k) skip the OOB clamp on the right /
+                # center samples.
+                is_interior = (n_base + self.tile_mn[1]) <= seqlen_k
+                n_right = Int32(0)
+                n_mid = Int32(0)
+                if is_interior:
+                    n_right = n_base + self.tile_mn[1] - 1
+                    n_mid = n_base + self.tile_mn[1] // 2
+                else:
+                    n_right = cutlass.min(n_base + self.tile_mn[1] - 1, seqlen_k - 1)
+                    n_mid = n_base + cutlass.min(seqlen_k - n_base, self.tile_mn[1]) // 2
+                kv_idx = n_base
+                if tidx == 1 or tidx == 3:
+                    kv_idx = n_right
                 elif tidx == 4:
+                    kv_idx = n_mid
+                thread_result = Boolean(False)
+                thread_is_valid = Boolean(False)
+                if tidx < 5:
                     thread_is_valid = Boolean(True)
                     thread_result = ssa_to_scalar(
+                        call_mask_mod(
+                            self.mask_mod,
                             ssa(batch_idx),
                             ssa(head_idx),
+                            ssa(q_idx_sample),
+                            ssa(kv_idx),
                             seqlen,
+                            aux_data,
                         )
                     )
                 has_unmasked = cute.arch.vote_any_sync(thread_result & thread_is_valid)
+                has_masked = cute.arch.vote_any_sync(Boolean(not thread_result) & thread_is_valid)
             else:
+                # Full path. Interior blocks (n_base + tile_n <= seqlen_k) drop the
+                # per-element bound check; the boundary block (at most one) keeps it.
                 thread_has_unmasked = Boolean(False)
                 thread_has_masked = Boolean(False)
                 kv_idx = Int32(0)
+                is_interior = (n_base + self.tile_mn[1]) <= seqlen_k
+                if is_interior:
+                    if thread_in_bounds:
+                        for c in cutlass.range(self.tile_mn[1], unroll_full=True):
                             mask_val = ssa_to_scalar(
+                                call_mask_mod(
+                                    self.mask_mod,
                                     ssa(batch_idx),
                                     ssa(head_idx),
+                                    ssa(q_idx_thread),
+                                    ssa(n_base + c),
                                     seqlen,
+                                    aux_data,
                                 )
                             )
+                            thread_has_unmasked |= Boolean(mask_val)
+                            thread_has_masked |= Boolean(not mask_val)
+                else:
+                    if thread_in_bounds:
+                        for c in cutlass.range(self.tile_mn[1], unroll_full=True):
+                            kv_idx = n_base + c
+                            if kv_idx < seqlen_k:
+                                mask_val = ssa_to_scalar(
+                                    call_mask_mod(
+                                        self.mask_mod,
+                                        ssa(batch_idx),
+                                        ssa(head_idx),
+                                        ssa(q_idx_thread),
+                                        ssa(kv_idx),
+                                        seqlen,
+                                        aux_data,
+                                    )
+                                )
+                                thread_has_unmasked |= Boolean(mask_val)
+                                thread_has_masked |= Boolean(not mask_val)
+                warp_unmasked = cute.arch.vote_any_sync(thread_has_unmasked & thread_in_bounds)
+                warp_masked = cute.arch.vote_any_sync(thread_has_masked & thread_in_bounds)
                 if lane_id == 0:
+                    reduction_buffer[warp_idx, 0] = Int8(1) if warp_unmasked else Int8(0)
+                    reduction_buffer[warp_idx, 1] = Int8(1) if warp_masked else Int8(0)
                 cute.arch.sync_threads()
+                # Cross-warp OR via warp 0; thread 0 (lane 0 of warp 0) holds the result.
                 has_unmasked = Boolean(False)
                 has_masked = Boolean(False)
+                if warp_idx == 0:
+                    lane_unmasked = Boolean(False)
+                    lane_masked = Boolean(False)
+                    if lane_id < self.num_warps:
+                        lane_unmasked = reduction_buffer[lane_id, 0] != Int8(0)
+                        lane_masked = reduction_buffer[lane_id, 1] != Int8(0)
+                    has_unmasked = cute.arch.vote_any_sync(lane_unmasked)
+                    has_masked = cute.arch.vote_any_sync(lane_masked)
             # Only thread 0 updates the output arrays (common to both paths)
             if tidx == 0:
                 is_full = Boolean(has_unmasked and (not has_masked))
                 if is_partial:
+                    curr_mask_idx[num_mask_blocks] = n_block
                     num_mask_blocks += 1
                 elif is_full and const_expr(self.compute_full_blocks):
+                    curr_full_idx[num_full_blocks] = n_block
                     num_full_blocks += 1
         # Only thread 0 writes back the counts
         if tidx == 0:
+            mask_cnt, _, full_cnt, *_ = blocksparse_tensors
+            if const_expr(self.is_varlen_q):
+                mask_cnt[head_idx, global_m_block] = num_mask_blocks
+                if const_expr(self.compute_full_blocks):
+                    full_cnt[head_idx, global_m_block] = num_full_blocks
+            else:
+                mask_cnt[batch_idx, head_idx, m_block] = num_mask_blocks
+                if const_expr(self.compute_full_blocks):
+                    full_cnt[batch_idx, head_idx, m_block] = num_full_blocks
 def compute_block_sparsity(
     seqlen_q,
     seqlen_k,
     mask_mod: Callable,
+    aux_tensors: Optional[list],
     device,
+    aux_scalars: Optional[tuple] = None,
+    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,
+    cu_total_m_blocks: Optional[torch.Tensor] = None,
+    cu_block_idx_offsets: Optional[torch.Tensor] = None,
     compute_full_blocks: bool = True,
     use_fast_sampling: bool = False,
+) -> BlockSparseTensorsTorch:
     """
     Computes block sparsity for a given `mask_mod`.
         mask_mod: The `mask_mod` callable to use.
         aux_tensors: A list of auxiliary tensors.
         device: The device to use.
+        cu_seqlens_q: Cumulative q sequence lengths for varlen
+        cu_seqlens_k: Cumulative k sequence lengths for varlen
+        seqused_q: Per-batch effective q sequence lengths
+        seqused_k: Per-batch effective k sequence lengths
+        cu_total_m_blocks: Cumulative total m blocks tensor for varlen q
+        cu_block_idx_offsets: Cumulative offsets into the packed mask_block_idx /
+            full_block_idx tensors per batch (== cumsum of M_b * N_b).
         compute_full_blocks: Whether to compute full blocks. If False, only partially-masked blocks are computed.
         use_fast_sampling: Whether to use 5-point sampling (4 corners + center). This is much faster, but only suitable for masks where this check is sufficient.
     Returns:
+        BlockSparseTensorsTorch
     """
+    aux_scalars = tuple(aux_scalars) if aux_scalars else None
+    # Check if mask_mod is marked as suitable for 5-point sampling
     use_fast_sampling = getattr(mask_mod, "use_fast_sampling", use_fast_sampling)
     num_m_blocks = (seqlen_q + tile_m - 1) // tile_m
     num_n_blocks = (seqlen_k + tile_n - 1) // tile_n
+    if cu_seqlens_q is not None:
+        assert cu_total_m_blocks is not None, "total m blocks must be provided when varlen q"
+        total_m_blocks = cu_total_m_blocks[-1].item()
+        if cu_block_idx_offsets is None and (cu_seqlens_k is not None or seqused_k is not None):
+            # Derive cu_block_idx_offsets from per-batch K seqlens.
+            cu_block_idx_offsets_list = [0]
+            for batch_idx in range(batch_size):
+                batch_seqlen_q = cu_seqlens_q[batch_idx + 1].item() - cu_seqlens_q[batch_idx].item()
+                if cu_seqlens_k is not None:
+                    batch_seqlen_k = (
+                        cu_seqlens_k[batch_idx + 1].item() - cu_seqlens_k[batch_idx].item()
+                    )
+                else:
+                    batch_seqlen_k = seqused_k[batch_idx].item()
+                num_m_blocks_batch = (batch_seqlen_q + tile_m - 1) // tile_m
+                num_n_blocks_batch = (batch_seqlen_k + tile_n - 1) // tile_n
+                cu_block_idx_offsets_list.append(
+                    cu_block_idx_offsets_list[-1] + num_m_blocks_batch * num_n_blocks_batch
+                )
+            cu_block_idx_offsets = torch.tensor(
+                cu_block_idx_offsets_list, dtype=torch.int32, device=device
+            )
+        if cu_block_idx_offsets is not None:
+            total_n_blocks = cu_block_idx_offsets[-1].item()
+        else:
+            # Uniform-K varlen-Q: every batch has the same K seqlen.
+            total_n_blocks = total_m_blocks * num_n_blocks
+        mask_block_cnt = torch.zeros((num_heads, total_m_blocks), device=device, dtype=torch.int32)
+        mask_block_idx = torch.zeros((num_heads, total_n_blocks), device=device, dtype=torch.int32)
+        full_block_cnt = (
+            torch.zeros((num_heads, total_m_blocks), device=device, dtype=torch.int32)
+            if compute_full_blocks
+            else None
+        )
+        full_block_idx = (
+            torch.zeros((num_heads, total_n_blocks), device=device, dtype=torch.int32)
+            if compute_full_blocks
+            else None
+        )
+    else:
+        total_m_blocks = batch_size * num_m_blocks
+        total_n_blocks = batch_size * num_m_blocks * num_n_blocks
+        mask_block_cnt = torch.zeros(
+            (batch_size, num_heads, num_m_blocks), device=device, dtype=torch.int32
+        )
+        mask_block_idx = torch.zeros(
             (batch_size, num_heads, num_m_blocks, num_n_blocks), device=device, dtype=torch.int32
         )
+        full_block_cnt = (
+            torch.zeros((batch_size, num_heads, num_m_blocks), device=device, dtype=torch.int32)
+            if compute_full_blocks
+            else None
+        )
+        full_block_idx = (
+            torch.zeros(
+                (batch_size, num_heads, num_m_blocks, num_n_blocks),
+                device=device,
+                dtype=torch.int32,
+            )
+            if compute_full_blocks
+            else None
+        )
     blocksparse_tensors_torch = BlockSparseTensorsTorch(
         mask_block_cnt=mask_block_cnt,
         mask_block_idx=mask_block_idx,
         full_block_cnt=full_block_cnt,
         full_block_idx=full_block_idx,
+        cu_total_m_blocks=cu_total_m_blocks,
+        cu_block_idx_offsets=cu_block_idx_offsets,
         block_size=(tile_m, tile_n),
     )
     mask_mod_hash = hash_callable(mask_mod)
+    if aux_tensors is not None:
+        aux_tensor_metadata = get_aux_tensor_metadata(aux_tensors)
+    else:
+        aux_tensor_metadata = None
+    aux_scalar_metadata = tuple(type(s) for s in aux_scalars) if aux_scalars is not None else None
     compile_key = (
         tile_m,
         tile_n,
         mask_mod_hash,
+        aux_tensor_metadata,
+        aux_scalar_metadata,
         compute_full_blocks,
+        cu_seqlens_q is None,
+        cu_seqlens_k is None,
+        seqused_q is None,
+        seqused_k is None,
         aux_tensors is not None,
         use_fast_sampling,
     )
     if compile_key not in compute_block_sparsity.compile_cache:
+        (
+            cu_seqlens_q_tensor,
+            cu_seqlens_k_tensor,
+            seqused_q_tensor,
+            seqused_k_tensor,
+        ) = [
+            to_cute_tensor(t, assumed_align=4, leading_dim=0) if t is not None else None
+            for t in (
+                cu_seqlens_q,
+                cu_seqlens_k,
+                seqused_q,
+                seqused_k,
+            )
+        ]
+        blocksparse_tensors = to_cute_block_sparse_tensors(
+            blocksparse_tensors_torch, enable_tvm_ffi=True
+        )
+        if aux_tensors is not None:
+            cute_aux_tensors = [to_cute_aux_tensor(buf) for buf in aux_tensors]
+        else:
+            cute_aux_tensors = None
         kernel = BlockSparsityKernel(
             mask_mod,
             tile_mn=(tile_m, tile_n),
         )
         compute_block_sparsity.compile_cache[compile_key] = cute.compile(
+            kernel,
+            blocksparse_tensors,
+            seqlen_q,
+            seqlen_k,
+            cu_seqlens_q_tensor,
+            cu_seqlens_k_tensor,
+            seqused_q_tensor,
+            seqused_k_tensor,
+            AuxData(cute_aux_tensors, aux_scalars),
+            options="--enable-tvm-ffi",
         )
+    if not is_fake_mode():
+        compute_block_sparsity.compile_cache[compile_key](
+            (
+                blocksparse_tensors_torch.mask_block_cnt,
+                blocksparse_tensors_torch.mask_block_idx,
+                blocksparse_tensors_torch.full_block_cnt,
+                blocksparse_tensors_torch.full_block_idx,
+                blocksparse_tensors_torch.cu_total_m_blocks,
+                blocksparse_tensors_torch.cu_block_idx_offsets,
+                blocksparse_tensors_torch.dq_write_order,
+                blocksparse_tensors_torch.dq_write_order_full,
+            ),
+            seqlen_q,
+            seqlen_k,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqused_q,
+            seqused_k,
+            AuxData(aux_tensors, aux_scalars),
+        )
+    return blocksparse_tensors_torch
 compute_block_sparsity.compile_cache = {}

build/torch-cuda/cu_blocks_kernel.py ADDED Viewed

	@@ -0,0 +1,181 @@

+from typing import Callable, Optional
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32, const_expr
+from .utils import get_batch_from_cu_tensor
+class CuSeqlensToBlocksKernel:
+    def __init__(
+        self,
+        tile: int = 128,
+        num_threads: int = 1024,
+        seqlen_q_multiplier: int = 1,
+    ):
+        self.tile = tile
+        self.num_threads = num_threads
+        assert num_threads % 32 == 0
+        self.num_warps = num_threads // cute.arch.WARP_SIZE
+        self.seqlen_q_multiplier = seqlen_q_multiplier
+    @cute.jit
+    def __call__(
+        self,
+        mCuBlocks: cute.Tensor,
+        mCuSplitsBlocks: Optional[cute.Tensor],
+        mCuSeqlens: Optional[cute.Tensor],
+        mSeqUsed: Optional[cute.Tensor] = None,
+        mNumSplits: Optional[cute.Tensor] = None,
+        mVirtualBatchIdx: Optional[cute.Tensor] = None,
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
+        stream: cuda.CUstream = None,
+    ):
+        assert const_expr((mNumSplits is None) == (mCuSplitsBlocks is None))
+        assert const_expr(mCuSeqlens is not None or mSeqUsed is not None)
+        @cute.struct
+        class SharedStorage:
+            warp_block_count: cute.struct.MemRange[Int32, self.num_warps]
+            warp_split_count: cute.struct.MemRange[Int32, self.num_warps]
+        self.kernel(
+            mCuBlocks,
+            mCuSplitsBlocks,
+            mCuSeqlens,
+            mSeqUsed,
+            mNumSplits,
+            mVirtualBatchIdx,
+            SharedStorage,
+        ).launch(
+            grid=[1, 1, 1],
+            block=[self.num_threads, 1, 1],
+            stream=stream,
+        )
+    @cute.kernel
+    def kernel(
+        self,
+        mCuBlocks: cute.Tensor,
+        mCuSplitsBlocks: Optional[cute.Tensor],
+        mCuSeqlens: Optional[cute.Tensor],
+        mSeqUsed: Optional[cute.Tensor],
+        mNumSplits: Optional[cute.Tensor],
+        mVirtualBatchIdx: Optional[cute.Tensor],
+        SharedStorage: cutlass.Constexpr[Callable],
+    ):
+        has_splits = mNumSplits is not None
+        batch_size = mCuBlocks.shape[0] - 1
+        tidx = cute.arch.thread_idx()[0]
+        lane_idx = cute.arch.lane_idx()
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+        smem = cutlass.utils.SmemAllocator()
+        storage = smem.allocate(SharedStorage)
+        warp_block_count = storage.warp_block_count.get_tensor(cute.make_layout(self.num_warps))
+        warp_split_count = storage.warp_split_count.get_tensor(cute.make_layout(self.num_warps))
+        if tidx == 0:
+            mCuBlocks[0] = 0
+            if const_expr(has_splits):
+                mCuSplitsBlocks[0] = 0
+        # Process the batch in chunks of num_threads, carrying the running totals in the bases.
+        base = Int32(0)
+        base_splits = Int32(0)
+        num_chunks = (batch_size + self.num_threads - 1) // self.num_threads
+        for chunk in cutlass.range(num_chunks):
+            batch_idx = chunk * self.num_threads + tidx
+            seqlen = Int32(0)
+            batch_splits = Int32(0)
+            if batch_idx < batch_size:
+                if const_expr(mVirtualBatchIdx is not None):
+                    batch_idx = Int32(mVirtualBatchIdx[batch_idx])
+                if const_expr(mSeqUsed is not None):
+                    seqlen = mSeqUsed[batch_idx]
+                else:
+                    seqlen = mCuSeqlens[batch_idx + 1] - mCuSeqlens[batch_idx]
+                if const_expr(has_splits):
+                    batch_splits = mNumSplits[batch_idx]
+            seqlen *= self.seqlen_q_multiplier
+            num_blocks = (seqlen + self.tile - 1) // self.tile
+            num_split_blocks = num_blocks * batch_splits
+            total_blocks_for_batch = num_blocks
+            total_split_blocks_for_batch = num_split_blocks
+            for delta in (1, 2, 4, 8, 16):
+                other = cute.arch.shuffle_sync_up(total_blocks_for_batch, delta, mask_and_clamp=0)
+                if const_expr(has_splits):
+                    other_splits = cute.arch.shuffle_sync_up(
+                        total_split_blocks_for_batch, delta, mask_and_clamp=0
+                    )
+                    if lane_idx >= delta:
+                        total_split_blocks_for_batch += other_splits
+                if lane_idx >= delta:
+                    total_blocks_for_batch += other
+            if lane_idx == 31:
+                warp_block_count[warp_idx] = total_blocks_for_batch
+                if const_expr(has_splits):
+                    warp_split_count[warp_idx] = total_split_blocks_for_batch
+            cute.arch.sync_threads()
+            total_blocks_for_batch += base
+            total_split_blocks_for_batch += base_splits
+            for idx in cutlass.range(warp_idx):
+                total_blocks_for_batch += warp_block_count[idx]
+                if const_expr(has_splits):
+                    total_split_blocks_for_batch += warp_split_count[idx]
+            if batch_idx < batch_size:
+                mCuBlocks[chunk * self.num_threads + tidx + 1] = total_blocks_for_batch
+                if const_expr(has_splits):
+                    mCuSplitsBlocks[chunk * self.num_threads + tidx + 1] = (
+                        total_split_blocks_for_batch
+                    )
+            for idx in cutlass.range(self.num_warps):
+                base += warp_block_count[idx]
+                if const_expr(has_splits):
+                    base_splits += warp_split_count[idx]
+            # warp_block_count / warp_split_count are reused by the next chunk.
+            cute.arch.sync_threads()
+class CuBlocksToBatchKernel:
+    """Inverts a cumulative block count into a flat block -> batch lookup."""
+    def __init__(self, num_threads: int = 128):
+        self.num_threads = num_threads
+    @cute.jit
+    def __call__(
+        self,
+        mCuTotalBlocks: cute.Tensor,
+        mBlocksToBatchIdx: cute.Tensor,
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
+        stream: cuda.CUstream = None,
+    ):
+        num_blocks = mBlocksToBatchIdx.shape[0]
+        self.kernel(mCuTotalBlocks, mBlocksToBatchIdx).launch(
+            grid=[(num_blocks + self.num_threads - 1) // self.num_threads, 1, 1],
+            block=[self.num_threads, 1, 1],
+            stream=stream,
+        )
+    @cute.kernel
+    def kernel(
+        self,
+        mCuTotalBlocks: cute.Tensor,
+        mBlocksToBatchIdx: cute.Tensor,
+    ):
+        block_idx = cute.arch.block_idx()[0] * self.num_threads + cute.arch.thread_idx()[0]
+        if block_idx < mBlocksToBatchIdx.shape[0]:
+            mBlocksToBatchIdx[block_idx] = get_batch_from_cu_tensor(block_idx, mCuTotalBlocks)

build/torch-cuda/cute_dsl_utils.py CHANGED Viewed

@@ -1,12 +1,11 @@
 # 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
@@ -15,7 +14,6 @@ except ImportError:
 import cutlass
 import cutlass.cute as cute
-from cutlass.base_dsl.typing import JitArgument
 from cutlass.cutlass_dsl import NumericMeta
 from cutlass.cute.runtime import from_dlpack
@@ -30,6 +28,8 @@ torch2cute_dtype_map = {
     torch.float16: cutlass.Float16,
     torch.bfloat16: cutlass.BFloat16,
     torch.float32: cutlass.Float32,
 }
@@ -43,61 +43,73 @@ def get_device_capacity(device: torch.device = None) -> Tuple[int, int]:
     return torch.cuda.get_device_capability(device)
-@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
 def assume_strides_aligned(t):
@@ -120,7 +132,22 @@ def assume_tensor_aligned(t):
 def to_cute_tensor(t, assumed_align=16, leading_dim=-1, fully_dynamic=False, enable_tvm_ffi=True):
     """Convert torch tensor to cute tensor for TVM FFI. leading_dim=-1 defaults to t.ndim-1."""
-    tensor = from_dlpack(t.detach(), assumed_align=assumed_align, enable_tvm_ffi=enable_tvm_ffi)
     if fully_dynamic:
         return tensor.mark_layout_dynamic()
     if leading_dim == -1:
@@ -146,15 +173,37 @@ def to_cute_aux_tensor(t, enable_tvm_ffi=True):
     )
 def get_aux_tensor_metadata(aux_tensors):
-    return tuple(
-        (
-            getattr(t, "__assumed_align__", 0),
-            getattr(t, "__leading_dim__", -1),
-            hasattr(t, "__leading_dim__"),
         )
-        for t in aux_tensors
-    )
 def get_broadcast_dims(tensor: torch.Tensor) -> Tuple[bool, ...]:
@@ -164,4 +213,43 @@ def get_broadcast_dims(tensor: torch.Tensor) -> Tuple[bool, ...]:
     stride=0 as static, meaning kernels compiled with different broadcast
     patterns are not interchangeable.
     """
-    return tuple(s == 0 for s in tensor.stride())

 # Copyright (c) 2025, Tri Dao.
 import os
 from typing import Tuple
+from functools import lru_cache
 import torch
+from torch._subclasses.fake_tensor import FakeTensor
 try:
     from triton.tools.disasm import extract
 import cutlass
 import cutlass.cute as cute
 from cutlass.cutlass_dsl import NumericMeta
 from cutlass.cute.runtime import from_dlpack
     torch.float16: cutlass.Float16,
     torch.bfloat16: cutlass.BFloat16,
     torch.float32: cutlass.Float32,
+    torch.float8_e4m3fn: cutlass.Float8E4M3FN,
+    torch.float8_e5m2: cutlass.Float8E5M2,
 }
     return torch.cuda.get_device_capability(device)
+@lru_cache
+def _get_device_arch_and_num_sms(device_index: int) -> tuple[int, int]:
+    properties = torch.cuda.get_device_properties(device_index)
+    return properties.major * 10 + properties.minor, properties.multi_processor_count
+def get_num_sms_for_selection(device_index: int, arch: int) -> int:
+    """Return the SM count of the matching local GPU or cross-compilation target."""
+    override = os.getenv("FLASH_ATTENTION_NUM_SMS")
+    if override is not None:
+        num_sms = int(override)
+        if num_sms <= 0:
+            raise ValueError("FLASH_ATTENTION_NUM_SMS must be positive")
+        return num_sms
+    if torch.cuda.is_available():
+        device_arch, num_sms = _get_device_arch_and_num_sms(device_index)
+        if device_arch == arch:
+            return num_sms
+    raise RuntimeError(
+        "Cannot determine the target GPU's SM count; set FLASH_ATTENTION_NUM_SMS "
+        "when cross-compiling without a matching local GPU"
+    )
+def _has_aligned_pointer(tensor: torch.Tensor, align_bytes: int) -> bool:
+    address = (
+        tensor.storage_offset() * tensor.element_size()
+        if isinstance(tensor, FakeTensor)
+        else tensor.data_ptr()
+    )
+    return address % align_bytes == 0
+def _is_aligned_layout(tensor: torch.Tensor, align_bytes: int) -> bool:
+    """Return whether a tensor satisfies the pointer and stride ABI kernels assume."""
+    if tensor.stride(-1) != 1 or not _has_aligned_pointer(tensor, align_bytes):
+        return False
+    stride_alignment = max(1, align_bytes // tensor.element_size())
+    return all(stride == 0 or stride % stride_alignment == 0 for stride in tensor.stride()[:-1])
+def maybe_contiguous(tensor: torch.Tensor | None, align_bytes: int = 16):
+    """Canonicalize inputs to the pointer and stride alignment kernels assume."""
+    if tensor is None:
+        return None
+    if tensor.is_contiguous():
+        return (
+            tensor
+            if _has_aligned_pointer(tensor, align_bytes)
+            else tensor.clone(memory_format=torch.contiguous_format)
         )
+    if not _has_aligned_pointer(tensor, align_bytes):
+        return tensor.clone(memory_format=torch.contiguous_format)
+    return tensor if _is_aligned_layout(tensor, align_bytes) else tensor.contiguous()
+def validate_output_layout(tensor: torch.Tensor, name: str, align_bytes: int) -> None:
+    """Validate a caller-provided output or SplitKV workspace."""
+    assert 0 not in tensor.stride(), f"{name} must not have broadcast dimensions"
+    if tensor.is_contiguous():
+        assert _has_aligned_pointer(tensor, align_bytes), (
+            f"{name} must have aligned strides and a contiguous last dimension"
+        )
+        return
+    assert _is_aligned_layout(tensor, align_bytes), (
+        f"{name} must have aligned strides and a contiguous last dimension"
+    )
 def assume_strides_aligned(t):
 def to_cute_tensor(t, assumed_align=16, leading_dim=-1, fully_dynamic=False, enable_tvm_ffi=True):
     """Convert torch tensor to cute tensor for TVM FFI. leading_dim=-1 defaults to t.ndim-1."""
+    if t is None:
+        return None
+    # NOTE: torch 2.9.1 doesn't support fp8 via DLPack but 2.11.0 nightly does
+    # currently export raw bytes as uint8 and tell cutlass correct type
+    # can directly export as fp8 when torch supports it
+    if t.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
+        tensor = from_dlpack(
+            t.view(torch.uint8).detach(),
+            assumed_align=assumed_align,
+            enable_tvm_ffi=enable_tvm_ffi,
+        )
+        tensor.element_type = (
+            cutlass.Float8E4M3FN if t.dtype == torch.float8_e4m3fn else cutlass.Float8E5M2
+        )
+    else:
+        tensor = from_dlpack(t.detach(), assumed_align=assumed_align, enable_tvm_ffi=enable_tvm_ffi)
     if fully_dynamic:
         return tensor.mark_layout_dynamic()
     if leading_dim == -1:
     )
+def _resolve_aux_leading_dim(tensor: torch.Tensor) -> int | None:
+    """Pick the mode CuTe keeps as a static stride-1 leading dimension."""
+    leading_dim = getattr(tensor, "__leading_dim__", None)
+    if leading_dim is not None:
+        if tensor.ndim == 0:
+            raise ValueError("Scalar aux tensors cannot declare __leading_dim__")
+        leading_dim %= tensor.ndim
+        if tensor.stride(leading_dim) != 1:
+            raise ValueError("Aux tensor __leading_dim__ must identify a stride-1 dimension")
+        return leading_dim
+    unit_stride_dims = [dim for dim, stride in enumerate(tensor.stride()) if stride == 1]
+    if len(unit_stride_dims) <= 1:
+        return unit_stride_dims[0] if unit_stride_dims else None
+    nontrivial_dims = [dim for dim in unit_stride_dims if tensor.shape[dim] > 1]
+    if len(nontrivial_dims) != 1:
+        raise ValueError("Aux tensor layout has no unique stride-1 leading dimension")
+    return nontrivial_dims[0]
 def get_aux_tensor_metadata(aux_tensors):
+    """Return the static aux-tensor ABI facts that must key the compile cache."""
+    metadata = []
+    for tensor in aux_tensors:
+        leading_dim = _resolve_aux_leading_dim(tensor)
+        static_strides = tuple(
+            0 if stride == 0 else 1 if dim == leading_dim else None
+            for dim, stride in enumerate(tensor.stride())
         )
+        metadata.append((tensor.dtype, getattr(tensor, "__assumed_align__", None), static_strides))
+    return tuple(metadata)
 def get_broadcast_dims(tensor: torch.Tensor) -> Tuple[bool, ...]:
     stride=0 as static, meaning kernels compiled with different broadcast
     patterns are not interchangeable.
     """
+    strides = tensor.stride()
+    # Written this way for speed.
+    if 0 not in strides:
+        return (False,) * len(strides)
+    return tuple(stride == 0 for stride in strides)
+# credit: monellz (https://github.com/NVIDIA/cutlass/issues/2658#issuecomment-3630564264)
+def dump_kernel_attributes(compiled_kernel):
+    from cuda.bindings import driver
+    from cutlass.utils import HardwareInfo
+    import torch
+    device_id = torch.cuda.current_device()
+    hardware_info = HardwareInfo(device_id=device_id)
+    cubin_data = compiled_kernel.artifacts.CUBIN
+    assert cubin_data is not None, "cubin_data is None, need '--keep-cubin' option when compiling"
+    cuda_library = hardware_info._checkCudaErrors(
+        driver.cuLibraryLoadData(cubin_data, None, None, 0, None, None, 0)
+    )
+    kernels = hardware_info._checkCudaErrors(driver.cuLibraryEnumerateKernels(1, cuda_library))
+    kernel = hardware_info._checkCudaErrors(driver.cuKernelGetFunction(kernels[0]))
+    # more metrics: https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__EXEC.html#group__CUDA__EXEC_1g5e92a1b0d8d1b82cb00dcfb2de15961b
+    local_size_bytes = hardware_info._checkCudaErrors(
+        driver.cuFuncGetAttribute(
+            driver.CUfunction_attribute.CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES,
+            kernel,
+        )
+    )
+    num_regs = hardware_info._checkCudaErrors(
+        driver.cuFuncGetAttribute(
+            driver.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NUM_REGS,
+            kernel,
+        )
+    )
+    print("--- Kernel Info ---")
+    print(f"local_size_bytes: {local_size_bytes}")
+    print(f"num_regs: {num_regs}")
+    print("--- End Kernel Info ---")

build/torch-cuda/fa_logging.py ADDED Viewed

	@@ -0,0 +1,97 @@

+# Copyright (c) 2025, Tri Dao.
+"""Unified FlashAttention logging controlled by a single ``FA_LOG_LEVEL`` env var.
+Host-side messages go through Python ``logging`` (logger name ``flash_attn``).
+A default ``StreamHandler`` is attached automatically when ``FA_LOG_LEVEL >= 1``
+so that standalone scripts get output without extra setup; applications that
+configure their own logging can remove or replace it via the standard API.
+FA_LOG_LEVEL mapping::
+    0  off       nothing logged
+    1  host      host-side summaries only (no kernel printf)
+    2  kernel    host + curated kernel traces
+    3  max       host + all kernel traces (noisy, perf hit)
+Set via environment variable::
+    FA_LOG_LEVEL=1 python train.py
+Device-side ``cute.printf`` calls are compile-time eliminated via
+``cutlass.const_expr`` when the log level is below the callsite threshold,
+so there is zero performance cost when device logging is off.
+Changing the log level after kernel compilation requires a recompile
+(the level participates in the forward compile key).
+"""
+import logging
+import os
+import sys
+import cutlass.cute as cute
+from cutlass import const_expr
+_LOG_LEVEL_NAMES = {"off": 0, "host": 1, "kernel": 2, "max": 3}
+def _parse_log_level(raw: str) -> int:
+    if raw in _LOG_LEVEL_NAMES:
+        return _LOG_LEVEL_NAMES[raw]
+    try:
+        level = int(raw)
+    except ValueError:
+        return 0
+    return max(0, min(level, 3))
+_fa_log_level: int = _parse_log_level(os.environ.get("FA_LOG_LEVEL", "0"))
+_logger = logging.getLogger("flash_attn")
+_logger.addHandler(logging.NullHandler())
+_default_handler: logging.Handler | None = None
+def _configure_default_handler() -> None:
+    global _default_handler
+    if _fa_log_level >= 1:
+        if _default_handler is None:
+            _default_handler = logging.StreamHandler(sys.stdout)
+            _default_handler.setFormatter(logging.Formatter("[FA] %(message)s"))
+            _logger.addHandler(_default_handler)
+        _logger.setLevel(logging.DEBUG)
+    else:
+        if _default_handler is not None:
+            _logger.removeHandler(_default_handler)
+            _default_handler = None
+        _logger.setLevel(logging.WARNING)
+_configure_default_handler()
+def get_fa_log_level() -> int:
+    return _fa_log_level
+def set_fa_log_level(level: int | str) -> None:
+    """Set the FA log level programmatically.
+    Host logging takes effect immediately.  Device logging changes only
+    affect kernels compiled after this call (new compile-key selection).
+    """
+    global _fa_log_level
+    if isinstance(level, str):
+        level = _parse_log_level(level)
+    _fa_log_level = max(0, min(int(level), 3))
+    _configure_default_handler()
+def fa_log(level: int, msg: str):
+    if _fa_log_level >= level:
+        _logger.info(msg)
+def fa_printf(level: int, fmt, *args):
+    if const_expr(_fa_log_level >= level):
+        cute.printf(fmt, *args)

build/torch-cuda/flash_attn4/__init__.py DELETED Viewed

@@ -1,26 +0,0 @@
-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-cuda/flash_bwd.py CHANGED Viewed

@@ -11,7 +11,7 @@ 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 .quack import layout_utils
@@ -19,9 +19,13 @@ from . import ampere_helpers as sm80_utils
 from .cute_dsl_utils import assume_tensor_aligned
 from . import utils
 from .mask import AttentionMask
 from .seqlen_info import SeqlenInfoQK
 from .quack.cute_dsl_utils import ParamsBase
 from .tile_scheduler import SingleTileScheduler, SingleTileVarlenScheduler, TileSchedulerArguments
 class FlashAttentionBackwardSm80:
@@ -38,6 +42,7 @@ class FlashAttentionBackwardSm80:
         num_threads: int = 256,
         pack_gqa: bool = False,
         is_causal: bool = False,
         SdP_swapAB: bool = False,
         dKV_swapAB: bool = False,
         dQ_swapAB: bool = False,
@@ -45,6 +50,8 @@ class FlashAttentionBackwardSm80:
         AtomLayoutNdKV: int = 8,
         AtomLayoutMdQ: int = 1,
         V_in_regs: bool = False,
     ):
         """Initializes the configuration for a flash attention v2 kernel.
@@ -62,7 +69,8 @@ class FlashAttentionBackwardSm80:
         :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
@@ -77,6 +85,7 @@ class FlashAttentionBackwardSm80:
         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
@@ -89,6 +98,8 @@ class FlashAttentionBackwardSm80:
         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(
@@ -161,13 +172,13 @@ class FlashAttentionBackwardSm80:
         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")
@@ -372,17 +383,24 @@ class FlashAttentionBackwardSm80:
         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)))
@@ -417,12 +435,15 @@ class FlashAttentionBackwardSm80:
             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,
@@ -439,6 +460,8 @@ class FlashAttentionBackwardSm80:
             mSeqUsedK,
             softmax_scale,
             softmax_scale_log2,
             self.sQ_layout,
             self.sK_layout,
             self.sV_layout,
@@ -458,6 +481,7 @@ class FlashAttentionBackwardSm80:
             SharedStorage,
             tile_sched_params,
             TileScheduler,
         ).launch(
             grid=grid_dim,
             block=[self.num_threads, 1, 1],
@@ -483,6 +507,8 @@ class FlashAttentionBackwardSm80:
         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,
@@ -502,6 +528,7 @@ class FlashAttentionBackwardSm80:
         SharedStorage: cutlass.Constexpr,
         tile_sched_params: ParamsBase,
         TileScheduler: cutlass.Constexpr[Callable],
     ):
         # Thread index, block index
         tidx, _, _ = cute.arch.thread_idx()
@@ -512,15 +539,28 @@ class FlashAttentionBackwardSm80:
         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
             # ///////////////////////////////////////////////////////////////////////////////
@@ -538,7 +578,7 @@ class FlashAttentionBackwardSm80:
                 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])
@@ -615,8 +655,8 @@ class FlashAttentionBackwardSm80:
             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)
@@ -755,63 +795,73 @@ class FlashAttentionBackwardSm80:
                 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
@@ -842,7 +892,9 @@ class FlashAttentionBackwardSm80:
         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():
@@ -860,7 +912,7 @@ class FlashAttentionBackwardSm80:
         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()
@@ -871,13 +923,30 @@ class FlashAttentionBackwardSm80:
             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 = layout_utils.reshape_acc_to_mn(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)
@@ -887,7 +956,7 @@ class FlashAttentionBackwardSm80:
         # 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()
@@ -907,7 +976,20 @@ class FlashAttentionBackwardSm80:
         # 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))
@@ -944,7 +1026,7 @@ class FlashAttentionBackwardSm80:
             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,
@@ -968,7 +1050,7 @@ class FlashAttentionBackwardSm80:
         # MMA dK
         if cutlass.const_expr(self.Mma_dKV_is_RS):
-            tdVrP = layout_utils.reshape_acc_to_frgA(rdS)
         else:
             tdKrdS = mma_params.tdKrdS
         sm80_utils.gemm(

 import cutlass
 import cutlass.cute as cute
 from cutlass.cute.nvgpu import cpasync, warp
+from cutlass import Int32
 import cutlass.utils as utils_basic
 from .quack import layout_utils
 from .cute_dsl_utils import assume_tensor_aligned
 from . import utils
 from .mask import AttentionMask
+from .softmax import call_score_mod, call_score_mod_bwd
 from .seqlen_info import SeqlenInfoQK
+from .block_info import BlockInfo
 from .quack.cute_dsl_utils import ParamsBase
 from .tile_scheduler import SingleTileScheduler, SingleTileVarlenScheduler, TileSchedulerArguments
+from .block_sparsity import BlockSparseTensors
+from .utils import AuxData
 class FlashAttentionBackwardSm80:
         num_threads: int = 256,
         pack_gqa: bool = False,
         is_causal: bool = False,
+        is_local: bool = False,
         SdP_swapAB: bool = False,
         dKV_swapAB: bool = False,
         dQ_swapAB: bool = False,
         AtomLayoutNdKV: int = 8,
         AtomLayoutMdQ: int = 1,
         V_in_regs: bool = False,
+        score_mod: cutlass.Constexpr | None = None,
+        score_mod_bwd: cutlass.Constexpr | None = None,
     ):
         """Initializes the configuration for a flash attention v2 kernel.
         :param is_causal: is causal
         """
         self.dtype = dtype
+        # padding head_dim to a multiple of 32 (stricter than fwd's 16) due to
+        # backward kernel register layout requirements for dQ/dK/dV accumulation
         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.num_threads = num_threads
         self.pack_gqa = pack_gqa
         self.is_causal = is_causal
+        self.is_local = is_local
         self.num_stages_Q = num_stages_Q
         self.num_stages_dO = num_stages_dO
         self.SdP_swapAB = SdP_swapAB
         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
+        self.score_mod = score_mod
+        self.score_mod_bwd = score_mod_bwd
     @staticmethod
     def can_implement(
         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(mQ_type not in [cutlass.Float16, cutlass.BFloat16]):
             raise TypeError("Only Float16 or BFloat16 is supported")
+        if cutlass.const_expr(mLSE_type not in [cutlass.Float32]):
             raise TypeError("LSE tensor must be Float32")
+        if cutlass.const_expr(mdPsum_type not in [cutlass.Float32]):
             raise TypeError("dPsum tensor must be Float32")
+        if cutlass.const_expr(mdQaccum_type not 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")
         mdK: cute.Tensor,
         mdV: cute.Tensor,
         softmax_scale: cutlass.Float32,
         mCuSeqlensQ: Optional[cute.Tensor] = None,
         mCuSeqlensK: Optional[cute.Tensor] = None,
         mSeqUsedQ: Optional[cute.Tensor] = None,
         mSeqUsedK: Optional[cute.Tensor] = None,
         window_size_left: Int32 | int | None = None,
         window_size_right: Int32 | int | None = None,
         mdQ_semaphore: Optional[cute.Tensor] = None,
+        mdK_semaphore: Optional[cute.Tensor] = None,
+        mdV_semaphore: Optional[cute.Tensor] = None,
+        aux_data: AuxData = AuxData(),
+        blocksparse_tensors: Optional[BlockSparseTensors] = None,
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
+        stream: cuda.CUstream = None,
     ):
+        assert mdQ_semaphore is None and mdK_semaphore is None and mdV_semaphore is None, (
+            "determinism not supported yet for Sm80"
+        )
         # 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)))
             qhead_per_kvhead_packgqa=self.qhead_per_kvhead if cutlass.const_expr(self.pack_gqa) else 1,
             mCuSeqlensQ=mCuSeqlensK,
             mSeqUsedQ=mSeqUsedK,
+            cu_total_m_blocks_ptr=mCuTotalMBlocks,
         )
         tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
         grid_dim = TileScheduler.get_grid_shape(tile_sched_params)
+        softmax_scale_log2, _ = utils.compute_softmax_scale_log2(
+            softmax_scale, self.score_mod
+        )
         self.kernel(
             mQ,
             mK,
             mSeqUsedK,
             softmax_scale,
             softmax_scale_log2,
+            window_size_left,
+            window_size_right,
             self.sQ_layout,
             self.sK_layout,
             self.sV_layout,
             SharedStorage,
             tile_sched_params,
             TileScheduler,
+            aux_data,
         ).launch(
             grid=grid_dim,
             block=[self.num_threads, 1, 1],
         mSeqUsedK: Optional[cute.Tensor],
         softmax_scale: cutlass.Float32,
         softmax_scale_log2: cutlass.Float32,
+        window_size_left: Optional[Int32],
+        window_size_right: Optional[Int32],
         sQ_layout: cute.ComposedLayout,
         sK_layout: cute.ComposedLayout,
         sV_layout: cute.ComposedLayout,
         SharedStorage: cutlass.Constexpr,
         tile_sched_params: ParamsBase,
         TileScheduler: cutlass.Constexpr[Callable],
+        aux_data: AuxData = AuxData(),
     ):
         # Thread index, block index
         tidx, _, _ = cute.arch.thread_idx()
         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,
+                tile_m=self.m_block_size,
+                tile_n=self.n_block_size,
+            )
+            block_info = BlockInfo(
+                self.m_block_size,
+                self.n_block_size,
+                self.is_causal,
+                self.is_local,
+                False,
+                window_size_left,
+                window_size_right,
+            )
+            m_block_min, m_block_max = block_info.get_m_block_min_max(seqlen, n_block)
             # TODO: return early if m_block_max == 0
             # ///////////////////////////////////////////////////////////////////////////////
                 mdPsum_cur = mdPsum[batch_idx, head_idx, None]
                 mdQaccum_cur = mdQaccum[batch_idx, head_idx, None]
             else:
+                padded_offset_q = seqlen.padded_offset_q
                 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])
             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_rmem_tensor(acc_shape_dK, cutlass.Float32)
+            acc_dV = cute.make_rmem_tensor(acc_shape_dV, cutlass.Float32)
             acc_dK.fill(0.0)
             acc_dV.fill(0.0)
                 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=softmax_scale,
                 softmax_scale_log2=softmax_scale_log2,
+                aux_data=aux_data,
             )
+            if m_block_min < m_block_max:
+                # ///////////////////////////////////////////////////////////////////////////////
+                # 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,
+                    window_size_left,
+                    window_size_right,
+                )
+                mask_fn = partial(
+                    mask.apply_mask, n_block=n_block, thr_mma=thr_mma_sdp,
+                    batch_idx=batch_idx, head_idx=head_idx,
+                    mask_seqlen=True, mask_causal=self.is_causal, mask_local=self.is_local
                 )
+                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
         load_Q_LSE: Callable,
         load_dO_dPsum: Callable,
         m_block_max: cutlass.Int32,
+        softmax_scale: cutlass.Float32,
         softmax_scale_log2: cutlass.Float32,
+        aux_data: AuxData = AuxData(),
         mask_fn: Optional[Callable] = None,
     ):
         def load_Q_next():
         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_rmem_tensor(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()
             smem_copy_params.smem_thr_copy_QdO, smem_copy_params.smem_thr_copy_KV,
             swap_AB=self.SdP_swapAB,
         )
+        acc_S_pre = cute.make_fragment_like(acc_S)
+        acc_S_pre.store(acc_S.load())
         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
         )
+        acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S)
+        acc_S_pre_mn = layout_utils.reshape_acc_to_mn(acc_S_pre)
+        if cutlass.const_expr(self.score_mod is not None):
+            for r in cutlass.range(cute.size(acc_S_mn, mode=[0]), unroll_full=True):
+                acc_S_mn[r, None].store(
+                    call_score_mod(
+                        self.score_mod,
+                        acc_S_mn[r, None].load() * softmax_scale,
+                        0,
+                        0,
+                        0,
+                        0,
+                        None,
+                        aux_data,
+                    )
+                )
         if cutlass.const_expr(mask_fn is not None):
             mask_fn(acc_S, m_block=m_block)
         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)
         # 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_rmem_tensor(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()
         # 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):
+            grad_val = acc_S_mn[r, None].load() * (acc_dP_mn[r, None].load() - tLSErdPsum[r])
+            if cutlass.const_expr(self.score_mod_bwd is not None):
+                grad_val = call_score_mod_bwd(
+                    self.score_mod_bwd,
+                    grad_val,
+                    acc_S_pre_mn[r, None].load() * softmax_scale,
+                    0,
+                    0,
+                    0,
+                    0,
+                    None,
+                    aux_data,
+                )
+            acc_dP_mn[r, None].store(grad_val)
         # 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))
             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_rmem_tensor(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,
         # MMA dK
         if cutlass.const_expr(self.Mma_dKV_is_RS):
+            tdKrdS = layout_utils.reshape_acc_to_frgA(rdS)
         else:
             tdKrdS = mma_params.tdKrdS
         sm80_utils.gemm(

build/torch-cuda/flash_bwd_mla_dk_sm100.py ADDED Viewed

	@@ -0,0 +1,1185 @@

+# Modified from CUTLASS example file, original copyright:
+# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+from typing import Callable, Optional, Tuple, Union
+import cuda.bindings.driver as cuda
+import cutlass
+from cutlass import Int32, const_expr
+from cutlass.cutlass_dsl import dsl_user_op
+import cutlass.cute as cute
+import cutlass.cute.testing as testing
+import cutlass.utils as utils
+from cutlass.utils.gemm.sm100 import (
+    transform_partitioned_tensor_layout,
+    epilogue_smem_copy_and_partition,
+    epilogue_tmem_copy_and_partition,
+)
+import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
+from cutlass.cute.nvgpu import cpasync, tcgen05
+from .utils import get_batch_from_cu_tensor
+@dsl_user_op
+def elem_pointer(x: cute.Tensor, coord, *, loc=None, ip=None) -> cute.Pointer:
+    """
+    Get a pointer to an element at the specified coordinate in a tensor.
+    Args:
+        x: The tensor (typically a shared memory tensor)
+        coord: The coordinate tuple, can be hierarchical like (row, (col, cluster_idx))
+    Returns:
+        Pointer to the element at the specified coordinate
+    """
+    return x.iterator + cute.crd2idx(coord, x.layout, loc=loc, ip=ip)
+class dKGemmKernel:
+    def __init__(
+        self,
+        topk: int,
+        heads: int,
+        dim: int,
+        varlen: bool,
+    ):
+        self.varlen = varlen
+        self.topk = topk
+        self.heads = heads
+        self.dim = dim
+        # A operand dS'^T: (total_q, heads, topk), topk-major
+        self.ab_dtype = cutlass.BFloat16
+        self.a_major_mode = cute.nvgpu.OperandMajorMode.MN
+        # B operand Q: (total_q, heads, dim), dim-major
+        self.b_major_mode = cute.nvgpu.OperandMajorMode.MN
+        # Index operand I: (total_q, seqlen_k)
+        self.idx_dtype = cutlass.Int32
+        # Output dKaccum: (total_q, seqlen_k, dim), dim-major
+        self.c_dtype = cutlass.Float32
+        self.c_layout = utils.LayoutEnum.ROW_MAJOR
+        self.acc_dtype = cutlass.Float32
+        if self.topk % 256 == 1:
+            # The kernel schedule requires 2CTA instructions with this tile shape
+            self.cluster_shape_mn = (2, 1)
+            self.use_2cta_instrs = True
+        else:
+            self.cluster_shape_mn = (1, 1)
+            self.use_2cta_instrs = False
+        self.cta_group = tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE
+        self.atom_thr_size = 2 if self.use_2cta_instrs else 1
+        self.cta_tile_shape_mnk = (128, 64, 1)
+        self.mma_tiler_dK = (
+            self.cta_tile_shape_mnk[0] * self.atom_thr_size,
+            self.cta_tile_shape_mnk[1],
+            1,
+        )
+        self.arch = "sm_100"
+        self.occupancy = 1
+        # Set specialized warp ids
+        self.epilogue_warp_id = [0, 1, 2, 3]
+        self.mma_warp_id = 4
+        self.tma_warp_id = 5
+        self.sched_warp_id = 6
+        self.threads_per_cta = 32 * len(
+            [
+                self.mma_warp_id,
+                self.tma_warp_id,
+                self.sched_warp_id,
+                *self.epilogue_warp_id,
+            ]
+        )
+        # Set barrier id for cta sync, epilogue sync and tmem ptr sync
+        self.epilog_sync_bar_id = 1
+        self.tmem_alloc_sync_bar_id = 2
+        self.I_load_bar_id = 3
+    @cute.jit
+    def __call__(
+        self,
+        dS: cute.Tensor,  # (batch, seqlen_q, heads, topk) or (total_q, heads, topk)
+        I: cute.Tensor,  # (batch, seqlen_q, heads, topk) or (total_q, heads, topk)
+        Q: cute.Tensor,  # (batch, seqlen_q, heads, dim) or (total_q, heads, dim)
+        dKaccum: cute.Tensor,  # (batch, seqlen_k, dim) or (total_k, dim)
+        cuSeqlensQ: Optional[cute.Tensor],  # (batch + 1,)
+        cuSeqlensK: Optional[cute.Tensor],  # (batch + 1,)
+        stream: cuda.CUstream,
+    ):
+        if const_expr(self.ab_dtype != dS.element_type):
+            raise TypeError(f"Type must match: {self.ab_dtype} != {dS.element_type}")
+        if const_expr(self.ab_dtype != Q.element_type):
+            raise TypeError(f"Type must match: {self.ab_dtype} != {Q.element_type}")
+        if const_expr(self.c_dtype != dKaccum.element_type):
+            raise TypeError(f"Type must match: {self.c_dtype} != {dKaccum.element_type}")
+        if const_expr(self.varlen):
+            assert cuSeqlensQ is not None
+            assert cuSeqlensK is not None
+        # For non-varlen, group batch and seqlen modes into token mode
+        if const_expr(not self.varlen):
+            batch_divmod = cute.FastDivmodDivisor(dS.shape[0])
+            dS = cute.group_modes(dS, 0, 2)
+            I = cute.group_modes(I, 0, 2)
+            Q = cute.group_modes(Q, 0, 2)
+            dKaccum = cute.group_modes(dKaccum, 0, 2)
+        else:
+            batch_divmod = cute.FastDivmodDivisor(0)
+        # Permute everything to (_, heads, tokens) for MMA
+        dS_mkl = cute.make_tensor(dS.iterator, cute.select(dS.layout, [2, 1, 0]))
+        I_ml = cute.make_tensor(I.iterator, cute.select(I.layout, [1, 0]))
+        Q_nkl = cute.make_tensor(Q.iterator, cute.select(Q.layout, [2, 1, 0]))
+        dKaccum_nl = cute.make_tensor(dKaccum.iterator, cute.select(dKaccum.layout, [1, 0]))
+        # Configure tiled mma
+        self.tiled_mma_dK = utils.sm100.make_trivial_tiled_mma(
+            self.ab_dtype,
+            self.ab_dtype,
+            self.a_major_mode,
+            self.b_major_mode,
+            self.acc_dtype,
+            self.cta_group,
+            self.mma_tiler_dK[:2],
+        )
+        # Compute mma/cluster/tile shapes
+        mma_inst_shape_k = cute.size(self.tiled_mma_dK.shape_mnk, mode=[2])
+        mma_inst_tile_k = 4 if self.heads == 64 else 8
+        self.mma_tiler_dK = (
+            self.mma_tiler_dK[0],
+            self.mma_tiler_dK[1],
+            mma_inst_shape_k * mma_inst_tile_k,
+        )
+        self.cta_tile_shape_mnk_dK = (
+            self.mma_tiler_dK[0] // self.atom_thr_size,
+            self.mma_tiler_dK[1],
+            self.mma_tiler_dK[2],
+        )
+        # Compute cluster layout
+        self.cluster_layout_vmnk = cute.tiled_divide(
+            cute.make_layout((*self.cluster_shape_mn, 1)),
+            (self.atom_thr_size,),
+        )
+        # Compute number of multicast CTAs for A/B
+        self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2])
+        self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1])
+        self.is_a_mcast = self.num_mcast_ctas_a > 1
+        self.is_b_mcast = self.num_mcast_ctas_b > 1
+        # Compute epilogue subtile shape for TMA store
+        # self.epi_tile_dK = utils.sm100.compute_epilogue_tile_shape(
+        #     self.cta_tile_shape_mnk_dK,
+        #     self.use_2cta_instrs,
+        #     self.c_layout,
+        #     self.c_dtype,
+        # )
+        self.epi_tile_dK = (cute.make_layout(128), cute.make_layout(32))
+        self.epi_tile_dK_width = cute.size(self.epi_tile_dK[1].shape)
+        self.dK_smem_layout = utils.sm100.make_smem_layout_epi(
+            self.c_dtype, self.c_layout, self.epi_tile_dK, 1
+        )
+        smem_capacity = utils.get_smem_capacity_in_bytes()
+        # TMEM pipeline stages
+        # 64 TMEM columns per stage, 8 stages can fit in TMEM
+        self.num_acc_stage = 8
+        self.num_tmem_alloc_cols = 512
+        # SMEM pipeline stages
+        # TMA load Q (persists for the whole token)
+        self.num_Q_stage = 1
+        assert const_expr(self.mma_tiler_dK[1] * self.num_Q_stage == self.dim)
+        Q_smem_layout_stage_one = utils.sm100.make_smem_layout_b(
+            self.tiled_mma_dK, self.mma_tiler_dK, self.ab_dtype, 1
+        )
+        Q_bytes = cute.size_in_bytes(self.ab_dtype, Q_smem_layout_stage_one) * self.num_Q_stage
+        # TMA store-reduce dK
+        self.num_c_stage = 2
+        dK_bytes_per_stage = cute.size_in_bytes(self.c_dtype, self.dK_smem_layout)
+        dK_bytes = dK_bytes_per_stage * self.num_c_stage
+        # cp.async load I
+        self.num_I_stage = 2
+        self.I_smem_layout_staged = cute.make_layout(
+            (self.mma_tiler_dK[0] // self.atom_thr_size, self.num_I_stage)
+        )
+        I_bytes = cute.size_in_bytes(self.idx_dtype, self.I_smem_layout_staged)
+        # TMA load dS
+        dS_smem_layout_stage_one = utils.sm100.make_smem_layout_a(
+            self.tiled_mma_dK, self.mma_tiler_dK, self.ab_dtype, 1
+        )
+        dS_bytes_per_stage = cute.size_in_bytes(self.ab_dtype, dS_smem_layout_stage_one)
+        mbar_helpers_bytes = 1024
+        # Increase dS stages to fill SMEM
+        self.num_dS_stage = (
+            smem_capacity // self.occupancy - (mbar_helpers_bytes + dK_bytes + Q_bytes + I_bytes)
+        ) // dS_bytes_per_stage
+        dS_bytes = self.num_dS_stage * dS_bytes_per_stage
+        # Increase dK stages to fill remainder
+        self.num_c_stage += (
+            smem_capacity // self.occupancy
+            - (mbar_helpers_bytes + dK_bytes + dS_bytes + Q_bytes + I_bytes)
+        ) // dK_bytes_per_stage
+        # Increase I stages to fill remainder
+        # Compute shared memory layout
+        self.dS_smem_layout_staged = utils.sm100.make_smem_layout_a(
+            self.tiled_mma_dK, self.mma_tiler_dK, self.ab_dtype, self.num_dS_stage
+        )
+        self.Q_smem_layout_staged = utils.sm100.make_smem_layout_b(
+            self.tiled_mma_dK, self.mma_tiler_dK, self.ab_dtype, self.num_Q_stage
+        )
+        self.dK_smem_layout_staged = utils.sm100.make_smem_layout_epi(
+            self.c_dtype, self.c_layout, self.epi_tile_dK, self.num_c_stage
+        )
+        # TMA load for dS
+        dS_op = utils.sm100.cluster_shape_to_tma_atom_A(
+            self.cluster_shape_mn, self.tiled_mma_dK.thr_id
+        )
+        self.dS_smem_layout = cute.slice_(self.dS_smem_layout_staged, (None, None, None, 0))
+        tma_atom_dS, tma_tensor_dS = cute.nvgpu.make_tiled_tma_atom_A(
+            dS_op,
+            dS_mkl,
+            self.dS_smem_layout,
+            self.mma_tiler_dK,
+            self.tiled_mma_dK,
+            self.cluster_layout_vmnk.shape,
+        )
+        # TMA load for Q
+        Q_op = utils.sm100.cluster_shape_to_tma_atom_B(
+            self.cluster_shape_mn, self.tiled_mma_dK.thr_id
+        )
+        self.Q_smem_layout = cute.slice_(self.Q_smem_layout_staged, (None, None, None, 0))
+        tma_atom_Q, tma_tensor_Q = cute.nvgpu.make_tiled_tma_atom_B(
+            Q_op,
+            Q_nkl,
+            self.Q_smem_layout,
+            self.mma_tiler_dK,
+            self.tiled_mma_dK,
+            self.cluster_layout_vmnk.shape,
+        )
+        self.Q_load_bytes = (
+            cute.size_in_bytes(self.ab_dtype, self.Q_smem_layout) * self.atom_thr_size
+        )
+        self.dS_load_bytes = (
+            cute.size_in_bytes(self.ab_dtype, self.dS_smem_layout) * self.atom_thr_size
+        )
+        # coalesced store for dKaccum: 1 warp copies 1 epi tile row
+        vector_width_dK = self.epi_tile_dK_width * self.c_dtype.width // cute.arch.WARP_SIZE
+        copy_atom_dK = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            self.c_dtype,
+            num_bits_per_copy=vector_width_dK,
+        )
+        vector_elts_dK = vector_width_dK // self.c_dtype.width
+        thr_layout_dK = cute.make_layout((cute.arch.WARP_SIZE,))
+        val_layout_dK = cute.make_layout((vector_elts_dK,))
+        tiled_copy_dK = cute.make_tiled_copy_tv(copy_atom_dK, thr_layout_dK, val_layout_dK)
+        # cp.async load I: 1 warp copies 32 values, 4 warps copy 128 values
+        copy_atom_I = cute.make_copy_atom(
+            cute.nvgpu.cpasync.CopyG2SOp(),
+            self.idx_dtype,
+            num_bits_per_copy=32,
+        )
+        vI = cute.make_layout((1,))
+        tI = cute.make_layout((128,))
+        tiled_copy_I = cute.make_tiled_copy_tv(copy_atom_I, tI, vI)
+        # Setup clc stage by default
+        self.num_clc_stage = 1
+        assert self.num_clc_stage == 1, "Only single-stage CLC pipeline is supported"
+        # Response size is 4B * 4 elements
+        self.num_clc_response_bytes = 16
+        # Compute grid size and set up tile scheduler
+        total_q = cute.size(dS_mkl.shape[2])
+        cluster_shape_mnl = (*self.cluster_shape_mn, 1)
+        num_ctas_mnl = (self.cluster_shape_mn[0] * total_q, self.cluster_shape_mn[1], 1)
+        self.tile_sched_params = utils.ClcDynamicPersistentTileSchedulerParams(
+            num_ctas_mnl, cluster_shape_mnl
+        )
+        grid = utils.ClcDynamicPersistentTileScheduler.get_grid_shape(self.tile_sched_params)
+        # Define shared storage for kernel
+        buffer_align_bytes = 1024
+        @cute.struct
+        class SharedStorage:
+            dS_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_dS_stage * 2]
+            Q_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_Q_stage * 2]
+            acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]
+            tmem_dealloc_mbar: cutlass.Int64
+            tmem_holding_buf: cutlass.Int32
+            clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2]
+            clc_response: cute.struct.MemRange[cutlass.Int32, 4]
+            sI: cute.struct.Align[
+                cute.struct.MemRange[self.idx_dtype, cute.cosize(self.I_smem_layout_staged)],
+                buffer_align_bytes,
+            ]
+            sdS: cute.struct.Align[
+                cute.struct.MemRange[self.ab_dtype, cute.cosize(self.dS_smem_layout_staged)],
+                buffer_align_bytes,
+            ]
+            sQ: cute.struct.Align[
+                cute.struct.MemRange[self.ab_dtype, cute.cosize(self.Q_smem_layout_staged)],
+                buffer_align_bytes,
+            ]
+            sdK: cute.struct.Align[
+                cute.struct.MemRange[self.c_dtype, cute.cosize(self.dK_smem_layout_staged)],
+                buffer_align_bytes,
+            ]
+        # Launch the kernel synchronously
+        self.kernel(
+            self.tiled_mma_dK,
+            tma_atom_dS,
+            tma_tensor_dS,
+            tma_atom_Q,
+            tma_tensor_Q,
+            tiled_copy_dK,
+            dKaccum_nl,
+            tiled_copy_I,
+            I_ml,
+            batch_divmod,
+            cuSeqlensQ,
+            cuSeqlensK,
+            self.cluster_layout_vmnk,
+            self.dS_smem_layout_staged,
+            self.Q_smem_layout_staged,
+            self.dK_smem_layout_staged,
+            self.I_smem_layout_staged,
+            self.epi_tile_dK,
+            self.tile_sched_params,
+            SharedStorage,
+        ).launch(
+            grid=grid,
+            block=[self.threads_per_cta, 1, 1],
+            cluster=(*self.cluster_shape_mn, 1),
+            stream=stream,
+        )
+    # GPU device kernel
+    @cute.kernel
+    def kernel(
+        self,
+        tiled_mma_dK: cute.TiledMma,
+        tma_atom_dS: cute.CopyAtom,
+        mdS_mkl: cute.Tensor,
+        tma_atom_Q: Optional[cute.CopyAtom],
+        mQ_nkl: Optional[cute.Tensor],
+        tiled_copy_dK: cute.TiledCopy,
+        mdKaccum_nl: Optional[cute.Tensor],
+        tiled_copy_I: cute.TiledCopy,
+        mI_ml: cute.Tensor,
+        batch_divmod: cute.FastDivmodDivisor,
+        cuSeqlensQ: Optional[cute.Tensor],
+        cuSeqlensK: Optional[cute.Tensor],
+        cluster_layout_vmnk: cute.Layout,
+        dS_smem_layout_staged: cute.ComposedLayout,
+        Q_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout],
+        dK_smem_layout_staged: Union[cute.ComposedLayout, cute.Layout],
+        I_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout],
+        epi_tile_dK: cute.Tile,
+        tile_sched_params: utils.ClcDynamicPersistentTileSchedulerParams,
+        SharedStorage: cutlass.Constexpr[Callable],
+    ):
+        warp_idx = cute.arch.warp_idx()
+        warp_idx = cute.arch.make_warp_uniform(warp_idx)
+        #
+        # Prefetch tma desc
+        #
+        if warp_idx == self.tma_warp_id:
+            cpasync.prefetch_descriptor(tma_atom_dS)
+            cpasync.prefetch_descriptor(tma_atom_Q)
+        #
+        # Setup cta/thread coordinates
+        #
+        # Coords inside cluster
+        bidx, bidy, bidz = cute.arch.block_idx()
+        mma_tile_coord_v = bidx % self.atom_thr_size
+        is_leader_cta = mma_tile_coord_v == 0
+        cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+        is_first_cta_in_cluster = cta_rank_in_cluster == 0
+        block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
+        # Coord inside cta
+        tidx, _, _ = cute.arch.thread_idx()
+        #
+        # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier
+        #
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(SharedStorage)
+        # Initialize mainloop ab_pipeline (barrier) and states
+        dS_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+        num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
+        dS_pipeline_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, num_tma_producer
+        )
+        dS_producer, dS_consumer = pipeline.PipelineTmaUmma.create(
+            barrier_storage=storage.dS_full_mbar_ptr.data_ptr(),
+            num_stages=self.num_dS_stage,
+            producer_group=dS_pipeline_producer_group,
+            consumer_group=dS_pipeline_consumer_group,
+            tx_count=self.dS_load_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        Q_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+        Q_pipeline_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, num_tma_producer
+        )
+        Q_pipeline = pipeline.PipelineTmaUmma.create(
+            barrier_storage=storage.Q_full_mbar_ptr.data_ptr(),
+            num_stages=self.num_Q_stage,
+            producer_group=Q_pipeline_producer_group,
+            consumer_group=Q_pipeline_consumer_group,
+            tx_count=self.Q_load_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        )
+        # Initialize acc_pipeline (barrier) and states
+        acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+        num_acc_consumer_threads = len(self.epilogue_warp_id) * (2 if self.use_2cta_instrs else 1)
+        acc_pipeline_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, num_acc_consumer_threads
+        )
+        acc_pipeline = pipeline.PipelineUmmaAsync.create(
+            barrier_storage=storage.acc_full_mbar_ptr.data_ptr(),
+            num_stages=self.num_acc_stage,
+            producer_group=acc_pipeline_producer_group,
+            consumer_group=acc_pipeline_consumer_group,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        )
+        # Initialize clc_pipeline (barrier) and states
+        clc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+        cluster_size = cute.size(self.cluster_shape_mn)
+        num_clc_consumer_threads = 32 * (1 + cluster_size * (1 + len(self.epilogue_warp_id) + 1))
+        clc_pipeline_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, num_clc_consumer_threads
+        )
+        clc_pipeline = pipeline.PipelineClcFetchAsync.create(
+            barrier_storage=storage.clc_mbar_ptr.data_ptr(),
+            num_stages=self.num_clc_stage,
+            producer_group=clc_pipeline_producer_group,
+            consumer_group=clc_pipeline_consumer_group,
+            tx_count=self.num_clc_response_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        )
+        tmem_alloc_barrier = pipeline.NamedBarrier(
+            barrier_id=self.tmem_alloc_sync_bar_id,
+            num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)),
+        )
+        # Tensor memory dealloc barrier init
+        tmem = utils.TmemAllocator(
+            storage.tmem_holding_buf.ptr,
+            barrier_for_retrieve=tmem_alloc_barrier,
+            allocator_warp_id=self.epilogue_warp_id[0],
+            is_two_cta=self.use_2cta_instrs,
+            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
+        )
+        # Cluster arrive after barrier init
+        pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True)
+        # Initial clc response pointer
+        clc_response_ptr = storage.clc_response.data_ptr()
+        clc_consumer_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_clc_stage
+        )
+        #
+        # Setup smem tensor A/B/C
+        #
+        # (MMA, MMA_M, MMA_K, STAGE)
+        sdS = storage.sdS.get_tensor(
+            dS_smem_layout_staged.outer,
+            swizzle=dS_smem_layout_staged.inner,
+        )
+        # (MMA_M, STAGE)
+        sI = storage.sI.get_tensor(I_smem_layout_staged)
+        sQ = storage.sQ.get_tensor(
+            Q_smem_layout_staged.outer,
+            swizzle=Q_smem_layout_staged.inner,
+        )
+        sdK = storage.sdK.get_tensor(
+            dK_smem_layout_staged.outer,
+            swizzle=dK_smem_layout_staged.inner,
+        )
+        #
+        # Compute multicast mask for A/B buffer full
+        #
+        a_full_mcast_mask = None
+        b_full_mcast_mask = None
+        if const_expr(self.is_a_mcast or self.is_b_mcast or self.use_2cta_instrs):
+            a_full_mcast_mask = cpasync.create_tma_multicast_mask(
+                cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
+            )
+            b_full_mcast_mask = cpasync.create_tma_multicast_mask(
+                cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1
+            )
+        #
+        # Local_tile partition global tensors
+        #
+        # (bM, bK, RestM, RestK, RestL)
+        gdS_mkl = cute.local_tile(
+            mdS_mkl, cute.slice_(self.mma_tiler_dK, (None, 0, None)), (None, None, None)
+        )
+        # (bM, RestM, RestL)
+        gI_ml = cute.local_tile(mI_ml, cute.slice_(self.mma_tiler_dK, (None, 0, 0)), (None, None))
+        # (bN, bK, RestN, RestK, RestL)
+        gQ_nkl = cute.local_tile(
+            mQ_nkl,
+            cute.slice_(self.mma_tiler_dK, (0, None, None)),
+            (None, None, None),
+        )
+        n_tile_cnt = cute.size(gQ_nkl, mode=[2])
+        # (bN, RestN, RestL)
+        gdKaccum_nl = cute.local_tile(
+            mdKaccum_nl,
+            cute.slice_(self.mma_tiler_dK, (0, None, 0)),
+            (None, None),
+        )
+        m_tile_cnt = cute.size(gdS_mkl, mode=[2])
+        k_tile_cnt = cute.size(gdS_mkl, mode=[3])
+        assert const_expr(k_tile_cnt == 1)
+        #
+        # Partition global tensor for TiledMMA_A/B/C
+        #
+        thr_mma_dK = tiled_mma_dK.get_slice(mma_tile_coord_v)
+        # (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
+        tCgQ = thr_mma_dK.partition_B(gQ_nkl)
+        # (MMA, MMA_M, MMA_N, RestM, RestN, RestL)
+        tCgdKaccum_fake = thr_mma_dK.partition_C(
+            cute.make_identity_tensor(
+                cute.append(cute.slice_(self.mma_tiler_dK, (None, None, 0)), 1, up_to_rank=5)
+            )
+        )
+        # (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
+        tCgdS = thr_mma_dK.partition_A(gdS_mkl)
+        #
+        # Partition global/shared tensor for TMA load A/B
+        #
+        # TMA load A partition_S/D
+        a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape)
+        # ((atom_v, rest_v), STAGE)
+        # ((atom_v, rest_v), RestM, RestK, RestL)
+        tAsdS, tAgdS = cpasync.tma_partition(
+            tma_atom_dS,
+            block_in_cluster_coord_vmnk[2],
+            a_cta_layout,
+            cute.group_modes(sdS, 0, 3),
+            cute.group_modes(tCgdS, 0, 3),
+        )
+        # TMA load B partition_S/D
+        b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape)
+        # ((atom_v, rest_v), RestM, RestK, RestL)
+        tBsQ, tBgQ = cpasync.tma_partition(
+            tma_atom_Q,
+            block_in_cluster_coord_vmnk[1],
+            b_cta_layout,
+            cute.group_modes(sQ, 0, 3),
+            cute.group_modes(tCgQ, 0, 3),
+        )
+        #
+        # Partition shared/tensor memory tensor for TiledMMA_A/B/C
+        #
+        # (MMA, MMA_M, MMA_K, STAGE)
+        tCrdS = tiled_mma_dK.make_fragment_A(sdS)
+        # (MMA, MMA_N, MMA_K, STAGE)
+        tCrQ = tiled_mma_dK.make_fragment_B(sQ)
+        # (MMA, MMA_M, MMA_N)
+        acc_shape_dK = tiled_mma_dK.partition_shape_C(self.mma_tiler_dK[:2])
+        # (MMA, MMA_M, MMA_N, STAGE)
+        tCtdK_fake = tiled_mma_dK.make_fragment_C(cute.append(acc_shape_dK, self.num_acc_stage))
+        #
+        # Cluster wait before tensor memory alloc
+        #
+        pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk)
+        #
+        # Construct the scheduler
+        #
+        tile_sched = utils.ClcDynamicPersistentTileScheduler.create(
+            tile_sched_params,
+            cute.arch.block_idx(),
+            cute.arch.grid_dim(),
+            clc_response_ptr,
+        )
+        work_tile = tile_sched.initial_work_tile_info()
+        #
+        # Specialized TMA load warp
+        #
+        if warp_idx == self.tma_warp_id:
+            Q_producer_state = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Producer, self.num_Q_stage
+            )
+            #
+            # Persistent tile scheduling loop
+            #
+            dS_producer.reset()
+            peek_dS_empty_status = dS_producer.try_acquire()
+            while work_tile.is_valid_tile:
+                # Get tile coord from tile scheduler
+                cur_tile_coord = work_tile.tile_idx
+                total_q_coord = cur_tile_coord[0] // self.cluster_shape_mn[0]
+                #
+                # Slice to per mma tile index
+                #
+                # ((atom_v, rest_v), RestM, RestK)
+                tAgdS_slice = tAgdS[(None, None, None, total_q_coord)]
+                # ((atom_v, rest_v), RestN, RestK)
+                tBgQ_slice = tBgQ[(None, None, None, total_q_coord)]
+                #
+                # Tma load loop -- fully unrolled as all sizes are static
+                #
+                for m_tile in cutlass.range_constexpr(m_tile_cnt):
+                    # Conditionally wait for AB buffer empty
+                    dS_handle = dS_producer.acquire_and_advance(peek_dS_empty_status)
+                    # TMA load dS
+                    cute.copy(
+                        tma_atom_dS,
+                        tAgdS_slice[(None, m_tile, 0)],
+                        tAsdS[(None, dS_handle.index)],
+                        tma_bar_ptr=dS_handle.barrier,
+                        mcast_mask=a_full_mcast_mask,
+                    )
+                    peek_dS_empty_status = dS_producer.try_acquire()
+                    if m_tile == 0:
+                        # TMA load Q only on first m-tile
+                        for n_tile in cutlass.range_constexpr(n_tile_cnt):
+                            Q_pipeline.producer_acquire(Q_producer_state)
+                            Q_load_barrier = Q_pipeline.producer_get_barrier(Q_producer_state)
+                            cute.copy(
+                                tma_atom_Q,
+                                tBgQ_slice[(None, n_tile, 0)],
+                                tBsQ[(None, n_tile)],
+                                tma_bar_ptr=Q_load_barrier,
+                                mcast_mask=b_full_mcast_mask,
+                            )
+                            Q_pipeline.producer_commit(Q_producer_state)
+                            Q_producer_state.advance()
+                #
+                # Advance to next tile
+                #
+                clc_pipeline.consumer_wait(clc_consumer_state)
+                work_tile = tile_sched.get_current_work()
+                clc_pipeline.consumer_release(clc_consumer_state)
+                clc_consumer_state.advance()
+            #
+            # Wait A/B buffer empty
+            #
+            dS_producer.tail()
+            Q_pipeline.producer_tail(Q_producer_state)
+        #
+        # Sched warp
+        #
+        elif warp_idx == self.sched_warp_id and is_first_cta_in_cluster:
+            #
+            # Persistent tile scheduling loop
+            #
+            clc_producer_state = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.ProducerConsumer, self.num_clc_stage
+            )
+            while work_tile.is_valid_tile:
+                #
+                # Advance to next tile
+                #
+                clc_pipeline.producer_acquire(clc_producer_state)
+                mbarrier_addr = clc_pipeline.producer_get_barrier(clc_producer_state)
+                tile_sched.advance_to_next_work(mbarrier_addr)
+                clc_producer_state.advance()
+                clc_pipeline.consumer_wait(clc_consumer_state)
+                work_tile = tile_sched.get_current_work()
+                clc_pipeline.consumer_release(clc_consumer_state)
+                clc_consumer_state.advance()
+            clc_pipeline.producer_tail(clc_producer_state)
+        #
+        # Specialized MMA warp
+        #
+        elif warp_idx == self.mma_warp_id:
+            #
+            # Retrieving tensor memory ptr and make accumulator tensor
+            #
+            tmem.wait_for_alloc()
+            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+            # (MMA, MMA_M, MMA_N, STAGE)
+            tCtdK_base = cute.make_tensor(tmem_ptr, tCtdK_fake.layout)
+            #
+            # Persistent tile scheduling loop
+            #
+            acc_producer_state = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Producer, self.num_acc_stage
+            )
+            Q_consumer_state = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Consumer, self.num_Q_stage
+            )
+            # Peek (try_wait) AB buffer full for k_tile = 0
+            dS_consumer.reset()
+            peek_dS_full_status = cutlass.Boolean(1)
+            if is_leader_cta:
+                peek_dS_full_status = dS_consumer.try_wait()
+            while work_tile.is_valid_tile:
+                # Get tile coord from tile scheduler
+                cur_tile_coord = work_tile.tile_idx
+                total_q_coord = cur_tile_coord[0] // self.cluster_shape_mn[0]
+                #
+                # Mma mainloop
+                #
+                Q_consumer_state_wait = Q_consumer_state.clone()
+                for m_tile in cutlass.range(m_tile_cnt):
+                    if is_leader_cta:
+                        # Conditionally wait for AB buffer full
+                        dS_handle = dS_consumer.wait_and_advance(peek_dS_full_status)
+                        for n_tile in cutlass.range_constexpr(n_tile_cnt):
+                            tCtdK = tCtdK_base[(None, None, None, acc_producer_state.index)]
+                            tiled_mma_dK.set(tcgen05.Field.ACCUMULATE, False)
+                            acc_pipeline.producer_acquire(acc_producer_state)
+                            if m_tile == 0:
+                                Q_pipeline.consumer_wait(Q_consumer_state_wait)
+                                Q_consumer_state_wait.advance()
+                            num_kblocks = cute.size(tCrdS, mode=[2])
+                            for kblk_idx in cutlass.range(num_kblocks, unroll_full=True):
+                                cute.gemm(
+                                    tiled_mma_dK,
+                                    tCtdK,
+                                    tCrdS[(None, None, kblk_idx, dS_handle.index)],
+                                    tCrQ[(None, None, kblk_idx, n_tile)],
+                                    tCtdK,
+                                )
+                                # Enable accumulate on tCtdK after first kblock
+                                tiled_mma_dK.set(tcgen05.Field.ACCUMULATE, True)
+                            if m_tile == m_tile_cnt - 1:
+                                Q_pipeline.consumer_release(Q_consumer_state)
+                                Q_consumer_state.advance()
+                            acc_pipeline.producer_commit(acc_producer_state)
+                            acc_producer_state.advance()
+                        # Async arrive AB buffer empty
+                        dS_handle.release()
+                        peek_dS_full_status = dS_consumer.try_wait()
+                    else:
+                        for n_tile in cutlass.range_constexpr(n_tile_cnt):
+                            acc_producer_state.advance()
+                #
+                # Advance to next tile
+                #
+                clc_pipeline.consumer_wait(clc_consumer_state)
+                work_tile = tile_sched.get_current_work()
+                clc_pipeline.consumer_release(clc_consumer_state)
+                clc_consumer_state.advance()
+            #
+            # Wait for accumulator buffer empty
+            #
+            acc_pipeline.producer_tail(acc_producer_state)
+        #
+        # Specialized epilogue warps
+        #
+        elif warp_idx < self.mma_warp_id:
+            #
+            # Alloc tensor memory buffer
+            #
+            tmem.allocate(self.num_tmem_alloc_cols)
+            #
+            # Retrieving tensor memory ptr and make accumulator tensor
+            #
+            tmem.wait_for_alloc()
+            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+            # (MMA, MMA_M, MMA_N, STAGE)
+            tCtdK_base = cute.make_tensor(tmem_ptr, tCtdK_fake.layout)
+            # Both gemms share accumulator and TMA store pipelines
+            acc_consumer_state = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Consumer, self.num_acc_stage
+            )
+            c_producer_group = pipeline.CooperativeGroup(
+                pipeline.Agent.Thread,
+                32 * len(self.epilogue_warp_id),
+            )
+            I_load_barrier = pipeline.NamedBarrier(
+                barrier_id=self.I_load_bar_id,
+                num_threads=32 * len(self.epilogue_warp_id),
+            )
+            # (EPI_TOPK, REST_TOPK, TOKENS)
+            gI_tile = cute.local_tile(
+                gI_ml,
+                (epi_tile_dK[0],),
+                (0 if is_leader_cta else 1, None, None),
+            )
+            thr_copy_I = tiled_copy_I.get_slice(tidx)
+            # (COPY_ATOM, EPI_TOPK, REST_TOPK, TOKENS)
+            tIgI = thr_copy_I.partition_S(gI_tile)
+            # (COPY_ATOM, EPI_TOPK, STAGE)
+            tIsI = thr_copy_I.partition_D(sI)
+            num_subtiles_executed = 0
+            while work_tile.is_valid_tile:
+                # Get tile coord from tile scheduler
+                cur_tile_coord = work_tile.tile_idx
+                token = cutlass.Int64(cur_tile_coord[0]) // self.cluster_shape_mn[0]
+                # Get batch index from token_coord and cuSeqlensQ
+                # Get seqlen_k offset from cuSeqlensK
+                if const_expr(self.varlen):
+                    batch = get_batch_from_cu_tensor(token, cuSeqlensQ)
+                    seqlen_k_offset = cuSeqlensK[batch]
+                else:
+                    _, batch = divmod(token, batch_divmod)
+                    seqlen_k_offset = Int32(0)  # unused
+                sI_read_stage = 0
+                sI_write_stage = 0
+                # Prefetch load I
+                I_load_barrier.arrive_and_wait()
+                for m_tile in cutlass.range_constexpr(min(m_tile_cnt, self.num_I_stage - 1)):
+                    cute.copy(
+                        tiled_copy_I,
+                        tIgI[(None, None, m_tile, token)],
+                        tIsI[(None, None, sI_write_stage)],
+                    )
+                    cute.arch.cp_async_commit_group()
+                    sI_write_stage = (sI_write_stage + 1) % self.num_I_stage
+                for m_tile in cutlass.range(m_tile_cnt):
+                    I_load_barrier.arrive_and_wait()
+                    # cp.async load I
+                    if m_tile < m_tile_cnt - 1:
+                        cute.copy(
+                            tiled_copy_I,
+                            tIgI[(None, None, m_tile + 1, token)],
+                            tIsI[(None, None, sI_write_stage)],
+                        )
+                    cute.arch.cp_async_commit_group()
+                    sI_write_stage = (sI_write_stage + 1) % self.num_I_stage
+                    cute.arch.cp_async_wait_group(self.num_I_stage - 1)
+                    I_load_barrier.arrive_and_wait()
+                    sI_tile = sI[(None, sI_read_stage)]
+                    for n_tile in cutlass.range_constexpr(n_tile_cnt):
+                        # (MMA, MMA_M, MMA_N, RestM, RestN, RestL)
+                        mma_tile_coord_mnl = (
+                            m_tile,
+                            n_tile,
+                            token,
+                        )
+                        acc_consumer_state = self.epilogue_scatter_reduce(
+                            tidx,
+                            0 if is_leader_cta else 1,
+                            tiled_copy_dK,
+                            tCtdK_base,
+                            sdK,
+                            sI_tile,
+                            gdKaccum_nl,
+                            tCgdKaccum_fake,
+                            batch,
+                            seqlen_k_offset,
+                            epi_tile_dK,
+                            num_subtiles_executed,
+                            mma_tile_coord_mnl,
+                            acc_consumer_state,
+                            acc_pipeline,
+                        )
+                        num_subtiles_executed += self.mma_tiler_dK[1] // cute.size(epi_tile_dK[1])
+                    # Advance I consumer pipeline
+                    sI_read_stage = (sI_read_stage + 1) % self.num_I_stage
+                cute.arch.cp_async_wait_group(0)
+                #
+                # Advance to next tile
+                #
+                clc_pipeline.consumer_wait(clc_consumer_state)
+                work_tile = tile_sched.get_current_work()
+                clc_pipeline.consumer_release(clc_consumer_state)
+                clc_consumer_state.advance()
+            # # Wait for C store complete
+            # dK_pipeline.producer_tail()
+            #
+            # Dealloc the tensor memory buffer
+            #
+            tmem.relinquish_alloc_permit()
+            tmem.free(tmem_ptr)
+    @cute.jit
+    def epilogue_scatter_reduce(
+        self,
+        epi_tidx: Int32,
+        cta_idx: Int32,
+        tiled_copy_c: cute.CopyAtom,
+        tCtAcc_base: cute.Tensor,
+        sC: cute.Tensor,
+        sI_tile: cute.Tensor,
+        gC_base: cute.Tensor,
+        tCgC_fake: cute.Tensor,
+        batch: Int32,
+        seqlen_k_offset: Int32,
+        epi_tile: cute.Tile,
+        num_subtiles_executed: Int32,
+        mma_tile_coord_mnl: Tuple[Int32, Int32, cutlass.Int64],
+        acc_consumer_state: pipeline.PipelineState,
+        acc_pipeline: pipeline.PipelineAsync,
+    ) -> pipeline.PipelineState:
+        warp_idx = cute.arch.make_warp_uniform(epi_tidx // 32)
+        # Layout transformation for tCgC_base
+        # ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, REST_M, REST_N, REST_L)
+        # -> ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), REST_M, REST_N, REST_L)
+        tCgC_fake = transform_partitioned_tensor_layout(tCgC_fake)
+        # Layout transformation for tCtAcc_base
+        # ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, STAGE)
+        # -> ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), STAGE)
+        tCtAcc = transform_partitioned_tensor_layout(tCtAcc_base)
+        tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = epilogue_tmem_copy_and_partition(
+            self,
+            epi_tidx,
+            tCtAcc,
+            tCgC_fake,
+            epi_tile,
+            self.use_2cta_instrs,
+        )
+        tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype)
+        tiled_copy_r2s, tRS_rC, tRS_sC = epilogue_smem_copy_and_partition(
+            self, tiled_copy_t2r, tTR_rC, epi_tidx, sC
+        )
+        # (EPI_TILE_N, EPI_N, RestN, RestL)
+        gC_epi = cute.flat_divide(gC_base, (epi_tile[1],))
+        # (EPI_TILE_N, EPI_N, SEQLEN_K)
+        gC_epi = gC_epi[None, None, mma_tile_coord_mnl[1], None]
+        # (EPI_TILE_N, MMA_M, STAGE)
+        sC_epi = cute.make_tensor(
+            sC.iterator,
+            cute.select(sC.layout, [1, 0, 2]),
+            # swizzle=sC.layout.inner,
+        )
+        thr_copy_c = tiled_copy_c.get_slice(epi_tidx % 32)
+        # (COPY_ATOM_N, COPY_N, MMA_M, STAGE)
+        tCsC = thr_copy_c.partition_S(sC_epi)
+        # (COPY_ATOM_N, COPY_N)
+        tCrC = cute.make_fragment_like(tCsC[None, None, 0, 0])
+        # (COPY_ATOM_N, COPY_N, EPI_N, SEQLEN_K)
+        tCgC = thr_copy_c.partition_D(gC_epi)
+        # Set tensor memory buffer for current tile
+        # (T2R, T2R_M, T2R_N, EPI_M, EPI_M)
+        tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_consumer_state.index)]
+        #
+        # Wait for accumulator buffer full
+        #
+        acc_pipeline.consumer_wait(acc_consumer_state)
+        tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
+        epilog_sync_barrier = pipeline.NamedBarrier(
+            barrier_id=self.epilog_sync_bar_id,
+            num_threads=32 * len(self.epilogue_warp_id),
+        )
+        #
+        # Store accumulator to global memory in subtiles
+        #
+        subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
+        for subtile_idx in range(subtile_cnt):
+            #
+            # Load accumulator from tensor memory buffer to register
+            #
+            tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)]
+            cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc)
+            #
+            # Convert to C type
+            #
+            acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load()
+            acc_vec = acc_vec.to(self.c_dtype)
+            tRS_rC.store(acc_vec)
+            #
+            # Store C to shared memory
+            #
+            c_buffer = (num_subtiles_executed + subtile_idx) % self.num_c_stage
+            cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)])
+            # Fence and barrier to make sure shared memory store is visible to TMA store
+            cute.arch.fence_proxy("async.shared", space="cta")
+            epilog_sync_barrier.arrive_and_wait()
+            #
+            # TMA store C to global memory (issued by lane 0 from all epi warps)
+            #
+            for topk_idx_in_warp in cutlass.range(32):
+                topk_idx = topk_idx_in_warp + warp_idx * 32
+                seqlen_k_idx_in_batch = sI_tile[topk_idx]
+                if const_expr(self.varlen):
+                    seqlen_k_idx = seqlen_k_idx_in_batch + seqlen_k_offset
+                else:
+                    seqlen_k_idx = (batch, seqlen_k_idx_in_batch)
+                cute.copy(tiled_copy_c, tCsC[(None, None, topk_idx, c_buffer)], tCrC)
+                # Skip -1 sentinel slots (invalid top-k entries)
+                if seqlen_k_idx_in_batch >= 0:
+                    for j in cutlass.range_constexpr(cute.size(tCrC, mode=[1])):
+                        for i in cutlass.range_constexpr(cute.size(tCrC, mode=[0])):
+                            ptr = elem_pointer(tCgC, (i, j, subtile_idx, seqlen_k_idx))
+                            cute.arch.atomic_add(
+                                ptr=ptr,
+                                val=tCrC[i, j],
+                            )
+            epilog_sync_barrier.arrive_and_wait()
+        epilog_sync_barrier.arrive_and_wait()
+        #
+        # Async arrive accumulator buffer empty
+        #
+        with cute.arch.elect_one():
+            acc_pipeline.consumer_release(acc_consumer_state)
+        acc_consumer_state.advance()
+        return acc_consumer_state
+    def check_can_implement(self):
+        """Check if parameters are valid.
+        :raises testing.CantImplementError: If the mma tiler, cluster shape, or alignments are invalid
+        """
+        if self.dim != 64:
+            raise testing.CantImplementError(f"Only dim = 64 supported for dK gemm, got {self.dim}")
+        # Check valid MMA tile shape and cluster shape
+        if not (
+            (not self.use_2cta_instrs and self.mma_tiler_dK[0] in [64, 128])
+            or (self.use_2cta_instrs and self.mma_tiler_dK[0] in [128, 256])
+        ):
+            raise testing.CantImplementError(
+                f"Invalid mma tiler & use_2cta_instrs: {self.mma_tiler_dK}, {self.use_2cta_instrs}"
+            )
+        if self.mma_tiler_dK[1] not in range(32, 257, 32):
+            raise testing.CantImplementError(f"Invalid mma tiler N: {self.mma_tiler_dK[1]}")
+        # Skip illegal cluster shape
+        if self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0:
+            raise testing.CantImplementError(f"Invalid cluster shape M: {self.cluster_shape_mn[0]}")
+        # Skip invalid cluster shape
+        is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0
+        if (
+            self.cluster_shape_mn[0] * self.cluster_shape_mn[1] > 16
+            or self.cluster_shape_mn[0] <= 0
+            or self.cluster_shape_mn[1] <= 0
+            or not is_power_of_2(self.cluster_shape_mn[0])
+            or not is_power_of_2(self.cluster_shape_mn[1])
+        ):
+            raise testing.CantImplementError(f"Invalid cluster shape: {self.cluster_shape_mn}")
+        # Check that all tensors are 16B aligned for TMA
+        def check_contiguous_16B_alignment(dtype, num_major_elements):
+            num_contiguous_elements = 16 * 8 // dtype.width
+            return num_major_elements % num_contiguous_elements == 0
+        if (
+            not check_contiguous_16B_alignment(self.ab_dtype, self.topk)
+            or not check_contiguous_16B_alignment(self.ab_dtype, self.dim)
+            or not check_contiguous_16B_alignment(self.c_dtype, self.dim)
+        ):
+            raise testing.CantImplementError(
+                f"Invalid tensor alignment: {self.ab_dtype=}, {self.c_dtype=}, {self.topk=}, {self.dim=}"
+            )

build/torch-cuda/flash_bwd_mla_dq_dqv_sm100.py ADDED Viewed

	@@ -0,0 +1,1297 @@

+# Copyright (c) 2026, Colfax International.
+"""
+CuTe DSL implementation of dQ+dQv gemm for DSA backward.
+Performs both dQ = dS @ K and dQv = dS @ V, where K and V are
+gathered according to index tensor mIdxTopK.
+This uses MQA with 128 heads.
+Inputs:
+    - dS:      [batch, seqlen_q, nheads, top_k] or [total_q, nheads, top_k]
+    - K:       [batch, seqlen_k, hdim]          or [total_k, hdim]
+    - V:       [batch, seqlen_k, hdim_v]        or [total_k, hdim_v]
+    - IdxTopK: [batch, seqlen_q, top_k]         or [total_q, top_k]
+Outputs:
+    - dQ:  [batch, seqlen_q, nheads, hdim]   or [total_q, nheads, hdim]
+    - dQv: [batch, seqlen_q, nheads, hdim_v] or [total_q, nheads, hdim_v]
+All sizes are known at compile time except seqlen_q, which is the batch dimension.
+Representative numbers are:
+    - nheads = 128
+    - hdim   = 64
+    - hdim_v = 512
+    - top_k  = 2048
+We launch a cluster of shape (1, 2) with mma tile 128x256, so that one cluster
+covers the full dQv mma. dS is loaded via TMA and multicast across the CTAs.
+Cluster 0 also performs the dQ mma with tile size 128x64.
+K and V are loaded via CpAsync, with logic according to CpasyncGatherKVManager.
+"""
+from functools import partial
+from typing import Optional, Tuple, Type
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+import cutlass.utils as utils
+from cutlass import Int32, const_expr
+from cutlass.cute import FastDivmodDivisor
+from cutlass.cute.nvgpu import cpasync, tcgen05
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
+from .topk_gather_kv import CpasyncGatherKVManager
+from .utils import get_batch_from_cu_tensor
+class dQdQvGemmKernel:
+    def __init__(
+        self,
+        acc_dtype: Type[cutlass.Numeric],
+        nheads: int,
+        head_dim_k: Optional[int],
+        head_dim_v: int,
+        top_k: int,
+    ):
+        self.acc_dtype: Type[cutlass.Numeric] = acc_dtype
+        self.nheads = nheads
+        assert self.nheads == 128, (
+            "only 128 heads supported; will expand to include 64 heads in a future PR."
+        )
+        self.head_dim_k = head_dim_k or 0  # when head_dim_k not provided, dQ is not computed
+        self.head_dim_v = head_dim_v
+        self.top_k = top_k
+        self.tile_k = 128
+        self.cluster_shape_mn = (1, 2)
+        self.mma_tiler_dQ = (self.nheads, self.head_dim_k, self.tile_k)
+        self.mma_tiler_dQv = (self.nheads, self.head_dim_v // 2, self.tile_k)
+        self.num_mainloop_iters = self.top_k // self.tile_k
+        self.arch = "sm_100"
+        self.cta_group = tcgen05.CtaGroup.ONE
+        self.occupancy = 1
+        self.threads_per_warp = cute.arch.WARP_SIZE
+        # ---- Set specialized warp ids ----
+        self.epilogue_warp_ids = (0, 1, 2, 3)
+        self.kv_load_warp_ids = (4, 5, 6, 7)
+        self.mma_warp_id = 8
+        self.tma_warp_id = 9
+        self.sched_warp_id = 10
+        self.threads_per_cta = 32 * len(
+            (
+                self.mma_warp_id,
+                self.tma_warp_id,
+                self.sched_warp_id,
+                *self.epilogue_warp_ids,
+                *self.kv_load_warp_ids,
+            )
+        )
+        # ---- Set barrier id for cta sync, epilogue sync and tmem ptr sync ----
+        self.epilog_sync_bar_id = 1
+        self.tmem_alloc_sync_bar_id = 2
+        self.tmem_dealloc_sync_bar_id = 3
+        self.kv_load_sync_bar_id = 4
+        self.epilog_sync_barrier = pipeline.NamedBarrier(
+            barrier_id=self.epilog_sync_bar_id,
+            num_threads=self.threads_per_warp * len(self.epilogue_warp_ids),
+        )
+        self.kv_load_sync_barrier = pipeline.NamedBarrier(
+            barrier_id=self.kv_load_sync_bar_id,
+            num_threads=self.threads_per_warp * len(self.kv_load_warp_ids),
+        )
+        self.is_persistent = False
+        # ---- pipeline stages ---- TODO: tune these
+        self.num_stages_dS = 2
+        self.num_stages_KV = 2
+        self.num_stages_acc = 1
+        self.num_stages_clc = 1
+        # ---- register allocation ----
+        self.num_regs_KV = 224
+        self.num_regs_epi = 128
+        self.num_regs_other = 112
+    @cute.jit
+    def __call__(
+        self,
+        mdS: cute.Tensor,
+        mK: Optional[cute.Tensor],
+        mV: cute.Tensor,
+        mdQ: cute.Tensor,
+        mdQv: cute.Tensor,
+        mIdxTopK: cute.Tensor,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
+        stream: cuda.CUstream = None,
+    ):
+        self.compute_dQ = const_expr(mK is not None)
+        # ---- dtype info ----
+        self.ds_dtype: Type[cutlass.Numeric] = mdS.element_type
+        self.kv_dtype: Type[cutlass.Numeric] = mV.element_type
+        self.dq_dtype: Type[cutlass.Numeric] = mdQv.element_type
+        if const_expr(self.compute_dQ):
+            assert self.kv_dtype == mV.element_type
+            assert self.dq_dtype == mdQ.element_type
+        varlen_q = const_expr(mCuSeqlensQ is not None)
+        varlen_k = const_expr(mCuSeqlensK is not None)
+        # ------------------------------------------------------------------ #
+        # Reshape GMEM layouts for static strides                            #
+        # ------------------------------------------------------------------ #
+        seqlen_q = Int32(0) if const_expr(varlen_q) else mdS.shape[1]
+        seqlen_q_divmod = FastDivmodDivisor(seqlen_q)
+        seqlen_k = Int32(0) if const_expr(varlen_k) else mV.shape[1]
+        # ---- group batch and seqlen modes in nonvarlen case ----
+        def group_batch_seqlen(t: cute.Tensor, varlen: bool) -> cute.Tensor:
+            if const_expr(not varlen):
+                t = cute.make_tensor(
+                    t.iterator,
+                    cute.make_layout(
+                        (t.shape[1], t.shape[0], *t.shape[2:]),
+                        stride=(t.stride[1], t.stride[0], *t.stride[2:]),
+                    ),
+                )
+                t = cute.group_modes(t, 0, 2)
+            return t
+        mdS = group_batch_seqlen(mdS, varlen_q)
+        mdQv = group_batch_seqlen(mdQv, varlen_q)
+        mV = group_batch_seqlen(mV, varlen_k)
+        mIdxTopK = group_batch_seqlen(mIdxTopK, varlen_q)
+        if const_expr(self.compute_dQ):
+            mdQ = group_batch_seqlen(mdQ, varlen_q)
+            mK = group_batch_seqlen(mK, varlen_k)
+        # ---- transpose and make static modes static ----
+        def static_reshape(t: cute.Tensor, *static_shapes) -> cute.Tensor:
+            static_modes = range(1, len(t.shape))
+            return cute.make_tensor(
+                t.iterator,
+                cute.make_layout(
+                    (*static_shapes, t.shape[0]),
+                    stride=(*(t.stride[i] for i in static_modes), t.stride[0]),
+                ),
+            )
+        mdS = static_reshape(mdS, self.nheads, self.top_k)
+        mdQv = static_reshape(mdQv, self.nheads, self.head_dim_v)
+        mV = static_reshape(mV, self.head_dim_v)
+        mIdxTopK = static_reshape(mIdxTopK, self.top_k)
+        if const_expr(self.compute_dQ):
+            mdQ = static_reshape(mdQ, self.nheads, self.head_dim_k)
+            mK = static_reshape(mK, self.head_dim_k)
+        # ---- layout info ----
+        self.ds_major_mode = utils.LayoutEnum.from_tensor(mdS).mma_major_mode()
+        self.kv_major_mode = utils.LayoutEnum.from_tensor(mV).mma_major_mode()
+        self.dq_layout = utils.LayoutEnum.from_tensor(mdQv)
+        if const_expr(self.compute_dQ):
+            assert self.dq_layout == utils.LayoutEnum.from_tensor(mdQ)
+        # ------------------------------------------------------------------ #
+        # Setup attributes that depend on kernel inputs                      #
+        # ------------------------------------------------------------------ #
+        tiled_mma_v = utils.sm100.make_trivial_tiled_mma(
+            self.ds_dtype,
+            self.ds_major_mode,
+            self.kv_major_mode,
+            self.acc_dtype,
+            self.cta_group,
+            self.mma_tiler_dQv[:2],
+        )
+        if const_expr(self.compute_dQ):
+            tiled_mma_k = utils.sm100.make_trivial_tiled_mma(
+                self.ds_dtype,
+                self.ds_major_mode,
+                self.kv_major_mode,
+                self.acc_dtype,
+                self.cta_group,
+                self.mma_tiler_dQ[:2],
+            )
+        self.cta_tile_shape_dQv = (
+            self.mma_tiler_dQv[0],
+            self.mma_tiler_dQv[1],
+            self.mma_tiler_dQv[2],
+        )
+        # ---- Compute cluster layout ----
+        self.cluster_layout_vmnk = cute.tiled_divide(
+            cute.make_layout((*self.cluster_shape_mn, 1)),
+            (tiled_mma_v.thr_id.shape,),
+        )
+        # ---- Compute number of multicast CTAs for A/B ----
+        self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2])
+        # ---- Compute epi tiles for dQ/dQv ----
+        self.epi_tile_dQv = utils.sm100.compute_epilogue_tile_shape(
+            self.cta_tile_shape_dQv,
+            False,  # use_2cta_instrs
+            self.dq_layout,
+            self.dq_dtype,
+        )
+        self.epi_tile_dQ = None
+        if const_expr(self.compute_dQ):
+            self.epi_tile_dQ = utils.sm100.compute_epilogue_tile_shape(
+                self.mma_tiler_dQ,
+                False,
+                self.dq_layout,
+                self.dq_dtype,
+            )
+        # ---- Device-specific attributes ----
+        self.smem_capacity = utils.get_smem_capacity_in_bytes()
+        self.num_tmem_alloc_cols = 512
+        # ------------------------------------------------------------------ #
+        # Make SMEM layouts                                                  #
+        # ------------------------------------------------------------------ #
+        sdS_layout = utils.sm100.make_smem_layout_a(
+            tiled_mma_v,
+            self.mma_tiler_dQ,
+            self.ds_dtype,
+            self.num_stages_dS,
+        )
+        sV_layout = utils.sm100.make_smem_layout_b(
+            tiled_mma_v,
+            self.mma_tiler_dQv,
+            self.kv_dtype,
+            self.num_stages_KV,
+        )
+        sdQv_layout = utils.sm100.make_smem_layout_epi(
+            self.dq_dtype,
+            self.dq_layout,
+            self.epi_tile_dQv,
+            self.num_stages_acc,
+        )
+        sK_layout, sdQ_layout = None, None
+        if const_expr(self.compute_dQ):
+            sK_layout = utils.sm100.make_smem_layout_b(
+                tiled_mma_k,
+                self.mma_tiler_dQ,
+                self.kv_dtype,
+                self.num_stages_KV,
+            )
+            sdQ_layout = utils.sm100.make_smem_layout_epi(
+                self.dq_dtype,
+                self.dq_layout,
+                self.epi_tile_dQ,
+                self.num_stages_acc,
+            )
+        # ------------------------------------------------------------------ #
+        # Set up TMA load/stores                                             #
+        # ------------------------------------------------------------------ #
+        atom_thr_size = cute.size(tiled_mma_v.thr_id.shape)
+        # ---- Setup TMA load for dS ----
+        dS_op = utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma_v.thr_id)
+        dS_smem_layout = cute.slice_(sdS_layout, (None, None, None, 0))
+        tma_atom_dS, tma_tensor_dS = cute.nvgpu.make_tiled_tma_atom_A(
+            dS_op,
+            mdS,
+            dS_smem_layout,
+            self.mma_tiler_dQv,
+            tiled_mma_v,
+            self.cluster_layout_vmnk.shape,
+        )
+        dS_copy_size = cute.size_in_bytes(self.ds_dtype, dS_smem_layout)
+        self.num_tma_load_bytes = dS_copy_size * atom_thr_size
+        # ---- Setup TMA store for dQ and dQV ----
+        dQv_epi_smem_layout = cute.select(sdQv_layout, mode=[0, 1])
+        tma_atom_dQv, tma_tensor_dQv = cpasync.make_tiled_tma_atom(
+            cpasync.CopyBulkTensorTileS2GOp(), mdQv, dQv_epi_smem_layout, self.epi_tile_dQv
+        )
+        tma_atom_dQ, tma_tensor_dQ = None, None
+        if const_expr(self.compute_dQ):
+            dQ_epi_smem_layout = cute.select(sdQ_layout, mode=[0, 1])
+            tma_atom_dQ, tma_tensor_dQ = cpasync.make_tiled_tma_atom(
+                cpasync.CopyBulkTensorTileS2GOp(), mdQ, dQ_epi_smem_layout, self.epi_tile_dQ
+            )
+        # ------------------------------------------------------------------ #
+        # Set up shared storage for SMEM                                     #
+        # ------------------------------------------------------------------ #
+        self.buffer_align_bytes = 1024
+        sdS_size = cute.cosize(sdS_layout)
+        sK_size = cute.cosize(sK_layout) if const_expr(self.compute_dQ) else 0
+        sV_size = cute.cosize(sV_layout)
+        sdQ_size = cute.cosize(sdQ_layout) if const_expr(self.compute_dQ) else 0
+        sdQv_size = cute.cosize(sdQv_layout)
+        assert sdQ_size <= sK_size, f"require {sdQ_size=} <= {sK_size=}"
+        assert sdQv_size <= sV_size, f"require {sdQv_size=} <= {sV_size=}"
+        self.overlap_kv_epi = self.compute_dQ
+        if const_expr(self.overlap_kv_epi):
+            sdQ_size = 0
+            sdQv_size = 0
+        @cute.struct
+        class SharedStorage:
+            mbar_ptr_dS: cute.struct.MemRange[cutlass.Int64, self.num_stages_dS * 2]
+            mbar_ptr_KV: cute.struct.MemRange[cutlass.Int64, self.num_stages_KV * 2]
+            mbar_ptr_dQ_dQv: cute.struct.MemRange[cutlass.Int64, self.num_stages_acc * 2]
+            mbar_ptr_KV_cpasync: cute.struct.MemRange[cutlass.Int64, self.num_stages_KV * 2]
+            mbar_ptr_load_kv_epi: cute.struct.MemRange[cutlass.Int64, 2]
+            # Tmem holding buffer
+            tmem_dealloc_mbar: cutlass.Int64
+            tmem_holding_buf: cutlass.Int32
+            # Clc pointers
+            clc_ptr: cute.struct.Align[
+                cute.struct.MemRange[cutlass.Int64, self.num_stages_clc * 2], 16
+            ]
+            clc_response_ptr: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, 4], 16]
+            # Smem tensors
+            sdS: cute.struct.Align[
+                cute.struct.MemRange[self.ds_dtype, sdS_size],
+                self.buffer_align_bytes,
+            ]
+            sK: cute.struct.Align[
+                cute.struct.MemRange[self.kv_dtype, sK_size],
+                self.buffer_align_bytes,
+            ]
+            sV: cute.struct.Align[
+                cute.struct.MemRange[self.kv_dtype, sV_size],
+                self.buffer_align_bytes,
+            ]
+            sdQ: cute.struct.Align[
+                cute.struct.MemRange[self.dq_dtype, sdQ_size],
+                self.buffer_align_bytes,
+            ]
+            sdQv: cute.struct.Align[
+                cute.struct.MemRange[self.dq_dtype, sdQv_size],
+                self.buffer_align_bytes,
+            ]
+        self.shared_storage = SharedStorage
+        # ---- Compute grid size ----
+        self.tile_sched_params, grid = self._compute_grid(
+            mdQv,
+            self.cta_tile_shape_dQv,
+            self.cluster_shape_mn,
+        )
+        self.num_clc_response_bytes = 16
+        # permute grid to conform to grid_dim_z <= 65536 constraint;
+        # this is undone in the kernel
+        grid = (grid[2], grid[1], grid[0])
+        # cute.printf("dQ/dQv grid: {}", grid)
+        # print("dQ/dQv SMEM: ", self.shared_storage.size_in_bytes())
+        # ---- Launch the kernel synchronously ----
+        self.kernel(
+            tiled_mma_k if const_expr(self.compute_dQ) else None,
+            tiled_mma_v,
+            tma_atom_dS,
+            tma_tensor_dS,
+            mK,
+            mV,
+            tma_atom_dQ,
+            tma_tensor_dQ,
+            tma_atom_dQv,
+            tma_tensor_dQv,
+            mIdxTopK,
+            mCuSeqlensQ,
+            mCuSeqlensK,
+            seqlen_q_divmod,
+            self.cluster_layout_vmnk,
+            sdS_layout,
+            sK_layout,
+            sV_layout,
+            sdQ_layout,
+            sdQv_layout,
+            self.epi_tile_dQ,
+            self.epi_tile_dQv,
+            self.tile_sched_params,
+            seqlen_k,
+        ).launch(
+            grid=grid,
+            block=[self.threads_per_cta, 1, 1],
+            cluster=(*self.cluster_shape_mn, 1),
+            smem=self.shared_storage.size_in_bytes(),
+            stream=stream,
+        )
+    # GPU device kernel
+    @cute.kernel
+    def kernel(
+        self,
+        tiled_mma_k: Optional[cute.TiledMma],
+        tiled_mma_v: cute.TiledMma,
+        tma_atom_dS: cute.CopyAtom,
+        mdS: cute.Tensor,
+        mK: Optional[cute.Tensor],
+        mV: cute.Tensor,
+        tma_atom_dQ: Optional[cute.CopyAtom],
+        mdQ: Optional[cute.Tensor],
+        tma_atom_dQv: cute.CopyAtom,
+        mdQv: cute.Tensor,
+        mIdxTopK: cute.Tensor,
+        mCuSeqlensQ: cute.Tensor,
+        mCuSeqlensK: cute.Tensor,
+        seqlen_q_divmod: FastDivmodDivisor,
+        cluster_layout_vmnk: cute.Layout,
+        sdS_layout: cute.ComposedLayout,
+        sK_layout: Optional[cute.ComposedLayout],
+        sV_layout: cute.ComposedLayout,
+        sdQ_layout: Optional[cute.ComposedLayout],
+        sdQv_layout: cute.ComposedLayout,
+        epi_tile_dQ: Optional[cute.Tile],
+        epi_tile_dQv: cute.Tile,
+        tile_sched_params: utils.ClcDynamicPersistentTileSchedulerParams,
+        seqlen_k_static: Int32,
+    ):
+        """
+        GPU device kernel performing the Persistent batched GEMM computation.
+        """
+        warp_idx = cute.arch.warp_idx()
+        warp_idx = cute.arch.make_warp_uniform(warp_idx)
+        # ------------------------------------------------------------------ #
+        # Prefetch TMA descriptors                                           #
+        # ------------------------------------------------------------------ #
+        if warp_idx == self.tma_warp_id:
+            cpasync.prefetch_descriptor(tma_atom_dS)
+            cpasync.prefetch_descriptor(tma_atom_dQv)
+            if const_expr(self.compute_dQ):
+                cpasync.prefetch_descriptor(tma_atom_dQ)
+        # ------------------------------------------------------------------ #
+        # Cluster coordinates                                                #
+        # ------------------------------------------------------------------ #
+        bidx, bidy, bidz = cute.arch.block_idx()
+        gridx, gridy, gridz = cute.arch.grid_dim()
+        mma_v_tile_coord_v = bidx % cute.size(tiled_mma_v.thr_id.shape)
+        if const_expr(self.compute_dQ):
+            mma_k_tile_coord_v = bidx % cute.size(tiled_mma_k.thr_id.shape)
+        cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+        block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
+        is_first_cta = cta_rank_in_cluster == 0
+        tidx, _, _ = cute.arch.thread_idx()
+        # ------------------------------------------------------------------ #
+        # Shared storage allocation                                          #
+        # ------------------------------------------------------------------ #
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(self.shared_storage)
+        # ------------------------------------------------------------------ #
+        # Initialize pipelines                                               #
+        # ------------------------------------------------------------------ #
+        ThreadCooperativeGroup = partial(pipeline.CooperativeGroup, pipeline.Agent.Thread)
+        kv_commit_group = ThreadCooperativeGroup(1)
+        clc_producer_group = ThreadCooperativeGroup(1)
+        num_clc_consumer_threads = 32 * (
+            1  # sched warp on CTA0 only
+            + cute.size(self.cluster_shape_mn)
+            * (1 + len(self.epilogue_warp_ids) + len(self.kv_load_warp_ids) + 1)
+            # tma + epi + kv_load + mma, on BOTH CTAs
+        )
+        clc_consumer_group = ThreadCooperativeGroup(num_clc_consumer_threads)
+        mma_warp = ThreadCooperativeGroup(1)
+        tma_warp = ThreadCooperativeGroup(cute.size(self.cluster_shape_mn))
+        tma_warp_local = ThreadCooperativeGroup(1)
+        epilogue_warps = ThreadCooperativeGroup(len(self.epilogue_warp_ids))
+        load_warps = ThreadCooperativeGroup(len(self.kv_load_warp_ids))
+        pipeline_dS = pipeline.PipelineTmaUmma.create(
+            barrier_storage=storage.mbar_ptr_dS.data_ptr(),
+            num_stages=self.num_stages_dS,
+            producer_group=mma_warp,
+            consumer_group=tma_warp,
+            tx_count=self.num_tma_load_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        )
+        pipeline_KV = pipeline.PipelineAsyncUmma.create(
+            barrier_storage=storage.mbar_ptr_KV.data_ptr(),
+            num_stages=self.num_stages_KV,
+            producer_group=kv_commit_group,
+            consumer_group=mma_warp,
+            defer_sync=True,
+        )
+        pipeline_dQ_dQv = pipeline.PipelineUmmaAsync.create(
+            barrier_storage=storage.mbar_ptr_dQ_dQv.data_ptr(),
+            num_stages=self.num_stages_acc,
+            producer_group=mma_warp,
+            consumer_group=epilogue_warps,
+            defer_sync=True,
+        )
+        pipeline_clc = pipeline.PipelineClcFetchAsync.create(
+            barrier_storage=storage.clc_ptr.data_ptr(),
+            num_stages=self.num_stages_clc,
+            producer_group=clc_producer_group,
+            consumer_group=clc_consumer_group,
+            tx_count=self.num_clc_response_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        )
+        pipeline_load_kv_epi = None
+        if const_expr(self.overlap_kv_epi):
+            pipeline_load_kv_epi = pipeline.PipelineAsync.create(
+                barrier_storage=storage.mbar_ptr_load_kv_epi.data_ptr(),
+                num_stages=1,
+                producer_group=epilogue_warps,
+                consumer_group=load_warps,
+                defer_sync=True,
+            )
+        # ------------------------------------------------------------------ #
+        # TMEM Allocation                                                    #
+        # ------------------------------------------------------------------ #
+        tmem_alloc_barrier = pipeline.NamedBarrier(
+            barrier_id=self.tmem_alloc_sync_bar_id,
+            num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_ids)),
+        )
+        # ---- Tensor memory dealloc barrier init ----
+        tmem = utils.TmemAllocator(
+            storage.tmem_holding_buf.ptr,
+            barrier_for_retrieve=tmem_alloc_barrier,
+            allocator_warp_id=self.epilogue_warp_ids[0],
+            is_two_cta=False,
+            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
+        )
+        # ---- Cluster arrive after barrier init ----
+        pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True)
+        # ---- Initial clc response pointer ----
+        clc_response_ptr = storage.clc_response_ptr.data_ptr()
+        clc_consumer_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_stages_clc
+        )
+        # ------------------------------------------------------------------ #
+        # SMEM tensors                                                       #
+        # ------------------------------------------------------------------ #
+        # (MMA, MMA_M, MMA_K, STAGE)
+        sdS = storage.sdS.get_tensor(sdS_layout.outer, swizzle=sdS_layout.inner)
+        if const_expr(self.compute_dQ):
+            sK = storage.sK.get_tensor(sK_layout.outer, swizzle=sK_layout.inner)
+        sV = storage.sV.get_tensor(sV_layout.outer, swizzle=sV_layout.inner)
+        # sdQ = storage.sdQ.get_tensor(sdQ_layout.outer, swizzle=sdQ_layout.inner)
+        if const_expr(self.compute_dQ):
+            sdQ = cute.make_tensor(
+                cute.recast_ptr(sK.iterator, sdQ_layout.inner, self.dq_dtype), sdQ_layout.outer
+            )
+        if const_expr(not self.compute_dQ):
+            sdQv = storage.sdQv.get_tensor(sdQv_layout.outer, swizzle=sdQv_layout.inner)
+        else:
+            sdQv = cute.make_tensor(
+                cute.recast_ptr(sV.iterator, sdQv_layout.inner, self.dq_dtype), sdQv_layout.outer
+            )
+        dS_full_mcast_mask = cpasync.create_tma_multicast_mask(
+            cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
+        )
+        # ------------------------------------------------------------------ #
+        # Global tile partitioning                                           #
+        # ------------------------------------------------------------------ #
+        # (bM, bK, RestM, RestK, RestL)
+        gdS = cute.local_tile(
+            mdS, cute.slice_(self.mma_tiler_dQv, (None, 0, None)), (None, None, None)
+        )
+        # (bM, bN, RestM, RestN, RestL)
+        if const_expr(self.compute_dQ):
+            gdQ = cute.local_tile(
+                mdQ, cute.slice_(self.mma_tiler_dQ, (None, None, 0)), (None, None, None)
+            )
+        gdQv = cute.local_tile(
+            mdQv, cute.slice_(self.mma_tiler_dQv, (None, None, 0)), (None, None, None)
+        )
+        k_tile_cnt = cute.size(gdS, mode=[3])
+        # ------------------------------------------------------------------ #
+        # TiledMMA partitioning                                              #
+        # ------------------------------------------------------------------ #
+        thr_mma_v = tiled_mma_v.get_slice(mma_v_tile_coord_v)
+        # (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
+        tdQvgdS = thr_mma_v.partition_A(gdS)
+        # (MMA, MMA_M, MMA_N, RestM, RestN, RestL)
+        tdQvgdQv = thr_mma_v.partition_C(gdQv)
+        if const_expr(self.compute_dQ):
+            thr_mma_k = tiled_mma_k.get_slice(mma_k_tile_coord_v)
+            tdQgdQ = thr_mma_k.partition_C(gdQ)
+        # ------------------------------------------------------------------ #
+        # TMA partition for dS                                               #
+        # ------------------------------------------------------------------ #
+        dS_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape)
+        # ((atom_v, rest_v), STAGE)
+        # ((atom_v, rest_v), RestM, RestK, RestL)
+        tdSsdS, tdSgdS = cpasync.tma_partition(
+            tma_atom_dS,
+            block_in_cluster_coord_vmnk[2],
+            dS_cta_layout,
+            cute.group_modes(sdS, 0, 3),
+            cute.group_modes(tdQvgdS, 0, 3),
+        )
+        # ------------------------------------------------------------------ #
+        # MMA fragments                                                      #
+        # ------------------------------------------------------------------ #
+        # (MMA, MMA_M, MMA_K, STAGE)
+        tdQvrdS = tiled_mma_v.make_fragment_A(sdS)
+        # (MMA, MMA_N, MMA_K, STAGE)
+        tdQvrV = tiled_mma_v.make_fragment_B(sV)
+        # (MMA, MMA_M, MMA_N)
+        acc_v_shape = tiled_mma_v.partition_shape_C(self.mma_tiler_dQv[:2])
+        # (MMA, MMA_M, MMA_N, STAGE)
+        tdQvtAcc_fake = tiled_mma_v.make_fragment_C(cute.append(acc_v_shape, self.num_stages_acc))
+        if const_expr(self.compute_dQ):
+            # (MMA, MMA_N, MMA_K, STAGE)
+            tdQrK = tiled_mma_k.make_fragment_B(sK)
+            # (MMA, MMA_M, MMA_N)
+            acc_k_shape = tiled_mma_k.partition_shape_C(self.mma_tiler_dQ[:2])
+            # (MMA, MMA_M, MMA_N, STAGE)
+            tdQtAcc_fake = tiled_mma_k.make_fragment_C(
+                cute.append(acc_k_shape, self.num_stages_acc)
+            )
+        # ------------------------------------------------------------------ #
+        # Cluster wait before tensor memory alloc                            #
+        # ------------------------------------------------------------------ #
+        pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk)
+        # ------------------------------------------------------------------ #
+        # Tile Scheduler                                                     #
+        # ------------------------------------------------------------------ #
+        tile_sched = utils.ClcDynamicPersistentTileScheduler.create(
+            tile_sched_params,
+            cute.arch.block_idx(),
+            cute.arch.grid_dim(),
+            clc_response_ptr,
+        )
+        work_tile = tile_sched.initial_work_tile_info()
+        # ------------------------------------------------------------------ #
+        # TMA load warp                                                      #
+        # ------------------------------------------------------------------ #
+        if warp_idx == self.tma_warp_id:
+            cute.arch.setmaxregister_decrease(self.num_regs_other)
+            producer_state_dS = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Producer, stages=self.num_stages_dS
+            )
+            while work_tile.is_valid_tile:
+                # ---- Get tile coord from tile scheduler ----
+                token, cta, _ = work_tile.tile_idx
+                # ((atom_v, rest_v), RestK)
+                tdSgdS_slice = tdSgdS[(None, 0, None, token)]
+                # ---- mainloop ----
+                for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
+                    pipeline_dS.producer_acquire(producer_state_dS)
+                    index_dS = producer_state_dS.index
+                    # ---- TMA load dS ----
+                    cute.copy(
+                        tma_atom_dS,
+                        tdSgdS_slice[(None, k_tile)],
+                        tdSsdS[(None, index_dS)],
+                        tma_bar_ptr=pipeline_dS.producer_get_barrier(producer_state_dS),
+                        mcast_mask=dS_full_mcast_mask,
+                    )
+                    producer_state_dS.advance()
+                # ---- Advance to next tile ----
+                pipeline_clc.consumer_wait(clc_consumer_state)
+                work_tile = tile_sched.get_current_work()
+                pipeline_clc.consumer_release(clc_consumer_state)
+                clc_consumer_state.advance()
+            # ---- Wait dS buffer empty ----
+            pipeline_dS.producer_tail(producer_state_dS)
+        # ------------------------------------------------------------------ #
+        # Clc Scheduler warp                                                 #
+        # ------------------------------------------------------------------ #
+        if warp_idx == self.sched_warp_id and is_first_cta:
+            cute.arch.setmaxregister_decrease(self.num_regs_other)
+            clc_producer_state = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.ProducerConsumer, self.num_stages_clc
+            )
+            while work_tile.is_valid_tile:
+                pipeline_clc.producer_acquire(clc_producer_state)
+                mbar_addr = pipeline_clc.producer_get_barrier(clc_producer_state)
+                tile_sched.advance_to_next_work(mbar_addr)
+                clc_producer_state.advance()
+                pipeline_clc.consumer_wait(clc_consumer_state)
+                work_tile = tile_sched.get_current_work()
+                pipeline_clc.consumer_release(clc_consumer_state)
+                clc_consumer_state.advance()
+            pipeline_clc.producer_tail(clc_producer_state)
+        # ------------------------------------------------------------------ #
+        # CpAsync KV load warps                                              #
+        # ------------------------------------------------------------------ #
+        if warp_idx >= self.kv_load_warp_ids[0] and warp_idx <= self.kv_load_warp_ids[-1]:
+            cute.arch.setmaxregister_increase(self.num_regs_KV)
+            find_batch = partial(
+                self.find_batch_from_q, seqlen_q_divmod=seqlen_q_divmod, mCuSeqlensQ=mCuSeqlensQ
+            )
+            load_epi_consumer_state = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Consumer, 1
+            )
+            producer_state_KV = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Producer, stages=self.num_stages_KV
+            )
+            kv_tidx = tidx % (len(self.kv_load_warp_ids) * self.threads_per_warp)
+            kv_warp_idx = warp_idx % len(self.kv_load_warp_ids)
+            mV_cta = cute.domain_offset((cta_rank_in_cluster * (self.head_dim_v // 2), 0), mV)
+            while work_tile.is_valid_tile:
+                # ---- Get tile coord from tile scheduler ----
+                token, cta, _ = work_tile.tile_idx
+                batch_idx = find_batch(token)
+                k_batch_offset = (
+                    mCuSeqlensK[batch_idx] if const_expr(mCuSeqlensK is not None) else Int32(0)
+                )
+                seqlen_k = (
+                    mCuSeqlensK[batch_idx + 1] - k_batch_offset
+                    if const_expr(mCuSeqlensK is not None)
+                    else seqlen_k_static
+                )
+                if const_expr(mCuSeqlensK is not None):
+                    if const_expr(self.compute_dQ):
+                        mK_cur = cute.domain_offset((0, k_batch_offset), mK)[None, None]
+                    mV_cur = cute.domain_offset((0, k_batch_offset), mV_cta)[None, None]
+                else:
+                    if const_expr(self.compute_dQ):
+                        mK_cur = cute.domain_offset((0, (0, batch_idx)), mK)[None, None]
+                    mV_cur = cute.domain_offset((0, (0, batch_idx)), mV_cta)[None, None]
+                mIdxTopK_cur = mIdxTopK[None, token]
+                cpasync_gather_kv_manager = CpasyncGatherKVManager.create(
+                    mIdxTopK_cur,
+                    0,
+                    kv_tidx,
+                    kv_warp_idx,
+                    self.top_k,
+                    seqlen_k,
+                    self.mma_tiler_dQv[2],
+                    self.head_dim_k,
+                    self.mma_tiler_dQv[1],
+                    1,
+                    len(self.kv_load_warp_ids) * self.threads_per_warp,
+                    mV.element_type,
+                    1,
+                )
+                # ---- K/V load mainloop ----
+                for k_tile in cutlass.range_constexpr(k_tile_cnt):
+                    # ---- Load top-k index tensor ----
+                    cpasync_gather_kv_manager.load_index_topk(k_tile, transpose=True)
+                    stage = producer_state_KV.index
+                    pipeline_KV.producer_acquire(producer_state_KV)
+                    # ---- Load V (and optionally load K) ----
+                    cpasync_gather_kv_manager.load_X(mV_cur, sV[None, None, None, stage], True, "V")
+                    if const_expr(self.compute_dQ):
+                        if is_first_cta:
+                            cpasync_gather_kv_manager.load_X(
+                                mK_cur, sK[None, None, None, stage], True, "K"
+                            )
+                    cute.arch.cp_async_commit_group()
+                    cute.arch.cp_async_wait_group(0)
+                    self.kv_load_sync_barrier.arrive_and_wait()
+                    cute.arch.fence_proxy("async.shared", space="cta")
+                    if kv_warp_idx == 0:
+                        with cute.arch.elect_one():
+                            pipeline_KV.producer_commit(producer_state_KV)
+                    producer_state_KV.advance()
+                if const_expr(self.overlap_kv_epi):
+                    pipeline_load_kv_epi.consumer_wait(load_epi_consumer_state)
+                    with cute.arch.elect_one():
+                        pipeline_load_kv_epi.consumer_release(load_epi_consumer_state)
+                    load_epi_consumer_state.advance()
+                # ---- Advance to next tile ----
+                pipeline_clc.consumer_wait(clc_consumer_state)
+                work_tile = tile_sched.get_current_work()
+                pipeline_clc.consumer_release(clc_consumer_state)
+                clc_consumer_state.advance()
+            pipeline_KV.producer_tail(producer_state_KV)
+        # ------------------------------------------------------------------ #
+        # MMA warp                                                           #
+        # ------------------------------------------------------------------ #
+        if warp_idx == self.mma_warp_id:
+            cute.arch.setmaxregister_decrease(self.num_regs_other)
+            # --- Retrieve TMEM ptr and make accumulator tensors
+            tmem.wait_for_alloc()
+            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+            # (MMA, MMA_M, MMA_N, STAGE)
+            if const_expr(self.compute_dQ):
+                tdQtAcc_base = cute.make_tensor(tmem_ptr, tdQtAcc_fake.layout)
+            tdQvtAcc_ptr = tmem_ptr + (
+                tcgen05.find_tmem_tensor_col_offset(tdQtAcc_base)
+                if const_expr(self.compute_dQ)
+                else 0
+            )
+            tdQvtAcc_base = cute.make_tensor(tdQvtAcc_ptr, tdQvtAcc_fake.layout)
+            consumer_state_dS = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Consumer, stages=self.num_stages_dS
+            )
+            consumer_state_KV = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Consumer, stages=self.num_stages_KV
+            )
+            producer_state_dQ_dQv = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Producer, self.num_stages_acc
+            )
+            while work_tile.is_valid_tile:
+                # ---- Get tile coord from tile scheduler ----
+                token, cta, _ = work_tile.tile_idx
+                # ---- Set tensor memory buffer for current tile ----
+                # (MMA, MMA_M, MMA_N)
+                tdQvtAcc = tdQvtAcc_base[(None, None, None, producer_state_dQ_dQv.index)]
+                if const_expr(self.compute_dQ):
+                    tdQtAcc = tdQtAcc_base[(None, None, None, producer_state_dQ_dQv.index)]
+                # ---- Wait for accumulator buffer empty ----
+                pipeline_dQ_dQv.producer_acquire(producer_state_dQ_dQv)
+                # ---- Reset the ACCUMULATE field for each tile ----
+                tiled_mma_v.set(tcgen05.Field.ACCUMULATE, False)
+                if const_expr(self.compute_dQ):
+                    tiled_mma_k.set(tcgen05.Field.ACCUMULATE, False)
+                # ---- Mma mainloop ----
+                for k_tile in cutlass.range_constexpr(k_tile_cnt):
+                    pipeline_dS.consumer_wait(consumer_state_dS)
+                    pipeline_KV.consumer_wait(consumer_state_KV)
+                    dS_stage = consumer_state_dS.index
+                    KV_stage = consumer_state_KV.index
+                    num_kblocks = cute.size(tdQvrdS, mode=[2])
+                    for kblk_idx in cutlass.range(num_kblocks, unroll_full=True):
+                        # dQv += dS @ V
+                        cute.gemm(
+                            tiled_mma_v,
+                            tdQvtAcc,
+                            tdQvrdS[(None, None, kblk_idx, dS_stage)],
+                            tdQvrV[(None, None, kblk_idx, KV_stage)],
+                            tdQvtAcc,
+                        )
+                        # Enable accumulate on tdQvtAcc after first kblock
+                        tiled_mma_v.set(tcgen05.Field.ACCUMULATE, True)
+                        if const_expr(self.compute_dQ):
+                            if is_first_cta:
+                                # dQ += dS @ K
+                                cute.gemm(
+                                    tiled_mma_k,
+                                    tdQtAcc,
+                                    tdQvrdS[(None, None, kblk_idx, dS_stage)],
+                                    tdQrK[(None, None, kblk_idx, KV_stage)],
+                                    tdQtAcc,
+                                )
+                                # Enable accumulate on tdQtAcc after first kblock
+                                tiled_mma_k.set(tcgen05.Field.ACCUMULATE, True)
+                    pipeline_dS.consumer_release(consumer_state_dS)
+                    pipeline_KV.consumer_release(consumer_state_KV)
+                    consumer_state_dS.advance()
+                    consumer_state_KV.advance()
+                pipeline_dQ_dQv.producer_commit(producer_state_dQ_dQv)
+                producer_state_dQ_dQv.advance()
+                # ---- Advance to next tile ----
+                pipeline_clc.consumer_wait(clc_consumer_state)
+                work_tile = tile_sched.get_current_work()
+                pipeline_clc.consumer_release(clc_consumer_state)
+                clc_consumer_state.advance()
+            pipeline_dQ_dQv.producer_tail(producer_state_dQ_dQv)
+        # ------------------------------------------------------------------ #
+        # Epilogue warps                                                     #
+        # ------------------------------------------------------------------ #
+        if warp_idx >= self.epilogue_warp_ids[0] and warp_idx <= self.epilogue_warp_ids[-1]:
+            cute.arch.setmaxregister_increase(self.num_regs_epi)
+            # ---- Alloc tensor memory buffer ----
+            tmem.allocate(self.num_tmem_alloc_cols)
+            # ---- Retrieving tensor memory ptr and make accumulator tensor ----
+            tmem.wait_for_alloc()
+            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+            # (MMA, MMA_M, MMA_N, STAGE)
+            if const_expr(self.compute_dQ):
+                tdQtAcc_base = cute.make_tensor(tmem_ptr, tdQtAcc_fake.layout)
+            tdQvtAcc_ptr = tmem_ptr + (
+                tcgen05.find_tmem_tensor_col_offset(tdQtAcc_base)
+                if const_expr(self.compute_dQ)
+                else 0
+            )
+            tdQvtAcc_base = cute.make_tensor(tdQvtAcc_ptr, tdQvtAcc_fake.layout)
+            epi_idx = tidx
+            # print(f"tdQvtAcc_base.layout = {tdQvtAcc_base.layout}")
+            # ---- TMEM -> RMEM -> SMEM -> GMEM copies + partitions ----
+            tiled_copy_dQv_t2r, tTR_dQvtAcc_base, tTR_dQvrAcc = (
+                self.epilogue_tmem_copy_and_partition(
+                    epi_idx,
+                    tdQvtAcc_base,
+                    tdQvgdQv,
+                    epi_tile_dQv,
+                    self.mma_tiler_dQv,
+                    self.dq_layout,
+                    self.dq_dtype,
+                )
+            )
+            # print(f"tTR_tdQvtAcc_base.layout = {tTR_dQvtAcc_base.layout}")
+            tTR_rdQv = cute.make_rmem_tensor(tTR_dQvrAcc.shape, self.dq_dtype)
+            (
+                tiled_copy_dQv_r2s,
+                tRS_rdQv,
+                tRS_sdQv,
+            ) = self.epilogue_smem_copy_and_partition(
+                self.dq_layout,
+                self.dq_dtype,
+                tiled_copy_dQv_t2r,
+                tTR_rdQv,
+                epi_idx,
+                sdQv,
+            )
+            bSG_sdQv, bSG_gdQv_partitioned = self.epilogue_gmem_copy_and_partition(
+                tma_atom_dQv,
+                tdQvgdQv,
+                epi_tile_dQv,
+                sdQv,
+            )
+            if const_expr(self.compute_dQ):
+                (tiled_copy_dQ_t2r, tTR_dQtAcc_base, tTR_dQrAcc) = (
+                    self.epilogue_tmem_copy_and_partition(
+                        epi_idx,
+                        tdQtAcc_base,
+                        tdQgdQ,
+                        epi_tile_dQ,
+                        self.mma_tiler_dQ,
+                        self.dq_layout,
+                        self.dq_dtype,
+                    )
+                )
+                tTR_rdQ = cute.make_rmem_tensor(tTR_dQrAcc.shape, self.dq_dtype)
+                (
+                    tiled_copy_dQ_r2s,
+                    tRS_rdQ,
+                    tRS_sdQ,
+                ) = self.epilogue_smem_copy_and_partition(
+                    self.dq_layout,
+                    self.dq_dtype,
+                    tiled_copy_dQ_t2r,
+                    tTR_rdQ,
+                    epi_idx,
+                    sdQ,
+                )
+                bSG_sdQ, bSG_gdQ_partitioned = self.epilogue_gmem_copy_and_partition(
+                    tma_atom_dQ,
+                    tdQgdQ,
+                    epi_tile_dQ,
+                    sdQ,
+                )
+            acc_consumer_state = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Consumer, self.num_stages_acc
+            )
+            load_epi_producer_state = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Producer, 1
+            )
+            epi_producer_group = pipeline.CooperativeGroup(
+                pipeline.Agent.Thread,
+                cute.arch.WARP_SIZE,
+            )
+            pipeline_epi = pipeline.PipelineTmaStore.create(
+                num_stages=self.num_stages_acc,
+                producer_group=epi_producer_group,
+            )
+            # ---- Persistent tile scheduling loop for epilogue ----
+            while work_tile.is_valid_tile:
+                # ---- Get current work tile ----
+                token, cta, bid_x = work_tile.tile_idx
+                bSG_gdQv = bSG_gdQv_partitioned[
+                    (None, None, None, bid_x, cta_rank_in_cluster, token)
+                ]
+                tTR_dQvtAcc = tTR_dQvtAcc_base[
+                    (None, None, None, None, None, acc_consumer_state.index)
+                ]
+                tTR_dQvtAcc = cute.group_modes(tTR_dQvtAcc, 3, cute.rank(tTR_dQvtAcc))
+                bSG_gdQv = cute.group_modes(bSG_gdQv, 1, cute.rank(bSG_gdQv))
+                subtile_cnt_v = cute.size(tTR_dQvtAcc.shape, mode=[3])
+                epi_subtile_counter_v = 0
+                if const_expr(self.compute_dQ):
+                    bSG_gdQ = bSG_gdQ_partitioned[(None, None, None, bid_x, cta, token)]
+                    tTR_dQtAcc = tTR_dQtAcc_base[
+                        (None, None, None, None, None, acc_consumer_state.index)
+                    ]
+                    tTR_dQtAcc = cute.group_modes(tTR_dQtAcc, 3, cute.rank(tTR_dQtAcc))
+                    bSG_gdQ = cute.group_modes(bSG_gdQ, 1, cute.rank(bSG_gdQ))
+                    subtile_cnt_k = cute.size(tTR_dQtAcc.shape, mode=[3])
+                    epi_subtile_counter_k = 0
+                pipeline_dQ_dQv.consumer_wait(acc_consumer_state)
+                for subtile_idx in cutlass.range(subtile_cnt_v, unroll_full=True):
+                    store_dQ = (
+                        const_expr(self.compute_dQ)
+                        and is_first_cta
+                        and (subtile_idx < subtile_cnt_k)
+                    )
+                    if not store_dQ:
+                        tTR_dQvtAcc_mn = tTR_dQvtAcc[(None, None, None, subtile_idx)]
+                        cute.copy(tiled_copy_dQv_t2r, tTR_dQvtAcc_mn, tTR_dQvrAcc)
+                        cute.arch.fence_view_async_tmem_load()
+                        # convert to output dtype
+                        tRS_rdQv.store(
+                            tiled_copy_dQv_r2s.retile(tTR_dQvrAcc).load().to(self.dq_dtype)
+                        )
+                        epi_buffer = epi_subtile_counter_v % self.num_stages_acc
+                        cute.copy(
+                            tiled_copy_dQv_r2s, tRS_rdQv, tRS_sdQv[(None, None, None, epi_buffer)]
+                        )
+                        cute.arch.fence_proxy("async.shared", space="cta")
+                        self.epilog_sync_barrier.arrive_and_wait()
+                        if warp_idx == self.epilogue_warp_ids[0]:
+                            cute.copy(
+                                tma_atom_dQv,
+                                bSG_sdQv[(None, epi_buffer)],
+                                bSG_gdQv[(None, subtile_idx)],
+                            )
+                            pipeline_epi.producer_commit()
+                            pipeline_epi.producer_acquire()
+                        self.epilog_sync_barrier.arrive_and_wait()
+                        epi_subtile_counter_v += 1
+                    elif const_expr(self.compute_dQ):
+                        tTR_dQtAcc_mn = tTR_dQtAcc[(None, None, None, subtile_idx)]
+                        tTR_dQvtAcc_mn = tTR_dQvtAcc[(None, None, None, subtile_idx)]
+                        cute.copy(tiled_copy_dQ_t2r, tTR_dQtAcc_mn, tTR_dQrAcc)
+                        cute.copy(tiled_copy_dQv_t2r, tTR_dQvtAcc_mn, tTR_dQvrAcc)
+                        # convert to output dtype
+                        tRS_rdQ.store(tiled_copy_dQ_r2s.retile(tTR_dQrAcc).load().to(self.dq_dtype))
+                        tRS_rdQv.store(
+                            tiled_copy_dQv_r2s.retile(tTR_dQvrAcc).load().to(self.dq_dtype)
+                        )
+                        epi_buffer = epi_subtile_counter_v % self.num_stages_acc
+                        epi_buffer_k = epi_subtile_counter_k % self.num_stages_acc
+                        cute.copy(
+                            tiled_copy_dQv_r2s, tRS_rdQv, tRS_sdQv[(None, None, None, epi_buffer)]
+                        )
+                        cute.copy(
+                            tiled_copy_dQ_r2s, tRS_rdQ, tRS_sdQ[(None, None, None, epi_buffer_k)]
+                        )
+                        cute.arch.fence_proxy("async.shared", space="cta")
+                        self.epilog_sync_barrier.arrive_and_wait()
+                        if warp_idx == self.epilogue_warp_ids[0]:
+                            cute.copy(
+                                tma_atom_dQv,
+                                bSG_sdQv[(None, epi_buffer)],
+                                bSG_gdQv[(None, subtile_idx)],
+                            )
+                            cute.copy(
+                                tma_atom_dQ,
+                                bSG_sdQ[(None, epi_buffer_k)],
+                                bSG_gdQ[(None, subtile_idx)],
+                            )
+                            pipeline_epi.producer_commit()
+                            pipeline_epi.producer_acquire()
+                        self.epilog_sync_barrier.arrive_and_wait()
+                        epi_subtile_counter_v += 1
+                        epi_subtile_counter_k += 1
+                if const_expr(self.overlap_kv_epi):
+                    pipeline_load_kv_epi.producer_acquire(load_epi_producer_state)
+                    with cute.arch.elect_one():
+                        pipeline_load_kv_epi.producer_commit(load_epi_producer_state)
+                    load_epi_producer_state.advance()
+                with cute.arch.elect_one():
+                    pipeline_dQ_dQv.consumer_release(acc_consumer_state)
+                acc_consumer_state.advance()
+                # ---- Advance to next tile ----
+                pipeline_clc.consumer_wait(clc_consumer_state)
+                work_tile = tile_sched.get_current_work()
+                pipeline_clc.consumer_release(clc_consumer_state)
+                clc_consumer_state.advance()
+            # ---- Dealloc the tensor memory buffer ----
+            tmem.relinquish_alloc_permit()
+            self.epilog_sync_barrier.arrive_and_wait()
+            tmem.free(tmem_ptr)
+            pipeline_epi.producer_tail()
+    def epilogue_tmem_copy_and_partition(
+        self,
+        tidx: cutlass.Int32,
+        tAcc: cute.Tensor,
+        gC_mnl: cute.Tensor,
+        epi_tile: cute.Tile,
+        mma_tiler_mnk,
+        c_layout,
+        c_dtype,
+    ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
+        copy_atom_t2r = utils.sm100.get_tmem_load_op(
+            mma_tiler_mnk,
+            c_layout,
+            c_dtype,
+            self.acc_dtype,
+            epi_tile,
+            use_2cta_instrs=False,
+        )
+        # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE)
+        tAcc_epi = cute.flat_divide(tAcc[((None, None), 0, 0, None)], epi_tile)
+        # (EPI_TILE_M, EPI_TILE_N)
+        tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)])
+        thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
+        # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, STAGE)
+        tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi)
+        # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL)
+        gC_mnl_epi = cute.flat_divide(gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile)
+        # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL)
+        tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi)
+        # (T2R, T2R_M, T2R_N)
+        rAcc_shape = tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape
+        tTR_rAcc = cute.make_rmem_tensor(rAcc_shape, self.acc_dtype)
+        return (tiled_copy_t2r, tTR_tAcc, tTR_rAcc)
+    def epilogue_smem_copy_and_partition(
+        self,
+        c_layout,
+        c_dtype: Type[cutlass.Numeric],
+        tiled_copy_t2r: cute.TiledCopy,
+        tTR_rC: cute.Tensor,
+        tidx: cutlass.Int32,
+        sC: cute.Tensor,
+    ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
+        copy_atom_r2s = utils.sm100.get_smem_store_op(
+            c_layout, c_dtype, self.acc_dtype, tiled_copy_t2r
+        )
+        tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r)
+        # (R2S, R2S_M, R2S_N, PIPE_D)
+        thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
+        tRS_sC = thr_copy_r2s.partition_D(sC)
+        # (R2S, R2S_M, R2S_N)
+        tRS_rC = tiled_copy_r2s.retile(tTR_rC)
+        return tiled_copy_r2s, tRS_rC, tRS_sC
+    def epilogue_gmem_copy_and_partition(
+        self,
+        tma_atom,
+        gC,
+        epi_tile,
+        sC,
+    ) -> Tuple[cute.Tensor, cute.Tensor]:
+        gC_epi = cute.flat_divide(gC[((None, None), 0, 0, None, None, None)], epi_tile)
+        sC_for_tma_partition = cute.group_modes(sC, 0, 2)
+        gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2)
+        bSG_sC, bSG_gC = cpasync.tma_partition(
+            tma_atom,
+            0,
+            cute.make_layout(1),
+            sC_for_tma_partition,
+            gC_for_tma_partition,
+        )
+        return bSG_sC, bSG_gC
+    @cute.jit
+    def find_batch_from_q(
+        self,
+        token: Int32,
+        seqlen_q_divmod: FastDivmodDivisor,
+        mCuSeqlensQ: Optional[cute.Tensor],
+    ) -> Int32:
+        """Find batch index from q token (binary search for varlen, divmod otherwise)"""
+        if const_expr(mCuSeqlensQ is not None):
+            return get_batch_from_cu_tensor(token, mCuSeqlensQ)
+        else:
+            batch, _ = divmod(token, seqlen_q_divmod)
+            return batch
+    @staticmethod
+    def _compute_grid(c, cta_tile_shape_mnk, cluster_shape_mn):
+        c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0))
+        gc = cute.zipped_divide(c, tiler=c_shape)
+        c_logical_shape = gc[(0, (None, None, None))].shape
+        num_ctas_mnl = (
+            cute.size(c_logical_shape[0]),
+            cute.size(c_logical_shape[1]),
+            cute.size(c_logical_shape[2]),
+        )
+        tile_sched_params = utils.ClcDynamicPersistentTileSchedulerParams(
+            num_ctas_mnl, (*cluster_shape_mn, 1)
+        )
+        grid = utils.ClcDynamicPersistentTileScheduler.get_grid_shape(tile_sched_params)
+        return tile_sched_params, grid

build/torch-cuda/flash_bwd_mla_sm100.py ADDED Viewed

	@@ -0,0 +1,2135 @@

+# Copyright (c) 2026, Colfax International.
+import math
+from functools import partial
+from typing import Callable, Optional
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Int64, Int32, Boolean, const_expr
+import cutlass.pipeline as pipeline
+from cutlass.cute.nvgpu import cpasync, tcgen05
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass.utils import ClcDynamicPersistentTileScheduler
+from .quack import copy_utils, layout_utils
+from .pack_gqa import pack_gqa_layout
+from .seqlen_info import SeqlenInfoQK
+from .block_info import BlockInfo
+from . import blackwell_helpers as fa_sm100_utils
+from .tile_scheduler import (
+    SchedulerState,
+    SchedulingMode,
+    TileSchedulerArguments,
+    TileSchedulerProtocol,
+    SingleTileScheduler,
+    SingleTileLPTScheduler,
+    SingleTileVarlenScheduler,
+    ParamsBase,
+)
+from .fa_logging import fa_log, fa_printf
+from .utils import smid, elem_pointer, get_batch_from_cu_tensor
+from .copy_utils import tiled_copy_2d, atomic_add_fp32x4
+from .topk_gather_kv import CpasyncGatherKVManager
+from .named_barrier import NamedBarrierBwdSm100_MLA2CTA
+class FlashAttentionSparseMLABackwardSm100:
+    def __init__(
+        self,
+        is_causal: bool = False,
+        topk_length: int = 2048,
+        qhead_per_kvhead: int = 1,
+        nheads_kv: int = 1,
+        hdim: int = 64,
+        hdimv: int = 512,
+        has_seqused_q: bool = False,
+        disable_bitmask: bool = False,
+        use_clc_scheduler: bool = True,
+    ):
+        use_cpasync_load_KV = True
+        self.is_causal = is_causal
+        self.is_local = False
+        self.pack_gqa = True
+        self.qhead_per_kvhead = qhead_per_kvhead
+        self.nheads_kv = nheads_kv
+        self.has_seqused_q = has_seqused_q
+        self.use_tma_O = True
+        self.use_cpasync_load_KV = True
+        self.use_tma_KV = False
+        self.topk_length = topk_length
+        self.is_topk_gather = True
+        assert qhead_per_kvhead == 128 or qhead_per_kvhead == 64
+        # user-provided option if topk indices guaranteed in bounds
+        self.disable_bitmask = disable_bitmask
+        # ==== tile scheduler ====
+        self.static_persistent = False
+        self.use_clc_scheduler = use_clc_scheduler
+        self.sched_stages = 1
+        self.scheduling_mode = (
+            SchedulingMode.CLC if self.use_clc_scheduler else SchedulingMode.STATIC
+        )
+        if const_expr(has_seqused_q):
+            self.TileScheduler = SingleTileVarlenScheduler
+        elif self.use_clc_scheduler:
+            self.TileScheduler = SingleTileLPTScheduler
+        else:
+            self.TileScheduler = SingleTileScheduler
+        fa_log(
+            1,
+            f"TileScheduler={self.TileScheduler.__name__}, scheduling_mode={self.scheduling_mode.name}",
+        )
+        # ==== thread info ====
+        self.num_softmax_threads = 128
+        self.num_epilogue_threads = 128
+        self.num_load_threads = 32
+        self.num_mma_threads = 32
+        self.num_empty_threads = 0
+        self.num_relay_threads = 32
+        self.num_cpasync_load_threads = 128
+        self.num_threads = 512
+        self.num_warps = self.num_threads // 32
+        self.softmax_warp_indices = (0, 1, 2, 3)
+        self.epilogue_warp_indices = (4, 5, 6, 7)
+        self.load_warp_id = 8
+        self.mma_warp_id = 9
+        self.clc_scheduler_warp_id = 10
+        self.relay_warp_id = 11
+        self.cpasync_load_warp_indices = (12, 13, 14, 15)
+        self.empty_warp_ids = ()
+        # ==== register usage ====
+        assert self.num_warps == 16
+        self.num_regs_load = 128
+        self.num_regs_mma = 128
+        self.num_regs_softmax = 128
+        self.num_regs_epilogue = 128
+        self.num_regs_cpasync = 128
+        self.num_regs_other = 128
+        # self.num_regs_load = 128 - 32
+        # self.num_regs_mma = 128 - 32
+        # self.num_regs_softmax = 128 + 32
+        # self.num_regs_epilogue = 128 + 32
+        # self.num_regs_cpasync = 128 - 32
+        # self.num_regs_other = 48
+        self.num_regs_per_thread = 128
+        self.num_regs_total = 512
+        assert (
+            self.num_regs_mma
+            + self.num_regs_softmax
+            + self.num_regs_epilogue
+            + self.num_regs_cpasync
+            <= self.num_regs_total
+        )
+        # ==== 2cta info ====
+        self.use_2cta_instrs = True
+        self.cta_group = tcgen05.CtaGroup.TWO
+        self.cta_group_size = 2
+        self.cluster_shape_mn = (2, 1)
+        self.cluster_shape_mnk = (2, 1, 1)
+        # ==== problem shape info ====
+        self.hdim = hdim  # ignored
+        self.hdimv = hdimv
+        self.tile_m = qhead_per_kvhead
+        self.tile_n = 64
+        self.cta_tiler_mn = (self.tile_m // self.cta_group_size, self.tile_n)
+        self.cluster_tile_n = self.cta_group_size * self.tile_n
+        self.num_hdimv_splits = 2  # split hdimv in half for our Qv @ V^T and P @ V mmas.
+        self.tile_P = (self.tile_m, self.tile_n)
+        self.tile_Pt = (self.tile_n, self.tile_m)
+        self.tile_dS = (self.tile_m, self.tile_n)
+        self.tile_dSt = (self.tile_n, self.tile_m)
+        self.tile_dV = (self.tile_n, 32)
+        # ==== MMA info ====
+        # dP.T = V    @ dO.T , N x M x dv
+        # dV  += P.T  @ dO   , N x dv x M
+        # dV  += dS.T @ Qv   , N x dv x M
+        self.mma_tiler_VdO = (
+            self.cluster_tile_n,
+            self.tile_m,
+            self.hdimv // self.num_hdimv_splits,
+        )
+        self.mma_tiler_PtdOt = (
+            self.cluster_tile_n,
+            self.hdimv // self.num_hdimv_splits,
+            self.tile_m,
+        )
+        self.mma_tiler_dStQvt = (
+            self.cluster_tile_n,
+            self.hdimv // self.num_hdimv_splits,
+            self.tile_m,
+        )
+        # note: store P.T, dS.T as tile_n major (i.e., as P and dS)
+        self.major_mode_V = tcgen05.OperandMajorMode.K
+        self.major_mode_dO = tcgen05.OperandMajorMode.K
+        self.major_mode_Pt = tcgen05.OperandMajorMode.MN
+        self.major_mode_dOt = tcgen05.OperandMajorMode.MN
+        self.major_mode_dSt = tcgen05.OperandMajorMode.MN
+        self.major_mode_Qvt = tcgen05.OperandMajorMode.MN
+        self.operand_source_V = tcgen05.OperandSource.SMEM
+        self.operand_source_Pt = tcgen05.OperandSource.SMEM
+        self.operand_source_dSt = tcgen05.OperandSource.SMEM
+        # ==== pipeline info ====
+        # stationary: dOi
+        # mainloop:
+        # *) P, scaleP => Pt
+        # *) dSt
+        # *) Vi, i = {0, 1}
+        # *) dOti => Qvi => dVi, i = {0, 1}
+        # redundant names for ease-of-use
+        self.num_stages_V = 2
+        self.num_stages_dO = 2
+        self.num_stages_P = 1
+        self.num_stages_Pt = 1
+        self.num_stages_dS = 1
+        self.num_stages_dSt = 1
+        self.num_stages_dOt = 2
+        self.num_stages_Qv = 2
+        self.num_stages_Qvt = 2
+        self.num_stages_dP = 1
+        self.num_stages_dPt = 1
+        self.num_stages_dV = 2  # == hdimv splits, for Umma <-> Async
+        self.num_epi_stages_dV = 8  # == 2 splits x 4 slots/split
+        self.num_stages_scaleP = 1
+        self.num_stages_dPsum = 1
+        # ==== dtype info ====
+        self.dtype_acc = Float32
+        # ==== TMEM info ====
+        SM100_TMEM_CAPACITY_COLUMNS = 512
+        self.tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS
+        self.tmem_cols_dP = self.tile_m // self.cta_group_size
+        self.tmem_cols_dVi = (self.hdimv // self.num_hdimv_splits) // self.cta_group_size
+        self.tmem_offset_dV0 = 0
+        self.tmem_offset_dV1 = self.tmem_offset_dV0 + self.tmem_cols_dVi
+        self.tmem_offsets_dV = [self.tmem_offset_dV0, self.tmem_offset_dV1]
+        self.tmem_offset_dP = self.tmem_offset_dV1 + self.tmem_cols_dVi
+        self.total_tmem = self.tmem_offset_dP + self.tmem_cols_dP
+        assert self.total_tmem <= self.tmem_alloc_cols, (
+            f"Total TMEM columns allocated {self.total_tmem} exceeds capacity {self.tmem_alloc_cols}"
+        )
+    def _get_shared_storage_cls(self):
+        self.buffer_align_bytes = 1024
+        def smem_struct_align(dtype, staged_layout):
+            return cute.struct.Align[
+                cute.struct.MemRange[dtype, cute.cosize(staged_layout)],
+                self.buffer_align_bytes,
+            ]
+        def mbar_struct(num_stages):
+            return cute.struct.MemRange[Int64, 2 * num_stages]
+        # sV, sdO, sP = sPt, sdSt = sdS, sdOt = sQvt = sdV
+        (
+            sV_struct,
+            sdO_struct,
+            sP_struct,
+            sdS_struct,
+            sQvt_struct,
+            sScaleP_struct,
+            sdPsum_struct,
+        ) = (
+            smem_struct_align(dtype, layout)
+            for dtype, layout in [
+                (self.dtype, self.sV_layout_staged),
+                (self.dtype, self.sdO_layout_staged),
+                (self.dtype, self.sPt_layout_staged),
+                (self.dtype, self.sdSt_layout_staged),
+                (self.dtype, self.sQvt_layout_staged),
+                (self.dtype_scale, self.sScaleP_layout_staged),
+                (self.dtype_scale, self.sdPsum_layout_staged),
+            ]
+        )
+        (
+            mbar_ptr_V_struct,  # load V
+            mbar_ptr_dO_struct,  # load dO
+            mbar_ptr_dOt_Qvt_struct,  # load dOt => Qvt
+            mbar_ptr_dSt_struct,  # store dS
+            mbar_ptr_P_struct,  # load P
+            mbar_ptr_Pt_struct,  # store Pt
+            mbar_ptr_dPt_struct,  # dP mma
+            mbar_ptr_dV_struct,  # dV mma
+            mbar_ptr_scaleP_struct,  # load scaleP
+            mbar_ptr_dPsum_struct,  # load dPsum
+        ) = (
+            mbar_struct(n)
+            for n in [
+                self.num_stages_V,
+                self.num_stages_dO,
+                self.num_stages_Qvt,
+                self.num_stages_dSt,
+                self.num_stages_P,
+                self.num_stages_Pt,
+                self.num_stages_dPt,
+                self.num_stages_dV,
+                self.num_stages_scaleP,
+                self.num_stages_dPsum,
+            ]
+        )
+        tmem_dealloc_mbar_struct = Int64
+        tmem_holding_buf_struct = Int32
+        self.sched_stages = 1
+        clc_response_size = self.sched_stages * 4 if self.use_clc_scheduler else 0
+        clc_mbar_size = self.sched_stages * 2 if self.use_clc_scheduler else 0
+        @cute.struct
+        class SharedStorage:
+            mbar_ptr_V: mbar_ptr_V_struct
+            mbar_ptr_V_cpasync: mbar_ptr_V_struct
+            mbar_ptr_dO: mbar_ptr_dO_struct
+            mbar_ptr_dOt_Qvt: mbar_ptr_dOt_Qvt_struct
+            mbar_ptr_P: mbar_ptr_P_struct
+            mbar_ptr_Pt: mbar_ptr_Pt_struct
+            mbar_ptr_dSt: mbar_ptr_dSt_struct
+            mbar_ptr_dPt: mbar_ptr_dPt_struct
+            mbar_ptr_dV: mbar_ptr_dV_struct
+            mbar_ptr_dV_epi: mbar_ptr_dV_struct
+            mbar_ptr_scaleP: mbar_ptr_scaleP_struct
+            mbar_ptr_dPsum: mbar_ptr_dPsum_struct
+            tmem_dealloc_mbar: tmem_dealloc_mbar_struct
+            tmem_holding_buf: tmem_holding_buf_struct
+            clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, clc_mbar_size]
+            clc_response: cute.struct.MemRange[Int32, clc_response_size]
+            sScaleP: sScaleP_struct
+            sdPsum: sdPsum_struct
+            sV: sV_struct
+            sdO: sdO_struct
+            sP: sP_struct
+            sdS: sdS_struct
+            sQv: sQvt_struct
+        # print("smem bytes = ", SharedStorage.size_in_bytes())
+        return SharedStorage
+    # fmt: off
+    @cute.jit
+    def __call__(
+        self,
+        mdO: cute.Tensor,  # (b, s_q, h, dv) or (total_q, h, dv) if there is cu_seqlens_q
+        mV: cute.Tensor,   # (b_k, s_k, h_k, dv) or (total_k, h_k, dv) if there is cu_seqlens_k
+        mQv: cute.Tensor,  # == mdO
+        mP: cute.Tensor,   # (b, s_q, h, topk) or (total_q, h, topk)
+        mdV: cute.Tensor,  # == mV
+        mdS: cute.Tensor,  # == mP
+        mIndexTopk: cute.Tensor,  # (b, s_q, topk) or (total_q, topk) if there is cu_seqlens_q
+        softmax_scale: Float32,
+        mScaleP: Optional[cute.Tensor] = None,      # (b, s_q, topk//128, h) or (total_q, topk//128, h)
+        mdPsum: Optional[cute.Tensor] = None,       # (b, s_q, h) or (total_q, h) if there is cu_seqlens_q
+        mCuSeqlensQ: Optional[cute.Tensor] = None,  # (b + 1)
+        mCuSeqlensK: Optional[cute.Tensor] = None,  # (b + 1)
+        mSeqUsedQ: Optional[cute.Tensor] = None,    # (b)
+        mSeqUsedK: Optional[cute.Tensor] = None,    # (b)
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
+        stream: cuda.CUstream = None,
+    ):
+        # fmt: on
+        # ==== dtype info ====
+        self.dtype = mdO.element_type
+        self.dtype_dV = mdV.element_type
+        self.dtype_scale = Float32
+        self.dtype_index = mIndexTopk.element_type
+        assert self.dtype.width == 16
+        assert self.dtype_dV.width == 32
+        assert self.dtype_index == Int32
+        if const_expr(mScaleP is not None):
+            assert mScaleP.element_type == self.dtype_scale
+        if const_expr(mdPsum is not None):
+            assert mdPsum.element_type == self.dtype_scale
+        # ==== Prepare Tensors ====
+        new_stride = lambda mX: (
+            *(cute.assume(s, divby=128 // mX.element_type.width) for s in mX.stride[:-1]),
+            mX.stride[-1],
+        )
+        mQv, mV, mdV, mdO, mP, mdS, mScaleP, mdPsum = [
+            cute.make_tensor(mX.iterator, cute.make_layout(mX.shape, stride=new_stride(mX)))
+            if mX is not None
+            else None
+            for mX in (mQv, mV, mdV, mdO, mP, mdS, mScaleP, mdPsum)
+        ]
+        # (b, s, h, d)  -> (s, d, h, b)  or
+        # (total, h, d) -> (total, d, h)
+        QO_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1]
+        KV_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensK is None) else [0, 2, 1]
+        mQv, mdO, mP, mdS = [
+            cute.make_tensor(mX.iterator, cute.select(mX.layout, mode=QO_layout_transpose))
+            if mX is not None
+            else None
+            for mX in (mQv, mdO, mP, mdS)
+        ]
+        mV, mdV = [
+            cute.make_tensor(mX.iterator, cute.select(mX.layout, mode=KV_layout_transpose))
+            if mX is not None
+            else None
+            for mX in (mV, mdV)
+        ]
+        # (b, s, topk//128, h) -> (s, topk//128, h, b) or
+        # (total, topk//128, h) -> (total, topk//128, h)
+        ScaleP_layout_transpose = [1, 2, 3, 0] if const_expr(mCuSeqlensQ is None) else [0, 1, 2]
+        mScaleP = cute.make_tensor(
+            mScaleP.iterator, cute.select(mScaleP.layout, mode=ScaleP_layout_transpose)
+        )
+        # (b, s, h) -> (s, h, b) or
+        # (total, h) -> (total, h)
+        dPsum_layout_transpose = [1, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 1]
+        mdPsum = cute.make_tensor(
+            mdPsum.iterator, cute.select(mdPsum.layout, mode=dPsum_layout_transpose)
+        )
+        # (b, s_q, topk) -> (topk, s_q, b) or (total_q, topk) -> (topk, total_q)
+        topk_layout_transpose = [2, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 0]
+        mIndexTopk = cute.make_tensor(
+            mIndexTopk.iterator, cute.select(mIndexTopk.layout, mode=topk_layout_transpose)
+        )
+        topk_length_dynamic = mIndexTopk.shape[0]
+        if const_expr(self.pack_gqa):
+            mQv, mdO, mP, mdS, mScaleP = [
+                pack_gqa_layout(mX, self.qhead_per_kvhead, self.nheads_kv, head_idx=2)
+                if mX is not None
+                else None
+                for mX in (mQv, mdO, mP, mdS, mScaleP)
+            ]
+            if const_expr(mdPsum is not None):
+                mdPsum = pack_gqa_layout(mdPsum, self.qhead_per_kvhead, self.nheads_kv, head_idx=1)
+        # ((h/h_k, s_q), dv, h_k, b) -> (dv, (h/h_k, s_q), h_k, b)
+        # or ((h/h_k, total_q), dv, h_k) -> (dv, (h/h_k, total_q), h_k)
+        mma_operand_layout_transpose = (
+            [1, 0, 2, 3] if const_expr(mCuSeqlensQ is None) else [1, 0, 2]
+        )
+        mQvt, mdOt = [
+            cute.make_tensor(mX.iterator, cute.select(mX.layout, mode=mma_operand_layout_transpose))
+            for mX in (mQv, mdO)
+        ]
+        # fmt: off
+        # ==== Prepare MMAs ====
+        # (local_var, dtype_a, major_a, major_b, mma_tiler, operand_source_a)
+        _mma_specs = [
+            ("tiled_mma_VdO",    self.dtype, self.major_mode_V,   self.major_mode_dO,  self.mma_tiler_VdO,    self.operand_source_V),
+            ("tiled_mma_PtdOt",  self.dtype, self.major_mode_Pt,  self.major_mode_dOt, self.mma_tiler_PtdOt,  self.operand_source_Pt),
+            ("tiled_mma_dStQvt", self.dtype, self.major_mode_dSt, self.major_mode_Qvt, self.mma_tiler_dStQvt, self.operand_source_dSt),
+        ]
+        tiled_mma_VdO, tiled_mma_PtdOt, tiled_mma_dStQvt = (
+            sm100_utils.make_trivial_tiled_mma(
+                dtype_a, major_a, major_b, self.dtype_acc, self.cta_group, mma_tiler[:2], operand_source_a,
+            )
+            for _, dtype_a, major_a, major_b, mma_tiler, operand_source_a in _mma_specs
+        )
+        # ==== Prepare SMEM layouts and TMAs ====
+        # (attr, make_fn, tiled_mma, mma_tiler, dtype, num_stages)
+        _smem_layout_specs = [
+            ("sV_layout",   sm100_utils.make_smem_layout_a, tiled_mma_VdO,    self.mma_tiler_VdO,    self.dtype, self.num_stages_V),
+            ("sdO_layout",  sm100_utils.make_smem_layout_b, tiled_mma_VdO,    self.mma_tiler_VdO,    self.dtype, self.num_stages_dO),
+            ("sPt_layout",  sm100_utils.make_smem_layout_a, tiled_mma_PtdOt,  self.mma_tiler_PtdOt,  self.dtype, self.num_stages_Pt),
+            ("sdOt_layout", sm100_utils.make_smem_layout_b, tiled_mma_PtdOt,  self.mma_tiler_PtdOt,  self.dtype, self.num_stages_dOt),
+            ("sdSt_layout", sm100_utils.make_smem_layout_a, tiled_mma_dStQvt, self.mma_tiler_dStQvt, self.dtype, self.num_stages_dSt),
+            ("sQvt_layout", sm100_utils.make_smem_layout_b, tiled_mma_dStQvt, self.mma_tiler_dStQvt, self.dtype, self.num_stages_Qvt),
+        ]
+        for attr, make_fn, tiled_mma, mma_tiler, dtype, num_stages in _smem_layout_specs:
+            ab_kwarg = "a_dtype" if make_fn is sm100_utils.make_smem_layout_a else "b_dtype"
+            staged = make_fn(
+                tiled_mma=tiled_mma,
+                mma_tiler_mnk=mma_tiler,
+                num_stages=num_stages,
+                **{ab_kwarg: dtype},
+            )
+            setattr(self, f"{attr}_staged", staged)
+            setattr(self, attr, cute.select(staged, mode=[0, 1, 2]))
+        # Prepare additional SMEM load layouts
+        self.P_layout_major = cutlass.utils.LayoutEnum.from_tensor(mP)
+        self.sP_layout_staged = sm100_utils.make_smem_layout_epi(
+            self.dtype, self.P_layout_major, self.tile_P, self.num_stages_P
+        )
+        self.sP_layout = cute.select(self.sP_layout_staged, mode=[0, 1])
+        self.sScaleP_layout_staged = cute.make_layout((self.tile_m, self.num_stages_scaleP))
+        self.sScaleP_layout = cute.select(self.sScaleP_layout_staged, mode=[0])
+        self.sdPsum_layout_staged = cute.make_layout((self.tile_m, self.num_stages_dPsum))
+        self.sdPsum_layout = cute.select(self.sdPsum_layout_staged, mode=[0])
+        # ==== TMA load ====
+        for attr, dtype, layout in [
+            ("tma_copy_bytes_V",   self.dtype, self.sV_layout),
+            ("tma_copy_bytes_dO",  self.dtype, self.sdO_layout),
+            ("tma_copy_bytes_dOt", self.dtype, self.sdOt_layout),
+            ("tma_copy_bytes_Qvt", self.dtype, self.sQvt_layout),
+        ]:
+            setattr(self, attr, cute.size_in_bytes(dtype, layout) * self.cta_group_size)
+        assert self.tma_copy_bytes_dOt == self.tma_copy_bytes_Qvt
+        self.tma_copy_bytes_P = cute.size_in_bytes(self.dtype, self.sP_layout)
+        self.tma_copy_bytes_scaleP = cute.size_in_bytes(self.dtype_scale, self.sScaleP_layout)
+        self.tma_copy_bytes_dPsum = cute.size_in_bytes(self.dtype_scale, self.sdPsum_layout)
+        tma_load_op = cpasync.CopyBulkTensorTileG2SOp(self.cta_group)
+        cta_layout_vmnk = cute.tiled_divide(
+            cute.make_layout(self.cluster_shape_mnk), (tiled_mma_VdO.thr_id.shape,)
+        )
+        cta_shape = cta_layout_vmnk.shape
+        def make_tma(make_fn, mX, smem_layout, mma_tiler, tiled_mma):
+            return make_fn(tma_load_op, mX, smem_layout, mma_tiler, tiled_mma, cta_shape)
+        A, B = cute.nvgpu.make_tiled_tma_atom_A, cute.nvgpu.make_tiled_tma_atom_B
+        # (atom_name, tensor_name, make_fn, m, smem_layout, mma_tiler, tiled_mma)
+        _tma_specs = [
+            ("tma_atom_dO",  "tma_tensor_dO",  B, mdO,  self.sdO_layout,  self.mma_tiler_VdO,    tiled_mma_VdO),
+            ("tma_atom_dOt", "tma_tensor_dOt", B, mdOt, self.sdOt_layout, self.mma_tiler_PtdOt,  tiled_mma_PtdOt),
+            ("tma_atom_Qvt", "tma_tensor_Qvt", B, mQvt, self.sQvt_layout, self.mma_tiler_dStQvt, tiled_mma_dStQvt),
+        ]
+        _tmas = {}
+        for atom_name, tensor_name, make_fn, m, smem_layout, mma_tiler, tiled_mma in _tma_specs:
+            _tmas[atom_name], _tmas[tensor_name] = (
+                make_tma(make_fn, m, smem_layout, mma_tiler, tiled_mma)
+            )
+        (tma_atom_dO,  tma_tensor_dO,
+         tma_atom_dOt, tma_tensor_dOt,
+         tma_atom_Qvt, tma_tensor_Qvt) = _tmas.values()
+        # Make TMA load for P separately
+        tma_atom_P, tma_tensor_P = cute.nvgpu.cpasync.make_tiled_tma_atom(
+            cpasync.CopyBulkTensorTileG2SOp(),
+            mP,
+            self.sP_layout,
+            self.tile_P,
+        )
+        # ==== TMA store ====
+        tma_store_op = cpasync.CopyBulkTensorTileS2GOp()
+        self.dS_layout_major = cutlass.utils.LayoutEnum.from_tensor(mdS)
+        self.dV_layout_major = cutlass.utils.LayoutEnum.from_tensor(mdV)
+        # (tile_m, tile_n, dS_stages) = (nheads, 64, dS_stage)
+        sdS_layout_staged = sm100_utils.make_smem_layout_epi(
+            self.dtype, self.dS_layout_major, self.tile_dS, self.num_stages_dSt
+        )
+        # (tile_n, 32, dV_epi_stages) = (64, 32, 4 x 2)
+        sdV_layout_staged = sm100_utils.make_smem_layout_epi(
+            self.dtype_dV, self.dV_layout_major, self.tile_dV, self.num_epi_stages_dV
+        )
+        tma_atom_dS, tma_tensor_dS = cpasync.make_tiled_tma_atom(
+            tma_store_op, mdS, cute.select(sdS_layout_staged, mode=[0, 1]), self.tile_dS
+        )
+        # fmt: on
+        # ==== Allocate shared memory ====
+        SharedStorage = self._get_shared_storage_cls()
+        # ==== Tile scheduler ====
+        TileScheduler = self.TileScheduler
+        fa_printf(1, "mdO = {}", mdO.layout)
+        batch_size_for_sched = cute.size(mdO.shape[3]) if const_expr(mCuSeqlensQ is None) else 1
+        tile_sched_args = TileSchedulerArguments(
+            num_block=cute.ceil_div(cute.size(mdO.shape[0]), self.tile_m),
+            num_head=cute.size(mdO.shape[2]),
+            num_batch=batch_size_for_sched,
+            num_splits=1,
+            seqlen_k=cute.size(mV.shape[0]),
+            headdim=self.hdim,
+            headdim_v=self.hdimv,
+            total_q=cute.size(mdO.shape[0])
+            if const_expr(mCuSeqlensQ is not None)
+            else cute.size(mdO.shape[0]) * cute.size(mdO.shape[3]),
+            tile_shape_mn=self.cta_tiler_mn,
+            mCuSeqlensQ=mCuSeqlensQ,
+            mSeqUsedQ=mSeqUsedQ,
+            qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
+            element_size=self.dtype.width // 8,
+            is_persistent=self.static_persistent,
+            lpt=False,
+            is_split_kv=False,
+            cluster_shape_mn=self.cluster_shape_mn,
+            use_cluster_idx=True,
+        )
+        tile_sched_params = TileScheduler.to_underlying_arguments(
+            tile_sched_args, scheduling_mode=self.scheduling_mode
+        )
+        self.tile_scheduler_cls = TileScheduler
+        grid_dim = TileScheduler.get_grid_shape(tile_sched_params)
+        fa_printf(1, "grid = {}", grid_dim)
+        # ==== Named Barrier ====
+        self.cpasync_barrier = cutlass.pipeline.NamedBarrier(
+            barrier_id=int(NamedBarrierBwdSm100_MLA2CTA.Cpasync),
+            num_threads=self.num_cpasync_load_threads,
+        )
+        self.softmax_barrier = cutlass.pipeline.NamedBarrier(
+            barrier_id=int(NamedBarrierBwdSm100_MLA2CTA.Softmax),
+            num_threads=self.num_softmax_threads,
+        )
+        self.epi_barrier = cutlass.pipeline.NamedBarrier(
+            barrier_id=int(NamedBarrierBwdSm100_MLA2CTA.Epilogue),
+            num_threads=self.num_epilogue_threads,
+        )
+        LOG2_E = math.log2(math.e)
+        softmax_scale_log2 = softmax_scale * LOG2_E
+        # ==== Launch kernel ====
+        block_dim = (self.num_threads, 1, 1)
+        self.kernel(
+            mV,
+            mdV,
+            tma_tensor_dO,
+            tma_tensor_dOt,
+            tma_tensor_Qvt,
+            tma_tensor_P,
+            tma_tensor_dS,
+            mScaleP,
+            mdPsum,
+            mCuSeqlensQ,
+            mCuSeqlensK,
+            mSeqUsedQ,
+            mSeqUsedK,
+            mIndexTopk,
+            tma_atom_dO,
+            tma_atom_dOt,
+            tma_atom_Qvt,
+            tma_atom_P,
+            tma_atom_dS,
+            self.sV_layout_staged,
+            self.sdO_layout_staged,
+            self.sdOt_layout_staged,
+            self.sQvt_layout_staged,
+            self.sP_layout_staged,  # load P
+            sdS_layout_staged,  # store dS
+            sdV_layout_staged,
+            self.sPt_layout_staged,  # mma Pt
+            self.sdSt_layout_staged,  # mma dSt
+            self.sScaleP_layout_staged,
+            self.sdPsum_layout_staged,
+            tiled_mma_VdO,
+            tiled_mma_PtdOt,
+            tiled_mma_dStQvt,
+            softmax_scale,
+            softmax_scale_log2,
+            topk_length_dynamic,
+            tile_sched_params,
+            SharedStorage,
+        ).launch(
+            grid=grid_dim,
+            block=block_dim,
+            cluster=self.cluster_shape_mnk,
+            smem=SharedStorage.size_in_bytes(),
+            stream=stream,
+        )
+    @cute.kernel
+    def kernel(
+        self,
+        mV: cute.Tensor,
+        mdV: cute.Tensor,
+        mdO: cute.Tensor,
+        mdOt: cute.Tensor,
+        mQvt: cute.Tensor,
+        mP: cute.Tensor,
+        mdS: cute.Tensor,
+        mScaleP: Optional[cute.Tensor],
+        mdPsum: Optional[cute.Tensor],
+        mCuSeqlensQ: Optional[cute.Tensor],
+        mCuSeqlensK: Optional[cute.Tensor],
+        mSeqUsedQ: Optional[cute.Tensor],
+        mSeqUsedK: Optional[cute.Tensor],
+        mIndexTopk: Optional[cute.Tensor],
+        tma_atom_dO: cute.CopyAtom,
+        tma_atom_dOt: cute.CopyAtom,
+        tma_atom_Qvt: cute.CopyAtom,
+        tma_atom_P: cute.CopyAtom,
+        tma_atom_dS: cute.CopyAtom,
+        sV_layout_staged: cute.ComposedLayout,
+        sdO_layout_staged: cute.ComposedLayout,
+        sdOt_layout_staged: cute.ComposedLayout,
+        sQvt_layout_staged: cute.ComposedLayout,
+        sP_layout_staged: cute.ComposedLayout,
+        sdS_layout_staged: cute.ComposedLayout,
+        sdV_layout_staged: cute.ComposedLayout,
+        sPt_layout_staged: cute.ComposedLayout,
+        sdSt_layout_staged: cute.ComposedLayout,
+        sScaleP_layout_staged: cute.Layout,
+        sdPsum_layout_staged: cute.Layout,
+        tiled_mma_VdO: cute.TiledMma,
+        tiled_mma_PtdOt: cute.TiledMma,
+        tiled_mma_dStQvt: cute.TiledMma,
+        softmax_scale: Float32,
+        softmax_scale_log2: Float32,
+        topk_length_dynamic: Optional[Int32],
+        tile_sched_params: ParamsBase,
+        SharedStorage: cutlass.Constexpr[Callable],
+    ):
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+        cta_layout_vmnk = cute.tiled_divide(
+            cute.make_layout(self.cluster_shape_mnk), (tiled_mma_VdO.thr_id.shape,)
+        )
+        mma_tile_coord_v = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+        is_leader_cta = mma_tile_coord_v == 0
+        # ==== Allocate SMEM ====
+        smem = cutlass.utils.SmemAllocator()
+        storage = smem.allocate(SharedStorage)
+        # ==== Prepare TMEM allocator ====
+        tmem_alloc_barrier = pipeline.NamedBarrier(
+            barrier_id=int(NamedBarrierBwdSm100_MLA2CTA.TmemPtr),
+            num_threads=self.num_mma_threads + self.num_softmax_threads + self.num_epilogue_threads,
+        )
+        tmem = cutlass.utils.TmemAllocator(
+            storage.tmem_holding_buf.ptr,
+            barrier_for_retrieve=tmem_alloc_barrier,
+            allocator_warp_id=self.mma_warp_id,
+            is_two_cta=self.use_2cta_instrs,
+            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
+        )
+        # ==== Prefetch TMA descriptors ====
+        if warp_idx == self.load_warp_id:
+            cpasync.prefetch_descriptor(tma_atom_dO)
+            cpasync.prefetch_descriptor(tma_atom_dOt)
+            cpasync.prefetch_descriptor(tma_atom_Qvt)
+            cpasync.prefetch_descriptor(tma_atom_P)
+            cpasync.prefetch_descriptor(tma_atom_dS)
+        # ==== Construct pipelines ====
+        tma_warp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1)
+        mma_warp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1)
+        sm_warps = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_softmax_threads // 32)
+        store_warp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1)
+        sm_threads = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_softmax_threads)
+        epi_threads = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_epilogue_threads)
+        sm_threads_cluster = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, self.num_softmax_threads * self.cta_group_size
+        )
+        epi_threads_cluster = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, self.num_epilogue_threads * self.cta_group_size
+        )
+        cpasync_load_threads = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, self.num_cpasync_load_threads
+        )
+        relay_warps_cluster = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.cta_group_size)
+        relay_threads = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_relay_threads)
+        TmaUmma = pipeline.PipelineTmaUmma
+        TmaAsync = pipeline.PipelineTmaAsync
+        AsyncUmma = pipeline.PipelineAsyncUmma
+        UmmaAsync = pipeline.PipelineUmmaAsync
+        Async = pipeline.PipelineAsync
+        def make_pipeline(cls, mbar_ptr, num_stages, producer, consumer, tx_count=None):
+            return cls.create(
+                barrier_storage=mbar_ptr.data_ptr(),
+                num_stages=num_stages,
+                producer_group=producer,
+                consumer_group=consumer,
+                defer_sync=True,
+                **(
+                    {"cta_layout_vmnk": cta_layout_vmnk}
+                    if cls is not Async and cls is not TmaAsync
+                    else {}
+                ),
+                **({"tx_count": tx_count} if tx_count is not None else {}),
+            )
+        # Unconditional pipelines
+        # fmt: off
+        # TmaUmma: dO, dOt & Qvt
+        pipeline_dO = make_pipeline(TmaUmma, storage.mbar_ptr_dO, self.num_stages_dO, tma_warp, mma_warp, self.tma_copy_bytes_dO)
+        pipeline_dOt_Qvt = make_pipeline(TmaUmma, storage.mbar_ptr_dOt_Qvt, self.num_stages_Qvt, tma_warp, mma_warp, self.tma_copy_bytes_dOt)
+        # TmaAsync: P, scaleP
+        pipeline_P = make_pipeline(TmaAsync, storage.mbar_ptr_P, self.num_stages_P, tma_warp, sm_warps, self.tma_copy_bytes_P)
+        pipeline_scaleP = make_pipeline(TmaAsync, storage.mbar_ptr_scaleP, self.num_stages_scaleP, tma_warp, sm_warps, self.tma_copy_bytes_scaleP)
+        pipeline_dPsum = make_pipeline(TmaAsync, storage.mbar_ptr_dPsum, self.num_stages_dPsum, tma_warp, sm_warps, self.tma_copy_bytes_dPsum)
+        # AsyncUmma: Pt => dV mma, dSt => dV mma
+        pipeline_Pt = make_pipeline(AsyncUmma, storage.mbar_ptr_Pt, self.num_stages_Pt, sm_threads_cluster, mma_warp)
+        pipeline_dSt = make_pipeline(AsyncUmma, storage.mbar_ptr_dSt, self.num_stages_dSt, sm_threads_cluster, mma_warp)
+        # UmmaAsync: dPt, dV
+        pipeline_dPt = make_pipeline(UmmaAsync, storage.mbar_ptr_dPt, self.num_stages_dPt, mma_warp, sm_threads_cluster)
+        pipeline_dV = make_pipeline(UmmaAsync, storage.mbar_ptr_dV, self.num_stages_dV, mma_warp, epi_threads_cluster)
+        # Async: dV_epi
+        pipeline_dV_epi = make_pipeline(Async, storage.mbar_ptr_dV_epi, self.num_stages_dV, tma_warp, store_warp)
+        pipeline_V         = make_pipeline(AsyncUmma, storage.mbar_ptr_V,         self.num_stages_V, relay_warps_cluster,  mma_warp)
+        pipeline_V_cpasync = make_pipeline(Async,     storage.mbar_ptr_V_cpasync, self.num_stages_V, cpasync_load_threads, relay_threads)
+        # fmt: on
+        pipeline.pipeline_init_arrive(cluster_shape_mn=cta_layout_vmnk, is_relaxed=True)
+        # ==== Get SMEM tensors ====
+        # fmt: off
+        sV, sdO, sP, sPt, sdOt, sdS, sdSt, sQvt = (
+            store.get_tensor(layout.outer, swizzle=layout.inner)
+            for store, layout in [
+                (storage.sV,  sV_layout_staged),
+                (storage.sdO, sdO_layout_staged),
+                (storage.sP,  sP_layout_staged),    # P & Pt overlap
+                (storage.sP,  sPt_layout_staged),   # P & Pt overlap
+                (storage.sQv, sdOt_layout_staged),  # {dOt, Qvt, dV} overlap
+                (storage.sdS, sdS_layout_staged),   # dS & dSt overlap
+                (storage.sdS, sdSt_layout_staged),  # dS & dSt overlap
+                (storage.sQv, sQvt_layout_staged),  # {dOt, Qvt, dV} overlap
+            ]
+        )
+        sdV = cute.make_tensor(
+            cute.recast_ptr(sdOt.iterator, sdV_layout_staged.inner, self.dtype_acc), sdV_layout_staged.outer
+        )
+        assert cute.cosize(sdV) * self.dtype_acc.width // self.dtype.width == cute.cosize(sdOt)
+        sScaleP = storage.sScaleP.get_tensor(sScaleP_layout_staged)
+        sdPsum = storage.sdPsum.get_tensor(sdPsum_layout_staged)
+        # fmt: on
+        # ==== Get thread MMAs and accumulator fragments ====
+        thr_mma_VdO = tiled_mma_VdO.get_slice(mma_tile_coord_v)
+        thr_mma_PtdOt = tiled_mma_PtdOt.get_slice(mma_tile_coord_v)
+        thr_mma_dStQvt = tiled_mma_dStQvt.get_slice(mma_tile_coord_v)
+        acc_shape_dPt = thr_mma_VdO.partition_shape_C(self.mma_tiler_VdO[:2])
+        acc_shape_dVi = thr_mma_PtdOt.partition_shape_C(self.mma_tiler_PtdOt[:2])
+        tdPtdP_fake = thr_mma_VdO.make_fragment_C(acc_shape_dPt)
+        tdVtdV0_fake = thr_mma_PtdOt.make_fragment_C(acc_shape_dVi)
+        tdVtdV1_fake = thr_mma_PtdOt.make_fragment_C(acc_shape_dVi)
+        # tdPtdP = cute.make_tensor(tdPtdP.iterator + self.tmem_offset_dP, tdPtdP.layout)
+        # tdVtdV0 = cute.make_tensor(tdVtdV0.iterator + self.tmem_offset_dV0, tdVtdV0.layout)
+        # tdVtdV1 = cute.make_tensor(tdVtdV1.iterator + self.tmem_offset_dV1, tdVtdV1.layout)
+        block_info = BlockInfo(
+            self.tile_m * self.cta_group_size,
+            self.tile_n,
+            is_causal=self.is_causal,
+            qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
+        )
+        SeqlenInfoCls = partial(
+            SeqlenInfoQK.create,
+            seqlen_q_static=mdO.shape[0] if const_expr(not self.pack_gqa) else mdO.shape[0][1],
+            seqlen_k_static=mV.shape[0],
+            tile_m=self.tile_m,
+            tile_n=self.tile_n,
+            mCuSeqlensQ=mCuSeqlensQ,
+            mCuSeqlensK=mCuSeqlensK,
+            mSeqUsedQ=mSeqUsedQ,
+            mSeqUsedK=mSeqUsedK,
+        )
+        if const_expr(self.use_clc_scheduler):
+            clc_response_ptr = storage.clc_response.data_ptr()
+            clc_mbar_ptr = storage.clc_mbar_ptr.data_ptr()
+            clc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+            num_clc_consumer_warps_per_cta = self.num_threads // cute.arch.WARP_SIZE
+            num_clc_consumer_warps = num_clc_consumer_warps_per_cta * self.cta_group_size
+            clc_pipeline_consumer_group = pipeline.CooperativeGroup(
+                pipeline.Agent.Thread, cute.arch.WARP_SIZE * num_clc_consumer_warps
+            )
+            sched_ctx = SchedulerState.create_clc(
+                hw_scheduler=ClcDynamicPersistentTileScheduler.create(
+                    self.tile_scheduler_cls.clc_problem_shape(tile_sched_params),
+                    cute.arch.block_idx(),
+                    cute.arch.grid_dim(),
+                    clc_response_ptr,
+                ),
+                pipeline=pipeline.PipelineClcFetchAsync.create(
+                    barrier_storage=clc_mbar_ptr,
+                    num_stages=self.sched_stages,
+                    producer_group=clc_pipeline_producer_group,
+                    consumer_group=clc_pipeline_consumer_group,
+                    tx_count=16,
+                    cta_layout_vmnk=cta_layout_vmnk,
+                ),
+                consumer_state=pipeline.make_pipeline_state(
+                    pipeline.PipelineUserType.Consumer, self.sched_stages
+                ),
+                producer_state=pipeline.make_pipeline_state(
+                    pipeline.PipelineUserType.Producer, self.sched_stages
+                ),
+            )
+            tile_scheduler = self.tile_scheduler_cls.create(tile_sched_params, ctx=sched_ctx)
+        else:
+            tile_scheduler = self.tile_scheduler_cls.create(tile_sched_params)
+        assert isinstance(tile_scheduler, TileSchedulerProtocol), (
+            f"tile_scheduler is not a TileSchedulerProtocol: {type(tile_scheduler)}"
+        )
+        pipeline.pipeline_init_wait(cluster_shape_mn=cta_layout_vmnk)
+        if const_expr(self.use_clc_scheduler):
+            if warp_idx == self.clc_scheduler_warp_id:
+                cute.arch.setmaxregister_decrease(self.num_regs_other)
+                if is_leader_cta:
+                    self.clc_scheduler_warp(tile_scheduler)
+                else:
+                    self.empty_warp(tile_scheduler)
+            for i in cutlass.range_constexpr(len(self.empty_warp_ids)):
+                if warp_idx == self.empty_warp_ids[i] and warp_idx != self.clc_scheduler_warp_id:
+                    cute.arch.setmaxregister_decrease(self.num_regs_other)
+                    self.empty_warp(tile_scheduler)
+        else:
+            for i in cutlass.range_constexpr(len(self.empty_warp_ids)):
+                if warp_idx == self.empty_warp_ids[i]:
+                    cute.arch.setmaxregister_decrease(self.num_regs_other)
+        if const_expr(self.use_cpasync_load_KV):
+            if warp_idx == self.relay_warp_id:
+                if const_expr(self.num_regs_load < self.num_regs_per_thread):
+                    cute.arch.setmaxregister_decrease(self.num_regs_load)
+                self.relay(
+                    pipeline_V,
+                    pipeline_V_cpasync,
+                    topk_length_dynamic,
+                    block_info,
+                    SeqlenInfoCls,
+                    tile_scheduler=tile_scheduler,
+                )
+            if warp_idx in self.cpasync_load_warp_indices:
+                if const_expr(self.num_regs_cpasync < self.num_regs_per_thread):
+                    cute.arch.setmaxregister_decrease(self.num_regs_cpasync)
+                self.load_cpasync(
+                    mIndexTopk,
+                    mV,
+                    sV,
+                    pipeline_V,
+                    pipeline_V_cpasync,
+                    topk_length_dynamic,
+                    block_info,
+                    SeqlenInfoCls,
+                    mCuSeqlensQ,
+                    tile_scheduler=tile_scheduler,
+                )
+        if warp_idx == self.load_warp_id:
+            if const_expr(self.num_regs_load < self.num_regs_per_thread):
+                cute.arch.setmaxregister_decrease(self.num_regs_load)
+            self.load(
+                mdO,
+                mP,
+                mdOt,
+                mQvt,
+                mScaleP,
+                mdPsum,
+                sdO,
+                sP,
+                sdOt,
+                sQvt,
+                sScaleP,
+                sdPsum,
+                tma_atom_dO,
+                tma_atom_P,
+                tma_atom_dOt,
+                tma_atom_Qvt,
+                pipeline_dO,
+                pipeline_P,
+                pipeline_dOt_Qvt,
+                pipeline_Pt,
+                pipeline_dV_epi,
+                pipeline_scaleP,
+                pipeline_dPsum,
+                thr_mma_VdO,
+                thr_mma_PtdOt,
+                thr_mma_dStQvt,
+                topk_length_dynamic,
+                block_info,
+                SeqlenInfoCls,
+                mCuSeqlensQ,
+                tile_scheduler=tile_scheduler,
+            )
+        if warp_idx == self.mma_warp_id:
+            if const_expr(self.num_regs_mma < self.num_regs_per_thread):
+                cute.arch.setmaxregister_decrease(self.num_regs_mma)
+            # ==== Allocate TMEM ====
+            tmem.allocate(self.tmem_alloc_cols)
+            tmem.wait_for_alloc()
+            tmem_ptr = tmem.retrieve_ptr(self.dtype_acc)
+            tdPtdP = cute.make_tensor(tmem_ptr + self.tmem_offset_dP, tdPtdP_fake.layout)
+            tdVtdV0 = cute.make_tensor(tmem_ptr + self.tmem_offset_dV0, tdVtdV0_fake.layout)
+            tdVtdV1 = cute.make_tensor(tmem_ptr + self.tmem_offset_dV1, tdVtdV1_fake.layout)
+            self.mma(
+                sV,
+                sdO,
+                sPt,
+                sdOt,
+                sdSt,
+                sQvt,
+                tdPtdP,
+                tdVtdV0,
+                tdVtdV1,
+                tiled_mma_VdO,
+                tiled_mma_PtdOt,
+                tiled_mma_dStQvt,
+                pipeline_V,
+                pipeline_dO,
+                pipeline_dPt,
+                pipeline_Pt,
+                pipeline_dOt_Qvt,
+                pipeline_dSt,
+                pipeline_dV,
+                is_leader_cta,
+                topk_length_dynamic,
+                block_info,
+                SeqlenInfoCls,
+                mCuSeqlensQ,
+                tile_scheduler=tile_scheduler,
+            )
+            tmem.relinquish_alloc_permit()
+            tmem_alloc_barrier.arrive_and_wait()
+            tmem.free(tmem_ptr)
+        if warp_idx in self.softmax_warp_indices:
+            cute.arch.setmaxregister_increase(self.num_regs_softmax)
+            tmem.wait_for_alloc()
+            tmem_ptr = tmem.retrieve_ptr(self.dtype_acc)
+            tdPtdP = cute.make_tensor(tmem_ptr + self.tmem_offset_dP, tdPtdP_fake.layout)
+            self.compute_loop(
+                softmax_scale,
+                softmax_scale_log2,
+                thr_mma_VdO,
+                tdPtdP,
+                sP,
+                sdS,
+                sScaleP,
+                sdPsum,
+                mdS,
+                tma_atom_dS,
+                pipeline_P,
+                pipeline_Pt,
+                pipeline_dPt,
+                pipeline_dSt,
+                pipeline_scaleP,
+                pipeline_dPsum,
+                topk_length_dynamic,
+                block_info,
+                SeqlenInfoCls,
+                mCuSeqlensQ,
+                tile_scheduler=tile_scheduler,
+            )
+            tmem_alloc_barrier.arrive()
+        if warp_idx in self.epilogue_warp_indices:
+            if const_expr(self.num_regs_epilogue < self.num_regs_per_thread):
+                cute.arch.setmaxregister_decrease(self.num_regs_epilogue)
+            elif const_expr(self.num_regs_epilogue > self.num_regs_per_thread):
+                cute.arch.setmaxregister_increase(self.num_regs_epilogue)
+            tmem.wait_for_alloc()
+            tmem_ptr = tmem.retrieve_ptr(self.dtype_acc)
+            tdVtdV0 = cute.make_tensor(tmem_ptr + self.tmem_offset_dV0, tdVtdV0_fake.layout)
+            tdVtdV1 = cute.make_tensor(tmem_ptr + self.tmem_offset_dV1, tdVtdV1_fake.layout)
+            self.dVacc_store(
+                mIndexTopk,
+                mdV,
+                sdV,
+                tdVtdV0,
+                tdVtdV1,
+                thr_mma_PtdOt,
+                pipeline_dV,
+                pipeline_dV_epi,
+                topk_length_dynamic,
+                block_info,
+                SeqlenInfoCls,
+                mCuSeqlensQ,
+                tile_scheduler=tile_scheduler,
+            )
+            tmem_alloc_barrier.arrive()
+    @cute.jit
+    def clc_scheduler_warp(
+        self,
+        tile_scheduler: TileSchedulerProtocol,
+    ):
+        work_tile = tile_scheduler.initial_work_tile_info()
+        while work_tile.is_valid_tile:
+            tile_scheduler.prefetch_next_work()
+            work_tile = tile_scheduler.advance_to_next_work()
+            cta_m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+            if cute.arch.thread_idx()[0] == self.clc_scheduler_warp_id * cute.arch.WARP_SIZE:
+                fa_printf(
+                    3,
+                    "[CLC] query sm={} cta={} (m_blk={},h={},b={},s={}) valid={}\n",
+                    smid(),
+                    cute.arch.block_idx()[0],
+                    work_tile.tile_idx[0],
+                    work_tile.tile_idx[1],
+                    work_tile.tile_idx[2],
+                    work_tile.tile_idx[3],
+                    work_tile.is_valid_tile,
+                )
+        tile_scheduler.producer_tail()
+    @cute.jit
+    def empty_warp(
+        self,
+        tile_scheduler: TileSchedulerProtocol,
+    ):
+        work_tile = tile_scheduler.initial_work_tile_info()
+        while work_tile.is_valid_tile:
+            work_tile = tile_scheduler.advance_to_next_work()
+    @cute.jit
+    def relay(
+        self,
+        pipeline_V: pipeline.PipelineAsyncUmma,
+        pipeline_V_cpasync: pipeline.PipelineAsync,
+        topk_length_dynamic: Optional[Int32],
+        block_info: BlockInfo,
+        SeqlenInfoCls: Callable,
+        tile_scheduler: TileSchedulerProtocol,
+    ):
+        # ==== Make pipeline states ====
+        # pipeline_V producer
+        # pipeline_V_cpasync consumer
+        producer_state_V = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, stages=self.num_stages_V
+        )
+        consumer_state_V = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, stages=self.num_stages_V
+        )
+        relay_V_fn = partial(self.relay_inner, pipeline_V_cpasync, pipeline_V)
+        work_tile = tile_scheduler.initial_work_tile_info()
+        while work_tile.is_valid_tile:
+            # m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+            # seqlen = SeqlenInfoCls(batch_idx)
+            num_n_block_groups = self.topk_length // self.cluster_tile_n
+            # num_n_block_groups = topk_length_dynamic // self.cluster_tile_n
+            # ==== Mainloop ====
+            for _ in cutlass.range(num_n_block_groups, unroll=1):
+                for _ in cutlass.range_constexpr(self.num_hdimv_splits):
+                    consumer_state_V, producer_state_V = relay_V_fn(
+                        consumer_state_V, producer_state_V
+                    )
+            # Advance to next tile
+            work_tile = tile_scheduler.advance_to_next_work()
+        pipeline_V.producer_tail(producer_state_V)
+    @cute.jit
+    def relay_inner(
+        self,
+        pipeline_cpasync: pipeline.PipelineAsync,
+        pipeline_mma: pipeline.PipelineAsyncUmma,
+        consumer_state: pipeline.PipelineState,
+        producer_state: pipeline.PipelineState,
+    ):
+        pipeline_cpasync.consumer_wait(consumer_state)
+        with cute.arch.elect_one():
+            pipeline_mma.producer_commit(producer_state)
+        consumer_state.advance()
+        producer_state.advance()
+        return consumer_state, producer_state
+    @cute.jit
+    def load_cpasync(
+        self,
+        mIndexTopk: cute.Tensor,
+        mV: cute.Tensor,
+        sV: cute.Tensor,
+        pipeline_V: pipeline.PipelineAsyncUmma,
+        pipeline_V_cpasync: pipeline.PipelineAsync,
+        topk_length_dynamic: Optional[Int32],
+        block_info: BlockInfo,
+        SeqlenInfoCls: Callable,
+        mCuSeqlensQ: Optional[cute.Tensor],
+        tile_scheduler: TileSchedulerProtocol,
+    ):
+        # ==== cpasync load warpgroup ====
+        # Description: loads tiles of V from gmem to smem using cpasync
+        # produces: V
+        # consumes: -
+        cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+        tidx = cute.arch.thread_idx()[0] % self.num_cpasync_load_threads
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % (
+            self.num_cpasync_load_threads // 32
+        )
+        # ==== Make pipeline states ====
+        # producer: acquire PipelineAsyncUmma <- mma
+        # producer: commit  PipelineAsync     -> relay
+        producer_state_V = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, stages=self.num_stages_V
+        )
+        work_tile = tile_scheduler.initial_work_tile_info()
+        while work_tile.is_valid_tile:
+            m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+            if const_expr(mCuSeqlensQ is not None):
+                batch_idx = get_batch_from_cu_tensor(m_block, mCuSeqlensQ)
+            seqlen = SeqlenInfoCls(batch_idx)
+            num_n_block_groups = self.topk_length // self.cluster_tile_n
+            # num_n_block_groups = topk_length_dynamic // self.cluster_tile_n
+            if const_expr(seqlen.has_cu_seqlens_q):
+                # m_block means absolute m_idx
+                mIndexTopk_cur = mIndexTopk[None, m_block]
+            else:
+                mIndexTopk_cur = mIndexTopk[None, m_block, batch_idx]
+            if const_expr(self.is_causal):
+                m_local_idx = (
+                    m_block - seqlen.offset_q if const_expr(seqlen.has_cu_seqlens_q) else m_block
+                )
+                seqlen_k_limit = m_local_idx + 1 + seqlen.seqlen_k - seqlen.seqlen_q
+            else:
+                seqlen_k_limit = seqlen.seqlen_k
+            cpasync_gather_kv_manager = CpasyncGatherKVManager.create(
+                mIndexTopk_cur,
+                cta_rank_in_cluster,
+                tidx,
+                warp_idx,
+                self.topk_length,
+                seqlen_k_limit,
+                self.cluster_tile_n,
+                self.hdim,
+                self.hdimv,
+                self.num_hdimv_splits,
+                self.num_cpasync_load_threads,
+                mV.element_type,
+                self.cta_group_size,
+                self.cpasync_barrier,
+                self.disable_bitmask,
+            )
+            # (seqlen_k, hdimv)
+            mV_cur = seqlen.offset_batch_K(mV, batch_idx, dim=3)[None, None, head_idx]
+            load_V = partial(
+                self.cpasync_gather_load_KV,
+                cpasync_gather_kv_manager,
+                pipeline_V,
+                pipeline_V_cpasync,
+                sV,
+                False,
+                "V",
+                mV_cur,
+            )
+            # ==== Mainloop ====
+            for n_block_group in cutlass.range(num_n_block_groups, unroll=1):
+                cpasync_gather_kv_manager.load_index_topk(n_block_group, transpose=False)
+                for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                    producer_state_V = load_V(producer_state_V, d_offset=split * self.hdimv // 2)
+            # Advance to next tile
+            work_tile = tile_scheduler.advance_to_next_work()
+        pipeline_V.producer_tail(producer_state_V)
+    @cute.jit
+    def cpasync_gather_load_KV(
+        self,
+        cpasync_gather_kv_manager: CpasyncGatherKVManager,
+        pipeline_mma: pipeline.PipelineAsyncUmma,
+        pipeline_cpasync: pipeline.PipelineAsync,
+        sX: cute.Tensor,
+        transpose: bool,
+        K_or_V: str,
+        mX: cute.Tensor,
+        producer_state: pipeline.PipelineState,
+        d_offset: int = 0,
+    ):
+        stage = producer_state.index
+        pipeline_mma.producer_acquire(producer_state)
+        cpasync_gather_kv_manager.load_X(
+            mX, sX[None, None, None, stage], transpose, K_or_V, d_offset
+        )
+        cute.arch.cp_async_commit_group()
+        pipeline_cpasync.sync_object_full.arrive_cp_async_mbarrier(stage)
+        producer_state.advance()
+        return producer_state
+    @cute.jit
+    def load(
+        self,
+        mdO: cute.Tensor,
+        mP: cute.Tensor,
+        mdOt: cute.Tensor,
+        mQvt: cute.Tensor,
+        mScaleP: Optional[cute.Tensor],
+        mdPsum: Optional[cute.Tensor],
+        sdO: cute.Tensor,
+        sP: cute.Tensor,
+        sdOt: cute.Tensor,
+        sQvt: cute.Tensor,
+        sScaleP: cute.Tensor,
+        sdPsum: cute.Tensor,
+        tma_atom_dO: cute.CopyAtom,
+        tma_atom_P: cute.CopyAtom,
+        tma_atom_dOt: cute.CopyAtom,
+        tma_atom_Qvt: cute.CopyAtom,
+        pipeline_dO: pipeline.PipelineAsync,  # TmaUmma
+        pipeline_P: pipeline.PipelineAsync,  # TmaAsync
+        pipeline_dOt_Qvt: pipeline.PipelineAsync,  # TmaUmma
+        pipeline_Pt: pipeline.PipelineAsync,  # AsyncUmma
+        pipeline_dV_epi: pipeline.PipelineAsync,  # Async
+        pipeline_scaleP: pipeline.PipelineAsync,  # TmaAsync
+        pipeline_dPsum: pipeline.PipelineAsync,  # TmaAsync
+        thr_mma_VdO: cute.ThrMma,
+        thr_mma_PtdOt: cute.ThrMma,
+        thr_mma_dStQvt: cute.ThrMma,
+        topk_length_dynamic: Optional[Int32],
+        block_info: BlockInfo,
+        SeqlenInfoCls: Callable,
+        mCuSeqlensQ: Optional[cute.Tensor],
+        tile_scheduler: TileSchedulerProtocol,
+    ):
+        # ==== Load warp ====
+        # Description: loads tiles of dO, P, dOt, Qvt from gmem to smem using TMA
+        # produces: dO, P, dOt, Qvt
+        # consumes: -
+        cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+        lane_idx = cute.arch.lane_idx()
+        # ==== Make pipeline states ====
+        Producer = pipeline.PipelineUserType.Producer
+        producer_state_dO = pipeline.make_pipeline_state(Producer, stages=self.num_stages_dO)
+        producer_state_P = pipeline.make_pipeline_state(Producer, stages=self.num_stages_P)
+        producer_state_dOt_Qvt = pipeline.make_pipeline_state(Producer, stages=self.num_stages_dOt)
+        producer_state_dV_epi = pipeline.make_pipeline_state(Producer, stages=self.num_stages_dV)
+        producer_state_scaleP = pipeline.make_pipeline_state(
+            Producer, stages=self.num_stages_scaleP
+        )
+        producer_state_dPsum = pipeline.make_pipeline_state(Producer, stages=self.num_stages_dPsum)
+        copy_atom_stats = cute.make_copy_atom(cpasync.CopyBulkG2SOp(), Float32)
+        work_tile = tile_scheduler.initial_work_tile_info()
+        while work_tile.is_valid_tile:
+            m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+            if const_expr(mCuSeqlensQ is not None):
+                batch_idx = get_batch_from_cu_tensor(m_block, mCuSeqlensQ)
+            seqlen = SeqlenInfoCls(batch_idx)
+            if const_expr(mCuSeqlensQ is not None):
+                m_block -= seqlen.offset_q
+            num_n_block_groups = self.topk_length // self.cluster_tile_n
+            # num_n_block_groups = topk_length_dynamic // self.cluster_tile_n
+            # ==== Partition GMEM tensors ====
+            # (seqlen_q, topk or hdimv)
+            mP_cur = seqlen.offset_batch_Q(mP, batch_idx, dim=3)[None, None, head_idx]
+            mdO_cur = seqlen.offset_batch_Q(mdO, batch_idx, dim=3)[None, None, head_idx]
+            # (hdimv, seqlen_q)
+            offset = (
+                (0, seqlen.offset_q) if const_expr(not self.pack_gqa) else (0, (0, seqlen.offset_q))
+            )
+            if const_expr(not seqlen.has_cu_seqlens_q):
+                mdOt_cur = mdOt[None, None, head_idx, batch_idx]
+                mQvt_cur = mQvt[None, None, head_idx, batch_idx]
+            else:
+                mdOt_cur = cute.domain_offset(offset, mdOt[None, None, head_idx])
+                mQvt_cur = cute.domain_offset(offset, mQvt[None, None, head_idx])
+            gScaleP = None
+            if const_expr(mScaleP is not None):
+                mScaleP_cur = seqlen.offset_batch_Q(mScaleP, batch_idx, dim=3)[None, None, head_idx]
+                # (tile_m, topk//128)
+                gScaleP = cute.local_tile(mScaleP_cur, (self.tile_m,), (m_block, None))
+            gdPsum = None
+            if const_expr(mdPsum is not None):
+                mdPsum_cur = seqlen.offset_batch_Q(mdPsum, batch_idx, dim=2)[None, head_idx]
+                # (tile_m)
+                gdPsum = cute.local_tile(mdPsum_cur, (self.tile_m,), (m_block,))
+            # (tile_m, tile_n, n_blocks)
+            gP = cute.local_tile(
+                mP_cur,
+                (self.tile_m, self.tile_n),
+                (m_block, None),
+            )
+            # (tile_m, hdimv//2, 2)
+            gdO = cute.local_tile(
+                mdO_cur,
+                (self.mma_tiler_VdO[1], self.mma_tiler_VdO[2]),
+                (m_block, None),
+            )
+            # (hdimv//2, tile_m, 2)
+            gdOt = cute.local_tile(
+                mdOt_cur,
+                (self.mma_tiler_PtdOt[1], self.mma_tiler_PtdOt[2]),
+                (None, m_block),
+            )
+            gQvt = cute.local_tile(
+                mQvt_cur,
+                (self.mma_tiler_dStQvt[1], self.mma_tiler_dStQvt[2]),
+                (None, m_block),
+            )
+            tdPgdO = thr_mma_VdO.partition_B(gdO)
+            tdVgdOt = thr_mma_PtdOt.partition_B(gdOt)
+            tdVgQvt = thr_mma_dStQvt.partition_B(gQvt)
+            # (V, REST)
+            tPsP, tPgP = cpasync.tma_partition(
+                atom=tma_atom_P,
+                cta_coord=0,
+                cta_layout=cute.make_layout(1),
+                smem_tensor=cute.group_modes(sP, 0, 2),
+                gmem_tensor=cute.group_modes(gP, 0, 2),
+            )
+            tdOsdO, tdOgdO = cpasync.tma_partition(
+                atom=tma_atom_dO,
+                cta_coord=0,
+                cta_layout=cute.make_layout(1),
+                smem_tensor=cute.group_modes(sdO, 0, 3),
+                gmem_tensor=cute.group_modes(tdPgdO, 0, 3),
+            )
+            tdOtsdOt, tdOtgdOt = cpasync.tma_partition(
+                atom=tma_atom_dOt,
+                cta_coord=0,
+                cta_layout=cute.make_layout(1),
+                smem_tensor=cute.group_modes(sdOt, 0, 3),
+                gmem_tensor=cute.group_modes(tdVgdOt, 0, 3),
+            )
+            tQvtsQvt, tQvtgQvt = cpasync.tma_partition(
+                atom=tma_atom_Qvt,
+                cta_coord=0,
+                cta_layout=cute.make_layout(1),
+                smem_tensor=cute.group_modes(sQvt, 0, 3),
+                gmem_tensor=cute.group_modes(tdVgQvt, 0, 3),
+            )
+            load_P = partial(self.load_inner, tma_atom_P, tPgP, tPsP, pipeline_P)
+            load_dO = partial(self.load_inner, tma_atom_dO, tdOgdO, tdOsdO, pipeline_dO)
+            load_dOt = partial(self.load_inner, tma_atom_dOt, tdOtgdOt, tdOtsdOt, pipeline_dOt_Qvt)
+            load_Qvt = partial(self.load_inner, tma_atom_Qvt, tQvtgQvt, tQvtsQvt, pipeline_dOt_Qvt)
+            load_scaleP = partial(
+                self.load_inner, copy_atom_stats, gScaleP, sScaleP, pipeline_scaleP, bulk_copy=True
+            )
+            load_dPsum = partial(
+                self.load_inner, copy_atom_stats, gdPsum, sdPsum, pipeline_dPsum, bulk_copy=True
+            )
+            # ==== Load stationary operands ====
+            for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                producer_state_dO = load_dO(producer_state_dO, block=split)
+            producer_state_dPsum = load_dPsum(producer_state_dPsum)
+            # ==== Mainloop ====
+            for n_block_group in cutlass.range(num_n_block_groups, unroll=1):
+                n_block = 2 * n_block_group + cta_rank_in_cluster
+                # load ScaleP
+                if const_expr(mScaleP is not None):
+                    producer_state_scaleP = load_scaleP(producer_state_scaleP, block=n_block_group)
+                pipeline_Pt.producer_acquire(producer_state_P)
+                producer_state_P = load_P(producer_state_P, block=n_block)
+                for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                    pipeline_dV_epi.producer_acquire(producer_state_dV_epi)
+                    producer_state_dV_epi.advance()
+                    producer_state_dOt_Qvt = load_dOt(producer_state_dOt_Qvt, block=split)
+                for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                    producer_state_dOt_Qvt = load_Qvt(producer_state_dOt_Qvt, block=split)
+            # Advance to next tile
+            work_tile = tile_scheduler.advance_to_next_work()
+        pipeline_P.producer_tail(producer_state_P)
+        pipeline_dO.producer_tail(producer_state_dO)
+        pipeline_dOt_Qvt.producer_tail(producer_state_dOt_Qvt)
+    @cute.jit
+    def load_inner(
+        self,
+        copy_atom: cute.CopyAtom,
+        tXgX: cute.Tensor,
+        tXsX: cute.Tensor,
+        load_pipeline: pipeline.PipelineAsync,
+        producer_state: pipeline.PipelineState,
+        block: Optional[Int32] = None,
+        bulk_copy: bool = False,
+    ):
+        if const_expr(block is not None):
+            tXgX = tXgX[(None, block)]
+        if const_expr(cute.rank(tXsX) != 1):
+            assert cute.rank(tXsX) == 2, f"wrong rank for tXsX, got {cute.rank(tXsX)}"
+            stage = producer_state.index
+            tXsX = tXsX[(None, stage)]
+        load_pipeline.producer_acquire(producer_state)
+        mbar_ptr = load_pipeline.producer_get_barrier(producer_state)
+        if const_expr(bulk_copy):
+            with cute.arch.elect_one():
+                cute.copy(copy_atom, tXgX, tXsX, mbar_ptr=mbar_ptr)
+        else:
+            cute.copy(copy_atom, tXgX, tXsX, tma_bar_ptr=mbar_ptr)
+        producer_state.advance()
+        return producer_state
+    @cute.jit
+    def mma(
+        self,
+        sV: cute.Tensor,
+        sdO: cute.Tensor,
+        sPt: cute.Tensor,
+        sdOt: cute.Tensor,
+        sdSt: cute.Tensor,
+        sQvt: cute.Tensor,
+        tdPtdP: cute.Tensor,
+        tdVtdV0: cute.Tensor,
+        tdVtdV1: cute.Tensor,
+        tiled_mma_VdO: cute.TiledMma,
+        tiled_mma_PtdOt: cute.TiledMma,
+        tiled_mma_dStQvt: cute.TiledMma,
+        pipeline_V: pipeline.PipelineAsync,  # AsyncUmma
+        pipeline_dO: pipeline.PipelineAsync,  # TmaUmma
+        pipeline_dPt: pipeline.PipelineAsync,  # UmmaAsync
+        pipeline_Pt: pipeline.PipelineAsync,  # AsyncUmma
+        pipeline_dOt_Qvt: pipeline.PipelineAsync,  # TmaUmma
+        pipeline_dSt: pipeline.PipelineAsync,  # AsyncUmma
+        pipeline_dV: pipeline.PipelineAsync,  # UmmaAsync
+        is_leader_cta: Boolean,
+        topk_length_dynamic: Optional[Int32],
+        block_info: BlockInfo,
+        SeqlenInfoCls: Callable,
+        mCuSeqlensQ: Optional[cute.Tensor],
+        tile_scheduler: TileSchedulerProtocol,
+    ):
+        # ==== mma warp ====
+        # Description: Computes dP = V @ dO^T, dV = P^T @ dO, and dV += dS^T @ Qv
+        # i.e. dP = gemm(V, dO), dV += gemm(P.T, dO.T), dV += gemm(dS.T, Qv.T)
+        # Produces: dP, dV
+        # Consumes: V, dO, P.T, dO.T, dS.T, Qv.T
+        lane_idx = cute.arch.lane_idx()
+        tdVtdVs = [tdVtdV0, tdVtdV1]
+        # Set accumulate = True for dS^T@Qv since we are accumulating on the P^T@dO result
+        tiled_mma_dStQvt.set(tcgen05.Field.ACCUMULATE, True)
+        # Operands for dP=V@dO^T
+        tdPrV = tiled_mma_VdO.make_fragment_A(sV)
+        tdPrdO = tiled_mma_VdO.make_fragment_B(sdO)
+        # Operands for dVi=P^T@dOi
+        tdVrPt = tiled_mma_PtdOt.make_fragment_A(sPt)
+        tdVrdOt = tiled_mma_PtdOt.make_fragment_B(sdOt)
+        # Operands for dVi+=dS^T@Qvi
+        tdVrdSt = tiled_mma_dStQvt.make_fragment_A(sdSt)
+        tdVrQvt = tiled_mma_dStQvt.make_fragment_B(sQvt)
+        use_ptx_gemm_VdO = False
+        use_ptx_gemm_PtdOt = False
+        use_ptx_gemm_dStQvt = False
+        # GEMM functions
+        if const_expr(use_ptx_gemm_VdO):
+            gemm_VdO = partial(
+                fa_sm100_utils.gemm_ptx_partial,
+                tiled_mma_VdO.op,
+                self.tmem_offset_dP,
+                zero_init=True,
+                cta_group=self.cta_group_size,
+            )
+        else:
+            gemm_VdO = partial(
+                fa_sm100_utils.gemm,
+                tiled_mma_VdO,
+                tdPtdP,
+            )
+        if const_expr(use_ptx_gemm_PtdOt):
+            gemm_PtdOt = [
+                partial(
+                    fa_sm100_utils.gemm_ptx_partial,
+                    tiled_mma_PtdOt.op,
+                    self.tmem_offsets_dV[split],
+                    zero_init=True,
+                    cta_group=self.cta_group_size,
+                )
+                for split in range(self.num_hdimv_splits)
+            ]
+        else:
+            gemm_PtdOt = [
+                partial(
+                    fa_sm100_utils.gemm,
+                    tiled_mma_PtdOt,
+                    tdVtdVs[split],
+                    zero_init=True,
+                )
+                for split in range(self.num_hdimv_splits)
+            ]
+        if const_expr(use_ptx_gemm_dStQvt):
+            gemm_dStQvt = [
+                partial(
+                    fa_sm100_utils.gemm_ptx_partial,
+                    tiled_mma_dStQvt.op,
+                    self.tmem_offsets_dV[split],
+                    zero_init=False,
+                    cta_group=self.cta_group_size,
+                )
+                for split in range(self.num_hdimv_splits)
+            ]
+        else:
+            gemm_dStQvt = [
+                partial(
+                    fa_sm100_utils.gemm,
+                    tiled_mma_dStQvt,
+                    tdVtdVs[split],
+                    zero_init=False,
+                )
+                for split in range(self.num_hdimv_splits)
+            ]
+        Consumer, Producer = pipeline.PipelineUserType.Consumer, pipeline.PipelineUserType.Producer
+        consumer_state_V = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_V)
+        consumer_state_dO = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_dO)
+        consumer_state_Pt = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_Pt)
+        consumer_state_dOt_Qvt = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_dOt)
+        consumer_state_dSt = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_dSt)
+        producer_state_dPt = pipeline.make_pipeline_state(Producer, stages=self.num_stages_dPt)
+        producer_state_dV = pipeline.make_pipeline_state(Producer, stages=self.num_stages_dV)
+        mma_VdO = partial(
+            self.mma_inner, gemm_VdO, pipeline_V, tdPrV, sV, tdPrdO, sdO, swap_AB_stage=True, use_ptx=use_ptx_gemm_VdO
+        )
+        mma_PtdOt = partial(
+            self.mma_inner, gemm_PtdOt, pipeline_dOt_Qvt, tdVrPt, sPt, tdVrdOt, sdOt, use_ptx=use_ptx_gemm_PtdOt
+        )
+        mma_dStQvt = partial(
+            self.mma_inner, gemm_dStQvt, pipeline_dOt_Qvt, tdVrdSt, sdSt, tdVrQvt, sQvt, use_ptx=use_ptx_gemm_dStQvt
+        )
+        work_tile = tile_scheduler.initial_work_tile_info()
+        while work_tile.is_valid_tile:
+            # m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+            # if const_expr(mCuSeqlensQ is not None):
+            #     batch_idx = get_batch_from_cu_tensor(m_block, mCuSeqlensQ)
+            # seqlen = SeqlenInfoCls(batch_idx)
+            # num_n_block_groups = self.topk_length // self.cluster_tile_n
+            num_n_block_groups = topk_length_dynamic // self.cluster_tile_n
+            if is_leader_cta:
+                # ==== Prologue ====
+                consumer_wait_state_dO = consumer_state_dO.clone()
+                for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                    pipeline_dO.consumer_wait(consumer_wait_state_dO)
+                    consumer_wait_state_dO.advance()
+                # ==== Mainloop ====
+                for _ in cutlass.range(num_n_block_groups, unroll=1):
+                    # 1. dP = V @ dO^T
+                    # mma inner waits for V
+                    pipeline_dPt.producer_acquire(producer_state_dPt)
+                    for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                        consumer_state_V = mma_VdO(
+                            consumer_state_V, a_stage=split, zero_init=split == 0
+                        )
+                    pipeline_dPt.producer_commit(producer_state_dPt)
+                    producer_state_dPt.advance()
+                    # 2. dV = P^T @ dO
+                    # mma inner waits for dOt
+                    pipeline_Pt.consumer_wait(consumer_state_Pt)
+                    producer_acquire_state_dV = producer_state_dV.clone()
+                    for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                        pipeline_dV.producer_acquire(producer_acquire_state_dV)
+                        producer_acquire_state_dV.advance()
+                        consumer_state_dOt_Qvt = mma_PtdOt(consumer_state_dOt_Qvt, acc_stage=split)
+                    pipeline_Pt.consumer_release(consumer_state_Pt)
+                    consumer_state_Pt.advance()
+                    # 3. dV += dS^T @ Qv
+                    # mma inner waits for Qvt
+                    pipeline_dSt.consumer_wait(consumer_state_dSt)
+                    for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                        consumer_state_dOt_Qvt = mma_dStQvt(consumer_state_dOt_Qvt, acc_stage=split)
+                        pipeline_dV.producer_commit(producer_state_dV)
+                        producer_state_dV.advance()
+                    pipeline_dSt.consumer_release(consumer_state_dSt)
+                    consumer_state_dSt.advance()
+                # ==== Epilogue ====
+                for _ in cutlass.range_constexpr(self.num_hdimv_splits):
+                    pipeline_dO.consumer_release(consumer_state_dO)
+                    consumer_state_dO.advance()
+            # Advance to next tile
+            work_tile = tile_scheduler.advance_to_next_work()
+        pipeline_dPt.producer_tail(producer_state_dPt)
+        pipeline_dV.producer_tail(producer_state_dV)
+    @cute.jit
+    def mma_inner(
+        self,
+        gemm,
+        load_pipeline,
+        tCrA,
+        sA,
+        tCrB,
+        sB,
+        consumer_state: pipeline.PipelineState,
+        acc_stage: Optional[Int32] = None,
+        a_stage: Int32 = 0,
+        zero_init: Optional[bool] = None,
+        swap_AB_stage: bool = False,
+        use_ptx: bool = True,
+    ):
+        if const_expr(acc_stage is not None):
+            gemm = gemm[acc_stage]
+        smem_stage = consumer_state.index
+        if const_expr(not swap_AB_stage):
+            a_stage = a_stage
+            b_stage = smem_stage
+        else:
+            a_stage = smem_stage
+            b_stage = a_stage
+        tCrA_cur = tCrA[None, None, None, a_stage]
+        sA_cur = sA[None, None, None, a_stage]
+        tCrB_cur = tCrB[None, None, None, b_stage]
+        sB_cur = sB[None, None, None, b_stage]
+        kwargs = dict(tCrA=tCrA_cur, tCrB=tCrB_cur)
+        if const_expr(use_ptx):
+            kwargs |= dict(sA=sA_cur, sB=sB_cur)
+        if const_expr(zero_init is not None):
+            kwargs["zero_init"] = zero_init
+        load_pipeline.consumer_wait(consumer_state)
+        gemm(**kwargs)
+        load_pipeline.consumer_release(consumer_state)
+        consumer_state.advance()
+        return consumer_state
+    @cute.jit
+    def compute_loop(
+        self,
+        softmax_scale: Float32,
+        softmax_scale_log2: Float32,
+        thr_mma_VdO: cute.ThrMma,
+        tdPtdP: cute.Tensor,
+        sP: cute.Tensor,
+        sdS: cute.Tensor,
+        sScaleP: cute.Tensor,
+        sdPsum: cute.Tensor,
+        mdS: cute.Tensor,
+        tma_atom_dS: cute.CopyAtom,
+        pipeline_P: pipeline.PipelineAsync,  # TmaAsync
+        pipeline_Pt: pipeline.PipelineAsync,  # AsyncUmma
+        pipeline_dPt: pipeline.PipelineAsync,  # UmmaAsync
+        pipeline_dSt: pipeline.PipelineAsync,  # AsyncUmma
+        pipeline_scaleP: pipeline.PipelineAsync,  # TmaAsync
+        pipeline_dPsum: pipeline.PipelineAsync,  # TmaAsync
+        topk_length_dynamic: Optional[Int32],
+        block_info: BlockInfo,
+        SeqlenInfoCls: Callable,
+        mCuSeqlensQ: Optional[cute.Tensor],
+        tile_scheduler: TileSchedulerProtocol,
+    ):
+        tidx = cute.arch.thread_idx()[0] % self.num_softmax_threads
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % (
+            self.num_softmax_threads // 32
+        )
+        cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+        leader_warp = warp_idx == 0
+        # 256b // 32 = 8 values, 128 mqa // 2 => tmem_rep = 8
+        tmem_rep = self.tile_m // self.cta_group_size // 8
+        copy_atom_t2r = cute.make_copy_atom(
+            tcgen05.copy.Ld16x256bOp(tcgen05.copy.Repetition(tmem_rep)),
+            self.dtype_acc,
+        )
+        # ((64,(64,2)),1,1):((65536,(1,4194304 = 65536*64)),0,0)
+        tdPtdP = tdPtdP[(None, None), 0, 0]
+        tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tdPtdP)
+        thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
+        # (T2R, T2R_M, T2R_N)
+        # (((64,16),1),2,1):(((1,65536),0),1048576,0)>, 1048576/65536 = 16
+        tdPtdP_t2r = thr_copy_t2r.partition_S(tdPtdP)
+        cdP = cute.make_identity_tensor(self.mma_tiler_VdO[:2])  # (128, 128)
+        tdPcdP = thr_mma_VdO.partition_C(cdP)[(None, None), 0, 0]  # (64,128):(1@0,1@1)
+        # (((2,2,8),1),2,1):(((1@1,8@0,8@1),0),16@0,0)
+        tdPcdP_t2r = thr_copy_t2r.partition_D(tdPcdP)
+        assert tdPcdP_t2r.shape[0][1] == 1, f"unexpected tdPcdP_t2r shape, got {tdPcdP_t2r.shape}"
+        smem_load_op = cute.nvgpu.warp.LdMatrix8x8x16bOp(True, 4)  # ldsm x num_matrices = ldsm x 4
+        smem_store_op = cute.nvgpu.warp.StMatrix8x8x16bOp(True, 4)  # stsm x num_matrices = stsm x 4
+        smem_load_atom = cute.make_copy_atom(smem_load_op, self.dtype)
+        smem_store_atom = cute.make_copy_atom(smem_store_op, self.dtype)
+        tiled_copy_r2s = cute.make_tiled_copy_D(smem_store_atom, tiled_copy_t2r)
+        tiled_copy_s2r = cute.make_tiled_copy_D(smem_load_atom, tiled_copy_t2r)
+        thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
+        thr_copy_s2r = tiled_copy_s2r.get_slice(tidx)
+        sPt_load_layout = cute.make_ordered_layout(
+            self.tile_Pt + (self.num_stages_P,), order=(1, 0, 2)
+        )
+        # (tile_n, tile_m, stages_P)
+        sPt = cute.composition(sP, sPt_load_layout)
+        sdSt = cute.composition(sdS, sPt_load_layout)
+        # (R2S, R2S_M, R2S_N, PIPE_D)
+        # ((8,4),2,1,1):((1,1024),16,0,0)
+        tSR_sPt = thr_copy_s2r.partition_S(sPt)
+        tRS_sdSt = thr_copy_r2s.partition_D(sdSt)
+        # ((2,2),(2,8,1),stage):((0,0),(1,8,0),_)
+        tPsScaleP_nm = self.broadcast_tensor_nm_view(sScaleP, thr_mma_VdO, thr_copy_t2r)
+        tdPsdPsum_nm = self.broadcast_tensor_nm_view(sdPsum, thr_mma_VdO, thr_copy_t2r)
+        Consumer, Producer = pipeline.PipelineUserType.Consumer, pipeline.PipelineUserType.Producer
+        consumer_state_P = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_P)
+        consumer_state_dPt = pipeline.make_pipeline_state(Consumer, stages=1)
+        consumer_state_scaleP = pipeline.make_pipeline_state(
+            Consumer, stages=self.num_stages_scaleP
+        )
+        consumer_state_dPsum = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_dPsum)
+        producer_state_Pt = pipeline.make_pipeline_state(Producer, stages=self.num_stages_Pt)
+        producer_state_dSt = pipeline.make_pipeline_state(Producer, stages=self.num_stages_dSt)
+        work_tile = tile_scheduler.initial_work_tile_info()
+        while work_tile.is_valid_tile:
+            m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+            if const_expr(mCuSeqlensQ is not None):
+                batch_idx = get_batch_from_cu_tensor(m_block, mCuSeqlensQ)
+            seqlen = SeqlenInfoCls(batch_idx)
+            if const_expr(mCuSeqlensQ is not None):
+                m_block -= seqlen.offset_q
+            num_n_block_groups = self.topk_length // self.cluster_tile_n
+            # num_n_block_groups = topk_length_dynamic // self.cluster_tile_n
+            mdS_cur = seqlen.offset_batch_Q(mdS, batch_idx, dim=3)[None, None, head_idx]
+            gdS = cute.local_tile(mdS_cur, (self.tile_m, self.tile_n), (m_block, None))
+            store_dS, _, _ = copy_utils.tma_get_copy_fn(
+                tma_atom_dS,
+                0,
+                cute.make_layout(1),
+                sdS,
+                gdS,
+            )
+            pipeline_dPsum.consumer_wait(consumer_state_dPsum)
+            tdPsdPsum_cur = tdPsdPsum_nm[0, None, consumer_state_dPsum.index]
+            tdPrdPsum_cur_f32 = cute.make_rmem_tensor(tdPsdPsum_cur.shape, dtype=self.dtype_scale)
+            cute.autovec_copy(tdPsdPsum_cur, tdPrdPsum_cur_f32)
+            for n_block_group in cutlass.range(num_n_block_groups, unroll=1):
+                n_block = 2 * n_block_group + cta_rank_in_cluster
+                pipeline_P.consumer_wait(consumer_state_P)
+                # todo: ablate wait -> try_wait
+                pipeline_scaleP.consumer_wait(consumer_state_scaleP)
+                # (((2,2,8),1),2,1):(((1,2,4),0),32,0)
+                rPt = cute.make_rmem_tensor(tdPcdP_t2r.shape, self.dtype)
+                # (S2R, S2R_M, S2R_N)
+                rPt_copy_view = tiled_copy_s2r.retile(rPt)
+                tSR_sPt_cur = tSR_sPt[None, None, None, consumer_state_P.index]
+                cute.copy(tiled_copy_s2r, tSR_sPt_cur, rPt_copy_view)
+                # ((2,2),(2,8,1)):((2,32),(1,4,0))
+                rP_nm = layout_utils.reshape_acc_to_mn(rPt[(None, 0), None, None])
+                tPsScaleP_cur = tPsScaleP_nm[0, None, consumer_state_scaleP.index]
+                tPrScaleP_cur_f32 = cute.make_rmem_tensor(
+                    tPsScaleP_cur.shape, dtype=self.dtype_scale
+                )
+                tPrScaleP_cur = cute.make_rmem_tensor(tPsScaleP_cur.shape, dtype=self.dtype)
+                cute.autovec_copy(tPsScaleP_cur, tPrScaleP_cur_f32)
+                tPrScaleP_cur.store(tPrScaleP_cur_f32.load().to(self.dtype))
+                # scale P
+                for n in cutlass.range_constexpr(cute.size(rP_nm.shape[0])):
+                    rP_cur = rP_nm[n, None]
+                    rP_cur.store(rP_cur.load() * tPrScaleP_cur.load())
+                cute.arch.sync_warp()
+                cute.copy(tiled_copy_r2s, rPt_copy_view, tSR_sPt_cur)
+                cute.arch.fence_view_async_shared()
+                self.softmax_barrier.arrive_and_wait()
+                pipeline_scaleP.consumer_release(consumer_state_scaleP)
+                consumer_state_scaleP.advance()
+                pipeline_Pt.producer_commit(producer_state_Pt)
+                producer_state_Pt.advance()
+                # note: mma also signals Pt free, signal acquired in tma warp
+                pipeline_P.consumer_release(consumer_state_P)
+                consumer_state_P.advance()
+                pipeline_dPt.consumer_wait(consumer_state_dPt)
+                # (((2,2,8),1),2,1):(((1,2,4),0),32,0)
+                tdPrdP_t2r = cute.make_rmem_tensor(tdPcdP_t2r.shape, self.dtype_acc)
+                cute.copy(tiled_copy_t2r, tdPtdP_t2r, tdPrdP_t2r)
+                cute.arch.fence_view_async_tmem_load()
+                self.softmax_barrier.arrive_and_wait()
+                pipeline_dPt.consumer_release(consumer_state_dPt)
+                consumer_state_dPt.advance()
+                # dS = P o (dP - dPsum)
+                rdP_nm = layout_utils.reshape_acc_to_mn(tdPrdP_t2r[(None, 0), None, None])
+                for n in cutlass.range_constexpr(cute.size(rdP_nm.shape[0])):
+                    rdP_cur = rdP_nm[n, None]
+                    rdP_cur.store(rdP_cur.load() - tdPrdPsum_cur_f32.load())
+                rPt.store(rPt.load() * (tdPrdP_t2r.load() * softmax_scale).to(self.dtype))
+                # wait for tma store to free dSt buffer
+                if leader_warp:
+                    cute.arch.cp_async_bulk_wait_group(1 - self.num_stages_dSt, read=True)
+                self.softmax_barrier.arrive_and_wait()
+                # note: dS guaranteed free as mma operand
+                pipeline_dSt.producer_acquire(producer_state_dSt)
+                tRS_sdSt_cur = tRS_sdSt[None, None, None, producer_state_dSt.index]
+                cute.copy(tiled_copy_r2s, rPt_copy_view, tRS_sdSt_cur)
+                cute.arch.fence_view_async_shared()
+                self.softmax_barrier.arrive_and_wait()
+                pipeline_dSt.producer_commit(producer_state_dSt)
+                # tma store
+                if leader_warp:
+                    store_dS(src_idx=producer_state_dSt.index, dst_idx=n_block)
+                    cute.arch.cp_async_bulk_commit_group()
+                producer_state_dSt.advance()
+            pipeline_dPsum.consumer_release(consumer_state_dPsum)
+            consumer_state_dPsum.advance()
+            # Advance to next tile
+            work_tile = tile_scheduler.advance_to_next_work()
+        # producer tails
+    @cute.jit
+    def broadcast_tensor_nm_view(
+        self,
+        sX: cute.Tensor,  # (tile_m, num_stages)
+        thr_mma: cute.ThrMma,
+        thr_copy_t2r: cute.ThrCopy,
+    ):
+        assert cute.size(sX.shape[0]) == self.tile_m
+        num_stages = sX.shape[1] if const_expr(cute.rank(sX) > 1) else 1
+        sX_2D_cluster = cute.make_tensor(
+            sX.iterator,
+            cute.make_layout(
+                (self.tile_m, self.cluster_tile_n, num_stages),
+                stride=(1, 0, self.tile_m),
+            ),
+        )
+        sXt_2D_cluster = layout_utils.transpose_view(sX_2D_cluster)
+        sXt_2D = thr_mma.partition_C(sXt_2D_cluster)[(None, None), 0, 0, None]
+        tXsXt_2D = thr_copy_t2r.partition_D(sXt_2D)[(None, 0), None, None, None]
+        tXsXt_nm = layout_utils.make_acc_tensor_mn_view(tXsXt_2D)
+        return tXsXt_nm
+    @cute.jit
+    def dVacc_store(
+        self,
+        mIndexTopk: cute.Tensor,
+        mdV: cute.Tensor,
+        sdV: cute.Tensor,
+        tdVtdV0: cute.Tensor,
+        tdVtdV1: cute.Tensor,
+        thr_mma_PtdOt: cute.ThrMma,
+        pipeline_dV: pipeline.PipelineAsync,  # UmmaAsync
+        pipeline_dV_epi: pipeline.PipelineAsync,  # Async
+        topk_length_dynamic: Optional[Int32],
+        block_info: BlockInfo,
+        SeqlenInfoCls: Callable,
+        mCuSeqlensQ: Optional[cute.Tensor],
+        tile_scheduler: TileSchedulerProtocol,
+    ):
+        # ==== dVaccum store warpgroup ====
+        # produces: -
+        # consumes: dV
+        tdVtdV0 = tdVtdV0[(None, None), 0, 0]
+        tdVtdV1 = tdVtdV1[(None, None), 0, 0]
+        num_epi_warps = self.num_epilogue_threads // 32
+        cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+        tidx = cute.arch.thread_idx()[0] % self.num_epilogue_threads
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % num_epi_warps
+        leader_warp = warp_idx == 0
+        wg_half = warp_idx // 2
+        consumer_state_dV = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, stages=self.num_stages_dV
+        )
+        copy_atom_t2r = cute.make_copy_atom(
+            tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)),
+            self.dtype_acc,
+        )
+        tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tdVtdV0)
+        thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
+        tdVtdV0_t2r = thr_copy_t2r.partition_S(tdVtdV0)
+        tdVtdV1_t2r = thr_copy_t2r.partition_S(tdVtdV1)
+        tdVtdVs_t2r = [tdVtdV0_t2r, tdVtdV1_t2r]
+        cdVmma = cute.make_identity_tensor(self.mma_tiler_PtdOt[:2])
+        tdVcdVmma = thr_mma_PtdOt.partition_C(cdVmma)[(None, None), 0, 0]
+        tdVcdVmma_t2r = thr_copy_t2r.partition_D(tdVcdVmma)
+        # 64 threads x 4 values to tile over tile_dV = (64, 32)
+        tiled_copy_r2s = tiled_copy_2d(self.dtype_acc, 4, 64)
+        thr_copy_r2s = tiled_copy_r2s.get_slice(tidx % 64)
+        # ((4,1),1,8,(1,8)):((1,0),0,4,(0,2048))
+        tRS_sdV = thr_copy_r2s.partition_D(sdV)
+        tiled_copy_s2r = copy_utils.tiled_copy_2d(self.dtype_acc, 8, self.num_epilogue_threads, 4)
+        thr_copy_s2r = tiled_copy_s2r.get_slice(tidx)
+        # (V, M, N, STAGE)
+        tSR_sdV = thr_copy_s2r.partition_S(sdV)
+        cdV = cute.make_identity_tensor(cute.product_each(sdV.shape[:2]))
+        # (V, M, N)
+        tdVcdV = thr_copy_s2r.partition_S(cdV)
+        gmem_rows_per_thread = cute.size(tSR_sdV.shape[1])
+        work_tile = tile_scheduler.initial_work_tile_info()
+        while work_tile.is_valid_tile:
+            m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+            if const_expr(mCuSeqlensQ is not None):
+                batch_idx = get_batch_from_cu_tensor(m_block, mCuSeqlensQ)
+            seqlen = SeqlenInfoCls(batch_idx)
+            num_n_block_groups = self.topk_length // self.cluster_tile_n
+            # num_n_block_groups = topk_length_dynamic // self.cluster_tile_n
+            # (seqlen_k, hdimv)
+            mdV_cur = seqlen.offset_batch_K(mdV, batch_idx, dim=3)[None, None, head_idx]
+            # (topk, dv)
+            if const_expr(seqlen.has_cu_seqlens_q):
+                # m_block means absolute m_idx
+                mIndexTopk_cur = mIndexTopk[None, m_block]
+            else:
+                mIndexTopk_cur = mIndexTopk[None, m_block, batch_idx]
+            # ==== Mainloop ====
+            for n_block_group in cutlass.range(num_n_block_groups, unroll=1):
+                n_block = 2 * n_block_group + cta_rank_in_cluster
+                rIdxTopK = cute.make_rmem_tensor((gmem_rows_per_thread,), dtype=self.dtype_index)
+                for j in cutlass.range_constexpr(gmem_rows_per_thread):
+                    n_idx = n_block * self.tile_n + tdVcdV[0, j, 0][0]
+                    rIdxTopK[j] = mIndexTopk_cur[n_idx]
+                for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                    tdVtdV_t2r = tdVtdVs_t2r[split]
+                    pipeline_dV.consumer_wait(consumer_state_dV)
+                    # TODO: record meaning of hard-coded values
+                    num_cols_per_store = self.tile_dV[1] * 2
+                    num_epi_subtiles = (self.hdimv // self.num_hdimv_splits) // num_cols_per_store
+                    assert num_cols_per_store == 64
+                    assert num_epi_subtiles == 4
+                    assert cute.size(tdVtdV_t2r.shape[2]) == num_epi_subtiles
+                    tdVrdV_cur_shape = tdVcdVmma_t2r[None, None, 0].shape
+                    tRS_rdV_cur_shape = tRS_sdV[None, None, None, 0].shape
+                    assert cute.size(tdVrdV_cur_shape) == cute.size(tRS_rdV_cur_shape)
+                    tdVrdV_out_shape = tSR_sdV[None, None, None, 0].shape + (2,)
+                    for i in cutlass.range_constexpr(num_epi_subtiles):
+                        tdVrdV_cur = cute.make_rmem_tensor(tdVrdV_cur_shape, self.dtype_acc)
+                        cute.copy(tiled_copy_t2r, tdVtdV_t2r[None, None, i], tdVrdV_cur)
+                        tRS_rdV_cur = cute.make_tensor(tdVrdV_cur.iterator, tRS_rdV_cur_shape)
+                        stage = 4 * split + 2 * wg_half + (i % 2)
+                        cute.copy(tiled_copy_r2s, tRS_rdV_cur, tRS_sdV[None, None, None, stage])
+                        cute.arch.fence_view_async_shared()
+                        self.epi_barrier.arrive_and_wait()
+                        tSR_rdV = cute.make_rmem_tensor(tdVrdV_out_shape, dtype=self.dtype_acc)
+                        for w in cutlass.range_constexpr(2):
+                            stage_out = 4 * split + 2 * w + (i % 2)
+                            cute.copy(
+                                tiled_copy_s2r,
+                                tSR_sdV[None, None, None, stage_out],
+                                tSR_rdV[None, None, None, w],
+                            )
+                        for j in cutlass.range_constexpr(gmem_rows_per_thread):
+                            gmem_n_idx = rIdxTopK[j]
+                            # Skip -1 sentinel slots (invalid top-k entries)
+                            if gmem_n_idx >= 0:
+                                for w in cutlass.range_constexpr(2):
+                                    dv_offset = (
+                                        self.hdimv // self.num_hdimv_splits * split  # 256 * split
+                                        + (self.hdimv // self.num_hdimv_splits // 2) * w  # 128 * w
+                                        + 32 * i
+                                    )
+                                    dv_offset += tdVcdV[0, j, 0][1]
+                                    gmem_coord = (gmem_n_idx, dv_offset)
+                                    dV_gmem_ptr = elem_pointer(mdV_cur, gmem_coord)
+                                    a = tSR_rdV[0, j, 0, w]
+                                    b = tSR_rdV[1, j, 0, w]
+                                    c = tSR_rdV[2, j, 0, w]
+                                    d = tSR_rdV[3, j, 0, w]
+                                    atomic_add_fp32x4(a, b, c, d, dV_gmem_ptr)
+                    cute.arch.fence_view_async_tmem_load()
+                    self.epi_barrier.arrive_and_wait()
+                    pipeline_dV.consumer_release(consumer_state_dV)
+                    if leader_warp:
+                        with cute.arch.elect_one():
+                            pipeline_dV_epi.consumer_release(consumer_state_dV)
+                    consumer_state_dV.advance()
+            # Advance to next tile
+            work_tile = tile_scheduler.advance_to_next_work()

build/torch-cuda/flash_bwd_postprocess.py CHANGED Viewed

@@ -2,7 +2,8 @@
 # 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
@@ -31,39 +32,49 @@ from .tile_scheduler import (
 )
 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,
         use_2cta_instrs: bool = False,
         cluster_size: int = 1,  # for varlen offsets
     ):
         """
         :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 // 10 in [8, 9, 10, 11], (
-            "Only Ampere (8.x), Hopper (9.x), and Blackwell (10.x, 11.x) 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
-        self.use_2cta_instrs = use_2cta_instrs and arch == 100 and head_dim != 64
         self.cluster_size = cluster_size
     @staticmethod
@@ -89,7 +100,7 @@ class FlashAttentionBackwardPostprocess:
         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)
@@ -101,9 +112,9 @@ class FlashAttentionBackwardPostprocess:
                 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,
@@ -125,7 +136,7 @@ class FlashAttentionBackwardPostprocess:
                 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
@@ -148,22 +159,22 @@ class FlashAttentionBackwardPostprocess:
             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
@@ -188,14 +199,18 @@ class FlashAttentionBackwardPostprocess:
         # 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
@@ -211,7 +226,10 @@ class FlashAttentionBackwardPostprocess:
         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]):
@@ -219,6 +237,19 @@ class FlashAttentionBackwardPostprocess:
         if const_expr(mdQaccum is not None):
             if const_expr(mdQaccum.element_type not in [cutlass.Float32]):
                 raise TypeError("dQaccum tensor must be Float32")
         mdQaccum, mdQ = [assume_tensor_aligned(t) for t in (mdQaccum, mdQ)]
@@ -253,6 +284,7 @@ class FlashAttentionBackwardPostprocess:
             tile_shape_mn=(self.tile_m, 1),
             mCuSeqlensQ=mCuSeqlensQ,
             mSeqUsedQ=mSeqUsedQ,
         )
         tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
@@ -264,6 +296,7 @@ class FlashAttentionBackwardPostprocess:
             mdQ,
             mCuSeqlensQ,
             mSeqUsedQ,
             scale,
             self.tiled_mma,
             self.dQ_swapAB,
@@ -288,6 +321,7 @@ class FlashAttentionBackwardPostprocess:
         mdQ: cute.Tensor,
         mCuSeqlensQ: Optional[cute.Tensor],
         mSeqUsedQ: Optional[cute.Tensor],
         scale: cutlass.Float32,
         tiled_mma: cute.TiledMma,
         dQ_swapAB: cutlass.Constexpr,
@@ -305,7 +339,7 @@ class FlashAttentionBackwardPostprocess:
         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
@@ -323,6 +357,73 @@ class FlashAttentionBackwardPostprocess:
         m_block, head_idx, batch_idx, _ = work_tile.tile_idx
         if work_tile.is_valid_tile:
             # ///////////////////////////////////////////////////////////////////////////////
             # Get the appropriate tiles for this thread block.
@@ -343,10 +444,7 @@ class FlashAttentionBackwardPostprocess:
                 mdQaccum_cur = mdQaccum[batch_idx, head_idx, None]
                 head_dim = mdQ.shape[3]
             else:
-                if cutlass.const_expr(self.arch >= 90):
-                    padded_offset_q = seqlen.padded_offset_q
-                else:
-                    padded_offset_q = seqlen.offset_q + batch_idx * self.tile_m
                 mdQ_cur = cute.domain_offset((seqlen.offset_q, 0), mdQ[None, head_idx, None])
                 mdQaccum_cur = cute.domain_offset(
                     (padded_offset_q * self.tile_hdim,), mdQaccum[head_idx, None]
@@ -371,7 +469,7 @@ class FlashAttentionBackwardPostprocess:
             seqlen_q = seqlen.seqlen_q
             seqlen_q_rounded = cute.round_up(seqlen_q, self.tile_m)
-            if const_expr(self.arch == 100 and self.use_2cta_instrs):
                 # 2-CTA: remap dQaccum layout into TMEM view before writing sdQ
                 num_reduce_threads = self.num_threads
                 thr_mma_dsk = tiled_mma.get_slice(tidx)
@@ -394,7 +492,7 @@ class FlashAttentionBackwardPostprocess:
                 g2s_thr_copy = tiled_copy_accum.get_slice(tidx)
                 # S -> R
-                tdQrdQ_fp32 = cute.make_fragment(tdQrdQ.shape, cutlass.Float32)
                 tdQrdQ_s2r = cute.make_tensor(tdQrdQ_fp32.iterator, tdQrdQ_fp32.shape)
                 smem_copy_atom = sm100_utils_basic.get_smem_store_op(
@@ -406,7 +504,7 @@ class FlashAttentionBackwardPostprocess:
                     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_fp32, mode=[1])
                 stage_stride = self.dQ_reduce_ncol
@@ -502,11 +600,11 @@ class FlashAttentionBackwardPostprocess:
                 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
@@ -522,7 +620,7 @@ class FlashAttentionBackwardPostprocess:
                     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
@@ -531,7 +629,7 @@ class FlashAttentionBackwardPostprocess:
                 # 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
                     )
@@ -553,7 +651,7 @@ class FlashAttentionBackwardPostprocess:
                     )
                 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

 # 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
+import operator
+from typing import Callable, NamedTuple, Optional, Type
 import cuda.bindings.driver as cuda
 )
+class LearnableSinkBwdTensors(NamedTuple):
+    dpsum: cute.Tensor
+    lse: cute.Tensor
+    sink: cute.Tensor
+    dsink: cute.Tensor
+    def __new_from_mlir_values__(self, values):
+        return LearnableSinkBwdTensors(*values)
 class FlashAttentionBackwardPostprocess:
     def __init__(
         self,
         dtype: Type[cutlass.Numeric],
         head_dim: int,
+        arch: int,
         tile_m: int = 128,
         num_threads: int = 256,
         AtomLayoutMdQ: int = 1,
         dQ_swapAB: bool = False,
         use_2cta_instrs: bool = False,
         cluster_size: int = 1,  # for varlen offsets
+        hdim_multiple_of: int = 32,
     ):
         """
         :param head_dim: head dimension
         :type head_dim: int
         :param tile_m: m block size
         :type tile_m: int
+        :param hdim_multiple_of: accumulator alignment shared with the main backward kernel.
         """
         self.dtype = dtype
         self.tile_m = tile_m
+        assert arch // 10 in [8, 9, 10, 11, 12], (
+            "Only Ampere (8.x), Hopper (9.x), and Blackwell (10.x, 11.x, 12.x) are supported"
         )
         self.arch = arch
         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
+        self.use_2cta_instrs = use_2cta_instrs and arch // 10 in [10, 11] and head_dim != 64
         self.cluster_size = cluster_size
     @staticmethod
         return True
     def _get_tiled_mma(self):
+        if const_expr(self.arch // 10 in [8, 12]):
             num_mma_warps = self.num_threads // 32
             atom_layout_dQ = (
                 (self.AtomLayoutMdQ, num_mma_warps // self.AtomLayoutMdQ, 1)
                 atom_layout_dQ,
                 permutation_mnk=(atom_layout_dQ[0] * 16, atom_layout_dQ[1] * 16, 16),
             )
+        elif const_expr(self.arch // 10 == 9):
+            num_wg_mma = self.num_threads // 128
+            atom_layout_dQ = (self.AtomLayoutMdQ, num_wg_mma // 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,
                 cta_group,
                 (self.tile_m, self.tile_hdim),
             )
+        if const_expr(self.arch // 10 in [8, 9, 12]):
             assert self.num_threads == tiled_mma.size
         return tiled_mma
             cute.make_layout(self.num_threads),
             cute.make_layout(async_copy_elems_accum),
         )
+        num_s2r_copy_elems = 1 if const_expr(self.arch // 10 in [8, 12]) else 4
+        if const_expr(self.arch // 10 in [8, 12]):
             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 // 10 == 9):
             num_threads_per_warp_group = 128
+            num_wg_mma = 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_wg_mma)),  # thr_layout
                 cute.make_layout(128 // Float32.width),  # val_layout
             )
             self.sdQaccum_layout = cute.make_layout(
+                (self.tile_m * self.tile_hdim // num_wg_mma, num_wg_mma)
             )
         else:
             self.dQ_reduce_ncol = 32
         # 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 // 10 in [8, 12]):
             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 // 10 == 9):
+            wg_d_dQ = num_wg_mma // self.AtomLayoutMdQ
             self.sdQ_layout = sm90_utils.make_smem_layout(
+                self.dtype,
+                LayoutEnum.ROW_MAJOR,
+                (self.tile_m, self.tile_hdim),
+                major_mode_size=self.tile_hdim // wg_d_dQ,
             )
         else:
             # TODO: this is hard-coded for hdim 128
         scale: cutlass.Float32,
         mCuSeqlensQ: Optional[cute.Tensor],
         mSeqUsedQ: Optional[cute.Tensor],
+        sink_tensors: LearnableSinkBwdTensors | None,
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
+        stream: cuda.CUstream = None,
     ):
         # Get the data type and check if it is fp16 or bf16
         if const_expr(mdQ.element_type not in [cutlass.Float16, cutlass.BFloat16]):
         if const_expr(mdQaccum is not None):
             if const_expr(mdQaccum.element_type not in [cutlass.Float32]):
                 raise TypeError("dQaccum tensor must be Float32")
+        if const_expr(sink_tensors is not None):
+            mdPsum, mLSE, mLearnableSink, mdSink = sink_tensors
+            if const_expr(
+                mLearnableSink.element_type
+                not in [cutlass.Float16, cutlass.BFloat16, cutlass.Float32]
+            ):
+                raise TypeError("Learnable sink tensor must be Float16, BFloat16, or Float32")
+            if const_expr(mdPsum.element_type not in [cutlass.Float32]):
+                raise TypeError("dPsum must be Float32")
+            if const_expr(mLSE.element_type not in [cutlass.Float32]):
+                raise TypeError("LSE must be Float32")
+            if const_expr(mdSink.element_type != mLearnableSink.element_type):
+                raise TypeError("dSink must have the learnable sink dtype")
         mdQaccum, mdQ = [assume_tensor_aligned(t) for t in (mdQaccum, mdQ)]
             tile_shape_mn=(self.tile_m, 1),
             mCuSeqlensQ=mCuSeqlensQ,
             mSeqUsedQ=mSeqUsedQ,
+            cu_total_m_blocks_ptr=mCuTotalMBlocks,
         )
         tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
             mdQ,
             mCuSeqlensQ,
             mSeqUsedQ,
+            sink_tensors,
             scale,
             self.tiled_mma,
             self.dQ_swapAB,
         mdQ: cute.Tensor,
         mCuSeqlensQ: Optional[cute.Tensor],
         mSeqUsedQ: Optional[cute.Tensor],
+        sink_tensors: LearnableSinkBwdTensors | None,
         scale: cutlass.Float32,
         tiled_mma: cute.TiledMma,
         dQ_swapAB: cutlass.Constexpr,
         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 // 10 in [8, 9, 12]):
             sdQ = cute.make_tensor(cute.recast_ptr(sdQaccum.iterator, dtype=self.dtype), sdQ_layout)
         else:
             # extra stage dimension
         m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+        # Reuse one existing dQ postprocess CTA per head to reduce dSink from
+        # the per-row dPsum and LSE written by backward preprocess. This avoids
+        # both a global atomic accumulator and a separate zero-initialization.
+        if const_expr(sink_tensors is not None):
+            mdPsum, mLSE, mLearnableSink, mdSink = sink_tensors
+            block_x, block_y, block_z = cute.arch.block_idx()
+            sink_head_idx = head_idx if const_expr(mCuSeqlensQ is None) else block_x
+            # Varlen uses block_x to select one CTA per head. block_y and block_z
+            # are currently always zero, but check them defensively.
+            reduce_sink = (
+                m_block == 0 and batch_idx == 0
+                if const_expr(mCuSeqlensQ is None)
+                else block_x < mdSink.shape[0] and block_y == 0 and block_z == 0
+            )
+            if reduce_sink:
+                sink_sum = Float32(0.0)
+                num_batch = (
+                    mdQ.shape[0] if const_expr(mCuSeqlensQ is None) else mCuSeqlensQ.shape[0] - 1
+                )
+                sink_val = Float32(mLearnableSink[sink_head_idx])
+                sink_batch = 0
+                while sink_batch < num_batch:
+                    sink_seqlen = SeqlenInfoQK.create(
+                        sink_batch,
+                        mdQ.shape[1],
+                        0,
+                        mCuSeqlensQ=mCuSeqlensQ,
+                        mSeqUsedQ=mSeqUsedQ,
+                        tile_m=self.tile_m * self.cluster_size,
+                    )
+                    sink_row = tidx
+                    while sink_row < sink_seqlen.seqlen_q:
+                        if const_expr(mCuSeqlensQ is None):
+                            dpsum_val = mdPsum[sink_batch, sink_head_idx, sink_row]
+                            lse_val = mLSE[sink_batch, sink_head_idx, sink_row]
+                        else:
+                            dpsum_val = mdPsum[
+                                sink_head_idx, sink_seqlen.padded_offset_q + sink_row
+                            ]
+                            lse_val = mLSE[sink_head_idx, sink_seqlen.offset_q + sink_row]
+                        lse_val = Float32(lse_val)
+                        sink_prob = (
+                            Float32(1.0)
+                            if lse_val == -Float32.inf
+                            else cute.math.exp2(
+                                (sink_val - lse_val) * utils.LOG2_E,
+                                fastmath=True,
+                            )
+                        )
+                        sink_sum += -sink_prob * Float32(dpsum_val)
+                        sink_row += self.num_threads
+                    sink_batch += 1
+                sink_sum = utils.warp_reduce(sink_sum, operator.add)
+                lane_idx = cute.arch.lane_idx()
+                warp_idx = tidx // cute.arch.WARP_SIZE
+                num_warps = self.num_threads // cute.arch.WARP_SIZE
+                if lane_idx == 0:
+                    sdQaccum_flat[warp_idx] = sink_sum
+                cute.arch.sync_threads()
+                if warp_idx == 0:
+                    sink_sum = sdQaccum_flat[lane_idx] if lane_idx < num_warps else Float32(0.0)
+                    sink_sum = utils.warp_reduce(sink_sum, operator.add)
+                    if lane_idx == 0:
+                        mdSink[sink_head_idx] = sink_sum.to(mdSink.element_type)
+                cute.arch.sync_threads()
         if work_tile.is_valid_tile:
             # ///////////////////////////////////////////////////////////////////////////////
             # Get the appropriate tiles for this thread block.
                 mdQaccum_cur = mdQaccum[batch_idx, head_idx, None]
                 head_dim = mdQ.shape[3]
             else:
+                padded_offset_q = seqlen.padded_offset_q
                 mdQ_cur = cute.domain_offset((seqlen.offset_q, 0), mdQ[None, head_idx, None])
                 mdQaccum_cur = cute.domain_offset(
                     (padded_offset_q * self.tile_hdim,), mdQaccum[head_idx, None]
             seqlen_q = seqlen.seqlen_q
             seqlen_q_rounded = cute.round_up(seqlen_q, self.tile_m)
+            if const_expr(self.arch // 10 in [10, 11] and self.use_2cta_instrs):
                 # 2-CTA: remap dQaccum layout into TMEM view before writing sdQ
                 num_reduce_threads = self.num_threads
                 thr_mma_dsk = tiled_mma.get_slice(tidx)
                 g2s_thr_copy = tiled_copy_accum.get_slice(tidx)
                 # S -> R
+                tdQrdQ_fp32 = cute.make_rmem_tensor(tdQrdQ.shape, cutlass.Float32)
                 tdQrdQ_s2r = cute.make_tensor(tdQrdQ_fp32.iterator, tdQrdQ_fp32.shape)
                 smem_copy_atom = sm100_utils_basic.get_smem_store_op(
                     tiler_mn=tiled_tmem_ld.tiler_mn,
                 )
                 tdQsdQ_r2s = thr_tmem_ld.partition_D(thr_mma_dsk.partition_C(sdQ))
+                tdQrdQ_r2s = cute.make_rmem_tensor(tdQsdQ_r2s.shape, self.dtype)
                 num_stages = cute.size(tdQrdQ_fp32, mode=[1])
                 stage_stride = self.dQ_reduce_ncol
                 tile_shape = (self.tile_m, self.tile_hdim)
                 acc = None
                 tiled_copy_t2r = None
+                if const_expr(self.arch // 10 in [8, 9, 12]):
                     acc_shape = tiled_mma.partition_shape_C(
                         tile_shape if const_expr(not dQ_swapAB) else tile_shape[::-1]
                     )
+                    acc = cute.make_rmem_tensor(acc_shape, cutlass.Float32)
                     assert cute.size(acc) == cute.size(tdQsdQaccum)
                 else:
                     thr_mma = tiled_mma.get_slice(0)  # 1-CTA
                     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_rmem_tensor(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
                 # Step 3: Copy dQ from register to smem
                 cute.arch.barrier()  # make sure all threads have finished loading dQaccum
+                if const_expr(self.arch // 10 in [8, 9, 12]):
                     copy_atom_r2s_dQ = utils.get_smem_store_atom(
                         self.arch, self.dtype, transpose=self.dQ_swapAB
                     )
                     )
                 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 // 10 in [8, 9, 12]):
                     taccdQrdQ = thr_copy_r2s_dQ.retile(rdQ)
                 else:
                     taccdQcdQ_shape = thr_copy_r2s_dQ.partition_S(cdQ).shape

build/torch-cuda/flash_bwd_preprocess.py CHANGED Viewed

@@ -1,27 +1,38 @@
 # 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, Literal
 import cuda.bindings.driver as cuda
 import cutlass
 import cutlass.cute as cute
-from cutlass import Float32
-from .quack import copy_utils
 from . import utils
-from .cute_dsl_utils import assume_tensor_aligned
-from .seqlen_info import SeqlenInfoQK
 from .quack.cute_dsl_utils import ParamsBase
 from .tile_scheduler import (
     SingleTileScheduler,
     SingleTileVarlenScheduler,
     TileSchedulerArguments,
 )
 class FlashAttentionBackwardPreprocess:
@@ -30,9 +41,14 @@ class FlashAttentionBackwardPreprocess:
         dtype: Type[cutlass.Numeric],
         head_dim: int,
         head_dim_v: int,
-        arch: Literal[80, 90, 100],
-        m_block_size: int = 128,
-        num_threads: int = 128,
     ):
         """
         All contiguous dimensions must be at least 16 bytes aligned which indicates the head dimension
@@ -40,31 +56,35 @@ class FlashAttentionBackwardPreprocess:
         :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
-        self.arch = arch
-        # 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.head_dim_v_padded = int(math.ceil(head_dim_v / hdim_multiple_of) * hdim_multiple_of)
         self.check_hdim_v_oob = head_dim_v != self.head_dim_v_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
@@ -77,7 +97,7 @@ class FlashAttentionBackwardPreprocess:
             return False
         if num_threads % 32 != 0:
             return False
-        if num_threads < m_block_size:  # For multiplying lse with log2
             return False
         return True
@@ -105,7 +125,7 @@ class FlashAttentionBackwardPreprocess:
         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
@@ -114,72 +134,147 @@ class FlashAttentionBackwardPreprocess:
     @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")
-        mO, mdO, mdQaccum = [assume_tensor_aligned(t) 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,
@@ -188,6 +283,7 @@ class FlashAttentionBackwardPreprocess:
             grid=grid_dim,
             block=[self.num_threads, 1, 1],
             stream=stream,
         )
     @cute.kernel
@@ -195,12 +291,16 @@ class FlashAttentionBackwardPreprocess:
         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,
@@ -213,149 +313,159 @@ class FlashAttentionBackwardPreprocess:
         work_tile = tile_scheduler.initial_work_tile_info()
         m_block, head_idx, batch_idx, _ = work_tile.tile_idx
         if work_tile.is_valid_tile:
             # ///////////////////////////////////////////////////////////////////////////////
             # Get the appropriate tiles for this thread block.
             # ///////////////////////////////////////////////////////////////////////////////
-            seqlen = SeqlenInfoQK.create(
-                batch_idx,
-                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_idx, None, head_idx, None]
-                mdO_cur = mdO[batch_idx, None, head_idx, None]
-                mdPsum_cur = mdPsum[batch_idx, head_idx, None]
-                headdim_v = mO.shape[3]
-            else:
-                mO_cur = cute.domain_offset((seqlen.offset_q, 0), mO[None, head_idx, None])
-                mdO_cur = cute.domain_offset((seqlen.offset_q, 0), mdO[None, head_idx, None])
-                padded_offset_q = seqlen.offset_q + batch_idx * self.m_block_size
-                if cutlass.const_expr(self.arch >= 90):
-                    padded_offset_q = padded_offset_q // self.m_block_size * self.m_block_size
-                mdPsum_cur = cute.domain_offset((padded_offset_q,), mdPsum[head_idx, None])
-                headdim_v = mO.shape[2]
-            blkOdO_shape = (self.m_block_size, self.head_dim_v_padded)
-            # (m_block_size, head_dim_v)
-            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_v_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_idx, head_idx, None]
-                else:
-                    mLSE_cur = cute.domain_offset((seqlen.offset_q,), mLSE[head_idx, 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_v_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_v_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_idx, head_idx, None]
-                else:
-                    mdQaccum_cur = cute.domain_offset(
-                        (padded_offset_q * self.head_dim_padded,), mdQaccum[head_idx, 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_idx, head_idx, None]
-                else:
-                    mLSElog2_cur = cute.domain_offset((padded_offset_q,), mLSElog2[head_idx, 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

 # 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.
+#
+# Computes D_i = (dO_i * O_i).sum(dim=-1), optionally adjusted for LSE gradient:
+#   D'_i = D_i - dLSE_i
+# This works because in the backward pass:
+#   dS_ij = P_ij * (dP_ij - D_i)                     [standard]
+# When LSE is differentiable, d(loss)/d(S_ij) gets an extra term dLSE_i * P_ij
+# (since d(LSE_i)/d(S_ij) = P_ij), giving:
+#   dS_ij = P_ij * (dP_ij - D_i) + dLSE_i * P_ij
+#         = P_ij * (dP_ij - (D_i - dLSE_i))
+# So the main backward kernel is unchanged; we just replace D with D' = D - dLSE here.
 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, const_expr
+from cutlass.cutlass_dsl import Arch, BaseDSL
+from .quack import copy_utils, layout_utils
 from . import utils
+from .seqlen_info import SeqlenInfo
 from .quack.cute_dsl_utils import ParamsBase
 from .tile_scheduler import (
     SingleTileScheduler,
     SingleTileVarlenScheduler,
     TileSchedulerArguments,
 )
+from .pack_gqa import pack_gqa_layout
 class FlashAttentionBackwardPreprocess:
         dtype: Type[cutlass.Numeric],
         head_dim: int,
         head_dim_v: int,
+        tile_m: int = 128,
+        num_threads: int = 256,
+        use_padded_offsets: bool = True,
+        nheads_major: bool = False,
+        pack_gqa: bool = False,
+        qhead_per_kvhead: int = 1,
+        nheads_kv: int = 1,
+        hdim_multiple_of: int = 32,
     ):
         """
         All contiguous dimensions must be at least 16 bytes aligned which indicates the head dimension
         :param head_dim: head dimension
         :type head_dim: int
+        :param tile_m: m block size
+        :type tile_m: int
         :param num_threads: number of threads
         :type num_threads: int
+        :param hdim_multiple_of: accumulator alignment shared with the main backward kernel.
         """
+        self.use_pdl = BaseDSL._get_dsl().get_arch_enum() >= Arch.sm_90a
         self.dtype = dtype
+        self.tile_m = tile_m
         self.head_dim_padded = int(math.ceil(head_dim / hdim_multiple_of) * hdim_multiple_of)
         self.head_dim_v_padded = int(math.ceil(head_dim_v / hdim_multiple_of) * hdim_multiple_of)
         self.check_hdim_v_oob = head_dim_v != self.head_dim_v_padded
         self.num_threads = num_threads
+        self.use_padded_offsets = use_padded_offsets
+        self.nheads_major = nheads_major
+        self.pack_gqa = pack_gqa
+        self.qhead_per_kvhead = qhead_per_kvhead
+        self.nheads_kv = nheads_kv
     @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
         :param num_threads: number of threads
         :type num_threads: int
             return False
         if num_threads % 32 != 0:
             return False
+        if num_threads < tile_m:  # For multiplying lse with log2
             return False
         return True
         universal_copy_bits = 128
         num_copy_elems_dQaccum = universal_copy_bits // Float32.width
         assert (
+            self.tile_m * 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,  # (batch, seqlen, nheads, head_dim_v) or (total_q, nheads, head_dim_v)
+        mdO: cute.Tensor,  # same shape as mO
+        mPdPsum: cute.Tensor,  # (batch, nheads, seqlen_padded) or (nheads, total_q_padded)
+        mLSE: Optional[cute.Tensor],  # (batch, nheads, seqlen) or (nheads, total_q)
+        mLSElog2: Optional[cute.Tensor],  # same shape as mPdPsum
+        # (batch, nheads, seqlen_padded * head_dim_v) or (nheads, total_q_padded * head_dim_v)
         mdQaccum: Optional[cute.Tensor],
+        mCuSeqlensQ: Optional[cute.Tensor],  # (batch + 1,)
+        mSeqUsedQ: Optional[cute.Tensor],  # (batch,)
+        mdLSE: Optional[cute.Tensor],  # (batch, nheads, seqlen) or (nheads, total_q)
+        mRowMax: Optional[cute.Tensor],  # (b, s, n, h) or (t, n, h)
+        mScaleP: Optional[cute.Tensor],  # == mRowMax
+        softmax_scale: Float32,
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
+        stream: cuda.CUstream = None,
     ):
         # Get the data type and check if it is fp16 or bf16
+        if const_expr(not (mO.element_type == mdO.element_type)):
             raise TypeError("All tensors must have the same data type")
+        if const_expr(mO.element_type not in [cutlass.Float16, cutlass.BFloat16]):
             raise TypeError("Only Float16 or BFloat16 is supported")
+        if const_expr(mPdPsum.element_type not in [Float32]):
+            raise TypeError("PdPsum tensor must be Float32")
+        if const_expr(mdQaccum is not None):
+            assert self.nheads_major is False
+            assert self.pack_gqa is False
+            assert self.use_padded_offsets is True
+            if const_expr(mdQaccum.element_type not in [Float32]):
                 raise TypeError("dQaccum tensor must be Float32")
+        if const_expr(mLSE is not None):
+            if const_expr(mLSE.element_type not in [Float32]):
                 raise TypeError("LSE tensor must be Float32")
+        if const_expr(mLSElog2 is not None):
+            if const_expr(mLSElog2.element_type not in [Float32]):
                 raise TypeError("LSElog2 tensor must be Float32")
+        if const_expr(mdLSE is not None):
+            if const_expr(mdLSE.element_type not in [Float32]):
+                raise TypeError("dLSE tensor must be Float32")
+        if const_expr(mScaleP is not None):
+            assert self.nheads_major is True
+            assert self.pack_gqa is True
+            assert mRowMax is not None
+            if const_expr(mScaleP.element_type not in [Float32]):
+                raise TypeError("ScaleP tensor must be Float32")
+            if const_expr(mRowMax.element_type not in [Float32]):
+                raise TypeError("RowMax tensor must be Float32")
         self._setup_attributes()
+        # (b, s, h, d)  -> (s, d, h, b)  or
+        # (total, h, d) -> (total, d, h)
+        QO_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1]
+        mO, mdO = [
+            cute.make_tensor(mX.iterator, cute.select(mX.layout, mode=QO_layout_transpose))
+            for mX in (mO, mdO)
+        ]
+        if const_expr(not self.nheads_major):
+            # (batch, nheads, seqlen) -> (seqlen, nheads, batch) or
+            # (nheads, total_q) -> (total_q, nheads)
+            transpose = [2, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 0]
+        else:
+            # (batch, seqlen, nheads) -> (seqlen, nheads, batch) or
+            # (total_q, nheads) -> (total_q, nheads)
+            transpose = [1, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 1]
+        mPdPsum, mLSE, mLSElog2, mdLSE, mdQaccum = [
+            layout_utils.select(mX, transpose) if mX is not None else None
+            for mX in (mPdPsum, mLSE, mLSElog2, mdLSE, mdQaccum)
+        ]
+        # (b, s, n, h) => (s, n, h, b) or
+        # (total, n, h) == (total, n, h)
+        rowmax_layout_transpose = [1, 2, 3, 0] if const_expr(mCuSeqlensQ is None) else [0, 1, 2]
+        if const_expr(mRowMax is not None):
+            mRowMax = layout_utils.select(mRowMax, rowmax_layout_transpose)
+        if const_expr(mScaleP is not None):
+            mScaleP = layout_utils.select(mScaleP, rowmax_layout_transpose)
+        # pack gqa
+        if const_expr(self.pack_gqa):
+            mO, mdO, mRowMax, mScaleP = [
+                pack_gqa_layout(mX, self.qhead_per_kvhead, self.nheads_kv, head_idx=2)
+                if mX is not None
+                else None
+                for mX in (mO, mdO, mRowMax, mScaleP)
+            ]
+            mPdPsum, mLSE, mLSElog2, mdLSE = [
+                pack_gqa_layout(mX, self.qhead_per_kvhead, self.nheads_kv, head_idx=1)
+                if mX is not None
+                else None
+                for mX in (mPdPsum, mLSE, mLSElog2, mdLSE)
+            ]
+        # mO: (s, d, h, b) or (total, d, h)
+        if const_expr(mCuSeqlensQ is not None):
             TileScheduler = SingleTileVarlenScheduler
+            num_head = mO.shape[2]
             num_batch = mCuSeqlensQ.shape[0] - 1
         else:
             TileScheduler = SingleTileScheduler
             num_head = mO.shape[2]
+            num_batch = mO.shape[3]
         tile_sched_args = TileSchedulerArguments(
+            num_block=cute.ceil_div(mO.shape[0], self.tile_m),
             num_head=num_head,
             num_batch=num_batch,
             num_splits=1,
             seqlen_k=0,
             headdim=0,
+            headdim_v=mO.shape[1],
+            total_q=cute.size(mO.shape[0])
+            if const_expr(mCuSeqlensQ is not None)
+            else cute.size(mO.shape[0]) * cute.size(mO.shape[3]),
+            tile_shape_mn=(self.tile_m, 1),
             mCuSeqlensQ=mCuSeqlensQ,
             mSeqUsedQ=mSeqUsedQ,
+            qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
+            cu_total_m_blocks_ptr=mCuTotalMBlocks,
         )
         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(
             mO,
             mdO,
+            mPdPsum,
             mLSE,
             mLSElog2,
             mdQaccum,
             mCuSeqlensQ,
             mSeqUsedQ,
+            mdLSE,
+            mRowMax,
+            mScaleP,
+            softmax_scale_log2,
             self.gmem_tiled_copy_O,
             self.gmem_tiled_copy_dQaccum,
             tile_sched_params,
             grid=grid_dim,
             block=[self.num_threads, 1, 1],
             stream=stream,
+            use_pdl=self.use_pdl,
         )
     @cute.kernel
         self,
         mO: cute.Tensor,
         mdO: cute.Tensor,
+        mPdPsum: cute.Tensor,
         mLSE: Optional[cute.Tensor],
         mLSElog2: Optional[cute.Tensor],
         mdQaccum: Optional[cute.Tensor],
         mCuSeqlensQ: Optional[cute.Tensor],
         mSeqUsedQ: Optional[cute.Tensor],
+        mdLSE: Optional[cute.Tensor],
+        mRowMax: Optional[cute.Tensor],
+        mScaleP: Optional[cute.Tensor],
+        softmax_scale_log2: Float32,
         gmem_tiled_copy_O: cute.TiledCopy,
         gmem_tiled_copy_dQaccum: cute.TiledCopy,
         tile_sched_params: ParamsBase,
         work_tile = tile_scheduler.initial_work_tile_info()
         m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+        # This kernel is launched with use_pdl=True, so the GPU may start executing it in
+        # "prologue" mode while the previous stream kernel is still running. We must wait
+        # before touching any upstream GMEM outputs (mO, mdO, mLSE); otherwise we risk
+        # reading a partially-written dout, which silently corrupts dpsum = sum(O * dO) and
+        # propagates to dQ/dK via dS = P * (dP - dpsum).
+        if const_expr(self.use_pdl):
+            cute.arch.griddepcontrol_wait()
         if work_tile.is_valid_tile:
             # ///////////////////////////////////////////////////////////////////////////////
             # Get the appropriate tiles for this thread block.
             # ///////////////////////////////////////////////////////////////////////////////
+            seqlen_static = mO.shape[0] if const_expr(not self.pack_gqa) else mO.shape[0][1]
+            seqlen = SeqlenInfo.create(
+                batch_idx, seqlen_static, mCuSeqlensQ, mSeqUsedQ, tile=self.tile_m
+            )
+            # (seqlen, dv)
+            mO_cur, mdO_cur = [
+                seqlen.offset_batch(mX, batch_idx, dim=3)[None, None, head_idx] for mX in (mO, mdO)
+            ]
+            mPdPsum_cur = seqlen.offset_batch(
+                mPdPsum, batch_idx, dim=2, padded=self.use_padded_offsets
+            )[None, head_idx]
+            headdim_v = mO_cur.shape[1]
+            seqlen_q = (
+                seqlen.seqlen
+                if const_expr(not self.pack_gqa)
+                else seqlen.seqlen * self.qhead_per_kvhead
             )
+            seqlen_q_rounded = cute.round_up(seqlen_q, self.tile_m)
+            seqlen_limit = seqlen_q - m_block * self.tile_m
+            lse = None
+            if const_expr(mLSE is not None):
+                mLSE_cur = seqlen.offset_batch(mLSE, batch_idx, dim=2)[None, head_idx]
+                gLSE = cute.local_tile(mLSE_cur, (self.tile_m,), (m_block,))
+                lse = Float32.inf
+                if tidx < seqlen_limit:
+                    lse = gLSE[tidx]
+            blk_shape = (self.tile_m, self.head_dim_v_padded)
+            gO = cute.local_tile(mO_cur, blk_shape, (m_block, 0))
+            gdO = cute.local_tile(mdO_cur, blk_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)
+            cO = cute.make_identity_tensor(blk_shape)
             tOcO = gmem_thr_copy_O.partition_S(cO)
             t0OcO = gmem_thr_copy_O.get_slice(0).partition_S(cO)
+            tOpO = None
+            if const_expr(self.check_hdim_v_oob):
+                tOpO = copy_utils.predicate_k(tOcO, limit=headdim_v)
+            # Each copy uses the predicate sliced to the current m (see loop below).
+            copy = copy_utils.copy
+            tOrO = cute.make_rmem_tensor_like(tOgO)
+            tOrdO = cute.make_rmem_tensor_like(tOgdO)
+            if const_expr(self.check_hdim_v_oob):
+                tOrO.fill(0.0)
+                tOrdO.fill(0.0)
+            assert tOgO.shape == tOgdO.shape
             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.
+                # This is bc the entries of t0OcO are known at compile time.
+                if t0OcO[0, m, 0][0] < seqlen_limit - tOcO[0][0]:
+                    # The predicate carries a CPY_M mode (broadcast over m); slice it to the
+                    # current m so its shape matches the m-sliced copy source. Otherwise the
+                    # vectorized copy atom's predicate-shape verification fails when head_dim_v
+                    # is not a multiple of the copy-atom width (e.g. 72, 104).
+                    tOpO_cur = tOpO[None, m, None] if const_expr(tOpO is not None) else None
+                    copy(tOgO[None, m, None], tOrO[None, m, None], pred=tOpO_cur)
+                    copy(tOgdO[None, m, None], tOrdO[None, m, None], pred=tOpO_cur)
+            # O and dO loads are done; signal that the next kernel can start.
+            # Correctness is ensured by griddepcontrol_wait() in bwd_sm90 before it reads our outputs.
+            if const_expr(self.use_pdl):
+                cute.arch.griddepcontrol_launch_dependents()
             # Sum across the "k" dimension
+            pdpsum = (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
+            pdpsum = utils.warp_reduce(pdpsum, operator.add, width=threads_per_row)
+            PdP_sum = cute.make_rmem_tensor(cute.size(tOrO, mode=[1]), Float32)
+            PdP_sum.store(pdpsum)
+            # If dLSE is provided, compute D' = D - dLSE (see module docstring for derivation).
+            gdLSE = None
+            if const_expr(mdLSE is not None):
+                mdLSE_cur = seqlen.offset_batch(mdLSE, batch_idx, dim=2)[None, head_idx]
+                gdLSE = cute.local_tile(mdLSE_cur, (self.tile_m,), (m_block,))
+            # Write PdPsum from rmem -> gmem
+            gPdPsum = cute.local_tile(mPdPsum_cur, (self.tile_m,), (m_block,))
+            # Only the thread corresponding to column 0 writes out the PdPsum to gmem
             if tOcO[0, 0, 0][1] == 0:
+                for m in cutlass.range(cute.size(PdP_sum), unroll_full=True):
                     row = tOcO[0, m, 0][0]
+                    PdPsum_val = 0.0
+                    if row < seqlen_limit:
+                        PdPsum_val = PdP_sum[m]
+                        if const_expr(mdLSE is not None):
+                            PdPsum_val -= gdLSE[row]
+                    gPdPsum[row] = PdPsum_val
             # Clear dQaccum
+            if const_expr(mdQaccum is not None):
+                mdQaccum_cur = seqlen.offset_batch(
+                    mdQaccum,
+                    batch_idx,
+                    dim=2,
+                    padded=self.use_padded_offsets,
+                    multiple=self.head_dim_padded,
+                )[None, head_idx]
+                blkdQaccum_shape = (self.tile_m * 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_rmem_tensor_like(tdQgdQaccum)
                 zero.fill(0.0)
                 cute.copy(gmem_tiled_copy_dQaccum, zero, tdQgdQaccum)
+            LOG2_E = math.log2(math.e)
+            lse_log2 = lse * LOG2_E if lse != -Float32.inf else 0.0
+            if const_expr(mLSElog2 is not None):
+                mLSElog2_cur = seqlen.offset_batch(
+                    mLSElog2, batch_idx, dim=2, padded=self.use_padded_offsets
+                )[None, head_idx]
+                gLSElog2 = cute.local_tile(mLSElog2_cur, (self.tile_m,), (m_block,))
                 LOG2_E = math.log2(math.e)
+                if tidx < seqlen_q_rounded - m_block * self.tile_m:
+                    gLSElog2[tidx] = lse_log2
+            if const_expr(mRowMax is not None):
+                assert mLSE is not None
+                # (s, n)
+                mRowMax_cur, mScaleP_cur = [
+                    seqlen.offset_batch(mX, batch_idx, dim=3)[None, None, head_idx]
+                    for mX in (mRowMax, mScaleP)
+                ]
+                # (tile_m, n)
+                gRowMax, gScaleP = [
+                    cute.local_tile(mX, (self.tile_m,), (m_block, None))
+                    for mX in (mRowMax_cur, mScaleP_cur)
+                ]
+                assert self.tile_m <= self.num_threads
+                if const_expr(self.tile_m == self.num_threads) or tidx < self.tile_m:
+                    for n in cutlass.range(gRowMax.shape[1], unroll=4):
+                        row_max = gRowMax[tidx, n]
+                        scale = 0.0
+                        if row_max != -Float32.inf and lse != -Float32.inf:
+                            scale = softmax_scale_log2 * row_max - lse_log2
+                            scale = cute.math.exp2(scale, fastmath=True)
+                        gScaleP[tidx, n] = scale

build/torch-cuda/flash_bwd_sm100.py CHANGED Viewed

@@ -14,7 +14,6 @@ from cutlass.cute.nvgpu import cpasync, tcgen05
 import cutlass.utils.blackwell_helpers as sm100_utils_basic
 from cutlass.pipeline import PipelineAsync
-from .quack import activation
 from .quack import layout_utils
 from . import utils
 from .cute_dsl_utils import assume_tensor_aligned
@@ -36,14 +35,35 @@ from . import barrier
 from .named_barrier import NamedBarrierBwdSm100
 from .softmax import apply_score_mod_inner, apply_score_mod_bwd_inner
 from .block_sparsity import BlockSparseTensors
 from .block_sparse_utils import (
     get_total_q_block_count_bwd,
     get_block_sparse_iteration_info_bwd,
     get_m_block_from_iter_bwd,
-    produce_block_sparse_q_loads_bwd_sm100,
 )
 class FlashAttentionBackwardSm100:
     arch = 100
@@ -58,22 +78,28 @@ class FlashAttentionBackwardSm100:
         tile_n: int = 128,
         is_persistent: bool = False,
         deterministic: bool = False,
         cluster_size: int = 1,
         use_2cta_instrs: bool = False,
         score_mod: cutlass.Constexpr | None = None,
         score_mod_bwd: cutlass.Constexpr | None = None,
         mask_mod: cutlass.Constexpr | None = None,
         has_aux_tensors: cutlass.Constexpr = False,
-        subtile_factor: cutlass.Constexpr[int] = 1,
     ):
-        # 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)
         self.check_hdim_oob = head_dim != self.tile_hdim
         self.check_hdim_v_oob = head_dim_v != self.tile_hdimv
         self.tile_m = tile_m
         self.tile_n = tile_n
@@ -81,14 +107,7 @@ class FlashAttentionBackwardSm100:
         assert self.tile_hdim <= 128 or (self.tile_hdim == 192 and self.tile_hdimv == 128)
         assert self.tile_hdimv <= 128
-        self.use_2cta_instrs = bool(
-            use_2cta_instrs
-            and cluster_size == 2
-            and not is_local
-            and score_mod is None
-            and score_mod_bwd is None
-            and mask_mod is None
-        )
         self.cta_group_size = 2 if self.use_2cta_instrs else 1
         assert self.tile_hdim != 192 or self.use_2cta_instrs, "Must use 2CTA for hdim 192"
@@ -111,19 +130,25 @@ class FlashAttentionBackwardSm100:
         assert cluster_size in (1, 2), "Only cluster_size=1 or 2 is supported"
         self.cluster_shape_mn = (cluster_size, 1)
         self.is_persistent = is_persistent
         self.is_causal = is_causal
         self.is_local = is_local
         self.qhead_per_kvhead = qhead_per_kvhead
         self.pack_gqa = False
         self.deterministic = deterministic
         # Score mod and mask mod support
         self.score_mod = score_mod
         self.score_mod_bwd = score_mod_bwd
         self.mask_mod = mask_mod
         self.has_aux_tensors = has_aux_tensors
-        self.subtile_factor = subtile_factor
         # For score_mod, use vec_size=1 (like forward) to handle per-element indices
         if cutlass.const_expr(has_aux_tensors):
             self.vec_size: cutlass.Constexpr = 1
@@ -137,6 +162,16 @@ class FlashAttentionBackwardSm100:
         # Generally slower to use store dS in smem for dK, and doesn't work for 2cta
         self.use_smem_dS_for_mma_dK = False
         self.reduce_warp_ids = (0, 1, 2, 3)
         self.compute_warp_ids = (4, 5, 6, 7, 8, 9, 10, 11)
         self.mma_warp_id = 12
@@ -191,7 +226,6 @@ class FlashAttentionBackwardSm100:
             self.tmem_dQ_offset = 512 - self.tile_hdim // 2
         else:
             self.tmem_S_offset = 0
-            self.tmem_P_offset = 0  # overlap with S
             self.tmem_dV_offset = self.tmem_S_offset + self.tile_n
             self.tmem_dP_offset = self.tmem_dV_offset + self.tile_hdimv
             self.tmem_dQ_offset = (
@@ -200,7 +234,14 @@ class FlashAttentionBackwardSm100:
                 else self.tmem_dP_offset
             )
             self.tmem_dK_offset = self.tmem_dP_offset + self.tile_m
-            self.tmem_dS_offset = self.tmem_dP_offset  # overlap with dP
         if (not is_causal and not is_local) or deterministic:
             self.num_regs_reduce = 136 if self.use_2cta_instrs else 152
@@ -238,6 +279,7 @@ class FlashAttentionBackwardSm100:
         self.Q_stage = 1 if self.use_2cta_instrs else 2
         self.dO_stage = 1
         self.single_stage = 1
         # LSE_stage = Q_stage and dPsum_stage = dO_stage
         self.sdKVaccum_stage = 2
         # number of tma reduce adds per dQacc mma
@@ -453,20 +495,21 @@ class FlashAttentionBackwardSm100:
         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,
         mdQ_semaphore: Optional[cute.Tensor] = None,
         mdK_semaphore: Optional[cute.Tensor] = None,
         mdV_semaphore: Optional[cute.Tensor] = None,
-        aux_tensors: Optional[list] = None,
         # Block-sparse tensors (Q direction - for iterating m_blocks per n_block):
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
     ):
         self.q_dtype = mQ.element_type
         self.k_dtype = mK.element_type
@@ -489,7 +532,9 @@ class FlashAttentionBackwardSm100:
             assert self.dk_dtype.width == 32, "Must accumulate dK in float precision for GQA"
             assert self.dv_dtype.width == 32, "Must accumulate dV in float precision for GQA"
-        mdQaccum, mdK, mdV = [assume_tensor_aligned(t) for t in (mdQaccum, mdK, mdV)]
         # (b, s, n, h) --> (s, h, n, b) or (t, n, h) -> (t, h, n)
         QO_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1]
@@ -706,7 +751,11 @@ class FlashAttentionBackwardSm100:
             TileScheduler = SingleTileLPTBwdScheduler
         else:
             TileScheduler = SingleTileScheduler
-        self.spt = (self.is_causal or self.is_local) and self.deterministic
         tile_sched_args = TileSchedulerArguments(
             cute.ceil_div(cute.size(mK.shape[0]), self.cta_tiler[0]),  # num_blocks
             cute.size(mQ.shape[2]),  # num_heads = num_query_heads
@@ -727,6 +776,7 @@ class FlashAttentionBackwardSm100:
             qhead_per_kvhead_packgqa=1,  # pack_gqa disabled for bwd
             element_size=self.k_dtype.width // 8,
             is_persistent=self.is_persistent,  # persistent mode not tested
             lpt=self.spt,
             head_swizzle=self.deterministic,
         )
@@ -782,7 +832,7 @@ class FlashAttentionBackwardSm100:
                     cutlass.Int64, self.dQaccum_reduce_stage // 2
                 ]
                 tmem_holding_buf: Int32
-                tmem_dealloc_mbar_ptr: cutlass.Int64
                 # 2-CTA
                 Qt_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.Q_stage]
@@ -852,6 +902,7 @@ class FlashAttentionBackwardSm100:
                 S_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.single_stage]
                 dP_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.single_stage]
                 dS_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.single_stage]
                 dKV_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.sdKVaccum_stage]
                 dQ_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2]
                 dQ_cluster_full_mbar_ptr: cute.struct.MemRange[
@@ -861,7 +912,7 @@ class FlashAttentionBackwardSm100:
                     cutlass.Int64, self.dQaccum_reduce_stage // 2
                 ]
                 tmem_holding_buf: Int32
-                tmem_dealloc_mbar_ptr: Int64
                 sQ: cute.struct.Align[
                     cute.struct.MemRange[cute.Uint8, sQ_alloc_bytes],
@@ -912,7 +963,7 @@ class FlashAttentionBackwardSm100:
             window_size_right = Int32(window_size_right)
         fastdiv_mods = None
-        if const_expr(aux_tensors is not None):
             seqlen_q = cute.size(mQ.shape[0]) // (
                 self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1
             )
@@ -921,21 +972,17 @@ class FlashAttentionBackwardSm100:
             seqlen_k_divmod = FastDivmodDivisor(seqlen_k)
             fastdiv_mods = (seqlen_q_divmod, seqlen_k_divmod)
         self.use_block_sparsity = cutlass.const_expr(blocksparse_tensors is not None)
-        if const_expr(self.use_2cta_instrs):
-            assert blocksparse_tensors is None, (
-                "2-CTA mode does not support block sparsity. "
-                "Please create kernel with use_2cta_instrs=False for block sparse attention."
-            )
-            assert window_size_left is None and window_size_right is None, (
-                "2-CTA mode does not support window attention. "
-                "Please create kernel with use_2cta_instrs=False for window attention."
             )
         # 2-CTA: 231424 and 1-CTA: 232448
         # print("SMEM: ", self.shared_storage.size_in_bytes())
-        if const_expr(self.use_block_sparsity or aux_tensors is not None):
             assert all(x is None for x in (mCuSeqlensQ, mCuSeqlensK, mSeqUsedQ, mSeqUsedK)), (
-                "Variable sequence length is not supported yet for blocksparse or aux tensors in bwd"
             )
         self.kernel(
@@ -997,7 +1044,7 @@ class FlashAttentionBackwardSm100:
             window_size_left,
             window_size_right,
             tile_sched_params,
-            aux_tensors,
             fastdiv_mods,
             blocksparse_tensors,
         ).launch(
@@ -1009,6 +1056,16 @@ class FlashAttentionBackwardSm100:
             min_blocks_per_mp=1,
         )
     @cute.kernel
     def kernel(
         self,
@@ -1070,7 +1127,7 @@ class FlashAttentionBackwardSm100:
         window_size_left: Optional[Int32],
         window_size_right: Optional[Int32],
         tile_sched_params: ParamsBase,
-        aux_tensors: Optional[list] = None,
         fastdiv_mods=(None, None),
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
     ):
@@ -1146,11 +1203,11 @@ class FlashAttentionBackwardSm100:
             * len((self.mma_warp_id, *self.compute_warp_ids, *self.reduce_warp_ids)),
         )
         tmem = cutlass.utils.TmemAllocator(
-            storage.tmem_holding_buf,
             barrier_for_retrieve=tmem_alloc_barrier,
             allocator_warp_id=self.mma_warp_id,
             is_two_cta=self.use_2cta_instrs,
-            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
         )
         # UMMA producers and AsyncThread consumers
@@ -1160,6 +1217,7 @@ class FlashAttentionBackwardSm100:
         pipeline_consumer_group_MMA_AsyncThread = cutlass.pipeline.CooperativeGroup(
             cutlass.pipeline.Agent.Thread, len(self.compute_warp_ids) * self.cta_group_size
         )
         pipeline_S_P = cutlass.pipeline.PipelineUmmaAsync.create(
             num_stages=1,
             producer_group=pipeline_producer_group_MMA_AsyncThread,
@@ -1209,6 +1267,16 @@ class FlashAttentionBackwardSm100:
             barrier_storage=storage.dS_mbar_ptr.data_ptr(),
             cta_layout_vmnk=cluster_layout_vmnk,
         )
         # TMA producer and UMMA consumers
         pipeline_producer_group = cutlass.pipeline.CooperativeGroup(
@@ -1412,14 +1480,7 @@ class FlashAttentionBackwardSm100:
         )
         TileSchedulerCls = partial(self.tile_scheduler_cls.create, tile_sched_params)
-        AttentionMaskCls = partial(
-            AttentionMask,
-            self.tile_m,
-            self.tile_n * self.cta_group_size,
-            swap_AB=True,
-            window_size_left=window_size_left,
-            window_size_right=window_size_right,
-        )
         #  EMPTY
         # (15)
         if warp_idx == self.empty_warp_id:
@@ -1440,6 +1501,7 @@ class FlashAttentionBackwardSm100:
                     block_info,
                     SeqlenInfoCls,
                     TileSchedulerCls,
                 )
         #  LOAD
@@ -1536,6 +1598,7 @@ class FlashAttentionBackwardSm100:
                 pipeline_dKV,
                 pipeline_dP,
                 pipeline_dQ,
                 block_info,
                 SeqlenInfoCls,
                 TileSchedulerCls,
@@ -1574,6 +1637,7 @@ class FlashAttentionBackwardSm100:
                 pipeline_dS,
                 pipeline_dKV,
                 pipeline_dP,
                 dS_cluster_empty_mbar_ptr,
                 dS_cluster_full_mbar_ptr,
                 dQaccum_empty_mbar_ptr,
@@ -1592,7 +1656,7 @@ class FlashAttentionBackwardSm100:
                 tiled_copy_r2s_dKV,
                 mdK_semaphore,
                 mdV_semaphore,
-                aux_tensors,
                 fastdiv_mods,
                 blocksparse_tensors,
             )
@@ -1631,6 +1695,7 @@ class FlashAttentionBackwardSm100:
         block_info: BlockInfo,
         SeqlenInfoCls: Callable,
         TileSchedulerCls: Callable,
     ):
         cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
         dS_cluster_phase = Int32(0)
@@ -1648,9 +1713,20 @@ class FlashAttentionBackwardSm100:
             process_tile = (
                 const_expr(not self.is_local and not self.is_varlen_q) or m_block_min < m_block_max
             )
             if process_tile:
-                num_iters = m_block_max - m_block_min
                 for _ in cutlass.range(num_iters, unroll=1):
                     # Wait for dS_xchg from peer CTA
                     cute.arch.mbarrier_wait(dS_cluster_full_mbar_ptr, phase=dS_cluster_phase)
@@ -1668,11 +1744,11 @@ class FlashAttentionBackwardSm100:
     @cute.jit
     def load(
         self,
-        thr_mma_S: cute.core.ThrMma,
-        thr_mma_dP: cute.core.ThrMma,
-        thr_mma_dV: cute.core.ThrMma,
-        thr_mma_dK: cute.core.ThrMma,
-        thr_mma_dQ: cute.core.ThrMma,
         mQ: cute.Tensor,
         mK: cute.Tensor,
         mKt: Optional[cute.Tensor],
@@ -1756,6 +1832,7 @@ class FlashAttentionBackwardSm100:
             )
             head_idx_kv = head_idx // self.qhead_per_kvhead
             n_block_cta_group = n_block // self.cta_group_size
             # GMEM tensors (varlen-aware)
             mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx]
@@ -1835,6 +1912,7 @@ class FlashAttentionBackwardSm100:
                 single_stage=True,
             )
             if const_expr(tma_atom_dOt is not None):
                 gdOt = cute.local_tile(
                     mdOt_cur, cute.select(self.mma_tiler_vdo, mode=[1, 2]), (None, 0)
@@ -1864,6 +1942,7 @@ class FlashAttentionBackwardSm100:
             load_dO = copy_utils.tma_producer_copy_fn(load_dO, pipeline_dO)
             # (4) dK += dS.T @ Q (2-CTA: needs separate Qt load)
             if const_expr(tma_atom_Qt is not None):
                 gQt = cute.local_tile(
                     mQt_cur, cute.select(self.mma_tiler_dsq, mode=[1, 2]), (0, None)
@@ -1880,6 +1959,7 @@ class FlashAttentionBackwardSm100:
                 load_Qt = copy_utils.tma_producer_copy_fn(load_Qt, pipeline_Qt)
             # (5) dQ = dS @ K
             if const_expr(self.use_2cta_instrs):
                 gKt = cute.local_tile(
                     mKt_cur, cute.select(self.mma_tiler_dsk, mode=[1, 2]), (0, n_block_cta_group)
@@ -1904,59 +1984,18 @@ class FlashAttentionBackwardSm100:
             # gdPsum = cute.logical_divide(gdPsum, (64,))[(None, block_in_cluster_coord_vmnk[1]), None]
             # copy_stats = partial(cute.copy, copy_atom_stats, mcast_mask=q_do_mcast_mask)
-            # some tiles might be empty due to block sparsity
-            if const_expr(self.use_block_sparsity):
-                total_m_block_cnt = get_total_q_block_count_bwd(
-                    blocksparse_tensors,
-                    batch_idx,
-                    head_idx,
-                    n_block,
-                    subtile_factor=self.subtile_factor,
-                    m_block_max=m_block_max,
-                )
-                process_tile = total_m_block_cnt > Int32(0)
-            else:
                 process_tile = (
                     const_expr(not self.is_local and not self.is_varlen_q)
                     or m_block_min < m_block_max
                 )
-            if process_tile:
-                if const_expr(self.use_block_sparsity):
-                    producer_state_Q_LSE, producer_state_dO_dPsum = (
-                        produce_block_sparse_q_loads_bwd_sm100(
-                            blocksparse_tensors,
-                            batch_idx,
-                            head_idx,
-                            n_block,
-                            producer_state_Q_LSE,
-                            producer_state_dO_dPsum,
-                            pipeline_Q,
-                            pipeline_LSE,
-                            pipeline_dO,
-                            pipeline_dPsum,
-                            load_K,
-                            load_V,
-                            load_Q,
-                            load_dO,
-                            copy_stats,
-                            gLSE,
-                            sLSE,
-                            gdPsum,
-                            sdPsum,
-                            self.tma_copy_bytes["K"],
-                            self.tma_copy_bytes["V"],
-                            should_load_Q=should_load_Q,
-                            should_load_dO=should_load_dO,
-                            subtile_factor=self.subtile_factor,
-                            m_block_max=m_block_max,
-                        )
-                    )
-                else:
                     first_m_block = m_block_min
                     if const_expr(self.use_2cta_instrs and self.tile_hdim == 192):
                         #### Prologue ####
                         assert should_load_Q and should_load_dO
                         # K & Q (for S)
                         pipeline_Q.producer_acquire(
                             producer_state_Q_Qt,
@@ -2060,6 +2099,10 @@ class FlashAttentionBackwardSm100:
                             pipeline_dO.producer_commit(producer_state_O_Ot)
                             producer_state_O_Ot.advance()
                     else:
                         #### Prologue ####
                         if const_expr(should_load_Q):
@@ -2089,7 +2132,7 @@ class FlashAttentionBackwardSm100:
                             pipeline_dO.producer_acquire(
                                 producer_state_dO_dPsum,
                                 extra_tx_count=self.tma_copy_bytes["V"] + self.tma_copy_bytes["dO"]
-                                if const_expr(tma_atom_dOt is not None)
                                 else self.tma_copy_bytes["V"],
                             )
                             load_V(
@@ -2098,7 +2141,7 @@ class FlashAttentionBackwardSm100:
                                 )
                             )
                             load_dO(first_m_block, producer_state=producer_state_dO_dPsum)
-                            if const_expr(tma_atom_dOt is not None):
                                 load_dOt(first_m_block, producer_state=producer_state_dO_dPsum)
                             pipeline_dO.producer_commit(producer_state_dO_dPsum)
@@ -2122,7 +2165,7 @@ class FlashAttentionBackwardSm100:
                         #### Main Loop ####
                         for m_block in cutlass.range(m_block_min + 1, m_block_max, unroll=1):
                             if const_expr(should_load_Q):
-                                if const_expr(tma_atom_Qt is not None):
                                     pipeline_Qt.producer_acquire(producer_state_Qt)
                                     load_Qt(m_block - 1, producer_state=producer_state_Qt)
                                     pipeline_Qt.producer_commit(producer_state_Qt)
@@ -2149,11 +2192,11 @@ class FlashAttentionBackwardSm100:
                                 pipeline_dO.producer_acquire(
                                     producer_state_dO_dPsum,
                                     extra_tx_count=self.tma_copy_bytes["dO"]
-                                    if const_expr(tma_atom_dOt is not None)
                                     else 0,
                                 )
                                 load_dO(m_block, producer_state=producer_state_dO_dPsum)
-                                if const_expr(tma_atom_dOt is not None):
                                     load_dOt(m_block, producer_state=producer_state_dO_dPsum)
                                 pipeline_dO.producer_commit(producer_state_dO_dPsum)
@@ -2171,27 +2214,104 @@ class FlashAttentionBackwardSm100:
                         #### Tail ####
                         if const_expr(should_load_Q):
-                            if const_expr(tma_atom_Qt is not None):
                                 pipeline_Qt.producer_acquire(producer_state_Qt)
                                 load_Qt(m_block_max - 1, producer_state=producer_state_Qt)
                                 pipeline_Qt.producer_commit(producer_state_Qt)
                                 producer_state_Qt.advance()
                 if const_expr(self.use_2cta_instrs and self.tile_hdim == 192):
-                    pipeline_Q.producer_tail(producer_state_Q_Qt)
-                    pipeline_LSE.producer_tail(producer_state_LSE)
-                    pipeline_dO.producer_tail(producer_state_O_Ot)
-                    pipeline_dPsum.producer_tail(producer_state_dPsum)
                 else:
-                    if const_expr(should_load_Q):
-                        pipeline_Q.producer_tail(producer_state_Q_LSE.clone())
-                        pipeline_LSE.producer_tail(producer_state_Q_LSE)
-                        if const_expr(tma_atom_Qt is not None):
-                            pipeline_Qt.producer_tail(producer_state_Qt)
-                    if const_expr(should_load_dO):
-                        pipeline_dO.producer_tail(producer_state_dO_dPsum.clone())
-                        pipeline_dPsum.producer_tail(producer_state_dO_dPsum)
             tile_scheduler.prefetch_next_work()
             tile_scheduler.advance_to_next_work()
             work_tile = tile_scheduler.get_current_work()
@@ -2232,6 +2352,7 @@ class FlashAttentionBackwardSm100:
         pipeline_dKV: PipelineAsync,
         pipeline_dP: PipelineAsync,
         pipeline_dQ: PipelineAsync,
         block_info: BlockInfo,
         SeqlenInfoCls: Callable,
         TileSchedulerCls: Callable,
@@ -2343,6 +2464,10 @@ class FlashAttentionBackwardSm100:
         )
         producer_phase_acc = Int32(1)  # For S & P, dP, dQ
         producer_phase_dQ = Int32(1)  # 2-CTA: separate phase for dQ pipeline
         consumer_state_dS = cutlass.pipeline.make_pipeline_state(
             cutlass.pipeline.PipelineUserType.Consumer, 1
         )
@@ -2366,8 +2491,8 @@ class FlashAttentionBackwardSm100:
                     blocksparse_tensors,
                     batch_idx,
                     head_idx,
-                    n_block,
-                    subtile_factor=self.subtile_factor,
                     m_block_max=m_block_max,
                 )
                 process_tile = block_iter_count > Int32(0)
@@ -2392,7 +2517,7 @@ class FlashAttentionBackwardSm100:
                     # 4. dV   = P.T  @ dO
                     # 5. dQ   = dS   @ K
-                    main_loop_iters = m_block_max - m_block_min
                     # empty waits
                     # pipeline_S_P.sync_object_empty.wait(0, producer_phase_acc)
@@ -2582,20 +2707,26 @@ class FlashAttentionBackwardSm100:
             else:
                 if is_leader_cta and process_tile:
                     accumulate_dK = False
                     # -----------------------------------------------------------
                     ###### Prologue
                     # -----------------------------------------------------------
-                    # 1. S  = Q0 @ K.T
-                    # 2. dP = V @ dOt.T
-                    # 3. dV = P @ dO
-                    # 1) S = K @ Q
                     handle_Q = pipeline_Q_consumer.wait_and_advance()
-                    pipeline_S_P.sync_object_empty.wait(0, producer_phase_acc)
                     mma_qk_fn(B_idx=handle_Q.index)
                     pipeline_S_P.sync_object_full.arrive(0, pipeline_S_P.producer_mask, cta_group)
-                    # 2) dP = V @ dOt.T
                     pipeline_dO.consumer_wait(consumer_state_dO)
                     pipeline_dP.sync_object_empty.wait(0, producer_phase_acc)
                     pipeline_dQ.sync_object_empty.wait(0, producer_phase_acc)
@@ -2604,19 +2735,22 @@ class FlashAttentionBackwardSm100:
                     producer_phase_acc ^= 1
                     # 3) dV = P.T @ dO
-                    pipeline_S_P.sync_object_empty.wait(0, producer_phase_acc)
-                    mma_pdo_fn(B_idx=consumer_state_dO.index, zero_init=True)
-                    pipeline_dO.consumer_release(consumer_state_dO)
-                    consumer_state_dO.advance()
                     # -----------------------------------------------------------
                     ###### MAIN LOOP
                     # -----------------------------------------------------------
-                    # 1. S  = K    @ Q.T
-                    # 2. dQ = dS   @ K
-                    # 3. dK = dS.T @ Q
-                    # 4. dP = V    @ dOt.T
-                    # 5. dV = P.T  @ dO
                     # For block sparsity, we use block_iter_count; for dense, use m_block range
                     # MMA doesn't need actual m_block indices, just the iteration count
@@ -2630,10 +2764,25 @@ class FlashAttentionBackwardSm100:
                     for _ in cutlass.range(main_loop_iters, unroll=1):
                         # (1) S.T = K @ Q.T
                         handle_Q_next = pipeline_Q_consumer.wait_and_advance()
                         mma_qk_fn(B_idx=handle_Q_next.index)
                         pipeline_S_P.sync_object_full.arrive(
                             0, pipeline_S_P.producer_mask, cta_group
                         )
                         # (2) dK += dS.T @ Q
                         pipeline_dS.consumer_wait(consumer_state_dS)
@@ -2647,22 +2796,34 @@ class FlashAttentionBackwardSm100:
                         pipeline_dS.consumer_release(consumer_state_dS)
                         consumer_state_dS.advance()
-                        # (4) dP = V @ dO.T
                         pipeline_dO.consumer_wait(consumer_state_dO)
                         pipeline_dQ.sync_object_empty.wait(0, producer_phase_acc)
                         mma_dov_fn(B_idx=consumer_state_dO.index)
                         pipeline_dP.sync_object_full.arrive(0, pipeline_dP.producer_mask, cta_group)
-                        # (5) dV += P.T @ dO
                         producer_phase_acc ^= 1
-                        pipeline_S_P.sync_object_empty.wait(0, producer_phase_acc)
-                        mma_pdo_fn(B_idx=consumer_state_dO.index, zero_init=False)
-                        pipeline_dO.consumer_release(consumer_state_dO)
-                        consumer_state_dO.advance()
                         handle_Q = handle_Q_next
-                    pipeline_S_P.sync_object_full.arrive(0, pipeline_S_P.producer_mask, cta_group)
                     # signal to the epilogue that dV is ready
                     # pipeline_dKV.producer_acquire(producer_state_dKV)
@@ -2753,13 +2914,17 @@ class FlashAttentionBackwardSm100:
         n_block,
         softmax_scale,
         seqlen_info,
-        aux_tensors=None,
         fastdiv_mods=(None, None),
     ):
         """Apply forward score modification for SM100 backward pass."""
-        # In bwd, S is computed as K @ Q.T so dimensions are (tile_n, tile_m)
-        cS = cute.make_identity_tensor((self.tile_n, self.tile_m))
-        cS = cute.domain_offset((n_block * self.tile_n, m_block * self.tile_m), cS)
         tScS = thr_mma_S.partition_C(cS)
         tScS_idx = thr_copy_t2r.partition_D(tScS)
@@ -2772,7 +2937,7 @@ class FlashAttentionBackwardSm100:
             softmax_scale,
             self.vec_size,
             self.qk_acc_dtype,
-            aux_tensors,
             fastdiv_mods,
             seqlen_info,
             constant_q_idx=None,
@@ -2790,7 +2955,7 @@ class FlashAttentionBackwardSm100:
         head_idx,
         softmax_scale,
         seqlen_info,
-        aux_tensors=None,
         fastdiv_mods=(None, None),
     ):
         """Apply backward score modification (joint graph) for SM100."""
@@ -2804,7 +2969,7 @@ class FlashAttentionBackwardSm100:
             softmax_scale,
             self.vec_size,
             self.qk_acc_dtype,
-            aux_tensors,
             fastdiv_mods,
             seqlen_info,
             constant_q_idx=None,
@@ -2815,10 +2980,10 @@ class FlashAttentionBackwardSm100:
     @cute.jit
     def compute_loop(
         self,
-        thr_mma_S: cute.core.ThrMma,
-        thr_mma_dP: cute.core.ThrMma,
-        thr_mma_dV: cute.core.ThrMma,
-        thr_mma_dK: cute.core.ThrMma,
         tStS: cute.Tensor,
         tdPtdP: cute.Tensor,
         tdVtdV: cute.Tensor,
@@ -2835,6 +3000,7 @@ class FlashAttentionBackwardSm100:
         pipeline_dS: PipelineAsync,
         pipeline_dKV: PipelineAsync,
         pipeline_dP: PipelineAsync,
         dS_cluster_empty_mbar_ptr: cute.Pointer,
         dS_cluster_full_mbar_ptr: cute.Pointer,
         dQaccum_empty_mbar_ptr: cute.Pointer,
@@ -2853,7 +3019,7 @@ class FlashAttentionBackwardSm100:
         tiled_copy_r2s_dKV: Optional[cute.TiledCopy],
         mdK_semaphore: Optional[cute.Tensor],
         mdV_semaphore: Optional[cute.Tensor],
-        aux_tensors: Optional[list] = None,
         fastdiv_mods=(None, None),
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
     ):
@@ -2888,13 +3054,17 @@ class FlashAttentionBackwardSm100:
         tileP_f32_like = self.cta_tiler[1] // 32 * self.v_dtype.width
         # tStS has shape ((128, 128), 1, 1), tStP has shape ((128, 64), 1, 1)
-        # tP overlap with tS
         tStP = cute.composition(tStS, (cute.make_layout((self.tile_n, tileP_f32_like)), 1, 1))
-        tStP = cute.make_tensor(tStS.iterator, tStP.layout)  # Otherwise the tmem address is wrong
         tScS = thr_mma_S.partition_C(cute.make_identity_tensor(self.mma_tiler_kq[:2]))
         tScP = cute.composition(tScS, (cute.make_layout((self.tile_n, tileP_f32_like)), 1, 1))
-        # tdS overlap with tdP
         tdPtdS = cute.composition(tdPtdP, (cute.make_layout((self.tile_n, tileP_f32_like)), 1, 1))
         tdPcdP = thr_mma_dP.partition_C(cute.make_identity_tensor(self.mma_tiler_vdo[:2]))
         tdPcdS = cute.composition(tdPcdP, (cute.make_layout((self.tile_n, tileP_f32_like)), 1, 1))
@@ -2953,6 +3123,12 @@ class FlashAttentionBackwardSm100:
             cutlass.pipeline.PipelineUserType.Consumer, 1
         )
         # consumer_phase_S_P_dP = Int32(0)
         producer_state_dS = pipeline.make_pipeline_state(  # Our impl has shortcut for stage==1
             cutlass.pipeline.PipelineUserType.Producer, 1
         )
@@ -2971,31 +3147,58 @@ class FlashAttentionBackwardSm100:
         while work_tile.is_valid_tile:
             n_block, head_idx, batch_idx, _ = work_tile.tile_idx
             seqlen = SeqlenInfoCls(batch_idx)
             m_block_min, m_block_max = block_info.get_m_block_min_max(
                 seqlen, n_block // self.cluster_shape_mnk[0]
             )
             mask = AttentionMaskCls(seqlen)
-            n_block_for_cluster = n_block // self.cta_group_size
             # TODO: condition mask_seqlen
             mask_fn = partial(
                 mask.apply_mask_sm100_transposed,
                 tScS_t2r=tScS_t2r,
                 t0ScS_t2r=t0ScS_t2r,
-                n_block=n_block_for_cluster,
                 mask_seqlen=True,
                 mask_causal=self.is_causal,
                 mask_local=self.is_local,
                 mask_mod=self.mask_mod,
                 batch_idx=batch_idx,
                 head_idx=head_idx,
-                aux_tensors=aux_tensors,
                 fastdiv_mods=fastdiv_mods,
             )
             # prefetch_LSE = not self.is_causal
             prefetch_LSE = False
-            # some tiles might be empty due to block sparsity
             if const_expr(self.use_block_sparsity):
                 (
                     curr_q_cnt,
                     curr_q_idx,
@@ -3006,22 +3209,19 @@ class FlashAttentionBackwardSm100:
                     blocksparse_tensors,
                     batch_idx,
                     head_idx,
-                    n_block,
-                    subtile_factor=self.subtile_factor,
                     m_block_max=m_block_max,
                 )
                 process_tile = loop_count > Int32(0)
-            else:
-                process_tile = (
-                    const_expr(not self.is_local and not self.is_varlen_q)
-                    or m_block_min < m_block_max
-                )
-                loop_count = m_block_max - m_block_min
             # Mainloop
             # Block sparsity: iterate over sparse m_block count and derive actual m_block
             # from Q_IDX/FULL_Q_IDX tensors. Dense: iterate m_block_min..m_block_max directly.
             for iter_idx in cutlass.range(loop_count, unroll=1):
                 if const_expr(self.use_block_sparsity):
                     m_block, is_full_block = get_m_block_from_iter_bwd(
                         iter_idx,
@@ -3029,27 +3229,34 @@ class FlashAttentionBackwardSm100:
                         curr_q_idx,
                         curr_full_cnt,
                         curr_full_idx,
-                        subtile_factor=self.subtile_factor,
                         m_block_max=m_block_max,
                     )
                     m_block_oob = m_block >= m_block_max
-                else:
-                    m_block = m_block_min + iter_idx
-                    m_block_oob = False
-                    is_full_block = False
                 # Prefetch 1 stage of LSE
                 pipeline_LSE.consumer_wait(consumer_state_LSE)
-                tSrLSE_s2r = cute.make_fragment(tScS_t2r[None, 0, 0, 0].shape, Float32)
                 if const_expr(prefetch_LSE and not self.shuffle_LSE):
                     cute.autovec_copy(tSsLSE[None, 0, 0, 0, consumer_state_LSE.index], tSrLSE_s2r)
-                pipeline_S_P.consumer_wait(consumer_state_S_P_dP)
                 # pipeline_S_P.sync_object_full.wait(0, consumer_phase_S_P_dP)
                 #### TMEM->RMEM (Load S from TMEM)
-                tSrS_t2r = cute.make_fragment(tScS_t2r.shape, Float32)
                 cute.copy(thr_copy_t2r, tStS_t2r, tSrS_t2r)
-                if const_expr(self.tile_hdim == 192):
                     # Signal S tmem load completion using pipeline_S_P when hdim 192
                     # dP is overlapped with S
                     cute.arch.fence_view_async_tmem_load()
@@ -3080,7 +3287,7 @@ class FlashAttentionBackwardSm100:
                         n_block,
                         softmax_scale,
                         seqlen,
-                        aux_tensors,
                         fastdiv_mods,
                     )
@@ -3097,7 +3304,7 @@ class FlashAttentionBackwardSm100:
                 #### P = exp(S - LSE)
                 # ---------------------------------------------
                 lane_idx = cute.arch.lane_idx()
-                tSrP_r2t_f32 = cute.make_fragment(tScP_r2t.shape, Float32)  # 64
                 tSrP_r2t = cute.recast_tensor(tSrP_r2t_f32, self.q_dtype)
                 for stage in cutlass.range_constexpr(num_stages):
                     tSrS_cur = tSrS_t2r[None, stage, 0, 0]
@@ -3125,10 +3332,14 @@ class FlashAttentionBackwardSm100:
                         tSrS_cur[2 * v + 1] = cute.math.exp2(tSrS_cur[2 * v + 1], fastmath=True)
                     utils.cvt_f16(tSrS_cur, tSrP_r2t[None, stage, 0, 0])
                     if const_expr(stage == 0):
-                        cute.arch.fence_view_async_tmem_load()
-                        # Without this barrier, we could have 1 warp writing to P in tmem while
-                        # another warp is still reading S from tmem.
-                        self.compute_sync_barrier.arrive_and_wait()
                     cute.copy(
                         thr_copy_r2t,
                         tSrP_r2t_f32[None, stage, None, None],
@@ -3137,12 +3348,22 @@ class FlashAttentionBackwardSm100:
                 cute.arch.fence_view_async_tmem_store()
                 cute.arch.fence_view_async_shared()
-                self.compute_sync_barrier.arrive_and_wait()
-                if const_expr(not self.tile_hdim == 192):
                     # Signal tmem store P completion with pipeline_S_P
                     with cute.arch.elect_one():
                         pipeline_S_P.consumer_release(consumer_state_S_P_dP)
                         # pipeline_S_P.sync_object_empty.arrive(0, pipeline_S_P.consumer_mask)
                 pipeline_LSE.consumer_release(consumer_state_LSE)
                 consumer_state_LSE.advance()
                 # ---------------------------------------------
@@ -3157,10 +3378,12 @@ class FlashAttentionBackwardSm100:
                 ##### dS.T = P.T * (dP.T - Psum)
                 for stage in cutlass.range_constexpr(num_stages):
-                    tdPrdP_t2r = cute.make_fragment(tScS_t2r[None, 0, None, None].shape, Float32)
                     cute.copy(thr_copy_t2r, tdPtdP_t2r[None, stage, None, None], tdPrdP_t2r)
                     cute.arch.fence_view_async_tmem_load()
-                    self.compute_sync_barrier.arrive_and_wait()
                     tdPrdP_cur = tdPrdP_t2r[None, 0, 0]
                     tSrS_cur = tSrS_t2r[None, stage, 0, 0]
                     tSsdPsum_cur = tSsdPsum[None, stage, 0, 0, consumer_state_dPsum.index]
@@ -3177,10 +3400,8 @@ class FlashAttentionBackwardSm100:
                                 utils.shuffle_sync(tSrdPsum, offset=2 * v),
                                 utils.shuffle_sync(tSrdPsum, offset=2 * v + 1),
                             )
-                        tdPrdP_cur[2 * v], tdPrdP_cur[2 * v + 1] = (
-                            activation.sub_packed_f32x2(
-                                (tdPrdP_cur[2 * v], tdPrdP_cur[2 * v + 1]), dPsum_pair
-                            )
                         )
                         tdPrdP_cur[2 * v], tdPrdP_cur[2 * v + 1] = cute.arch.mul_packed_f32x2(
                             (tSrS_cur[2 * v], tSrS_cur[2 * v + 1]),
@@ -3189,9 +3410,12 @@ class FlashAttentionBackwardSm100:
                     if const_expr(self.score_mod_bwd is not None):
                         tSrS_pre_cur = tSrS_pre[None, stage, 0, 0]
-                        cS_bwd = cute.make_identity_tensor((self.tile_n, self.tile_m))
                         cS_bwd = cute.domain_offset(
-                            (n_block * self.tile_n, m_block * self.tile_m), cS_bwd
                         )
                         tScS_bwd = thr_mma_S.partition_C(cS_bwd)
                         tScS_idx_bwd = thr_copy_t2r.partition_D(tScS_bwd)
@@ -3204,7 +3428,7 @@ class FlashAttentionBackwardSm100:
                             head_idx,
                             softmax_scale,
                             seqlen,
-                            aux_tensors,
                             fastdiv_mods,
                         )
                         # Zero out OOB positions (kv_idx >= seqlen_k) after score_mod_bwd
@@ -3252,7 +3476,13 @@ class FlashAttentionBackwardSm100:
                     cute.autovec_copy(tdPrdS_xchg, tRS_sdS_xchg[None, 0])
                 cute.arch.fence_view_async_shared()
-                self.compute_sync_barrier.arrive_and_wait()
                 pipeline_dPsum.consumer_release(consumer_state_dPsum)
                 consumer_state_dPsum.advance()
                 # when 2cta hdim 128, pipeline_dS also signals S tmem load completion so is deferred
@@ -3413,12 +3643,67 @@ class FlashAttentionBackwardSm100:
             tile_scheduler.advance_to_next_work()
             work_tile = tile_scheduler.get_current_work()
     @cute.jit
     def dQacc_reduce(
         self,
         mdQaccum: cute.Tensor,
         sdQaccum: cute.Tensor,
-        thr_mma_dQ: cute.core.ThrMma,
         tdQtdQ: cute.Tensor,
         pipeline_dQ: PipelineAsync,
         dQaccum_empty_mbar_ptr: Optional[cute.Pointer],
@@ -3470,6 +3755,7 @@ class FlashAttentionBackwardSm100:
         while work_tile.is_valid_tile:
             n_block, head_idx, batch_idx, _ = work_tile.tile_idx
             n_block_cta_group = n_block // self.cta_group_size  # for 2cta
             seqlen = SeqlenInfoCls(batch_idx)
             m_block_min, m_block_max = block_info.get_m_block_min_max(seqlen, n_block_cta_group)
             if const_expr(not seqlen.has_cu_seqlens_q):
@@ -3485,13 +3771,43 @@ class FlashAttentionBackwardSm100:
             )
             if const_expr(self.deterministic):
                 mdQ_semaphore_cur = mdQ_semaphore[None, None, head_idx, batch_idx]
-            # delay_semaphore_release = self.is_causal and not self.tile_hdim == 192
-            delay_semaphore_release = not self.tile_hdim == 192
-            # some tiles might be empty due to block sparsity
             if const_expr(self.use_block_sparsity):
                 (
                     curr_q_cnt,
                     curr_q_idx,
@@ -3502,22 +3818,30 @@ class FlashAttentionBackwardSm100:
                     blocksparse_tensors,
                     batch_idx,
                     head_idx,
-                    n_block,
-                    subtile_factor=self.subtile_factor,
                     m_block_max=m_block_max,
                 )
                 process_tile = loop_count > Int32(0)
-            else:
-                process_tile = (
-                    const_expr(not self.is_local and not self.is_varlen_q)
-                    or m_block_min < m_block_max
-                )
-                loop_count = m_block_max - m_block_min
             # dQacc_reduce mainloop
             # Block sparsity: iterate over sparse m_block count and derive actual m_block
             # from Q_IDX/FULL_Q_IDX tensors. Dense: iterate m_block_min..m_block_max directly.
             for iter_idx in cutlass.range(loop_count, unroll=1):
                 if const_expr(self.use_block_sparsity):
                     m_block, _ = get_m_block_from_iter_bwd(
                         iter_idx,
@@ -3525,16 +3849,13 @@ class FlashAttentionBackwardSm100:
                         curr_q_idx,
                         curr_full_cnt,
                         curr_full_idx,
-                        subtile_factor=self.subtile_factor,
                         m_block_max=m_block_max,
                     )
-                    if m_block_max > 0:
-                        m_block = cutlass.min(m_block, m_block_max - 1)
-                else:
-                    m_block = m_block_min + iter_idx
                 pipeline_dQ.consumer_wait(dQ_consumer_state)
                 # TMEM -> RMEM
-                tdQrdQ_t2r = cute.make_fragment(tdQrdQ_t2r_shape, Float32)
                 cute.copy(thr_copy_t2r, tdQtdQ_t2r, tdQrdQ_t2r)
                 cute.arch.fence_view_async_tmem_load()
                 cute.arch.sync_warp()
@@ -3542,6 +3863,8 @@ class FlashAttentionBackwardSm100:
                     pipeline_dQ.consumer_release(dQ_consumer_state)
                 dQ_consumer_state.advance()
                 gdQaccum_cur = gdQaccum[None, None, m_block]
                 tdQrdQ_shape = (
@@ -3559,22 +3882,27 @@ class FlashAttentionBackwardSm100:
                     cute.arch.fence_view_async_shared()
                     # semaphore acquire
                     if const_expr(self.deterministic and stage == 0):
-                        if const_expr(self.spt):
-                            _, n_block_max_for_m_block = block_info.get_n_block_min_max(
-                                seqlen, m_block
                             )
-                            lock_value = n_block_max_for_m_block - 1 - n_block_cta_group
-                        else:
-                            lock_value = n_block_cta_group
-                        barrier.wait_eq(
-                            mdQ_semaphore_cur[(m_block, None)].iterator,
-                            tidx,
-                            cta_rank_in_cluster,
-                            lock_value,
-                        )
                     self.reduce_sync_barrier.arrive_and_wait()
                     # Copy from shared memory to global memory
-                    if is_tma_warp:
                         with cute.arch.elect_one():
                             copy_utils.cpasync_reduce_bulk_add_f32(
                                 sdQaccum[None, smem_idx].iterator,
@@ -3583,20 +3911,12 @@ class FlashAttentionBackwardSm100:
                             )
                         cute.arch.cp_async_bulk_commit_group()
                         cute.arch.cp_async_bulk_wait_group(self.sdQaccum_stage - 1, read=read_flag)
                     self.reduce_sync_barrier.arrive_and_wait()
                     dQ_tma_store_producer_state.advance()
-                    # Directly add to gmem, much slower
-                    # tdQgdQ = thr_copy_dQaccum_r2s.partition_D(gdQaccum[None, stage, m_block])
-                    # assert cute.size(tdQrdQ_r2s) == cute.size(tdQgdQ)
-                    # for i in cutlass.range(cute.size(tdQrdQ_r2s) // 4, unroll_full=True):
-                    #     copy_utils.atomic_add_fp32x4(
-                    #         tdQrdQ_r2s[4 * i],
-                    #         tdQrdQ_r2s[4 * i + 1],
-                    #         tdQrdQ_r2s[4 * i + 2],
-                    #         tdQrdQ_r2s[4 * i + 3],
-                    #         utils.elem_pointer(tdQgdQ, 4 * i),
-                    #     )
-                    # semaphore release for prior m_block
                     if const_expr(self.deterministic and stage == 0 and delay_semaphore_release):
                         if m_block > m_block_min:
                             barrier.arrive_inc(
@@ -3618,12 +3938,16 @@ class FlashAttentionBackwardSm100:
                 # NOTE: arrive_inc calls red_release which issues membar
                 if const_expr(self.deterministic and not delay_semaphore_release):
                     if const_expr(self.sdQaccum_stage > 1 and not self.tile_hdim == 192):
-                        if is_tma_warp:
                             cute.arch.cp_async_bulk_wait_group(0, read=read_flag)
                         self.reduce_sync_barrier.arrive_and_wait()
-                    barrier.arrive_inc(
-                        mdQ_semaphore_cur[m_block, None].iterator, tidx, cta_rank_in_cluster, 1
-                    )
             if process_tile:
                 if is_tma_warp:
@@ -3639,7 +3963,10 @@ class FlashAttentionBackwardSm100:
                     )
             if const_expr(
-                self.deterministic and not self.spt and block_info.window_size_left is not None
             ):
                 m_block_global_max = cute.ceil_div(seqlen.seqlen_q, self.tile_m)
                 for m_block in cutlass.range(m_block_max, m_block_global_max, unroll=1):
@@ -3650,6 +3977,9 @@ class FlashAttentionBackwardSm100:
             tile_scheduler.advance_to_next_work()
             work_tile = tile_scheduler.get_current_work()
     @cute.jit
     def epilogue_dKV(
         self,
@@ -3659,8 +3989,8 @@ class FlashAttentionBackwardSm100:
         head_idx: Int32,
         n_block: Int32,
         seqlen,
-        thr_mma_dV: cute.core.ThrMma,
-        thr_mma_dK: cute.core.ThrMma,
         tdVtdV: cute.Tensor,
         tdKtdK: cute.Tensor,
         mdV: cute.Tensor,
@@ -3678,13 +4008,22 @@ class FlashAttentionBackwardSm100:
         mdV_cur = seqlen.offset_batch_K(mdV, batch_idx, dim=3)[None, None, head_idx]
         mdK_cur = seqlen.offset_batch_K(mdK, batch_idx, dim=3)[None, None, head_idx]
-        tmem_load_atom = cute.make_copy_atom(
-            tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(16)), Float32
         )
         # dV
         pipeline_dKV.consumer_wait(consumer_state_dKV)
-        tiled_tmem_ld_dV = tcgen05.make_tmem_copy(tmem_load_atom, tdVtdV)
         thr_tmem_ld_dV = tiled_tmem_ld_dV.get_slice(tidx)
         tdVtdV_t2r_p = thr_tmem_ld_dV.partition_S(tdVtdV)
@@ -3696,7 +4035,7 @@ class FlashAttentionBackwardSm100:
         tdVcdV_t2r_p = thr_tmem_ld_dV.partition_D(tdVcdV_tensor)
         tdVcdV_t2r = self.split_wg(tdVcdV_t2r_p, wg_idx, num_wg)
-        tdVrdV_t2r = cute.make_fragment(tdVcdV_t2r.shape, Float32)
         cute.copy(thr_tmem_ld_dV, tdVtdV_t2r, tdVrdV_t2r)
         cute.arch.fence_view_async_tmem_load()
@@ -3713,7 +4052,7 @@ class FlashAttentionBackwardSm100:
             tiler_mn=tiled_tmem_ld_dV.tiler_mn,
         )
-        tdVrdV_r2s = cute.make_fragment(tdVrdV_t2r.shape, self.dv_dtype)
         for i in cutlass.range_constexpr(cute.size(tdVrdV_t2r, mode=[1])):
             dV_vec = tdVrdV_t2r[(None, i, 0, 0)].load()
             tdVrdV_r2s[(None, i, 0, 0)].store(dV_vec.to(self.dv_dtype))
@@ -3725,8 +4064,11 @@ class FlashAttentionBackwardSm100:
         tdVgdV_r2g_p = thr_tmem_ld_dV.partition_D(tdVgdV)
         tdVgdV_r2g = self.split_wg(tdVgdV_r2g_p, wg_idx, num_wg)
         if tidx < seqlen.seqlen_k - self.tile_n * n_block:
-            cute.copy(tiled_gmem_store_dV, tdVrdV_r2s, tdVgdV_r2g)
         cute.arch.sync_warp()
         with cute.arch.elect_one():
@@ -3736,7 +4078,7 @@ class FlashAttentionBackwardSm100:
         # dK
         pipeline_dKV.consumer_wait(consumer_state_dKV)
-        tiled_tmem_ld_dK = tcgen05.make_tmem_copy(tmem_load_atom, tdKtdK)
         thr_tmem_ld_dK = tiled_tmem_ld_dK.get_slice(tidx)
         tdKtdK_t2r_p = thr_tmem_ld_dK.partition_S(tdKtdK)
@@ -3748,7 +4090,7 @@ class FlashAttentionBackwardSm100:
         tdKcdK_t2r_p = thr_tmem_ld_dK.partition_D(tdKcdK_tensor)
         tdKcdK_t2r = self.split_wg(tdKcdK_t2r_p, wg_idx, num_wg)
-        tdKrdK_t2r = cute.make_fragment(tdKcdK_t2r.shape, Float32)
         cute.copy(tiled_tmem_ld_dK, tdKtdK_t2r, tdKrdK_t2r)
         cute.arch.fence_view_async_tmem_load()
@@ -3766,7 +4108,7 @@ class FlashAttentionBackwardSm100:
             tiler_mn=tiled_tmem_ld_dK.tiler_mn,
         )
-        tdKrdK_r2s = cute.make_fragment(tdKrdK_t2r.shape, self.dk_dtype)
         for i in cutlass.range_constexpr(cute.size(tdKrdK_t2r, mode=[1])):
             dK_vec = tdKrdK_t2r[(None, i, 0, 0)].load() * softmax_scale
@@ -3779,14 +4121,33 @@ class FlashAttentionBackwardSm100:
         tdKgdK_r2g_p = thr_tmem_ld_dK.partition_D(tdKgdK)
         tdKgdK_r2g = self.split_wg(tdKgdK_r2g_p, wg_idx, num_wg)
         if tidx < seqlen.seqlen_k - self.tile_n * n_block:
-            cute.copy(tiled_gmem_store_dK, tdKrdK_r2s, tdKgdK_r2g)
         cute.arch.sync_warp()
         with cute.arch.elect_one():
             pipeline_dKV.consumer_release(consumer_state_dKV)
         return consumer_state_dKV
     @cute.jit
     def epilogue_dK_or_dV_tma(
         self,
@@ -3795,7 +4156,7 @@ class FlashAttentionBackwardSm100:
         head_idx: Int32,
         n_block: Int32,
         seqlen,
-        thr_mma: cute.core.ThrMma,
         tdKVtdKV: cute.Tensor,
         mdKV: cute.Tensor,
         sdKV: cute.Tensor,
@@ -3861,6 +4222,7 @@ class FlashAttentionBackwardSm100:
         deterministic_KV = self.deterministic and self.qhead_per_kvhead > 1
         if const_expr(deterministic_KV):
             mdKV_semaphore_cur = mdKV_semaphore[n_block, None, head_idx_kv, batch_idx]
         if const_expr(not self.dKV_postprocess):
@@ -3913,7 +4275,7 @@ class FlashAttentionBackwardSm100:
             if const_expr(num_epi_stages > 1):
                 tdKVcdKV_t2r = tdKVcdKV_t2r[None, epi_stage]
-            tdKVrdKV_t2r = cute.make_fragment(tdKVcdKV_t2r.shape, Float32)
             assert cute.size(tdKVrdKV_t2r) == cute.size(tdKVtdKV_t2r) // cute.arch.WARP_SIZE, (
                 "RMEM<->TMEM fragment size mismatch"
@@ -3929,7 +4291,7 @@ class FlashAttentionBackwardSm100:
                     tdKVrdKV_t2r[2 * i], tdKVrdKV_t2r[2 * i + 1] = cute.arch.mul_packed_f32x2(
                         (tdKVrdKV_t2r[2 * i], tdKVrdKV_t2r[2 * i + 1]), (scale, scale)
                     )
-            tdKVrdKV = cute.make_fragment(tdKVrdKV_t2r.shape, dtype)  # (32 columns)
             tdKVrdKV.store(tdKVrdKV_t2r.load().to(dtype))
             # RMEM -> SMEM -- copy, fence and barrier

 import cutlass.utils.blackwell_helpers as sm100_utils_basic
 from cutlass.pipeline import PipelineAsync
 from .quack import layout_utils
 from . import utils
 from .cute_dsl_utils import assume_tensor_aligned
 from .named_barrier import NamedBarrierBwdSm100
 from .softmax import apply_score_mod_inner, apply_score_mod_bwd_inner
 from .block_sparsity import BlockSparseTensors
+from .utils import AuxData
 from .block_sparse_utils import (
     get_total_q_block_count_bwd,
     get_block_sparse_iteration_info_bwd,
     get_m_block_from_iter_bwd,
+    produce_block_sparse_q_loads_bwd_sm100_2cta_hdim192,
+    produce_block_sparse_q_loads_bwd_sm100_default,
 )
+# NOTE [hdim64 dedicated P/dS TMEM slots]
+# The 1-CTA hdim64 backward uses TMEM columns [0, 384), leaving room for two 64-column
+# slots. With split_P_dS, P and dS live there instead of aliasing S and dP:
+#   * S is released as soon as the compute warps have loaded it, so the MMA warp issues
+#     QK_{t+1} while the softmax of tile t is still running; dV += P.T @ dO for tile t
+#     then trails QK_{t+1} and is gated by its own pipeline_P.
+#   * The compute-wide barriers that ordered the P-over-S and dS-over-dP overwrites are
+#     gone. They were cross-warp hazards: a warp's P (dS) columns overlap S (dP) lanes
+#     that another warp still has to load.
+#   * The remaining compute-wide barriers before pipeline arrivals become sync_warp: every
+#     mbarrier signalled from the compute warps expects one arrival per warp, so the
+#     mbarrier itself is the cross-warp join. This is 1-CTA only: at 2-CTA the barrier
+#     before the dS commit also orders every warp's sdS_xchg write ahead of the
+#     single-thread DSMEM copy to the peer CTA.
+# The S handshake keeps its own phase/state (producer_phase_S, consumer_state_S) so its
+# parity never depends on where the other pipelines flip. Perf evidence in PR #2804; the
+# aliased layout is what every other shape runs.
 class FlashAttentionBackwardSm100:
     arch = 100
         tile_n: int = 128,
         is_persistent: bool = False,
         deterministic: bool = False,
+        spt: Optional[bool] = None,
         cluster_size: int = 1,
         use_2cta_instrs: bool = False,
         score_mod: cutlass.Constexpr | None = None,
         score_mod_bwd: cutlass.Constexpr | None = None,
         mask_mod: cutlass.Constexpr | None = None,
         has_aux_tensors: cutlass.Constexpr = False,
+        q_subtile_factor: cutlass.Constexpr[int] = 1,
+        kv_subtile_factor: cutlass.Constexpr[int] = 1,
     ):
+        # Pad head_dim to a multiple of 32 (k_block_size). Must stay a multiple of 32 so
+        # that the dQ accumulator reduce (dQ_reduce_ncol up to 32) tiles evenly and matches
+        # both the preprocess kernel and the interface's head_dim_rounded (also mult of 32).
+        hdim_multiple_of = 32
         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)
         self.check_hdim_oob = head_dim != self.tile_hdim
         self.check_hdim_v_oob = head_dim_v != self.tile_hdimv
+        self.head_dim = head_dim
+        self.head_dim_v = head_dim_v
         self.tile_m = tile_m
         self.tile_n = tile_n
         assert self.tile_hdim <= 128 or (self.tile_hdim == 192 and self.tile_hdimv == 128)
         assert self.tile_hdimv <= 128
+        self.use_2cta_instrs = bool(use_2cta_instrs and cluster_size == 2)
         self.cta_group_size = 2 if self.use_2cta_instrs else 1
         assert self.tile_hdim != 192 or self.use_2cta_instrs, "Must use 2CTA for hdim 192"
         assert cluster_size in (1, 2), "Only cluster_size=1 or 2 is supported"
         self.cluster_shape_mn = (cluster_size, 1)
+        # The MMA warp's producer_phase_acc (shared by the S/P, dP and dQ empty waits in the
+        # aliased layout) is only consistent across one work tile per CTA.
+        assert not is_persistent, "persistent tile scheduling is not supported"
         self.is_persistent = is_persistent
         self.is_causal = is_causal
         self.is_local = is_local
         self.qhead_per_kvhead = qhead_per_kvhead
         self.pack_gqa = False
         self.deterministic = deterministic
+        self.spt_override = spt
         # Score mod and mask mod support
         self.score_mod = score_mod
         self.score_mod_bwd = score_mod_bwd
         self.mask_mod = mask_mod
         self.has_aux_tensors = has_aux_tensors
+        self.q_subtile_factor = q_subtile_factor
+        self.kv_subtile_factor = kv_subtile_factor
+        assert self.kv_subtile_factor == 1 or self.kv_subtile_factor % self.cta_group_size == 0
         # For score_mod, use vec_size=1 (like forward) to handle per-element indices
         if cutlass.const_expr(has_aux_tensors):
             self.vec_size: cutlass.Constexpr = 1
         # Generally slower to use store dS in smem for dK, and doesn't work for 2cta
         self.use_smem_dS_for_mma_dK = False
+        # See NOTE [hdim64 dedicated P/dS TMEM slots]
+        self.split_P_dS = (
+            cluster_size == 1
+            and self.tile_m == 128
+            and self.tile_n == 128
+            and self.tile_hdim == 64
+            and self.tile_hdimv == 64
+            and not self.use_smem_dS_for_mma_dK
+        )
         self.reduce_warp_ids = (0, 1, 2, 3)
         self.compute_warp_ids = (4, 5, 6, 7, 8, 9, 10, 11)
         self.mma_warp_id = 12
             self.tmem_dQ_offset = 512 - self.tile_hdim // 2
         else:
             self.tmem_S_offset = 0
             self.tmem_dV_offset = self.tmem_S_offset + self.tile_n
             self.tmem_dP_offset = self.tmem_dV_offset + self.tile_hdimv
             self.tmem_dQ_offset = (
                 else self.tmem_dP_offset
             )
             self.tmem_dK_offset = self.tmem_dP_offset + self.tile_m
+            if self.split_P_dS:
+                # P/dS are packed 16-bit over tile_m columns -> tile_m // 2 fp32 columns
+                self.tmem_P_offset = self.tmem_dK_offset + self.tile_hdim
+                self.tmem_dS_offset = self.tmem_P_offset + self.tile_m // 2
+                assert self.tmem_dS_offset + self.tile_m // 2 <= self.tmem_alloc_cols
+            else:
+                self.tmem_P_offset = self.tmem_S_offset  # overlap with S
+                self.tmem_dS_offset = self.tmem_dP_offset  # overlap with dP
         if (not is_causal and not is_local) or deterministic:
             self.num_regs_reduce = 136 if self.use_2cta_instrs else 152
         self.Q_stage = 1 if self.use_2cta_instrs else 2
         self.dO_stage = 1
         self.single_stage = 1
+        self.P_stage = 1 if self.split_P_dS else 0  # dedicated P pipeline, else P aliases S
         # LSE_stage = Q_stage and dPsum_stage = dO_stage
         self.sdKVaccum_stage = 2
         # number of tma reduce adds per dQacc mma
         mdK: cute.Tensor,
         mdV: cute.Tensor,
         softmax_scale: Float32,
         mCuSeqlensQ: Optional[cute.Tensor] = None,
         mCuSeqlensK: Optional[cute.Tensor] = None,
         mSeqUsedQ: Optional[cute.Tensor] = None,
         mSeqUsedK: Optional[cute.Tensor] = None,
         window_size_left: Int32 | int | None = None,
         window_size_right: Int32 | int | None = None,
         mdQ_semaphore: Optional[cute.Tensor] = None,
         mdK_semaphore: Optional[cute.Tensor] = None,
         mdV_semaphore: Optional[cute.Tensor] = None,
+        aux_data: AuxData = AuxData(),
         # Block-sparse tensors (Q direction - for iterating m_blocks per n_block):
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
+        stream: cuda.CUstream = None,
     ):
         self.q_dtype = mQ.element_type
         self.k_dtype = mK.element_type
             assert self.dk_dtype.width == 32, "Must accumulate dK in float precision for GQA"
             assert self.dv_dtype.width == 32, "Must accumulate dV in float precision for GQA"
+        mdQaccum, mdK, mdV, mLSE, mdPsum = [
+            assume_tensor_aligned(t) for t in (mdQaccum, mdK, mdV, mLSE, mdPsum)
+        ]
         # (b, s, n, h) --> (s, h, n, b) or (t, n, h) -> (t, h, n)
         QO_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1]
             TileScheduler = SingleTileLPTBwdScheduler
         else:
             TileScheduler = SingleTileScheduler
+        if const_expr(self.spt_override is None):
+            self.spt = (self.is_causal or self.is_local) and self.deterministic
+        else:
+            assert self.spt_override is not None
+            self.spt = self.spt_override and self.deterministic
         tile_sched_args = TileSchedulerArguments(
             cute.ceil_div(cute.size(mK.shape[0]), self.cta_tiler[0]),  # num_blocks
             cute.size(mQ.shape[2]),  # num_heads = num_query_heads
             qhead_per_kvhead_packgqa=1,  # pack_gqa disabled for bwd
             element_size=self.k_dtype.width // 8,
             is_persistent=self.is_persistent,  # persistent mode not tested
+            cu_total_m_blocks_ptr=mCuTotalMBlocks,
             lpt=self.spt,
             head_swizzle=self.deterministic,
         )
                     cutlass.Int64, self.dQaccum_reduce_stage // 2
                 ]
                 tmem_holding_buf: Int32
+                tmem_dealloc_mbar: cutlass.Int64
                 # 2-CTA
                 Qt_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.Q_stage]
                 S_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.single_stage]
                 dP_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.single_stage]
                 dS_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.single_stage]
+                P_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.P_stage]
                 dKV_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.sdKVaccum_stage]
                 dQ_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2]
                 dQ_cluster_full_mbar_ptr: cute.struct.MemRange[
                     cutlass.Int64, self.dQaccum_reduce_stage // 2
                 ]
                 tmem_holding_buf: Int32
+                tmem_dealloc_mbar: Int64
                 sQ: cute.struct.Align[
                     cute.struct.MemRange[cute.Uint8, sQ_alloc_bytes],
             window_size_right = Int32(window_size_right)
         fastdiv_mods = None
+        if const_expr(aux_data.tensors is not None):
             seqlen_q = cute.size(mQ.shape[0]) // (
                 self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1
             )
             seqlen_k_divmod = FastDivmodDivisor(seqlen_k)
             fastdiv_mods = (seqlen_q_divmod, seqlen_k_divmod)
         self.use_block_sparsity = cutlass.const_expr(blocksparse_tensors is not None)
+        if const_expr(self.use_block_sparsity and self.use_2cta_instrs):
+            # Both CTAs of a cluster must map to the same sparse KV column or they deadlock.
+            assert self.kv_subtile_factor % self.cta_group_size == 0, (
+                "2-CTA block-sparse backward requires kv_subtile_factor % cta_group_size == 0"
             )
         # 2-CTA: 231424 and 1-CTA: 232448
         # print("SMEM: ", self.shared_storage.size_in_bytes())
+        if const_expr(self.use_block_sparsity):
             assert all(x is None for x in (mCuSeqlensQ, mCuSeqlensK, mSeqUsedQ, mSeqUsedK)), (
+                "Variable sequence length is not supported yet for blocksparse in bwd"
             )
         self.kernel(
             window_size_left,
             window_size_right,
             tile_sched_params,
+            aux_data,
             fastdiv_mods,
             blocksparse_tensors,
         ).launch(
             min_blocks_per_mp=1,
         )
+    def _generate_attention_mask_cls(self, window_size_left, window_size_right):
+        return partial(
+            AttentionMask,
+            self.tile_m,
+            self.tile_n * self.cta_group_size,
+            swap_AB=True,
+            window_size_left=window_size_left,
+            window_size_right=window_size_right,
+        )
     @cute.kernel
     def kernel(
         self,
         window_size_left: Optional[Int32],
         window_size_right: Optional[Int32],
         tile_sched_params: ParamsBase,
+        aux_data: AuxData = AuxData(),
         fastdiv_mods=(None, None),
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
     ):
             * len((self.mma_warp_id, *self.compute_warp_ids, *self.reduce_warp_ids)),
         )
         tmem = cutlass.utils.TmemAllocator(
+            storage.tmem_holding_buf.ptr,
             barrier_for_retrieve=tmem_alloc_barrier,
             allocator_warp_id=self.mma_warp_id,
             is_two_cta=self.use_2cta_instrs,
+            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
         )
         # UMMA producers and AsyncThread consumers
         pipeline_consumer_group_MMA_AsyncThread = cutlass.pipeline.CooperativeGroup(
             cutlass.pipeline.Agent.Thread, len(self.compute_warp_ids) * self.cta_group_size
         )
+        # With split_P_dS, this pipeline guards only the S TMEM buffer (P has pipeline_P).
         pipeline_S_P = cutlass.pipeline.PipelineUmmaAsync.create(
             num_stages=1,
             producer_group=pipeline_producer_group_MMA_AsyncThread,
             barrier_storage=storage.dS_mbar_ptr.data_ptr(),
             cta_layout_vmnk=cluster_layout_vmnk,
         )
+        if const_expr(self.split_P_dS):
+            pipeline_P = cutlass.pipeline.PipelineAsyncUmma.create(
+                num_stages=1,
+                producer_group=pipeline_PdS_producer_group,
+                consumer_group=pipeline_PdS_consumer_group,
+                barrier_storage=storage.P_mbar_ptr.data_ptr(),
+                cta_layout_vmnk=cluster_layout_vmnk,
+            )
+        else:
+            pipeline_P = None
         # TMA producer and UMMA consumers
         pipeline_producer_group = cutlass.pipeline.CooperativeGroup(
         )
         TileSchedulerCls = partial(self.tile_scheduler_cls.create, tile_sched_params)
+        AttentionMaskCls = self._generate_attention_mask_cls(window_size_left, window_size_right)
         #  EMPTY
         # (15)
         if warp_idx == self.empty_warp_id:
                     block_info,
                     SeqlenInfoCls,
                     TileSchedulerCls,
+                    blocksparse_tensors,
                 )
         #  LOAD
                 pipeline_dKV,
                 pipeline_dP,
                 pipeline_dQ,
+                pipeline_P,
                 block_info,
                 SeqlenInfoCls,
                 TileSchedulerCls,
                 pipeline_dS,
                 pipeline_dKV,
                 pipeline_dP,
+                pipeline_P,
                 dS_cluster_empty_mbar_ptr,
                 dS_cluster_full_mbar_ptr,
                 dQaccum_empty_mbar_ptr,
                 tiled_copy_r2s_dKV,
                 mdK_semaphore,
                 mdV_semaphore,
+                aux_data,
                 fastdiv_mods,
                 blocksparse_tensors,
             )
         block_info: BlockInfo,
         SeqlenInfoCls: Callable,
         TileSchedulerCls: Callable,
+        blocksparse_tensors: Optional[BlockSparseTensors] = None,
     ):
         cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
         dS_cluster_phase = Int32(0)
             process_tile = (
                 const_expr(not self.is_local and not self.is_varlen_q) or m_block_min < m_block_max
             )
+            num_iters = m_block_max - m_block_min
+            if const_expr(self.use_block_sparsity):
+                assert blocksparse_tensors is not None
+                num_iters = get_total_q_block_count_bwd(
+                    blocksparse_tensors,
+                    batch_idx,
+                    head_idx,
+                    n_block // self.kv_subtile_factor,
+                    q_subtile_factor=self.q_subtile_factor,
+                    m_block_max=m_block_max,
+                )
+                process_tile = num_iters > Int32(0)
             if process_tile:
                 for _ in cutlass.range(num_iters, unroll=1):
                     # Wait for dS_xchg from peer CTA
                     cute.arch.mbarrier_wait(dS_cluster_full_mbar_ptr, phase=dS_cluster_phase)
     @cute.jit
     def load(
         self,
+        thr_mma_S: cute.ThrMma,
+        thr_mma_dP: cute.ThrMma,
+        thr_mma_dV: cute.ThrMma,
+        thr_mma_dK: cute.ThrMma,
+        thr_mma_dQ: cute.ThrMma,
         mQ: cute.Tensor,
         mK: cute.Tensor,
         mKt: Optional[cute.Tensor],
             )
             head_idx_kv = head_idx // self.qhead_per_kvhead
             n_block_cta_group = n_block // self.cta_group_size
+            n_block_sparse = n_block // self.kv_subtile_factor
             # GMEM tensors (varlen-aware)
             mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx]
                 single_stage=True,
             )
+            load_dOt = None
             if const_expr(tma_atom_dOt is not None):
                 gdOt = cute.local_tile(
                     mdOt_cur, cute.select(self.mma_tiler_vdo, mode=[1, 2]), (None, 0)
             load_dO = copy_utils.tma_producer_copy_fn(load_dO, pipeline_dO)
             # (4) dK += dS.T @ Q (2-CTA: needs separate Qt load)
+            load_Qt = None
             if const_expr(tma_atom_Qt is not None):
                 gQt = cute.local_tile(
                     mQt_cur, cute.select(self.mma_tiler_dsq, mode=[1, 2]), (0, None)
                 load_Qt = copy_utils.tma_producer_copy_fn(load_Qt, pipeline_Qt)
             # (5) dQ = dS @ K
+            load_Kt = None
             if const_expr(self.use_2cta_instrs):
                 gKt = cute.local_tile(
                     mKt_cur, cute.select(self.mma_tiler_dsk, mode=[1, 2]), (0, n_block_cta_group)
             # gdPsum = cute.logical_divide(gdPsum, (64,))[(None, block_in_cluster_coord_vmnk[1]), None]
             # copy_stats = partial(cute.copy, copy_atom_stats, mcast_mask=q_do_mcast_mask)
+            if const_expr(not self.use_block_sparsity):
                 process_tile = (
                     const_expr(not self.is_local and not self.is_varlen_q)
                     or m_block_min < m_block_max
                 )
+                if process_tile:
                     first_m_block = m_block_min
                     if const_expr(self.use_2cta_instrs and self.tile_hdim == 192):
                         #### Prologue ####
                         assert should_load_Q and should_load_dO
+                        assert load_dOt is not None and load_Qt is not None
                         # K & Q (for S)
                         pipeline_Q.producer_acquire(
                             producer_state_Q_Qt,
                             pipeline_dO.producer_commit(producer_state_O_Ot)
                             producer_state_O_Ot.advance()
+                        pipeline_Q.producer_tail(producer_state_Q_Qt)
+                        pipeline_LSE.producer_tail(producer_state_LSE)
+                        pipeline_dO.producer_tail(producer_state_O_Ot)
+                        pipeline_dPsum.producer_tail(producer_state_dPsum)
                     else:
                         #### Prologue ####
                         if const_expr(should_load_Q):
                             pipeline_dO.producer_acquire(
                                 producer_state_dO_dPsum,
                                 extra_tx_count=self.tma_copy_bytes["V"] + self.tma_copy_bytes["dO"]
+                                if const_expr(load_dOt is not None)
                                 else self.tma_copy_bytes["V"],
                             )
                             load_V(
                                 )
                             )
                             load_dO(first_m_block, producer_state=producer_state_dO_dPsum)
+                            if const_expr(load_dOt is not None):
                                 load_dOt(first_m_block, producer_state=producer_state_dO_dPsum)
                             pipeline_dO.producer_commit(producer_state_dO_dPsum)
                         #### Main Loop ####
                         for m_block in cutlass.range(m_block_min + 1, m_block_max, unroll=1):
                             if const_expr(should_load_Q):
+                                if const_expr(load_Qt is not None):
                                     pipeline_Qt.producer_acquire(producer_state_Qt)
                                     load_Qt(m_block - 1, producer_state=producer_state_Qt)
                                     pipeline_Qt.producer_commit(producer_state_Qt)
                                 pipeline_dO.producer_acquire(
                                     producer_state_dO_dPsum,
                                     extra_tx_count=self.tma_copy_bytes["dO"]
+                                    if const_expr(load_dOt is not None)
                                     else 0,
                                 )
                                 load_dO(m_block, producer_state=producer_state_dO_dPsum)
+                                if const_expr(load_dOt is not None):
                                     load_dOt(m_block, producer_state=producer_state_dO_dPsum)
                                 pipeline_dO.producer_commit(producer_state_dO_dPsum)
                         #### Tail ####
                         if const_expr(should_load_Q):
+                            if const_expr(load_Qt is not None):
                                 pipeline_Qt.producer_acquire(producer_state_Qt)
                                 load_Qt(m_block_max - 1, producer_state=producer_state_Qt)
                                 pipeline_Qt.producer_commit(producer_state_Qt)
                                 producer_state_Qt.advance()
+                            pipeline_Q.producer_tail(producer_state_Q_LSE.clone())
+                            pipeline_LSE.producer_tail(producer_state_Q_LSE)
+                            if const_expr(load_Qt is not None):
+                                pipeline_Qt.producer_tail(producer_state_Qt)
+                        if const_expr(should_load_dO):
+                            pipeline_dO.producer_tail(producer_state_dO_dPsum.clone())
+                            pipeline_dPsum.producer_tail(producer_state_dO_dPsum)
+            else:
+                assert blocksparse_tensors is not None
                 if const_expr(self.use_2cta_instrs and self.tile_hdim == 192):
+                    assert should_load_Q and should_load_dO
+                    assert load_dOt is not None and load_Qt is not None
+                    assert load_Kt is not None and pipeline_Qt is not None
+                    (
+                        producer_state_Q_Qt,
+                        producer_state_O_Ot,
+                        producer_state_LSE,
+                        producer_state_dPsum,
+                    ) = produce_block_sparse_q_loads_bwd_sm100_2cta_hdim192(
+                        blocksparse_tensors,
+                        batch_idx,
+                        head_idx,
+                        n_block_sparse,
+                        producer_state_Q_Qt,
+                        producer_state_O_Ot,
+                        producer_state_LSE,
+                        producer_state_dPsum,
+                        pipeline_Q,
+                        pipeline_LSE,
+                        pipeline_dO,
+                        pipeline_dPsum,
+                        pipeline_Qt,
+                        load_K,
+                        load_V,
+                        load_Q,
+                        load_dO,
+                        load_Qt,
+                        load_Kt,
+                        load_dOt,
+                        copy_stats,
+                        gLSE,
+                        sLSE,
+                        gdPsum,
+                        sdPsum,
+                        self.tma_copy_bytes["K"],
+                        self.tma_copy_bytes["V"],
+                        q_subtile_factor=self.q_subtile_factor,
+                        m_block_max=m_block_max,
+                    )
                 else:
+                    (
+                        producer_state_Q_LSE,
+                        producer_state_dO_dPsum,
+                        producer_state_Qt,
+                        producer_state_Kt,
+                    ) = produce_block_sparse_q_loads_bwd_sm100_default(
+                        blocksparse_tensors,
+                        batch_idx,
+                        head_idx,
+                        n_block_sparse,
+                        producer_state_Q_LSE,
+                        producer_state_dO_dPsum,
+                        pipeline_Q,
+                        pipeline_LSE,
+                        pipeline_dO,
+                        pipeline_dPsum,
+                        load_K,
+                        load_V,
+                        load_Q,
+                        load_dO,
+                        copy_stats,
+                        gLSE,
+                        sLSE,
+                        gdPsum,
+                        sdPsum,
+                        self.tma_copy_bytes["K"],
+                        self.tma_copy_bytes["V"],
+                        should_load_Q,
+                        should_load_dO,
+                        q_subtile_factor=self.q_subtile_factor,
+                        m_block_max=m_block_max,
+                        use_2cta_instrs=self.use_2cta_instrs,
+                        producer_state_Qt=producer_state_Qt,
+                        producer_state_Kt=producer_state_Kt,
+                        pipeline_Qt=pipeline_Qt,
+                        pipeline_Kt=pipeline_Kt,
+                        load_Qt=load_Qt,
+                        load_Kt=load_Kt,
+                        load_dOt=load_dOt,
+                        tma_copy_bytes_dO=self.tma_copy_bytes["dO"],
+                    )
             tile_scheduler.prefetch_next_work()
             tile_scheduler.advance_to_next_work()
             work_tile = tile_scheduler.get_current_work()
         pipeline_dKV: PipelineAsync,
         pipeline_dP: PipelineAsync,
         pipeline_dQ: PipelineAsync,
+        pipeline_P: Optional[PipelineAsync],
         block_info: BlockInfo,
         SeqlenInfoCls: Callable,
         TileSchedulerCls: Callable,
         )
         producer_phase_acc = Int32(1)  # For S & P, dP, dQ
         producer_phase_dQ = Int32(1)  # 2-CTA: separate phase for dQ pipeline
+        producer_phase_S = Int32(1)  # split_P_dS: separate phase for the S handshake
+        consumer_state_P = cutlass.pipeline.make_pipeline_state(
+            cutlass.pipeline.PipelineUserType.Consumer, 1
+        )
         consumer_state_dS = cutlass.pipeline.make_pipeline_state(
             cutlass.pipeline.PipelineUserType.Consumer, 1
         )
                     blocksparse_tensors,
                     batch_idx,
                     head_idx,
+                    n_block // self.kv_subtile_factor,
+                    q_subtile_factor=self.q_subtile_factor,
                     m_block_max=m_block_max,
                 )
                 process_tile = block_iter_count > Int32(0)
                     # 4. dV   = P.T  @ dO
                     # 5. dQ   = dS   @ K
+                    main_loop_iters = block_iter_count
                     # empty waits
                     # pipeline_S_P.sync_object_empty.wait(0, producer_phase_acc)
             else:
                 if is_leader_cta and process_tile:
                     accumulate_dK = False
+                    accumulate_dV = False
                     # -----------------------------------------------------------
                     ###### Prologue
                     # -----------------------------------------------------------
+                    # 1. S.T  = K @ Q.T
+                    # 2. dP.T = V @ dO.T
+                    # 3. dV   = P.T @ dO   (split_P_dS: issued in the main loop, after QK_1)
+                    # 1) S.T = K @ Q.T
                     handle_Q = pipeline_Q_consumer.wait_and_advance()
+                    if const_expr(self.split_P_dS):
+                        pipeline_S_P.sync_object_empty.wait(0, producer_phase_S)
+                    else:
+                        pipeline_S_P.sync_object_empty.wait(0, producer_phase_acc)
                     mma_qk_fn(B_idx=handle_Q.index)
                     pipeline_S_P.sync_object_full.arrive(0, pipeline_S_P.producer_mask, cta_group)
+                    if const_expr(self.split_P_dS):
+                        producer_phase_S ^= 1
+                    # 2) dP.T = V @ dO.T
                     pipeline_dO.consumer_wait(consumer_state_dO)
                     pipeline_dP.sync_object_empty.wait(0, producer_phase_acc)
                     pipeline_dQ.sync_object_empty.wait(0, producer_phase_acc)
                     producer_phase_acc ^= 1
                     # 3) dV = P.T @ dO
+                    if const_expr(not self.split_P_dS):
+                        pipeline_S_P.sync_object_empty.wait(0, producer_phase_acc)
+                        mma_pdo_fn(B_idx=consumer_state_dO.index, zero_init=True)
+                        pipeline_dO.consumer_release(consumer_state_dO)
+                        consumer_state_dO.advance()
                     # -----------------------------------------------------------
                     ###### MAIN LOOP
                     # -----------------------------------------------------------
+                    # Per trip t (the CTA owns a KV tile, so S and dP come out transposed):
+                    # 1.  S.T  = K    @ Q.T    (tile t+1)
+                    # 1b. dV  += P.T  @ dO     (tile t; split_P_dS only -- otherwise at 5)
+                    # 2.  dK  += dS.T @ Q
+                    # 3.  dQ   = dS   @ K
+                    # 4.  dP.T = V    @ dO.T   (tile t+1)
+                    # 5.  dV  += P.T  @ dO     (tile t+1; aliased layout only)
                     # For block sparsity, we use block_iter_count; for dense, use m_block range
                     # MMA doesn't need actual m_block indices, just the iteration count
                     for _ in cutlass.range(main_loop_iters, unroll=1):
                         # (1) S.T = K @ Q.T
                         handle_Q_next = pipeline_Q_consumer.wait_and_advance()
+                        if const_expr(self.split_P_dS):
+                            # S empty means "S_t loaded", not "P_t stored", so this QK
+                            # overlaps the softmax of tile t.
+                            pipeline_S_P.sync_object_empty.wait(0, producer_phase_S)
                         mma_qk_fn(B_idx=handle_Q_next.index)
                         pipeline_S_P.sync_object_full.arrive(
                             0, pipeline_S_P.producer_mask, cta_group
                         )
+                        if const_expr(self.split_P_dS):
+                            producer_phase_S ^= 1
+                            # (1b) dV += P_t.T @ dO_t, issued after QK_{t+1} so QK never
+                            # waits behind the P store.
+                            pipeline_P.consumer_wait(consumer_state_P)
+                            mma_pdo_fn(B_idx=consumer_state_dO.index, zero_init=not accumulate_dV)
+                            accumulate_dV = True
+                            pipeline_P.consumer_release(consumer_state_P)
+                            consumer_state_P.advance()
+                            pipeline_dO.consumer_release(consumer_state_dO)
+                            consumer_state_dO.advance()
                         # (2) dK += dS.T @ Q
                         pipeline_dS.consumer_wait(consumer_state_dS)
                         pipeline_dS.consumer_release(consumer_state_dS)
                         consumer_state_dS.advance()
+                        # (4) dP.T = V @ dO.T
                         pipeline_dO.consumer_wait(consumer_state_dO)
                         pipeline_dQ.sync_object_empty.wait(0, producer_phase_acc)
                         mma_dov_fn(B_idx=consumer_state_dO.index)
                         pipeline_dP.sync_object_full.arrive(0, pipeline_dP.producer_mask, cta_group)
+                        # (5) dV += P.T @ dO   (split_P_dS: done at (1b) of the next trip)
                         producer_phase_acc ^= 1
+                        if const_expr(not self.split_P_dS):
+                            pipeline_S_P.sync_object_empty.wait(0, producer_phase_acc)
+                            mma_pdo_fn(B_idx=consumer_state_dO.index, zero_init=False)
+                            pipeline_dO.consumer_release(consumer_state_dO)
+                            consumer_state_dO.advance()
                         handle_Q = handle_Q_next
+                    if const_expr(self.split_P_dS):
+                        # the last tile's dV += P.T @ dO; there is no further QK to overlap
+                        pipeline_P.consumer_wait(consumer_state_P)
+                        mma_pdo_fn(B_idx=consumer_state_dO.index, zero_init=not accumulate_dV)
+                        pipeline_P.consumer_release(consumer_state_P)
+                        consumer_state_P.advance()
+                        pipeline_dO.consumer_release(consumer_state_dO)
+                        consumer_state_dO.advance()
+                    else:
+                        pipeline_S_P.sync_object_full.arrive(
+                            0, pipeline_S_P.producer_mask, cta_group
+                        )
                     # signal to the epilogue that dV is ready
                     # pipeline_dKV.producer_acquire(producer_state_dKV)
         n_block,
         softmax_scale,
         seqlen_info,
+        aux_data: AuxData = AuxData(),
         fastdiv_mods=(None, None),
     ):
         """Apply forward score modification for SM100 backward pass."""
+        # In bwd, S is computed as K @ Q.T so dimensions are (tile_n, tile_m).
+        # With 2CTA, partition_C must see the full cluster tile so each CTA
+        # gets its own half of the tile.
+        cluster_tile_n = self.tile_n * self.cta_group_size
+        cluster_n_block = n_block // self.cta_group_size
+        cS = cute.make_identity_tensor((cluster_tile_n, self.tile_m))
+        cS = cute.domain_offset((cluster_n_block * cluster_tile_n, m_block * self.tile_m), cS)
         tScS = thr_mma_S.partition_C(cS)
         tScS_idx = thr_copy_t2r.partition_D(tScS)
             softmax_scale,
             self.vec_size,
             self.qk_acc_dtype,
+            aux_data,
             fastdiv_mods,
             seqlen_info,
             constant_q_idx=None,
         head_idx,
         softmax_scale,
         seqlen_info,
+        aux_data: AuxData = AuxData(),
         fastdiv_mods=(None, None),
     ):
         """Apply backward score modification (joint graph) for SM100."""
             softmax_scale,
             self.vec_size,
             self.qk_acc_dtype,
+            aux_data,
             fastdiv_mods,
             seqlen_info,
             constant_q_idx=None,
     @cute.jit
     def compute_loop(
         self,
+        thr_mma_S: cute.ThrMma,
+        thr_mma_dP: cute.ThrMma,
+        thr_mma_dV: cute.ThrMma,
+        thr_mma_dK: cute.ThrMma,
         tStS: cute.Tensor,
         tdPtdP: cute.Tensor,
         tdVtdV: cute.Tensor,
         pipeline_dS: PipelineAsync,
         pipeline_dKV: PipelineAsync,
         pipeline_dP: PipelineAsync,
+        pipeline_P: Optional[PipelineAsync],
         dS_cluster_empty_mbar_ptr: cute.Pointer,
         dS_cluster_full_mbar_ptr: cute.Pointer,
         dQaccum_empty_mbar_ptr: cute.Pointer,
         tiled_copy_r2s_dKV: Optional[cute.TiledCopy],
         mdK_semaphore: Optional[cute.Tensor],
         mdV_semaphore: Optional[cute.Tensor],
+        aux_data: AuxData = AuxData(),
         fastdiv_mods=(None, None),
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
     ):
         tileP_f32_like = self.cta_tiler[1] // 32 * self.v_dtype.width
         # tStS has shape ((128, 128), 1, 1), tStP has shape ((128, 64), 1, 1)
         tStP = cute.composition(tStS, (cute.make_layout((self.tile_n, tileP_f32_like)), 1, 1))
+        # Rebase the iterator explicitly (P may alias S), otherwise the tmem address is wrong
+        tStP = cute.make_tensor(
+            tStS.iterator + (self.tmem_P_offset - self.tmem_S_offset), tStP.layout
+        )
         tScS = thr_mma_S.partition_C(cute.make_identity_tensor(self.mma_tiler_kq[:2]))
         tScP = cute.composition(tScS, (cute.make_layout((self.tile_n, tileP_f32_like)), 1, 1))
         tdPtdS = cute.composition(tdPtdP, (cute.make_layout((self.tile_n, tileP_f32_like)), 1, 1))
+        tdPtdS = cute.make_tensor(
+            tdPtdP.iterator + (self.tmem_dS_offset - self.tmem_dP_offset), tdPtdS.layout
+        )
         tdPcdP = thr_mma_dP.partition_C(cute.make_identity_tensor(self.mma_tiler_vdo[:2]))
         tdPcdS = cute.composition(tdPcdP, (cute.make_layout((self.tile_n, tileP_f32_like)), 1, 1))
             cutlass.pipeline.PipelineUserType.Consumer, 1
         )
         # consumer_phase_S_P_dP = Int32(0)
+        consumer_state_S = pipeline.make_pipeline_state(
+            cutlass.pipeline.PipelineUserType.Consumer, 1
+        )
+        producer_state_P = pipeline.make_pipeline_state(
+            cutlass.pipeline.PipelineUserType.Producer, 1
+        )
         producer_state_dS = pipeline.make_pipeline_state(  # Our impl has shortcut for stage==1
             cutlass.pipeline.PipelineUserType.Producer, 1
         )
         while work_tile.is_valid_tile:
             n_block, head_idx, batch_idx, _ = work_tile.tile_idx
             seqlen = SeqlenInfoCls(batch_idx)
+            recompute_fastdiv_mods_q = const_expr(
+                aux_data.tensors is not None and (seqlen.has_cu_seqlens_q or seqlen.has_seqused_q)
+            )
+            recompute_fastdiv_mods_k = const_expr(
+                aux_data.tensors is not None and (seqlen.has_cu_seqlens_k or seqlen.has_seqused_k)
+            )
+            if const_expr(fastdiv_mods is not None and fastdiv_mods[0] is not None):
+                seqlen_q_divmod, seqlen_k_divmod = fastdiv_mods
+                fastdiv_mods = (
+                    seqlen_q_divmod
+                    if not recompute_fastdiv_mods_q
+                    else FastDivmodDivisor(seqlen.seqlen_q),
+                    seqlen_k_divmod
+                    if not recompute_fastdiv_mods_k
+                    else FastDivmodDivisor(seqlen.seqlen_k),
+                )
             m_block_min, m_block_max = block_info.get_m_block_min_max(
                 seqlen, n_block // self.cluster_shape_mnk[0]
             )
             mask = AttentionMaskCls(seqlen)
+            cluster_n_block = n_block // self.cta_group_size
             # TODO: condition mask_seqlen
             mask_fn = partial(
                 mask.apply_mask_sm100_transposed,
                 tScS_t2r=tScS_t2r,
                 t0ScS_t2r=t0ScS_t2r,
+                n_block=cluster_n_block,
                 mask_seqlen=True,
                 mask_causal=self.is_causal,
                 mask_local=self.is_local,
                 mask_mod=self.mask_mod,
                 batch_idx=batch_idx,
                 head_idx=head_idx,
+                aux_data=aux_data,
                 fastdiv_mods=fastdiv_mods,
             )
             # prefetch_LSE = not self.is_causal
             prefetch_LSE = False
+            curr_q_cnt = Int32(0)
+            curr_q_idx = None
+            curr_full_cnt = Int32(0)
+            curr_full_idx = None
+            loop_count = m_block_max - m_block_min
+            process_tile = (
+                const_expr(not self.is_local and not self.is_varlen_q) or m_block_min < m_block_max
+            )
             if const_expr(self.use_block_sparsity):
+                assert blocksparse_tensors is not None
                 (
                     curr_q_cnt,
                     curr_q_idx,
                     blocksparse_tensors,
                     batch_idx,
                     head_idx,
+                    n_block // self.kv_subtile_factor,
+                    q_subtile_factor=self.q_subtile_factor,
                     m_block_max=m_block_max,
                 )
                 process_tile = loop_count > Int32(0)
             # Mainloop
             # Block sparsity: iterate over sparse m_block count and derive actual m_block
             # from Q_IDX/FULL_Q_IDX tensors. Dense: iterate m_block_min..m_block_max directly.
             for iter_idx in cutlass.range(loop_count, unroll=1):
+                m_block = m_block_min + iter_idx
+                m_block_oob = False
+                is_full_block = False
                 if const_expr(self.use_block_sparsity):
                     m_block, is_full_block = get_m_block_from_iter_bwd(
                         iter_idx,
                         curr_q_idx,
                         curr_full_cnt,
                         curr_full_idx,
+                        q_subtile_factor=self.q_subtile_factor,
                         m_block_max=m_block_max,
                     )
                     m_block_oob = m_block >= m_block_max
                 # Prefetch 1 stage of LSE
                 pipeline_LSE.consumer_wait(consumer_state_LSE)
+                tSrLSE_s2r = cute.make_rmem_tensor(tScS_t2r[None, 0, 0, 0].shape, Float32)
                 if const_expr(prefetch_LSE and not self.shuffle_LSE):
                     cute.autovec_copy(tSsLSE[None, 0, 0, 0, consumer_state_LSE.index], tSrLSE_s2r)
+                if const_expr(self.split_P_dS):
+                    pipeline_S_P.consumer_wait(consumer_state_S)
+                else:
+                    pipeline_S_P.consumer_wait(consumer_state_S_P_dP)
                 # pipeline_S_P.sync_object_full.wait(0, consumer_phase_S_P_dP)
                 #### TMEM->RMEM (Load S from TMEM)
+                tSrS_t2r = cute.make_rmem_tensor(tScS_t2r.shape, Float32)
                 cute.copy(thr_copy_t2r, tStS_t2r, tSrS_t2r)
+                if const_expr(self.split_P_dS):
+                    # S is in registers: release it before the softmax so the MMA warp can
+                    # issue the next QK into the slot
+                    cute.arch.fence_view_async_tmem_load()
+                    cute.arch.sync_warp()
+                    with cute.arch.elect_one():
+                        pipeline_S_P.consumer_release(consumer_state_S)
+                    consumer_state_S.advance()
+                elif const_expr(self.tile_hdim == 192):
                     # Signal S tmem load completion using pipeline_S_P when hdim 192
                     # dP is overlapped with S
                     cute.arch.fence_view_async_tmem_load()
                         n_block,
                         softmax_scale,
                         seqlen,
+                        aux_data,
                         fastdiv_mods,
                     )
                 #### P = exp(S - LSE)
                 # ---------------------------------------------
                 lane_idx = cute.arch.lane_idx()
+                tSrP_r2t_f32 = cute.make_rmem_tensor(tScP_r2t.shape, Float32)  # 64
                 tSrP_r2t = cute.recast_tensor(tSrP_r2t_f32, self.q_dtype)
                 for stage in cutlass.range_constexpr(num_stages):
                     tSrS_cur = tSrS_t2r[None, stage, 0, 0]
                         tSrS_cur[2 * v + 1] = cute.math.exp2(tSrS_cur[2 * v + 1], fastmath=True)
                     utils.cvt_f16(tSrS_cur, tSrP_r2t[None, stage, 0, 0])
                     if const_expr(stage == 0):
+                        if const_expr(self.split_P_dS):
+                            # wait until the previous pdo UMMA has consumed the P slot
+                            pipeline_P.producer_acquire(producer_state_P)
+                        else:
+                            cute.arch.fence_view_async_tmem_load()
+                            # P overwrites S lanes another warp may still be loading, so
+                            # every warp must have loaded S before any stores P.
+                            self.compute_sync_barrier.arrive_and_wait()
                     cute.copy(
                         thr_copy_r2t,
                         tSrP_r2t_f32[None, stage, None, None],
                 cute.arch.fence_view_async_tmem_store()
                 cute.arch.fence_view_async_shared()
+                # The P (or S_P) and LSE mbarriers expect one arrival per compute warp
+                if const_expr(self.split_P_dS):
+                    cute.arch.sync_warp()
+                else:
+                    self.compute_sync_barrier.arrive_and_wait()
+                if const_expr(self.split_P_dS):
+                    with cute.arch.elect_one():
+                        pipeline_P.producer_commit(producer_state_P)
+                    producer_state_P.advance()
+                elif const_expr(not self.tile_hdim == 192):
                     # Signal tmem store P completion with pipeline_S_P
                     with cute.arch.elect_one():
                         pipeline_S_P.consumer_release(consumer_state_S_P_dP)
                         # pipeline_S_P.sync_object_empty.arrive(0, pipeline_S_P.consumer_mask)
+                # Normally we'd need syncwarp here since only 1 thread will signal in
+                # consumer_release, but we already have the barrier or sync_warp before this
                 pipeline_LSE.consumer_release(consumer_state_LSE)
                 consumer_state_LSE.advance()
                 # ---------------------------------------------
                 ##### dS.T = P.T * (dP.T - Psum)
                 for stage in cutlass.range_constexpr(num_stages):
+                    tdPrdP_t2r = cute.make_rmem_tensor(tScS_t2r[None, 0, None, None].shape, Float32)
                     cute.copy(thr_copy_t2r, tdPtdP_t2r[None, stage, None, None], tdPrdP_t2r)
                     cute.arch.fence_view_async_tmem_load()
+                    if const_expr(not self.split_P_dS):
+                        # dS overwrites dP with the same cross-warp lane sharing as P over S
+                        self.compute_sync_barrier.arrive_and_wait()
                     tdPrdP_cur = tdPrdP_t2r[None, 0, 0]
                     tSrS_cur = tSrS_t2r[None, stage, 0, 0]
                     tSsdPsum_cur = tSsdPsum[None, stage, 0, 0, consumer_state_dPsum.index]
                                 utils.shuffle_sync(tSrdPsum, offset=2 * v),
                                 utils.shuffle_sync(tSrdPsum, offset=2 * v + 1),
                             )
+                        tdPrdP_cur[2 * v], tdPrdP_cur[2 * v + 1] = cute.arch.sub_packed_f32x2(
+                            (tdPrdP_cur[2 * v], tdPrdP_cur[2 * v + 1]), dPsum_pair
                         )
                         tdPrdP_cur[2 * v], tdPrdP_cur[2 * v + 1] = cute.arch.mul_packed_f32x2(
                             (tSrS_cur[2 * v], tSrS_cur[2 * v + 1]),
                     if const_expr(self.score_mod_bwd is not None):
                         tSrS_pre_cur = tSrS_pre[None, stage, 0, 0]
+                        cluster_tile_n = self.tile_n * self.cta_group_size
+                        cluster_n_block = n_block // self.cta_group_size
+                        cS_bwd = cute.make_identity_tensor((cluster_tile_n, self.tile_m))
                         cS_bwd = cute.domain_offset(
+                            (cluster_n_block * cluster_tile_n, m_block * self.tile_m),
+                            cS_bwd,
                         )
                         tScS_bwd = thr_mma_S.partition_C(cS_bwd)
                         tScS_idx_bwd = thr_copy_t2r.partition_D(tScS_bwd)
                             head_idx,
                             softmax_scale,
                             seqlen,
+                            aux_data,
                             fastdiv_mods,
                         )
                         # Zero out OOB positions (kv_idx >= seqlen_k) after score_mod_bwd
                     cute.autovec_copy(tdPrdS_xchg, tRS_sdS_xchg[None, 0])
                 cute.arch.fence_view_async_shared()
+                # The dS and dPsum mbarriers expect one arrival per compute warp
+                if const_expr(self.split_P_dS):
+                    cute.arch.sync_warp()
+                else:
+                    self.compute_sync_barrier.arrive_and_wait()
+                # Normally we'd need syncwarp here since only 1 thread will signal in
+                # consumer_release, but we already have the barrier or sync_warp before this
                 pipeline_dPsum.consumer_release(consumer_state_dPsum)
                 consumer_state_dPsum.advance()
                 # when 2cta hdim 128, pipeline_dS also signals S tmem load completion so is deferred
             tile_scheduler.advance_to_next_work()
             work_tile = tile_scheduler.get_current_work()
+    @cute.jit
+    def _dq_semaphore_lock_value(
+        self,
+        iter_idx: Int32,
+        curr_q_cnt: Int32,
+        curr_dq_write_order: Optional[cute.Tensor],
+        curr_dq_write_order_full: Optional[cute.Tensor],
+        blocksparse_tensors: Optional[BlockSparseTensors],
+        block_info: BlockInfo,
+        seqlen,
+        m_block: Int32,
+        n_block: Int32,
+    ) -> Int32:
+        lock_value = n_block
+        if const_expr(self.spt):
+            n_block_max_for_m_block = block_info.get_n_block_max_for_m_block(seqlen, m_block)
+            lock_value = n_block_max_for_m_block - 1 - n_block
+        if const_expr(self.use_block_sparsity):
+            assert blocksparse_tensors is not None
+            if const_expr(blocksparse_tensors.dq_write_order is not None):
+                sparse_iter = iter_idx // self.q_subtile_factor
+                if sparse_iter < curr_q_cnt:
+                    assert curr_dq_write_order is not None
+                    lock_value = curr_dq_write_order[sparse_iter]
+                else:
+                    assert curr_dq_write_order_full is not None
+                    lock_value = curr_dq_write_order_full[sparse_iter - curr_q_cnt]
+                if const_expr(self.kv_subtile_factor > self.cta_group_size):
+                    groups_per_sparse_block = self.kv_subtile_factor // self.cta_group_size
+                    local_group = n_block % groups_per_sparse_block
+                    if const_expr(self.spt):
+                        # [NOTE] KV_subtile determ + spt
+                        # dq_write_order stores one rank per sparse block; each physical tile
+                        # derives its slot as rank * groups_per_sparse_block + (n_block % groups_per_sparse_block).
+                        # W/ kv_subtile the tail sparse column can have a physical tile that is
+                        # never scheduled; e.g. kv_tile = 128, seqlen = 1023, KV_Block = 384 -> 3
+                        # sparse blocks. The last block is covered in [768, 896), [896, 1024) and
+                        # then [1024, 1152) which no CTA ever runs. Since the highest tile gets
+                        # the lowest lock value under spt, we would hang!
+                        # In this case we locally reverse, [N+2, N+1, N*] where N* is not scheduled
+                        # -> [N+1, N, N+2*]. Ahh but won't the vacant N+2 slot stall the next sparse
+                        # block's CTAs? It will, so the writer holding N+1 bumps the increment by 2
+                        # instead of 1 (i.e. 1 + #unscheduled) :)
+                        total_groups = cute.ceil_div(
+                            seqlen.seqlen_k, self.tile_n * self.cta_group_size
+                        )
+                        groups_in_own_block = cutlass.min(
+                            groups_per_sparse_block,
+                            total_groups
+                            - (n_block // groups_per_sparse_block) * groups_per_sparse_block,
+                        )
+                        local_group = groups_in_own_block - 1 - local_group
+                    lock_value = lock_value * groups_per_sparse_block + local_group
+        return lock_value
     @cute.jit
     def dQacc_reduce(
         self,
         mdQaccum: cute.Tensor,
         sdQaccum: cute.Tensor,
+        thr_mma_dQ: cute.ThrMma,
         tdQtdQ: cute.Tensor,
         pipeline_dQ: PipelineAsync,
         dQaccum_empty_mbar_ptr: Optional[cute.Pointer],
         while work_tile.is_valid_tile:
             n_block, head_idx, batch_idx, _ = work_tile.tile_idx
             n_block_cta_group = n_block // self.cta_group_size  # for 2cta
+            n_block_sparse = n_block // self.kv_subtile_factor
             seqlen = SeqlenInfoCls(batch_idx)
             m_block_min, m_block_max = block_info.get_m_block_min_max(seqlen, n_block_cta_group)
             if const_expr(not seqlen.has_cu_seqlens_q):
             )
             if const_expr(self.deterministic):
+                assert mdQ_semaphore is not None
                 mdQ_semaphore_cur = mdQ_semaphore[None, None, head_idx, batch_idx]
+            delay_semaphore_release = not self.tile_hdim == 192 and not self.use_block_sparsity
+            dq_sem_release_inc = Int32(1)
+            if const_expr(
+                self.deterministic
+                and self.use_block_sparsity
+                and self.spt
+                and self.kv_subtile_factor > self.cta_group_size
+            ):
+                # A truncated tail block's last writer releases the missing increments,
+                # see: [NOTE] KV_subtile determ + spt
+                groups_per_sparse_block = self.kv_subtile_factor // self.cta_group_size
+                total_groups = cute.ceil_div(seqlen.seqlen_k, self.tile_n * self.cta_group_size)
+                tail_sparse_block_idx = (total_groups - 1) // groups_per_sparse_block
+                groups_in_tail = total_groups - tail_sparse_block_idx * groups_per_sparse_block
+                is_tail_bridge_group = (
+                    n_block_cta_group // groups_per_sparse_block == tail_sparse_block_idx
+                    and n_block_cta_group % groups_per_sparse_block == 0
+                )
+                if is_tail_bridge_group:
+                    dq_sem_release_inc = Int32(1) + groups_per_sparse_block - groups_in_tail
+            curr_q_cnt = Int32(0)
+            curr_q_idx = None
+            curr_full_cnt = Int32(0)
+            curr_full_idx = None
+            curr_dq_write_order = None
+            curr_dq_write_order_full = None
+            loop_count = m_block_max - m_block_min
+            process_tile = (
+                const_expr(not self.is_local and not self.is_varlen_q) or m_block_min < m_block_max
+            )
             if const_expr(self.use_block_sparsity):
+                assert blocksparse_tensors is not None
                 (
                     curr_q_cnt,
                     curr_q_idx,
                     blocksparse_tensors,
                     batch_idx,
                     head_idx,
+                    n_block_sparse,
+                    q_subtile_factor=self.q_subtile_factor,
                     m_block_max=m_block_max,
                 )
                 process_tile = loop_count > Int32(0)
+            if const_expr(self.deterministic and self.use_block_sparsity):
+                assert blocksparse_tensors is not None
+                if const_expr(blocksparse_tensors.dq_write_order is not None):
+                    assert blocksparse_tensors.dq_write_order is not None
+                    curr_dq_write_order = blocksparse_tensors.dq_write_order[
+                        batch_idx, head_idx, n_block_sparse, None
+                    ]
+                    if const_expr(blocksparse_tensors.dq_write_order_full is not None):
+                        assert blocksparse_tensors.dq_write_order_full is not None
+                        curr_dq_write_order_full = blocksparse_tensors.dq_write_order_full[
+                            batch_idx, head_idx, n_block_sparse, None
+                        ]
             # dQacc_reduce mainloop
             # Block sparsity: iterate over sparse m_block count and derive actual m_block
             # from Q_IDX/FULL_Q_IDX tensors. Dense: iterate m_block_min..m_block_max directly.
             for iter_idx in cutlass.range(loop_count, unroll=1):
+                m_block = m_block_min + iter_idx
+                m_block_oob_upper = False
                 if const_expr(self.use_block_sparsity):
                     m_block, _ = get_m_block_from_iter_bwd(
                         iter_idx,
                         curr_q_idx,
                         curr_full_cnt,
                         curr_full_idx,
+                        q_subtile_factor=self.q_subtile_factor,
                         m_block_max=m_block_max,
                     )
+                    m_block_oob_upper = m_block >= m_block_max
                 pipeline_dQ.consumer_wait(dQ_consumer_state)
                 # TMEM -> RMEM
+                tdQrdQ_t2r = cute.make_rmem_tensor(tdQrdQ_t2r_shape, Float32)
                 cute.copy(thr_copy_t2r, tdQtdQ_t2r, tdQrdQ_t2r)
                 cute.arch.fence_view_async_tmem_load()
                 cute.arch.sync_warp()
                     pipeline_dQ.consumer_release(dQ_consumer_state)
                 dQ_consumer_state.advance()
+                if m_block_max > 0:
+                    m_block = cutlass.min(m_block, m_block_max - 1)
                 gdQaccum_cur = gdQaccum[None, None, m_block]
                 tdQrdQ_shape = (
                     cute.arch.fence_view_async_shared()
                     # semaphore acquire
                     if const_expr(self.deterministic and stage == 0):
+                        if not m_block_oob_upper:
+                            lock_value = self._dq_semaphore_lock_value(
+                                iter_idx,
+                                curr_q_cnt,
+                                curr_dq_write_order,
+                                curr_dq_write_order_full,
+                                blocksparse_tensors,
+                                block_info,
+                                seqlen,
+                                m_block,
+                                n_block_cta_group,
+                            )
+                            barrier.wait_eq(
+                                mdQ_semaphore_cur[(m_block, None)].iterator,
+                                tidx,
+                                cta_rank_in_cluster,
+                                lock_value,
                             )
                     self.reduce_sync_barrier.arrive_and_wait()
                     # Copy from shared memory to global memory
+                    if is_tma_warp and not m_block_oob_upper:
                         with cute.arch.elect_one():
                             copy_utils.cpasync_reduce_bulk_add_f32(
                                 sdQaccum[None, smem_idx].iterator,
                             )
                         cute.arch.cp_async_bulk_commit_group()
                         cute.arch.cp_async_bulk_wait_group(self.sdQaccum_stage - 1, read=read_flag)
+                    elif is_tma_warp:
+                        # Drain pending TMA stores so SMEM buffers are safe to reuse
+                        cute.arch.cp_async_bulk_wait_group(0, read=read_flag)
                     self.reduce_sync_barrier.arrive_and_wait()
                     dQ_tma_store_producer_state.advance()
                     if const_expr(self.deterministic and stage == 0 and delay_semaphore_release):
                         if m_block > m_block_min:
                             barrier.arrive_inc(
                 # NOTE: arrive_inc calls red_release which issues membar
                 if const_expr(self.deterministic and not delay_semaphore_release):
                     if const_expr(self.sdQaccum_stage > 1 and not self.tile_hdim == 192):
+                        if is_tma_warp and not m_block_oob_upper:
                             cute.arch.cp_async_bulk_wait_group(0, read=read_flag)
                         self.reduce_sync_barrier.arrive_and_wait()
+                    if not m_block_oob_upper:
+                        barrier.arrive_inc(
+                            mdQ_semaphore_cur[m_block, None].iterator,
+                            tidx,
+                            cta_rank_in_cluster,
+                            dq_sem_release_inc,
+                        )
             if process_tile:
                 if is_tma_warp:
                     )
             if const_expr(
+                self.deterministic
+                and not self.spt
+                and not self.use_block_sparsity
+                and block_info.window_size_left is not None
             ):
                 m_block_global_max = cute.ceil_div(seqlen.seqlen_q, self.tile_m)
                 for m_block in cutlass.range(m_block_max, m_block_global_max, unroll=1):
             tile_scheduler.advance_to_next_work()
             work_tile = tile_scheduler.get_current_work()
+        if const_expr(not self.deterministic):
+            cute.arch.cp_async_bulk_wait_group(0, read=True)
     @cute.jit
     def epilogue_dKV(
         self,
         head_idx: Int32,
         n_block: Int32,
         seqlen,
+        thr_mma_dV: cute.ThrMma,
+        thr_mma_dK: cute.ThrMma,
         tdVtdV: cute.Tensor,
         tdKtdK: cute.Tensor,
         mdV: cute.Tensor,
         mdV_cur = seqlen.offset_batch_K(mdV, batch_idx, dim=3)[None, None, head_idx]
         mdK_cur = seqlen.offset_batch_K(mdK, batch_idx, dim=3)[None, None, head_idx]
+        # The gmem store below has no head-dim OOB protection from TMA, so when head_dim is
+        # not a multiple of the tile we predicate each 128-bit store vector (8 bf16 columns).
+        # The store's predicate granularity is one tmem-load repetition, so shrink the
+        # repetition to 8 columns for that case; 16 keeps fewer tmem loads otherwise.
+        tmem_load_atom_dV = cute.make_copy_atom(
+            tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(8 if self.check_hdim_v_oob else 16)),
+            Float32,
+        )
+        tmem_load_atom_dK = cute.make_copy_atom(
+            tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(8 if self.check_hdim_oob else 16)),
+            Float32,
         )
         # dV
         pipeline_dKV.consumer_wait(consumer_state_dKV)
+        tiled_tmem_ld_dV = tcgen05.make_tmem_copy(tmem_load_atom_dV, tdVtdV)
         thr_tmem_ld_dV = tiled_tmem_ld_dV.get_slice(tidx)
         tdVtdV_t2r_p = thr_tmem_ld_dV.partition_S(tdVtdV)
         tdVcdV_t2r_p = thr_tmem_ld_dV.partition_D(tdVcdV_tensor)
         tdVcdV_t2r = self.split_wg(tdVcdV_t2r_p, wg_idx, num_wg)
+        tdVrdV_t2r = cute.make_rmem_tensor(tdVcdV_t2r.shape, Float32)
         cute.copy(thr_tmem_ld_dV, tdVtdV_t2r, tdVrdV_t2r)
         cute.arch.fence_view_async_tmem_load()
             tiler_mn=tiled_tmem_ld_dV.tiler_mn,
         )
+        tdVrdV_r2s = cute.make_rmem_tensor(tdVrdV_t2r.shape, self.dv_dtype)
         for i in cutlass.range_constexpr(cute.size(tdVrdV_t2r, mode=[1])):
             dV_vec = tdVrdV_t2r[(None, i, 0, 0)].load()
             tdVrdV_r2s[(None, i, 0, 0)].store(dV_vec.to(self.dv_dtype))
         tdVgdV_r2g_p = thr_tmem_ld_dV.partition_D(tdVgdV)
         tdVgdV_r2g = self.split_wg(tdVgdV_r2g_p, wg_idx, num_wg)
+        tdVpdV = None
+        if const_expr(self.check_hdim_v_oob):
+            tdVpdV = self.predicate_hdim(tdVcdV_t2r, self.head_dim_v)
         if tidx < seqlen.seqlen_k - self.tile_n * n_block:
+            cute.copy(tiled_gmem_store_dV, tdVrdV_r2s, tdVgdV_r2g, pred=tdVpdV)
         cute.arch.sync_warp()
         with cute.arch.elect_one():
         # dK
         pipeline_dKV.consumer_wait(consumer_state_dKV)
+        tiled_tmem_ld_dK = tcgen05.make_tmem_copy(tmem_load_atom_dK, tdKtdK)
         thr_tmem_ld_dK = tiled_tmem_ld_dK.get_slice(tidx)
         tdKtdK_t2r_p = thr_tmem_ld_dK.partition_S(tdKtdK)
         tdKcdK_t2r_p = thr_tmem_ld_dK.partition_D(tdKcdK_tensor)
         tdKcdK_t2r = self.split_wg(tdKcdK_t2r_p, wg_idx, num_wg)
+        tdKrdK_t2r = cute.make_rmem_tensor(tdKcdK_t2r.shape, Float32)
         cute.copy(tiled_tmem_ld_dK, tdKtdK_t2r, tdKrdK_t2r)
         cute.arch.fence_view_async_tmem_load()
             tiler_mn=tiled_tmem_ld_dK.tiler_mn,
         )
+        tdKrdK_r2s = cute.make_rmem_tensor(tdKrdK_t2r.shape, self.dk_dtype)
         for i in cutlass.range_constexpr(cute.size(tdKrdK_t2r, mode=[1])):
             dK_vec = tdKrdK_t2r[(None, i, 0, 0)].load() * softmax_scale
         tdKgdK_r2g_p = thr_tmem_ld_dK.partition_D(tdKgdK)
         tdKgdK_r2g = self.split_wg(tdKgdK_r2g_p, wg_idx, num_wg)
+        tdKpdK = None
+        if const_expr(self.check_hdim_oob):
+            tdKpdK = self.predicate_hdim(tdKcdK_t2r, self.head_dim)
         if tidx < seqlen.seqlen_k - self.tile_n * n_block:
+            cute.copy(tiled_gmem_store_dK, tdKrdK_r2s, tdKgdK_r2g, pred=tdKpdK)
         cute.arch.sync_warp()
         with cute.arch.elect_one():
             pipeline_dKV.consumer_release(consumer_state_dKV)
         return consumer_state_dKV
+    @cute.jit
+    def predicate_hdim(self, tcoord: cute.Tensor, limit: cutlass.Constexpr[int]) -> cute.Tensor:
+        """Per-store-vector head-dim predicate for the non-TMA dK/dV epilogue store.
+        ``tcoord`` is the (tile_n, hdim) identity tensor partitioned like the register data,
+        shape ((rep, 1), chunks, 1, 1) with one contiguous ``rep``-wide column chunk per copy
+        vector; the predicate drops the vector mode and marks each chunk whose first column
+        is inside ``limit``. Valid because ``limit`` is a multiple of ``rep``.
+        """
+        tpred = cute.make_rmem_tensor(
+            cute.make_layout((1, cute.size(tcoord, mode=[1]), 1, 1)), cutlass.Boolean
+        )
+        for i in cutlass.range_constexpr(cute.size(tcoord, mode=[1])):
+            tpred[0, i, 0, 0] = tcoord[0, i, 0, 0][1] < limit
+        return tpred
     @cute.jit
     def epilogue_dK_or_dV_tma(
         self,
         head_idx: Int32,
         n_block: Int32,
         seqlen,
+        thr_mma: cute.ThrMma,
         tdKVtdKV: cute.Tensor,
         mdKV: cute.Tensor,
         sdKV: cute.Tensor,
         deterministic_KV = self.deterministic and self.qhead_per_kvhead > 1
         if const_expr(deterministic_KV):
+            assert mdKV_semaphore is not None
             mdKV_semaphore_cur = mdKV_semaphore[n_block, None, head_idx_kv, batch_idx]
         if const_expr(not self.dKV_postprocess):
             if const_expr(num_epi_stages > 1):
                 tdKVcdKV_t2r = tdKVcdKV_t2r[None, epi_stage]
+            tdKVrdKV_t2r = cute.make_rmem_tensor(tdKVcdKV_t2r.shape, Float32)
             assert cute.size(tdKVrdKV_t2r) == cute.size(tdKVtdKV_t2r) // cute.arch.WARP_SIZE, (
                 "RMEM<->TMEM fragment size mismatch"
                     tdKVrdKV_t2r[2 * i], tdKVrdKV_t2r[2 * i + 1] = cute.arch.mul_packed_f32x2(
                         (tdKVrdKV_t2r[2 * i], tdKVrdKV_t2r[2 * i + 1]), (scale, scale)
                     )
+            tdKVrdKV = cute.make_rmem_tensor(tdKVrdKV_t2r.shape, dtype)  # (32 columns)
             tdKVrdKV.store(tdKVrdKV_t2r.load().to(dtype))
             # RMEM -> SMEM -- copy, fence and barrier

build/torch-cuda/flash_bwd_sm120.py ADDED Viewed

	@@ -0,0 +1,55 @@

+# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+# SM120 (Blackwell GeForce / DGX Spark) backward pass.
+#
+# SM120 uses the same SM80-era MMA instructions (mma.sync.aligned.m16n8k16) but has
+# a smaller shared memory capacity (99 KB vs 163 KB on SM80). This module subclasses
+# FlashAttentionBackwardSm80 and overrides the SMEM capacity check accordingly.
+import cutlass
+import cutlass.utils as utils_basic
+from .flash_bwd import FlashAttentionBackwardSm80
+class FlashAttentionBackwardSm120(FlashAttentionBackwardSm80):
+    @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 on SM120.
+        Same logic as SM80 but uses SM120's shared memory capacity (99 KB).
+        """
+        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
+        # Shared memory usage: Q tile + dO tile + K tile + V tile
+        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
+        # SM120 has 99 KB shared memory (vs 163 KB on SM80)
+        smem_capacity = utils_basic.get_smem_capacity_in_bytes("sm_120")
+        if smem_usage > smem_capacity:
+            return False
+        return True

build/torch-cuda/flash_bwd_sm90.py CHANGED Viewed

@@ -24,10 +24,17 @@ from .seqlen_info import SeqlenInfoQK
 from .block_info import BlockInfo
 from . import pipeline
 from .quack.cute_dsl_utils import ParamsBase
-from .tile_scheduler import TileSchedulerArguments, SingleTileScheduler
 from .named_barrier import NamedBarrierBwd
 from .softmax import apply_score_mod_inner, apply_score_mod_bwd_inner
 from .block_sparsity import BlockSparseTensors
 from .block_sparse_utils import (
     get_total_q_block_count_bwd,
     produce_block_sparse_q_loads_bwd_sm90,
@@ -46,6 +53,8 @@ class FlashAttentionBackwardSm90:
         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,
@@ -63,11 +72,12 @@ class FlashAttentionBackwardSm90:
         score_mod_bwd: cutlass.Constexpr | None = None,
         mask_mod: cutlass.Constexpr | None = None,
         has_aux_tensors: cutlass.Constexpr = False,
-        subtile_factor: cutlass.Constexpr[int] = 1,
     ):
         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
@@ -77,7 +87,8 @@ class FlashAttentionBackwardSm90:
         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
@@ -92,23 +103,23 @@ class FlashAttentionBackwardSm90:
         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
         if qhead_per_kvhead > 1:
             assert self.same_hdim_kv, "GQA backward requires head_dim == head_dim_v"
-            assert self.num_mma_warp_groups == 2, "GQA backward assumes 2 warp groups"
         # 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
@@ -118,12 +129,18 @@ class FlashAttentionBackwardSm90:
         self.score_mod_bwd = score_mod_bwd
         self.mask_mod = mask_mod
         self.has_aux_tensors = has_aux_tensors
-        self.subtile_factor = subtile_factor
         if cutlass.const_expr(has_aux_tensors):
             self.vec_size: cutlass.Constexpr = 1
         else:
             self.vec_size: cutlass.Constexpr = 4
         self.qk_acc_dtype = Float32
     @staticmethod
     def can_implement(
@@ -182,32 +199,64 @@ class FlashAttentionBackwardSm90:
         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
         )
         # dKVaccum for GQA epilogue - reuses sV+sK memory recast as f32
         # TODO: assert that sVaccum and sKaccum don't overflow smem
     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,
@@ -215,12 +264,11 @@ class FlashAttentionBackwardSm90:
             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 = [
@@ -232,9 +280,8 @@ class FlashAttentionBackwardSm90:
                 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,
@@ -242,7 +289,8 @@ class FlashAttentionBackwardSm90:
             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,
@@ -250,8 +298,8 @@ class FlashAttentionBackwardSm90:
             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
@@ -305,23 +353,25 @@ class FlashAttentionBackwardSm90:
         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,
         mdQ_semaphore: Optional[cute.Tensor] = None,
         mdK_semaphore: Optional[cute.Tensor] = None,
         mdV_semaphore: Optional[cute.Tensor] = None,
-        aux_tensors: Optional[list] = None,
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
     ):
-        assert mdQ_semaphore is None and mdK_semaphore is None and mdV_semaphore is None, (
-            "determinism not supported yet for Sm90"
-        )
         self._check_type(
             *(
@@ -330,23 +380,45 @@ class FlashAttentionBackwardSm90:
             )
         )
         mQ, mK, mV, mdO, mLSE, mdPsum, mdQaccum, mdK, mdV = [
             assume_tensor_aligned(t) 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, mdO = [layout_utils.select(t, layout_transpose) for t in (mQ, mK, mV, mdO)]
         if const_expr(self.qhead_per_kvhead == 1):
-            mdK, mdV = [layout_utils.select(t, layout_transpose) for t in (mdK, mdV)]
         else:
-            accum_transpose = [2, 1, 0]  # (b, n, s*h) -> (s*h, n, b)
             mdK, mdV = [layout_utils.select(t, accum_transpose) for t in (mdK, mdV)]
-        LSE_dPsum_dQaccum_transpose = [2, 1, 0]  # (b, n, s) -> (s, n, b)
         mLSE, mdPsum, mdQaccum = [
             layout_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
@@ -354,10 +426,25 @@ class FlashAttentionBackwardSm90:
         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()
@@ -374,7 +461,7 @@ class FlashAttentionBackwardSm90:
         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
         )
         self.tma_copy_bytes["dKacc"] = self.tile_n * self.tile_hdim * Float32.width // 8
         self.tma_copy_bytes["dVacc"] = self.tile_n * self.tile_hdimv * Float32.width // 8
@@ -404,38 +491,60 @@ class FlashAttentionBackwardSm90:
             (self.tile_m, self.tile_hdimv),
         )
         if const_expr(self.qhead_per_kvhead == 1):
             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),
             )
         else:
             tma_atom_dK = tma_atom_dV = tma_tensor_dK = tma_tensor_dV = None
-        TileScheduler = SingleTileScheduler
         tile_sched_args = TileSchedulerArguments(
             cute.ceil_div(cute.size(mK.shape[0]), self.tile_n),
             cute.size(mQ.shape[2]),
-            cute.size(mQ.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)
@@ -448,7 +557,7 @@ class FlashAttentionBackwardSm90:
             softmax_scale_log2 = LOG2_E
         fastdiv_mods = None
-        if const_expr(aux_tensors is not None):
             seqlen_q = cute.size(mQ.shape[0])
             seqlen_k = cute.size(mK.shape[0])
             seqlen_q_divmod = FastDivmodDivisor(seqlen_q)
@@ -461,6 +570,11 @@ class FlashAttentionBackwardSm90:
         self.use_block_sparsity = cutlass.const_expr(blocksparse_tensors is not None)
         self.kernel(
             tma_tensor_Q,
             tma_tensor_K,
@@ -477,6 +591,10 @@ class FlashAttentionBackwardSm90:
             mLSE,
             mdPsum,
             mdQaccum,
             self.sQ_layout,
             self.sK_layout,
             self.sV_layout,
@@ -493,15 +611,21 @@ class FlashAttentionBackwardSm90:
             tile_sched_params,
             TileScheduler,
             SharedStorage,
-            aux_tensors,
             fastdiv_mods,
             blocksparse_tensors,
             qhead_per_kvhead_divmod,
         ).launch(
             grid=grid_dim,
             block=[self.num_threads, 1, 1],
             stream=stream,
             min_blocks_per_mp=1,
         )
     @cute.kernel
@@ -522,6 +646,10 @@ class FlashAttentionBackwardSm90:
         mLSE: cute.Tensor,
         mdPsum: cute.Tensor,
         mdQaccum: cute.Tensor,
         sQ_layout: cute.ComposedLayout,
         sK_layout: cute.ComposedLayout,
         sV_layout: cute.ComposedLayout,
@@ -538,19 +666,23 @@ class FlashAttentionBackwardSm90:
         tile_sched_params: ParamsBase,
         TileScheduler: cutlass.Constexpr[Callable],
         SharedStorage: cutlass.Constexpr[Callable],
-        aux_tensors: Optional[list] = None,
         fastdiv_mods=(None, None),
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
         qhead_per_kvhead_divmod: Optional[FastDivmodDivisor] = None,
     ):
         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)
@@ -604,25 +736,27 @@ class FlashAttentionBackwardSm90:
             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,
             swap_AB=self.SdP_swapAB,
         )
         TileSchedulerCls = partial(TileScheduler.create, tile_sched_params)
@@ -663,18 +797,20 @@ class FlashAttentionBackwardSm90:
                     TileSchedulerCls,
                     SeqlenInfoCls,
                     blocksparse_tensors,
                 )
         else:
-            cute.arch.setmaxregister_increase(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,
@@ -697,11 +833,24 @@ class FlashAttentionBackwardSm90:
                 SeqlenInfoCls,
                 AttentionMaskCls,
                 TileSchedulerCls,
-                aux_tensors,
                 fastdiv_mods,
                 blocksparse_tensors,
                 qhead_per_kvhead_divmod,
             )
     @cute.jit
     def load(
@@ -749,18 +898,22 @@ class FlashAttentionBackwardSm90:
                     if const_expr(self.qhead_per_kvhead == 1)
                     else head_idx // qhead_per_kvhead_divmod
                 )
-                mK_cur = mK[None, None, head_idx_kv, batch_idx]
                 gK = cute.local_tile(mK_cur, (self.tile_n, self.tile_hdim), (n_block, 0))
-                mV_cur = mV[None, None, head_idx_kv, 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(
@@ -786,14 +939,17 @@ class FlashAttentionBackwardSm90:
                 if const_expr(not self.use_block_sparsity):
                     total_m_block_cnt = m_block_max - m_block_min
-                    process_tile = const_expr(not self.is_local) or m_block_min < m_block_max
                 else:
                     total_m_block_cnt = get_total_q_block_count_bwd(
                         blocksparse_tensors,
                         batch_idx,
                         head_idx,
                         n_block,
-                        subtile_factor=self.subtile_factor,
                         m_block_max=m_block_max,
                     )
                     process_tile = total_m_block_cnt > Int32(0)
@@ -806,6 +962,8 @@ class FlashAttentionBackwardSm90:
                         )
                         load_K(tma_bar_ptr=pipeline_Q.producer_get_barrier(producer_state_Q))
                         load_Q(first_m_block, producer_state=producer_state_Q)
                         load_LSE(first_m_block, producer_state=producer_state_Q)
                         producer_state_dO_cur = (
                             producer_state_dO
@@ -836,6 +994,8 @@ class FlashAttentionBackwardSm90:
                             producer_state_Q.advance()
                             producer_state_dO.advance()
                     else:
                         producer_state_Q, producer_state_dO = produce_block_sparse_q_loads_bwd_sm90(
                             blocksparse_tensors,
                             batch_idx,
@@ -854,7 +1014,7 @@ class FlashAttentionBackwardSm90:
                             self.tma_copy_bytes["K"],
                             self.tma_copy_bytes["V"],
                             Q_stage_eq_dO_stage=(self.Q_stage == self.dO_stage),
-                            subtile_factor=self.subtile_factor,
                             m_block_max=m_block_max,
                         )
@@ -866,14 +1026,14 @@ class FlashAttentionBackwardSm90:
     def apply_score_mod(
         self,
         acc_S: cute.Tensor,
-        thr_mma_SdP: cute.core.ThrMma,
         batch_idx,
         head_idx,
         m_block,
         n_block,
         softmax_scale,
         seqlen_info: SeqlenInfoQK,
-        aux_tensors=None,
         fastdiv_mods=(None, None),
     ):
         # [NOTE] SdP_swapAB: swapAB transposes the tile, so use (n, m) indexing
@@ -897,7 +1057,7 @@ class FlashAttentionBackwardSm90:
             softmax_scale,
             self.vec_size,
             self.qk_acc_dtype,
-            aux_tensors,
             fastdiv_mods,
             seqlen_info,
             constant_q_idx=None,
@@ -910,14 +1070,14 @@ class FlashAttentionBackwardSm90:
         self,
         grad_tensor: cute.Tensor,
         score_tensor: cute.Tensor,
-        thr_mma_SdP: cute.core.ThrMma,
         batch_idx,
         head_idx,
         m_block,
         n_block,
         softmax_scale,
         seqlen_info: SeqlenInfoQK,
-        aux_tensors=None,
         fastdiv_mods=(None, None),
     ):
         cS = cute.make_identity_tensor(
@@ -941,7 +1101,7 @@ class FlashAttentionBackwardSm90:
             softmax_scale,
             self.vec_size,
             self.qk_acc_dtype,
-            aux_tensors,
             fastdiv_mods,
             seqlen_info,
             constant_q_idx=None,
@@ -958,6 +1118,8 @@ class FlashAttentionBackwardSm90:
         tiled_mma_dQ: cute.TiledMma,
         mdK: cute.Tensor,
         mdV: cute.Tensor,
         mdQaccum: cute.Tensor,
         sQ: cute.Tensor,
         sK: cute.Tensor,
@@ -980,20 +1142,24 @@ class FlashAttentionBackwardSm90:
         SeqlenInfoCls: Callable,
         AttentionMaskCls: Callable,
         TileSchedulerCls: Callable,
-        aux_tensors: Optional[list] = None,
         fastdiv_mods=(None, None),
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
         qhead_per_kvhead_divmod: Optional[FastDivmodDivisor] = None,
     ):
         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
         shape_mnk_S = (self.tile_m, self.tile_n, self.tile_hdim)
         _, tSrQ, tSrK = sm90_utils.partition_fragment_ABC(
@@ -1039,23 +1205,41 @@ class FlashAttentionBackwardSm90:
         # dQ = dS @ K
         sKt = layout_utils.transpose_view(sK)
         shape_mnk_dQ = (self.tile_m, self.tile_hdim, self.tile_n)
-        _, tdQrdS, tdQrKt = sm90_utils.partition_fragment_ABC(
-            wg_mma_dQ, shape_mnk_dQ, sdS, sKt, swap_AB=self.dQ_swapAB
-        )
-        mma_dsk_fn = partial(
-            gemm_zero_init, tiled_mma_dQ, shape_mnk_dQ[:2], tdQrdS, tdQrKt, swap_AB=self.dQ_swapAB
-        )
-        # Smem copy atom tiling
         copy_P_r2s = None
         if const_expr(sP is not None):
             sP_cpy = sP if const_expr(not self.SdP_swapAB) else sPt
             copy_P_r2s, _, _ = copy_utils.get_smem_store_C(
-                tiled_mma_SdP, sP_cpy, tidx, self.arch, transpose=self.SdP_swapAB
             )
         sdS_cpy = sdS if const_expr(not self.SdP_swapAB) else sdSt
         copy_dS_r2s, _, _ = copy_utils.get_smem_store_C(
-            tiled_mma_SdP, sdS_cpy, tidx, self.arch, transpose=self.SdP_swapAB
         )
         tLSEsLSE = layout_utils.mma_partition_C_vec(
@@ -1064,27 +1248,25 @@ class FlashAttentionBackwardSm90:
         tLSEsdPsum = layout_utils.mma_partition_C_vec(
             sdPsum, thr_mma_SdP, expand_shape=self.tile_n, is_colvec=not self.SdP_swapAB
         )
-        smem_thr_copy_dQaccum = r2s_tiled_copy_dQaccum.get_slice(tidx)
-        tdQsdQaccum = smem_thr_copy_dQaccum.partition_D(sdQaccum)
         PdS_barrier = cutlass.pipeline.NamedBarrier(
             barrier_id=int(NamedBarrierBwd.PdS), num_threads=self.num_mma_threads
         )
-        score_mod_fn = partial(
-            self.apply_score_mod,
-            thr_mma_SdP=thr_mma_SdP,
-            softmax_scale=softmax_scale,
-            aux_tensors=aux_tensors,
-            fastdiv_mods=fastdiv_mods,
-        )
-        score_mod_bwd_fn = partial(
-            self.apply_score_mod_bwd,
-            thr_mma_SdP=thr_mma_SdP,
-            softmax_scale=softmax_scale,
-            aux_tensors=aux_tensors,
-            fastdiv_mods=fastdiv_mods,
-        )
         mma_one_m_block_all = partial(
             self.mma_one_m_block,
@@ -1105,6 +1287,7 @@ class FlashAttentionBackwardSm90:
             PdS_barrier=PdS_barrier,
             # acc_dV=acc_dV,
             # acc_dK=acc_dK,
         )
         consumer_state_Q = cutlass.pipeline.make_pipeline_state(
@@ -1118,7 +1301,40 @@ class FlashAttentionBackwardSm90:
         while work_tile.is_valid_tile:
             n_block, head_idx, batch_idx, _ = work_tile.tile_idx
             seqlen = SeqlenInfoCls(batch_idx)
             mask = AttentionMaskCls(seqlen)
             score_mod_fn_cur = partial(
                 score_mod_fn,
                 batch_idx=batch_idx,
@@ -1136,14 +1352,17 @@ class FlashAttentionBackwardSm90:
             m_block_min, m_block_max = block_info.get_m_block_min_max(seqlen, n_block)
             if const_expr(not self.use_block_sparsity):
-                process_tile = const_expr(not self.is_local) or m_block_min < m_block_max
             else:
                 total_m_block_cnt = get_total_q_block_count_bwd(
                     blocksparse_tensors,
                     batch_idx,
                     head_idx,
                     n_block,
-                    subtile_factor=self.subtile_factor,
                     m_block_max=m_block_max,
                 )
                 process_tile = total_m_block_cnt > Int32(0)
@@ -1160,7 +1379,7 @@ class FlashAttentionBackwardSm90:
                         mask_causal=self.is_causal,
                         mask_local=self.is_local,
                         mask_mod=self.mask_mod,
-                        aux_tensors=aux_tensors,
                         fastdiv_mods=fastdiv_mods,
                     )
                     dKV_accumulate = False
@@ -1191,9 +1410,9 @@ class FlashAttentionBackwardSm90:
                         thr_mma_SdP=thr_mma_SdP,
                         score_mod_fn=score_mod_fn_cur,
                         score_mod_bwd_fn=score_mod_bwd_fn_cur,
-                        subtile_factor=self.subtile_factor,
                         m_block_max=m_block_max,
-                        aux_tensors=aux_tensors,
                         fastdiv_mods=fastdiv_mods,
                     )
@@ -1216,10 +1435,12 @@ class FlashAttentionBackwardSm90:
                     head_idx,
                     batch_idx,
                     qhead_per_kvhead_divmod,
                 )
             else:
-                # Block sparsity: KV tile with zero Q blocks produces no dK/dV; write zeros.
-                if const_expr(self.use_block_sparsity):
                     acc_dK.fill(0.0)
                     acc_dV.fill(0.0)
                     self.epilogue_dKV(
@@ -1239,6 +1460,8 @@ class FlashAttentionBackwardSm90:
                         head_idx,
                         batch_idx,
                         qhead_per_kvhead_divmod,
                     )
             tile_scheduler.advance_to_next_work()
@@ -1248,6 +1471,22 @@ class FlashAttentionBackwardSm90:
         if warp_idx == 4:
             cute.arch.cp_async_bulk_wait_group(0, read=True)
     @cute.jit
     def mma_one_m_block(
         self,
@@ -1266,16 +1505,17 @@ class FlashAttentionBackwardSm90:
         pipeline_dO: cutlass.pipeline.PipelineAsync,
         tLSEsLSE: cute.Tensor,
         tLSEsdPsum: cute.Tensor,
-        tdQsdQaccum: cute.Tensor,
         softmax_scale_log2: Float32,
         PdS_barrier: cutlass.pipeline.NamedBarrier,
         mask_fn: Optional[Callable] = None,
         score_mod_fn: Optional[Callable] = None,
         score_mod_bwd_fn: Optional[Callable] = None,
         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
@@ -1283,6 +1523,7 @@ class FlashAttentionBackwardSm90:
         # (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(
@@ -1301,10 +1542,12 @@ class FlashAttentionBackwardSm90:
         if cutlass.const_expr(mask_fn is not None):
             mask_fn(acc_S, m_block=m_block)
         acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S, transpose=self.SdP_swapAB)
         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
                 )
         tLSErdPsum = copy_utils.load_s2r(tLSEsdPsum[None, smem_idx_dO])
@@ -1321,8 +1564,9 @@ class FlashAttentionBackwardSm90:
         warpgroup.wait_group(0)
         acc_dP_mn = layout_utils.reshape_acc_to_mn(acc_dP, transpose=self.SdP_swapAB)
         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 const_expr(self.score_mod_bwd is not None):
             score_mod_bwd_fn(acc_dP, acc_S_pre, m_block=m_block)
@@ -1354,36 +1598,50 @@ class FlashAttentionBackwardSm90:
         # smem fence to make sure sdS is written before it's read by WGMMA
         cute.arch.fence_view_async_shared()
         PdS_barrier.arrive_and_wait()
-        # (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_view_async_shared()
-        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()
@@ -1408,6 +1666,8 @@ class FlashAttentionBackwardSm90:
         head_idx: Int32,
         batch_idx: Int32,
         qhead_per_kvhead_divmod: Optional[FastDivmodDivisor] = None,
     ):
         epi_barrier = cutlass.pipeline.NamedBarrier(
             barrier_id=int(NamedBarrierBwd.Epilogue), num_threads=self.num_mma_threads
@@ -1415,8 +1675,12 @@ class FlashAttentionBackwardSm90:
         warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
         if const_expr(self.qhead_per_kvhead == 1):
-            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(
@@ -1428,10 +1692,18 @@ class FlashAttentionBackwardSm90:
             sdV = sV if const_expr(not self.dKV_swapAB) else layout_utils.transpose_view(sV)
             sdK = sK if const_expr(not self.dKV_swapAB) else layout_utils.transpose_view(sK)
             copy_dV_r2s, _, _ = copy_utils.get_smem_store_C(
-                tiled_mma_dV, sdV, tidx, self.arch, transpose=self.dKV_swapAB
             )
             copy_dK_r2s, _, _ = copy_utils.get_smem_store_C(
-                tiled_mma_dK, sdK, tidx, self.arch, transpose=self.dKV_swapAB
             )
             cute.arch.cp_async_bulk_wait_group(1, read=True)
             epi_barrier.arrive_and_wait()
@@ -1450,15 +1722,26 @@ class FlashAttentionBackwardSm90:
                 store_dK()
                 cute.arch.cp_async_bulk_commit_group()
         else:
-            sdKaccum_shape0 = self.tile_n * self.tile_hdim // self.num_mma_warp_groups
-            sdVaccum_shape0 = self.tile_n * self.tile_hdimv // self.num_mma_warp_groups
-            sdKaccum_layout = cute.make_layout((sdKaccum_shape0, self.num_mma_warp_groups))
-            sdVaccum_layout = cute.make_layout((sdVaccum_shape0, self.num_mma_warp_groups))
             head_idx_kv = head_idx // qhead_per_kvhead_divmod
-            mdKaccum_cur = mdK[None, head_idx_kv, batch_idx]
             gdKaccum_ = cute.local_tile(mdKaccum_cur, (self.tile_n * self.tile_hdim,), (n_block,))
             gdKaccum = cute.flat_divide(gdKaccum_, (sdKaccum_shape0,))
-            mdVaccum_cur = mdV[None, head_idx_kv, batch_idx]
             gdVaccum_ = cute.local_tile(mdVaccum_cur, (self.tile_n * self.tile_hdimv,), (n_block,))
             gdVaccum = cute.flat_divide(gdVaccum_, (sdVaccum_shape0,))
             # These two overlap each other
@@ -1467,14 +1750,17 @@ class FlashAttentionBackwardSm90:
             sdVaccum = cute.make_tensor(sVaccum_ptr, sdVaccum_layout)
             tiled_copy_dKVaccum_r2s = cute.make_tiled_copy_tv(
                 cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128),
-                cute.make_layout((self.num_threads_per_warp_group, self.num_mma_warp_groups)),
                 cute.make_layout(128 // Float32.width),
             )
             thr_copy_dKVaccum_r2s = tiled_copy_dKVaccum_r2s.get_slice(tidx)
             tdKsdKaccum = thr_copy_dKVaccum_r2s.partition_D(sdKaccum)
             tdVsdVaccum = thr_copy_dKVaccum_r2s.partition_D(sdVaccum)
-            cute.arch.cp_async_bulk_wait_group(0, read=True)
             epi_barrier.arrive_and_wait()
             tdKrdKaccum_flat = cute.make_tensor(acc_dK.iterator, tdKsdKaccum.shape)
             cute.autovec_copy(tdKrdKaccum_flat, tdKsdKaccum)
@@ -1482,15 +1768,18 @@ class FlashAttentionBackwardSm90:
             epi_barrier.arrive_and_wait()
             if warp_idx == 4:
                 with cute.arch.elect_one():
-                    for wg_idx in cutlass.range_constexpr(self.num_mma_warp_groups):
                         copy_utils.cpasync_reduce_bulk_add_f32(
                             sdKaccum[None, wg_idx].iterator,
                             gdKaccum[None, wg_idx].iterator,
-                            self.tma_copy_bytes["dKacc"] // self.num_mma_warp_groups,
                         )
                 cute.arch.cp_async_bulk_commit_group()
-            cute.arch.cp_async_bulk_wait_group(0, read=True)
             epi_barrier.arrive_and_wait()
             tdVrdVaccum_flat = cute.make_tensor(acc_dV.iterator, tdVsdVaccum.shape)
             cute.autovec_copy(tdVrdVaccum_flat, tdVsdVaccum)
@@ -1498,13 +1787,16 @@ class FlashAttentionBackwardSm90:
             epi_barrier.arrive_and_wait()
             if warp_idx == 4:
                 with cute.arch.elect_one():
-                    for wg_idx in cutlass.range_constexpr(self.num_mma_warp_groups):
                         copy_utils.cpasync_reduce_bulk_add_f32(
                             sdVaccum[None, wg_idx].iterator,
                             gdVaccum[None, wg_idx].iterator,
-                            self.tma_copy_bytes["dVacc"] // self.num_mma_warp_groups,
                         )
                 cute.arch.cp_async_bulk_commit_group()
     @cute.jit
     def dQaccum_store(
@@ -1515,21 +1807,45 @@ class FlashAttentionBackwardSm90:
         TileSchedulerCls: cutlass.Constexpr[Callable],
         SeqlenInfoCls: cutlass.Constexpr[Callable],
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
     ):
         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)
             if const_expr(not self.use_block_sparsity):
-                process_tile = const_expr(not self.is_local) or m_block_min < m_block_max
                 loop_count = m_block_max - m_block_min
             else:
                 total_block_cnt = get_total_q_block_count_bwd(
@@ -1537,7 +1853,7 @@ class FlashAttentionBackwardSm90:
                     batch_idx,
                     head_idx,
                     n_block,
-                    subtile_factor=self.subtile_factor,
                     m_block_max=m_block_max,
                 )
                 process_tile = total_block_cnt > Int32(0)
@@ -1548,17 +1864,36 @@ class FlashAttentionBackwardSm90:
                         m_block = m_block_min + iter_idx
                         m_block_safe = m_block
-                        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,
                             )
-                        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
@@ -1567,11 +1902,24 @@ class FlashAttentionBackwardSm90:
                             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_safe].iterator,
                                     self.tma_copy_bytes["dQ"],
                                 )
                             cute.arch.cp_async_bulk_commit_group()
                 else:
                     dQaccum_store_block_sparse_bwd_sm90(
                         blocksparse_tensors,
                         batch_idx,
@@ -1579,13 +1927,29 @@ class FlashAttentionBackwardSm90:
                         n_block,
                         sdQaccum,
                         gdQaccum,
-                        subtile_factor=self.subtile_factor,
                         m_block_max=m_block_max,
-                        num_mma_warp_groups=self.num_mma_warp_groups,
                         num_threads_per_warp_group=self.num_threads_per_warp_group,
                         tma_copy_bytes_dQ=self.tma_copy_bytes["dQ"],
                     )
             tile_scheduler.advance_to_next_work()
             work_tile = tile_scheduler.get_current_work()
-        cute.arch.cp_async_bulk_wait_group(0, read=True)

 from .block_info import BlockInfo
 from . import pipeline
 from .quack.cute_dsl_utils import ParamsBase
+from .tile_scheduler import (
+    TileSchedulerArguments,
+    SingleTileScheduler,
+    SingleTileLPTBwdScheduler,
+    SingleTileVarlenScheduler,
+)
+from . import barrier
 from .named_barrier import NamedBarrierBwd
 from .softmax import apply_score_mod_inner, apply_score_mod_bwd_inner
 from .block_sparsity import BlockSparseTensors
+from .utils import AuxData
 from .block_sparse_utils import (
     get_total_q_block_count_bwd,
     produce_block_sparse_q_loads_bwd_sm90,
         head_dim_v: Optional[int] = None,
         qhead_per_kvhead: int = 1,
         is_causal: bool = False,
+        is_local: bool = False,
+        deterministic: bool = False,
         tile_m: int = 64,
         tile_n: int = 128,
         Q_stage: int = 2,
         score_mod_bwd: cutlass.Constexpr | None = None,
         mask_mod: cutlass.Constexpr | None = None,
         has_aux_tensors: cutlass.Constexpr = False,
+        q_subtile_factor: cutlass.Constexpr[int] = 1,
+        dQ_single_wg: bool = False,
     ):
         self.dtype = dtype
+        # Swapped dK/dV accumulators put head_dim on WGMMA's 64-row M axis.
+        hdim_multiple_of = 64 if dKV_swapAB else 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.check_hdim_v_oob = head_dim_v != self.tile_hdimv
         self.qhead_per_kvhead = qhead_per_kvhead
         self.is_causal = is_causal
+        self.is_local = is_local
+        self.deterministic = deterministic
         self.tile_m = tile_m
         self.tile_n = tile_n
         self.num_threads = num_threads
         self.AtomLayoutMSdP = AtomLayoutMSdP
         self.AtomLayoutNdKV = AtomLayoutNdKV
         self.AtomLayoutMdQ = AtomLayoutMdQ
+        self.num_wg_mma = (self.num_threads // 128) - 1
         self.mma_dkv_is_rs = (
             AtomLayoutMSdP == 1
+            and AtomLayoutNdKV == self.num_wg_mma
             and SdP_swapAB
             and not dKV_swapAB
         )
         self.V_in_regs = V_in_regs
+        # May be overridden in __call__ for varlen inputs.
         if qhead_per_kvhead > 1:
             assert self.same_hdim_kv, "GQA backward requires head_dim == head_dim_v"
+            assert self.num_wg_mma == 2, "GQA backward assumes 2 warp groups"
         # 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.
         self.shuffle_LSE = self.SdP_swapAB and self.tile_hdim <= 64
         self.shuffle_dPsum = self.SdP_swapAB and self.tile_hdim <= 64
         self.score_mod_bwd = score_mod_bwd
         self.mask_mod = mask_mod
         self.has_aux_tensors = has_aux_tensors
+        self.q_subtile_factor = q_subtile_factor
         if cutlass.const_expr(has_aux_tensors):
             self.vec_size: cutlass.Constexpr = 1
         else:
             self.vec_size: cutlass.Constexpr = 4
         self.qk_acc_dtype = Float32
+        # dQ_single_wg: WG0 computes the full dQ GEMM, WG1 skips it.
+        # Only valid for 2 MMA warp groups.
+        # Credit: Ben Spector
+        if dQ_single_wg:
+            assert self.num_wg_mma == 2, "dQ_single_wg only supports 2 warp groups"
+        self.num_wg_dQ = 1 if dQ_single_wg else self.num_wg_mma
     @staticmethod
     def can_implement(
         assert mQ_type == self.dtype
     def _setup_attributes(self):
+        # We need to accommodate both Q and Q^T (and dO and dO^T) in shared memory.
+        # Q & dO are used in the SdP Mma and Q^T and dO^T are used in the dKV Mma.
+        # The M dimension (tile_m) doesn't matter for the layout, only the K dimension
+        wg_d_dKV = self.num_wg_mma // self.AtomLayoutNdKV
+        self.sQ_layout, self.sdO_layout = [
+            # Need to set major_mode_size (mms) to accommodate Q and Q.T
+            sm90_utils.make_smem_layout(
+                self.dtype,
+                LayoutEnum.ROW_MAJOR,
+                shape,
+                stage,
+                major_mode_size=mms,
+            )
+            for shape, stage, mms in [
+                ((self.tile_m, self.tile_hdim), self.Q_stage, self.tile_hdim // wg_d_dKV),
+                ((self.tile_m, self.tile_hdimv), self.dO_stage, self.tile_hdim // wg_d_dKV),
             ]
         ]
+        wg_d_dQ = self.num_wg_dQ // self.AtomLayoutMdQ
+        # Accomodate both K and K.T
+        self.sK_layout = sm90_utils.make_smem_layout(
+            self.dtype,
+            LayoutEnum.ROW_MAJOR,
+            (self.tile_n, self.tile_hdim),
+            stage=None,
+            major_mode_size=self.tile_hdim // wg_d_dQ,
+        )
+        # There's only V, no V.T, so layout is normal
+        self.sV_layout = sm90_utils.make_smem_layout(
+            self.dtype, LayoutEnum.ROW_MAJOR, (self.tile_n, self.tile_hdimv), None
+        )
+        # Accomodate both S and S.T
+        wg_n_SdP = self.num_wg_mma // self.AtomLayoutMSdP
+        wg_n_dKV = self.AtomLayoutNdKV
+        self.sPdS_layout = sm90_utils.make_smem_layout(
+            self.dtype,
+            LayoutEnum.ROW_MAJOR,
+            (self.tile_m, self.tile_n),
+            stage=self.PdS_stage,
+            major_mode_size=math.gcd(self.tile_n // wg_n_SdP, self.tile_n // wg_n_dKV),
+        )
         self.sdQaccum_layout = cute.make_layout(
+            (self.tile_m * self.tile_hdim // self.num_wg_dQ, self.num_wg_dQ)
         )
         # 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_wg_dQ)),
             cute.make_layout(128 // Float32.width),  # val_layout
         )
         # dKVaccum for GQA epilogue - reuses sV+sK memory recast as f32
         # TODO: assert that sVaccum and sKaccum don't overflow smem
     def _get_tiled_mma(self):
+        maybe_swap_mn = lambda shape, swap: (shape[1], shape[0], *shape[2:]) if swap else shape
         # S = Q @ K.T, dP = dO @ V.T
+        atom_layout_SdP = (self.AtomLayoutMSdP, self.num_wg_mma // self.AtomLayoutMSdP, 1)
         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,
             warpgroup.OperandMajorMode.K,
             warpgroup.OperandMajorMode.K,
             Float32,
+            atom_layout_mnk=maybe_swap_mn(atom_layout_SdP, self.SdP_swapAB),
+            tiler_mn=(64, tiler_mn_SdP[1] if not self.SdP_swapAB else tiler_mn_SdP[0]),
         )
         # dV = P.T @ dO, dK = dS.T @ Q
+        atom_layout_dKV = (self.AtomLayoutNdKV, self.num_wg_mma // self.AtomLayoutNdKV, 1)
         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 = [
                 else warpgroup.OperandMajorMode.K,
                 warpgroup.OperandMajorMode.MN,
                 Float32,
+                atom_layout_mnk=maybe_swap_mn(atom_layout_dKV, self.dKV_swapAB),
+                tiler_mn=(64, tiler_mn_d[1] if not self.dKV_swapAB else tiler_mn_d[0]),
                 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
+        assert self.num_wg_dQ % self.AtomLayoutMdQ == 0
+        atom_layout_dQ = (self.AtomLayoutMdQ, self.num_wg_dQ // self.AtomLayoutMdQ, 1)
         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,
             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=maybe_swap_mn(atom_layout_dQ, self.dQ_swapAB),
+            tiler_mn=(64, tiler_mn_dQ[1] if not self.dQ_swapAB else tiler_mn_dQ[0]),
         )
         return tiled_mma_SdP, tiled_mma_dK, tiled_mma_dV, tiled_mma_dQ
         mdK: cute.Tensor,
         mdV: cute.Tensor,
         softmax_scale: Float32,
         mCuSeqlensQ: Optional[cute.Tensor] = None,
         mCuSeqlensK: Optional[cute.Tensor] = None,
         mSeqUsedQ: Optional[cute.Tensor] = None,
         mSeqUsedK: Optional[cute.Tensor] = None,
         window_size_left: Int32 | int | None = None,
         window_size_right: Int32 | int | None = None,
         mdQ_semaphore: Optional[cute.Tensor] = None,
         mdK_semaphore: Optional[cute.Tensor] = None,
         mdV_semaphore: Optional[cute.Tensor] = None,
+        aux_data: AuxData = AuxData(),
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
+        stream: cuda.CUstream = None,
     ):
+        # For GQA (qhead_per_kvhead > 1), multiple Q heads accumulate into the same dK/dV,
+        # so we need the float32 accum path + postprocess.
+        # For varlen_k with qhead_per_kvhead == 1, we use ragged TMA tensors.
+        self.varlen_k = mCuSeqlensK is not None or mSeqUsedK is not None
         self._check_type(
             *(
             )
         )
+        self.is_varlen_q = mCuSeqlensQ is not None or mSeqUsedQ is not None
         mQ, mK, mV, mdO, mLSE, mdPsum, mdQaccum, mdK, mdV = [
             assume_tensor_aligned(t) for t in (mQ, mK, mV, mdO, mLSE, mdPsum, mdQaccum, mdK, mdV)
         ]
+        # Non-varlen inputs are (b, s, n, h), varlen inputs are (s, n, h).
+        # We convert both to a seqlen-major view with head-dim second.
+        # Each tensor may have different rank when Q is padded (seqused_q) but K/V are unpadded (cu_seqlens_k).
+        def _qkv_transpose(t):
+            return layout_utils.select(t, [1, 3, 2, 0] if cute.rank(t.shape) == 4 else [0, 2, 1])
+        mQ, mK, mV, mdO = [_qkv_transpose(t) for t in (mQ, mK, mV, mdO)]
         if const_expr(self.qhead_per_kvhead == 1):
+            mdK, mdV = [_qkv_transpose(t) for t in (mdK, mdV)]
         else:
+            # Accum tensors are (b, n, s*h) for non-varlen and (n, s*h) for varlen.
+            accum_transpose = [2, 1, 0] if cute.rank(mdK.shape) == 3 else [1, 0]
             mdK, mdV = [layout_utils.select(t, accum_transpose) for t in (mdK, mdV)]
+        # Non-varlen stats are (b, n, s), varlen stats are (n, s).
+        LSE_dPsum_dQaccum_transpose = [2, 1, 0] if cute.rank(mLSE.shape) == 3 else [1, 0]
         mLSE, mdPsum, mdQaccum = [
             layout_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()
+        # (batch, num_head, num_m_blocks, cluster_size) -> (num_m_blocks, cluster_size, num_head, batch)
+        if const_expr(self.deterministic):
+            assert mdQ_semaphore is not None
+            mdQ_semaphore = layout_utils.select(mdQ_semaphore, mode=[2, 3, 1, 0])
+        if const_expr(self.deterministic and self.qhead_per_kvhead > 1):
+            assert mdK_semaphore is not None
+            assert mdV_semaphore is not None
+            mdK_semaphore, mdV_semaphore = [
+                layout_utils.select(t, mode=[2, 3, 1, 0]) for t in (mdK_semaphore, mdV_semaphore)
+            ]
+        else:
+            mdK_semaphore = None
+            mdV_semaphore = None
         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
+        REG_LIMIT = 504 if self.num_wg_mma == 2 else 512
+        if const_expr(self.num_wg_mma == 2):
+            if const_expr(self.num_wg_dQ == 1):
+                self.num_mma_regs_wg0 = 256
+                self.num_mma_regs_wg1 = 224
+            else:
+                self.num_mma_regs_wg0 = 240
+                self.num_mma_regs_wg1 = 240
+            self.num_mma_regs = self.num_mma_regs_wg0  # for backward compat
+            self.num_producer_regs = 24
+            assert (
+                self.num_mma_regs_wg0 + self.num_mma_regs_wg1 + self.num_producer_regs <= REG_LIMIT
+            )
+        else:  # 3 warp groups
+            self.num_mma_regs_wg0 = 160
+            self.num_mma_regs_wg1 = 160
+            self.num_mma_regs = 160
+            self.num_producer_regs = 32
+            assert self.num_mma_regs_wg0 * self.num_wg_mma + self.num_producer_regs <= REG_LIMIT
         self._setup_attributes()
         SharedStorage = self._get_shared_storage_cls()
         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_wg_dQ
         )
         self.tma_copy_bytes["dKacc"] = self.tile_n * self.tile_hdim * Float32.width // 8
         self.tma_copy_bytes["dVacc"] = self.tile_n * self.tile_hdimv * Float32.width // 8
             (self.tile_m, self.tile_hdimv),
         )
         if const_expr(self.qhead_per_kvhead == 1):
+            mdK_tma = (
+                copy_utils.create_ragged_tensor_for_tma(mdK, ragged_dim=0, ptr_shift=True)
+                if self.varlen_k
+                else mdK
+            )
+            mdV_tma = (
+                copy_utils.create_ragged_tensor_for_tma(mdV, ragged_dim=0, ptr_shift=True)
+                if self.varlen_k
+                else mdV
+            )
             tma_atom_dK, tma_tensor_dK = cpasync.make_tiled_tma_atom(
                 cpasync.CopyBulkTensorTileS2GOp(),
+                mdK_tma,
                 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_tma,
                 cute.select(self.sV_layout, mode=[0, 1]),
                 (self.tile_n, self.tile_hdimv),
             )
         else:
             tma_atom_dK = tma_atom_dV = tma_tensor_dK = tma_tensor_dV = None
+        if const_expr(mCuSeqlensK is not None or mSeqUsedK is not None):
+            TileScheduler = SingleTileVarlenScheduler
+        elif const_expr(self.deterministic):
+            TileScheduler = SingleTileLPTBwdScheduler
+        else:
+            TileScheduler = SingleTileScheduler
+        self.spt = (self.is_causal or self.is_local) and self.deterministic
         tile_sched_args = TileSchedulerArguments(
             cute.ceil_div(cute.size(mK.shape[0]), self.tile_n),
             cute.size(mQ.shape[2]),
+            cute.size(mK.shape[3])
+            if const_expr(mCuSeqlensK is None)
+            else cute.size(mCuSeqlensK.shape[0] - 1),  # num_batch
             1,  # num_splits
+            cute.size(mQ.shape[0]),  # pass seqlen_q or total_q for seqlen_k
+            mQ.shape[1],  # headdim
+            mV.shape[1],  # headdim_v
+            total_q=cute.size(mK.shape[0])
+            if const_expr(mCuSeqlensK is not None)
+            else cute.size(mK.shape[0]) * cute.size(mK.shape[3]),
+            tile_shape_mn=(self.tile_n, self.tile_m),  # Swapping the role of Q & K
+            mCuSeqlensQ=mCuSeqlensK,
+            mSeqUsedQ=mSeqUsedK,
             qhead_per_kvhead_packgqa=1,
             element_size=self.dtype.width // 8,
             is_persistent=False,
+            lpt=self.spt,
+            head_swizzle=self.deterministic,
+            cu_total_m_blocks_ptr=mCuTotalMBlocks,
         )
         tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
             softmax_scale_log2 = LOG2_E
         fastdiv_mods = None
+        if const_expr(aux_data.tensors is not None):
             seqlen_q = cute.size(mQ.shape[0])
             seqlen_k = cute.size(mK.shape[0])
             seqlen_q_divmod = FastDivmodDivisor(seqlen_q)
         self.use_block_sparsity = cutlass.const_expr(blocksparse_tensors is not None)
+        if const_expr(window_size_left is not None):
+            window_size_left = Int32(window_size_left)
+        if const_expr(window_size_right is not None):
+            window_size_right = Int32(window_size_right)
         self.kernel(
             tma_tensor_Q,
             tma_tensor_K,
             mLSE,
             mdPsum,
             mdQaccum,
+            mCuSeqlensQ,
+            mCuSeqlensK,
+            mSeqUsedQ,
+            mSeqUsedK,
             self.sQ_layout,
             self.sK_layout,
             self.sV_layout,
             tile_sched_params,
             TileScheduler,
             SharedStorage,
+            aux_data,
             fastdiv_mods,
             blocksparse_tensors,
             qhead_per_kvhead_divmod,
+            mdQ_semaphore,
+            mdK_semaphore,
+            mdV_semaphore,
+            window_size_left,
+            window_size_right,
         ).launch(
             grid=grid_dim,
             block=[self.num_threads, 1, 1],
             stream=stream,
             min_blocks_per_mp=1,
+            use_pdl=True,
         )
     @cute.kernel
         mLSE: cute.Tensor,
         mdPsum: cute.Tensor,
         mdQaccum: cute.Tensor,
+        mCuSeqlensQ: Optional[cute.Tensor],
+        mCuSeqlensK: Optional[cute.Tensor],
+        mSeqUsedQ: Optional[cute.Tensor],
+        mSeqUsedK: Optional[cute.Tensor],
         sQ_layout: cute.ComposedLayout,
         sK_layout: cute.ComposedLayout,
         sV_layout: cute.ComposedLayout,
         tile_sched_params: ParamsBase,
         TileScheduler: cutlass.Constexpr[Callable],
         SharedStorage: cutlass.Constexpr[Callable],
+        aux_data: AuxData = AuxData(),
         fastdiv_mods=(None, None),
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
         qhead_per_kvhead_divmod: Optional[FastDivmodDivisor] = None,
+        mdQ_semaphore: Optional[cute.Tensor] = None,
+        mdK_semaphore: Optional[cute.Tensor] = None,
+        mdV_semaphore: Optional[cute.Tensor] = None,
+        window_size_left: Optional[Int32] = None,
+        window_size_right: Optional[Int32] = None,
     ):
         warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
         # prefetch TMA descriptors
         if warp_idx == 0:
+            for atom in [tma_atom_Q, tma_atom_K, tma_atom_V, tma_atom_dO, tma_atom_dK, tma_atom_dV]:
+                if const_expr(atom is not None):
+                    cpasync.prefetch_descriptor(atom)
         smem = cutlass.utils.SmemAllocator()
         storage = smem.allocate(SharedStorage)
             self.is_causal,
             self.is_local,
             False,  # is_split_kv
+            window_size_left,
+            window_size_right,
             qhead_per_kvhead_packgqa=1,
         )
         SeqlenInfoCls = partial(
             SeqlenInfoQK.create,
             seqlen_q_static=mQ.shape[0],
             seqlen_k_static=mK.shape[0],
+            mCuSeqlensQ=mCuSeqlensQ,
+            mCuSeqlensK=mCuSeqlensK,
+            mSeqUsedQ=mSeqUsedQ,
+            mSeqUsedK=mSeqUsedK,
+            tile_m=self.tile_m,
+            tile_n=self.tile_n,
         )
         AttentionMaskCls = partial(
             AttentionMask,
             self.tile_m,
             self.tile_n,
+            window_size_left=window_size_left,
+            window_size_right=window_size_right,
             swap_AB=self.SdP_swapAB,
         )
         TileSchedulerCls = partial(TileScheduler.create, tile_sched_params)
                     TileSchedulerCls,
                     SeqlenInfoCls,
                     blocksparse_tensors,
+                    mdQ_semaphore,
                 )
         else:
             tidx, _, _ = cute.arch.thread_idx()
             tidx = tidx - 128
+            mma_args = (
                 tiled_mma_SdP,
                 tiled_mma_dK,
                 tiled_mma_dV,
                 tiled_mma_dQ,
                 mdK,
                 mdV,
+                mdK_semaphore,
+                mdV_semaphore,
                 mdQaccum,
                 sQ,
                 sK,
                 SeqlenInfoCls,
                 AttentionMaskCls,
                 TileSchedulerCls,
+                aux_data,
                 fastdiv_mods,
                 blocksparse_tensors,
                 qhead_per_kvhead_divmod,
             )
+            if const_expr(self.num_wg_dQ == self.num_wg_mma):
+                # Both WGs compute dQ
+                cute.arch.setmaxregister_increase(self.num_mma_regs_wg0)
+                self.mma(*mma_args, is_dQ_wg=True)
+            else:
+                # WG0 computes dQ, WG1 skips it
+                warp_idx_in_mma = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - 4
+                if warp_idx_in_mma < 4:
+                    cute.arch.setmaxregister_increase(self.num_mma_regs_wg0)
+                    self.mma(*mma_args, is_dQ_wg=True)
+                else:
+                    cute.arch.setmaxregister_increase(self.num_mma_regs_wg1)
+                    self.mma(*mma_args, is_dQ_wg=False)
     @cute.jit
     def load(
                     if const_expr(self.qhead_per_kvhead == 1)
                     else head_idx // qhead_per_kvhead_divmod
                 )
+                mK_cur = seqlen.offset_batch_K(mK, batch_idx, dim=3)[None, None, head_idx_kv]
+                mV_cur = seqlen.offset_batch_K(mV, batch_idx, dim=3)[None, None, head_idx_kv]
                 gK = cute.local_tile(mK_cur, (self.tile_n, self.tile_hdim), (n_block, 0))
                 gV = cute.local_tile(mV_cur, (self.tile_n, self.tile_hdimv), (n_block, 0))
+                mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx]
+                mLSE_cur = seqlen.offset_batch_Q(mLSE, batch_idx, dim=2, padded=True)[
+                    None, head_idx
+                ]
+                mdO_cur = seqlen.offset_batch_Q(mdO, batch_idx, dim=3)[None, None, head_idx]
+                mdPsum_cur = seqlen.offset_batch_Q(mdPsum, batch_idx, dim=2, padded=True)[
+                    None, head_idx
+                ]
                 gQ = cute.local_tile(mQ_cur, (self.tile_m, self.tile_hdim), (None, 0))
                 gdO = cute.local_tile(mdO_cur, (self.tile_m, self.tile_hdimv), (None, 0))
                 gLSE = cute.local_tile(mLSE_cur, (self.tile_m,), (None,))
                 gdPsum = cute.local_tile(mdPsum_cur, (self.tile_m,), (None,))
                 load_K, _, _ = copy_utils.tma_get_copy_fn(
                 if const_expr(not self.use_block_sparsity):
                     total_m_block_cnt = m_block_max - m_block_min
+                    process_tile = (
+                        const_expr(not self.is_local and not self.is_varlen_q)
+                        or m_block_min < m_block_max
+                    )
                 else:
                     total_m_block_cnt = get_total_q_block_count_bwd(
                         blocksparse_tensors,
                         batch_idx,
                         head_idx,
                         n_block,
+                        q_subtile_factor=self.q_subtile_factor,
                         m_block_max=m_block_max,
                     )
                     process_tile = total_m_block_cnt > Int32(0)
                         )
                         load_K(tma_bar_ptr=pipeline_Q.producer_get_barrier(producer_state_Q))
                         load_Q(first_m_block, producer_state=producer_state_Q)
+                        # Wait for bwd preprocess to finish writing LSE and dPsum
+                        cute.arch.griddepcontrol_wait()
                         load_LSE(first_m_block, producer_state=producer_state_Q)
                         producer_state_dO_cur = (
                             producer_state_dO
                             producer_state_Q.advance()
                             producer_state_dO.advance()
                     else:
+                        # The wait for bwd preprocess that the dense arm above does inline sits
+                        # in _load_q_do_block_sm90, at the same point in the load order.
                         producer_state_Q, producer_state_dO = produce_block_sparse_q_loads_bwd_sm90(
                             blocksparse_tensors,
                             batch_idx,
                             self.tma_copy_bytes["K"],
                             self.tma_copy_bytes["V"],
                             Q_stage_eq_dO_stage=(self.Q_stage == self.dO_stage),
+                            q_subtile_factor=self.q_subtile_factor,
                             m_block_max=m_block_max,
                         )
     def apply_score_mod(
         self,
         acc_S: cute.Tensor,
+        thr_mma_SdP: cute.ThrMma,
         batch_idx,
         head_idx,
         m_block,
         n_block,
         softmax_scale,
         seqlen_info: SeqlenInfoQK,
+        aux_data: AuxData = AuxData(),
         fastdiv_mods=(None, None),
     ):
         # [NOTE] SdP_swapAB: swapAB transposes the tile, so use (n, m) indexing
             softmax_scale,
             self.vec_size,
             self.qk_acc_dtype,
+            aux_data,
             fastdiv_mods,
             seqlen_info,
             constant_q_idx=None,
         self,
         grad_tensor: cute.Tensor,
         score_tensor: cute.Tensor,
+        thr_mma_SdP: cute.ThrMma,
         batch_idx,
         head_idx,
         m_block,
         n_block,
         softmax_scale,
         seqlen_info: SeqlenInfoQK,
+        aux_data: AuxData = AuxData(),
         fastdiv_mods=(None, None),
     ):
         cS = cute.make_identity_tensor(
             softmax_scale,
             self.vec_size,
             self.qk_acc_dtype,
+            aux_data,
             fastdiv_mods,
             seqlen_info,
             constant_q_idx=None,
         tiled_mma_dQ: cute.TiledMma,
         mdK: cute.Tensor,
         mdV: cute.Tensor,
+        mdK_semaphore: Optional[cute.Tensor],
+        mdV_semaphore: Optional[cute.Tensor],
         mdQaccum: cute.Tensor,
         sQ: cute.Tensor,
         sK: cute.Tensor,
         SeqlenInfoCls: Callable,
         AttentionMaskCls: Callable,
         TileSchedulerCls: Callable,
+        aux_data: AuxData = AuxData(),
         fastdiv_mods=(None, None),
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
         qhead_per_kvhead_divmod: Optional[FastDivmodDivisor] = None,
+        is_dQ_wg: cutlass.Constexpr[bool] = True,
     ):
         warp_group_idx = cute.arch.make_warp_uniform(tidx // self.num_threads_per_warp_group)
         warp_group_thread_layout = cute.make_layout(
+            self.num_wg_mma, 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 = None
+        if const_expr(is_dQ_wg):
+            wg_idx_dQ = warp_group_idx if const_expr(self.num_wg_dQ > 1) else 0
+            wg_mma_dQ = tiled_mma_dQ.get_slice(warp_group_thread_layout(wg_idx_dQ))
         # S = Q @ K.T
         shape_mnk_S = (self.tile_m, self.tile_n, self.tile_hdim)
         _, tSrQ, tSrK = sm90_utils.partition_fragment_ABC(
         # dQ = dS @ K
         sKt = layout_utils.transpose_view(sK)
         shape_mnk_dQ = (self.tile_m, self.tile_hdim, self.tile_n)
+        mma_dsk_fn = None
+        if const_expr(is_dQ_wg):
+            _, tdQrdS, tdQrKt = sm90_utils.partition_fragment_ABC(
+                wg_mma_dQ, shape_mnk_dQ, sdS, sKt, swap_AB=self.dQ_swapAB
+            )
+            mma_dsk_fn = partial(
+                gemm_zero_init,
+                tiled_mma_dQ,
+                shape_mnk_dQ[:2],
+                tdQrdS,
+                tdQrKt,
+                swap_AB=self.dQ_swapAB,
+            )
+        # Smem copy atom tiling for P/dS R2S
         copy_P_r2s = None
+        mms_PdS = self.tile_n // (self.num_wg_mma // self.AtomLayoutMSdP)
         if const_expr(sP is not None):
             sP_cpy = sP if const_expr(not self.SdP_swapAB) else sPt
             copy_P_r2s, _, _ = copy_utils.get_smem_store_C(
+                tiled_mma_SdP,
+                sP_cpy,
+                tidx,
+                transpose=self.SdP_swapAB,
+                position_independent=True,
+                major_mode_size=mms_PdS,
             )
         sdS_cpy = sdS if const_expr(not self.SdP_swapAB) else sdSt
         copy_dS_r2s, _, _ = copy_utils.get_smem_store_C(
+            tiled_mma_SdP,
+            sdS_cpy,
+            tidx,
+            transpose=self.SdP_swapAB,
+            position_independent=True,
+            major_mode_size=mms_PdS,
         )
         tLSEsLSE = layout_utils.mma_partition_C_vec(
         tLSEsdPsum = layout_utils.mma_partition_C_vec(
             sdPsum, thr_mma_SdP, expand_shape=self.tile_n, is_colvec=not self.SdP_swapAB
         )
+        # When shuffle=True, rows are distributed across 8 quads (4 threads each) within a warp.
+        # Each thread loads only ceil(num_rows/8) values;
+        shfl_copy = copy_utils.tiled_copy_1d(sLSE.element_type, num_threads=8, num_copy_elems=2)
+        if const_expr(self.shuffle_LSE):
+            tLSEsLSE = shfl_copy.get_slice(cute.arch.lane_idx() // 4).partition_S(tLSEsLSE)
+            # ((2, 1), 1, 2) -> (((2, 1), 1), 2)
+            tLSEsLSE = cute.group_modes(tLSEsLSE, 0, 2)
+        if const_expr(self.shuffle_dPsum):
+            tLSEsdPsum = shfl_copy.get_slice(cute.arch.lane_idx() // 4).partition_S(tLSEsdPsum)
+            tLSEsdPsum = cute.group_modes(tLSEsdPsum, 0, 2)
+        tdQsdQaccum = None
+        if const_expr(is_dQ_wg):
+            smem_thr_copy_dQaccum = r2s_tiled_copy_dQaccum.get_slice(tidx)
+            tdQsdQaccum = smem_thr_copy_dQaccum.partition_D(sdQaccum)
         PdS_barrier = cutlass.pipeline.NamedBarrier(
             barrier_id=int(NamedBarrierBwd.PdS), num_threads=self.num_mma_threads
         )
         mma_one_m_block_all = partial(
             self.mma_one_m_block,
             PdS_barrier=PdS_barrier,
             # acc_dV=acc_dV,
             # acc_dK=acc_dK,
+            is_dQ_wg=is_dQ_wg,
         )
         consumer_state_Q = cutlass.pipeline.make_pipeline_state(
         while work_tile.is_valid_tile:
             n_block, head_idx, batch_idx, _ = work_tile.tile_idx
             seqlen = SeqlenInfoCls(batch_idx)
+            recompute_fastdiv_mods_q = const_expr(
+                aux_data.tensors is not None and (seqlen.has_cu_seqlens_q or seqlen.has_seqused_q)
+            )
+            recompute_fastdiv_mods_k = const_expr(
+                aux_data.tensors is not None and (seqlen.has_cu_seqlens_k or seqlen.has_seqused_k)
+            )
+            if const_expr(fastdiv_mods is not None and fastdiv_mods[0] is not None):
+                seqlen_q_divmod, seqlen_k_divmod = fastdiv_mods
+                fastdiv_mods = (
+                    seqlen_q_divmod
+                    if not recompute_fastdiv_mods_q
+                    else FastDivmodDivisor(seqlen.seqlen_q),
+                    seqlen_k_divmod
+                    if not recompute_fastdiv_mods_k
+                    else FastDivmodDivisor(seqlen.seqlen_k),
+                )
             mask = AttentionMaskCls(seqlen)
+            score_mod_fn = partial(
+                self.apply_score_mod,
+                thr_mma_SdP=thr_mma_SdP,
+                softmax_scale=softmax_scale,
+                aux_data=aux_data,
+                fastdiv_mods=fastdiv_mods,
+            )
+            score_mod_bwd_fn = partial(
+                self.apply_score_mod_bwd,
+                thr_mma_SdP=thr_mma_SdP,
+                softmax_scale=softmax_scale,
+                aux_data=aux_data,
+                fastdiv_mods=fastdiv_mods,
+            )
             score_mod_fn_cur = partial(
                 score_mod_fn,
                 batch_idx=batch_idx,
             m_block_min, m_block_max = block_info.get_m_block_min_max(seqlen, n_block)
             if const_expr(not self.use_block_sparsity):
+                process_tile = (
+                    const_expr(not self.is_local and not self.is_varlen_q)
+                    or m_block_min < m_block_max
+                )
             else:
                 total_m_block_cnt = get_total_q_block_count_bwd(
                     blocksparse_tensors,
                     batch_idx,
                     head_idx,
                     n_block,
+                    q_subtile_factor=self.q_subtile_factor,
                     m_block_max=m_block_max,
                 )
                 process_tile = total_m_block_cnt > Int32(0)
                         mask_causal=self.is_causal,
                         mask_local=self.is_local,
                         mask_mod=self.mask_mod,
+                        aux_data=aux_data,
                         fastdiv_mods=fastdiv_mods,
                     )
                     dKV_accumulate = False
                         thr_mma_SdP=thr_mma_SdP,
                         score_mod_fn=score_mod_fn_cur,
                         score_mod_bwd_fn=score_mod_bwd_fn_cur,
+                        q_subtile_factor=self.q_subtile_factor,
                         m_block_max=m_block_max,
+                        aux_data=aux_data,
                         fastdiv_mods=fastdiv_mods,
                     )
                     head_idx,
                     batch_idx,
                     qhead_per_kvhead_divmod,
+                    mdK_semaphore,
+                    mdV_semaphore,
                 )
             else:
+                # KV tile with zero Q blocks produces no dK/dV; write zeros.
+                if const_expr(self.use_block_sparsity or self.is_local or self.is_varlen_q):
                     acc_dK.fill(0.0)
                     acc_dV.fill(0.0)
                     self.epilogue_dKV(
                         head_idx,
                         batch_idx,
                         qhead_per_kvhead_divmod,
+                        mdK_semaphore,
+                        mdV_semaphore,
                     )
             tile_scheduler.advance_to_next_work()
         if warp_idx == 4:
             cute.arch.cp_async_bulk_wait_group(0, read=True)
+    @staticmethod
+    @cute.jit
+    def _get_stat(tSrS: cute.Tensor, row: Int32, lane: Int32, shuffle: bool) -> Float32:
+        """Retrieve the statistic for a given accumulator row.
+        When shuffle=False, direct register indexing.
+        When shuffle=True, warp shuffle from the thread group that holds the value.
+        """
+        if const_expr(not shuffle):
+            return tSrS[row]
+        # tSrS: (((2, 1), 1), 1)), distributed across 8 threads in the warp
+        vecsize = cute.size(tSrS, mode=[0, 0])  # 2
+        idx0, off, idx1 = cute.idx2crd(row, (vecsize, 8, cute.shape(tSrS, mode=[0, 1])))
+        # register index: 0, 1, 0, 1, ..., 2, 3, 2, 3, ...
+        return utils.shuffle_sync(tSrS[idx0 + idx1 * vecsize], offset=off * 4 + (lane % 4))
     @cute.jit
     def mma_one_m_block(
         self,
         pipeline_dO: cutlass.pipeline.PipelineAsync,
         tLSEsLSE: cute.Tensor,
         tLSEsdPsum: cute.Tensor,
+        tdQsdQaccum: Optional[cute.Tensor],
         softmax_scale_log2: Float32,
         PdS_barrier: cutlass.pipeline.NamedBarrier,
+        is_dQ_wg: cutlass.Constexpr[bool] = True,
         mask_fn: Optional[Callable] = None,
         score_mod_fn: Optional[Callable] = None,
         score_mod_bwd_fn: Optional[Callable] = None,
         dKV_accumulate: Boolean = True,
     ):
         consumer_state_dO_cur = (
+            consumer_state_Q if const_expr(self.Q_stage == self.dO_stage) else consumer_state_dO
         )
         smem_idx_Q = consumer_state_Q.index
         smem_idx_dO = consumer_state_dO_cur.index if const_expr(self.dO_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)
+        # If shuffle_LSE, OOB reads are OK since sLSE is already padded
         tLSErLSE = copy_utils.load_s2r(tLSEsLSE[None, smem_idx_Q])
         # (2) [GEMM 2] dP = dO @ V.T
         pipeline_dO.consumer_wait(
         if cutlass.const_expr(mask_fn is not None):
             mask_fn(acc_S, m_block=m_block)
         acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S, transpose=self.SdP_swapAB)
+        lane_idx = cute.arch.lane_idx()
         for r in cutlass.range_constexpr(cute.size(acc_S_mn, mode=[0])):
+            lse_val = self._get_stat(tLSErLSE, r, lane_idx, shuffle=self.shuffle_LSE)
             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 - lse_val, fastmath=True
                 )
         tLSErdPsum = copy_utils.load_s2r(tLSEsdPsum[None, smem_idx_dO])
         warpgroup.wait_group(0)
         acc_dP_mn = layout_utils.reshape_acc_to_mn(acc_dP, transpose=self.SdP_swapAB)
         for r in cutlass.range_constexpr(cute.size(acc_dP_mn, mode=[0])):
+            dpsum_val = self._get_stat(tLSErdPsum, r, lane_idx, shuffle=self.shuffle_dPsum)
             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] - dpsum_val)
         if const_expr(self.score_mod_bwd is not None):
             score_mod_bwd_fn(acc_dP, acc_S_pre, m_block=m_block)
         # smem fence to make sure sdS is written before it's read by WGMMA
         cute.arch.fence_view_async_shared()
         PdS_barrier.arrive_and_wait()
+        if const_expr(is_dQ_wg):
+            # (6) [GEMM 4] dQ = dS @ K
+            acc_dQ = mma_dsk_fn(A_idx=smem_idx_PdS, wg_wait=1)
+            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)
+            # dQ R2S: wait for dQaccum_store to free the smem buffer, then write dQ to smem
+            # When dQ_single_wg, only WG0 enters here so warp_group_idx == 0
+            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_view_async_shared()
+            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)
+            pipeline_Q.consumer_release(consumer_state_Q)
+        else:
+            # dQ_single_wg: WG1 skips dQ, only does dV wait + dK
+            # (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)
+            pipeline_dO.consumer_release(consumer_state_dO_cur)
+            warpgroup.wait_group(0)
+            pipeline_Q.consumer_release(consumer_state_Q)
         consumer_state_Q.advance()
         consumer_state_dO.advance()
         head_idx: Int32,
         batch_idx: Int32,
         qhead_per_kvhead_divmod: Optional[FastDivmodDivisor] = None,
+        mdK_semaphore: Optional[cute.Tensor] = None,
+        mdV_semaphore: Optional[cute.Tensor] = None,
     ):
         epi_barrier = cutlass.pipeline.NamedBarrier(
             barrier_id=int(NamedBarrierBwd.Epilogue), num_threads=self.num_mma_threads
         warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
         if const_expr(self.qhead_per_kvhead == 1):
+            mdK_cur = seqlen.offset_batch_K(mdK, batch_idx, dim=3, ragged=self.varlen_k)[
+                None, None, head_idx
+            ]
+            mdV_cur = seqlen.offset_batch_K(mdV, batch_idx, dim=3, ragged=self.varlen_k)[
+                None, None, head_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(
             sdV = sV if const_expr(not self.dKV_swapAB) else layout_utils.transpose_view(sV)
             sdK = sK if const_expr(not self.dKV_swapAB) else layout_utils.transpose_view(sK)
             copy_dV_r2s, _, _ = copy_utils.get_smem_store_C(
+                tiled_mma_dV,
+                sdV,
+                tidx,
+                transpose=self.dKV_swapAB,
+                position_independent=True,
             )
             copy_dK_r2s, _, _ = copy_utils.get_smem_store_C(
+                tiled_mma_dK,
+                sdK,
+                tidx,
+                transpose=self.dKV_swapAB,
+                position_independent=True,
             )
             cute.arch.cp_async_bulk_wait_group(1, read=True)
             epi_barrier.arrive_and_wait()
                 store_dK()
                 cute.arch.cp_async_bulk_commit_group()
         else:
+            deterministic_KV = self.deterministic and self.qhead_per_kvhead > 1
+            sdKaccum_shape0 = self.tile_n * self.tile_hdim // self.num_wg_mma
+            sdVaccum_shape0 = self.tile_n * self.tile_hdimv // self.num_wg_mma
+            sdKaccum_layout = cute.make_layout((sdKaccum_shape0, self.num_wg_mma))
+            sdVaccum_layout = cute.make_layout((sdVaccum_shape0, self.num_wg_mma))
             head_idx_kv = head_idx // qhead_per_kvhead_divmod
+            if const_expr(deterministic_KV):
+                assert mdK_semaphore is not None
+                assert mdV_semaphore is not None
+                mdK_semaphore_cur = mdK_semaphore[n_block, None, head_idx_kv, batch_idx]
+                mdV_semaphore_cur = mdV_semaphore[n_block, None, head_idx_kv, batch_idx]
+                lock_value = head_idx % self.qhead_per_kvhead
+            mdKaccum_cur = seqlen.offset_batch_K(
+                mdK, batch_idx, dim=2, padded=True, multiple=self.tile_hdim
+            )[None, head_idx_kv]
+            mdVaccum_cur = seqlen.offset_batch_K(
+                mdV, batch_idx, dim=2, padded=True, multiple=self.tile_hdimv
+            )[None, head_idx_kv]
             gdKaccum_ = cute.local_tile(mdKaccum_cur, (self.tile_n * self.tile_hdim,), (n_block,))
             gdKaccum = cute.flat_divide(gdKaccum_, (sdKaccum_shape0,))
             gdVaccum_ = cute.local_tile(mdVaccum_cur, (self.tile_n * self.tile_hdimv,), (n_block,))
             gdVaccum = cute.flat_divide(gdVaccum_, (sdVaccum_shape0,))
             # These two overlap each other
             sdVaccum = cute.make_tensor(sVaccum_ptr, sdVaccum_layout)
             tiled_copy_dKVaccum_r2s = cute.make_tiled_copy_tv(
                 cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128),
+                cute.make_layout((self.num_threads_per_warp_group, self.num_wg_mma)),
                 cute.make_layout(128 // Float32.width),
             )
             thr_copy_dKVaccum_r2s = tiled_copy_dKVaccum_r2s.get_slice(tidx)
             tdKsdKaccum = thr_copy_dKVaccum_r2s.partition_D(sdKaccum)
             tdVsdVaccum = thr_copy_dKVaccum_r2s.partition_D(sdVaccum)
+            read_flag = const_expr(not deterministic_KV)
+            cute.arch.cp_async_bulk_wait_group(0, read=read_flag)
+            if const_expr(deterministic_KV):
+                barrier.wait_eq(mdK_semaphore_cur.iterator, tidx, 0, lock_value)
             epi_barrier.arrive_and_wait()
             tdKrdKaccum_flat = cute.make_tensor(acc_dK.iterator, tdKsdKaccum.shape)
             cute.autovec_copy(tdKrdKaccum_flat, tdKsdKaccum)
             epi_barrier.arrive_and_wait()
             if warp_idx == 4:
                 with cute.arch.elect_one():
+                    for wg_idx in cutlass.range_constexpr(self.num_wg_mma):
                         copy_utils.cpasync_reduce_bulk_add_f32(
                             sdKaccum[None, wg_idx].iterator,
                             gdKaccum[None, wg_idx].iterator,
+                            self.tma_copy_bytes["dKacc"] // self.num_wg_mma,
                         )
                 cute.arch.cp_async_bulk_commit_group()
+            cute.arch.cp_async_bulk_wait_group(0, read=read_flag)
+            if const_expr(deterministic_KV):
+                barrier.arrive_inc(mdK_semaphore_cur.iterator, tidx, 0, 1)
+                barrier.wait_eq(mdV_semaphore_cur.iterator, tidx, 0, lock_value)
             epi_barrier.arrive_and_wait()
             tdVrdVaccum_flat = cute.make_tensor(acc_dV.iterator, tdVsdVaccum.shape)
             cute.autovec_copy(tdVrdVaccum_flat, tdVsdVaccum)
             epi_barrier.arrive_and_wait()
             if warp_idx == 4:
                 with cute.arch.elect_one():
+                    for wg_idx in cutlass.range_constexpr(self.num_wg_mma):
                         copy_utils.cpasync_reduce_bulk_add_f32(
                             sdVaccum[None, wg_idx].iterator,
                             gdVaccum[None, wg_idx].iterator,
+                            self.tma_copy_bytes["dVacc"] // self.num_wg_mma,
                         )
                 cute.arch.cp_async_bulk_commit_group()
+            if const_expr(deterministic_KV):
+                cute.arch.cp_async_bulk_wait_group(0, read=read_flag)
+                barrier.arrive_inc(mdV_semaphore_cur.iterator, tidx, 0, 1)
     @cute.jit
     def dQaccum_store(
         TileSchedulerCls: cutlass.Constexpr[Callable],
         SeqlenInfoCls: cutlass.Constexpr[Callable],
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
+        mdQ_semaphore: Optional[cute.Tensor] = None,
     ):
+        tidx, _, _ = cute.arch.thread_idx()
+        # warp-local thread index (dQaccum_store runs on warp 1, global tidx 32-63)
+        warp_local_tidx = tidx % cute.arch.WARP_SIZE
+        read_flag = const_expr(not self.deterministic)
         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)
+            if const_expr(not seqlen.has_cu_seqlens_q):
+                mdQaccum_cur = mdQaccum[None, head_idx, batch_idx]
+            else:
+                mdQaccum_cur = cute.domain_offset(
+                    (seqlen.padded_offset_q * self.tile_hdim,), mdQaccum[None, head_idx]
+                )
+            # ((M * K / num_wg_dQ, num_wg_dQ), num_m_blocks)
+            gdQaccum = cute.local_tile(
+                mdQaccum_cur,
+                (
+                    cute.make_layout(
+                        (self.tile_m * self.tile_hdim // self.num_wg_dQ, self.num_wg_dQ)
+                    ),
+                ),
+                (None,),
             )
+            if const_expr(mdQ_semaphore is not None):
+                # mdQ_semaphore is (num_m_blocks, cluster_size, num_head, batch) after transpose
+                mdQ_semaphore_cur = mdQ_semaphore[None, None, head_idx, batch_idx]
             m_block_min, m_block_max = block_info.get_m_block_min_max(seqlen, n_block)
             if const_expr(not self.use_block_sparsity):
+                process_tile = (
+                    const_expr(not self.is_local and not self.is_varlen_q)
+                    or m_block_min < m_block_max
+                )
                 loop_count = m_block_max - m_block_min
             else:
                 total_block_cnt = get_total_q_block_count_bwd(
                     batch_idx,
                     head_idx,
                     n_block,
+                    q_subtile_factor=self.q_subtile_factor,
                     m_block_max=m_block_max,
                 )
                 process_tile = total_block_cnt > Int32(0)
                         m_block = m_block_min + iter_idx
                         m_block_safe = m_block
+                        num_dQ_chunks = self.num_wg_dQ
+                        for warp_group_idx in cutlass.range_constexpr(num_dQ_chunks):
+                            if const_expr(not self.deterministic):
+                                # If deterministic, we already waited at the end of the prev iter
+                                cute.arch.cp_async_bulk_wait_group(
+                                    num_dQ_chunks - 1 - warp_group_idx, read=read_flag
+                                )
                             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,
                             )
+                        # Semaphore acquire: wait for prior n_blocks to finish writing this m_block
+                        if const_expr(self.deterministic):
+                            if const_expr(self.spt):
+                                _, n_block_max_for_m_block = block_info.get_n_block_min_max(
+                                    seqlen, m_block_safe
+                                )
+                                lock_value = n_block_max_for_m_block - 1 - n_block
+                            else:
+                                lock_value = n_block
+                            barrier.wait_eq(
+                                mdQ_semaphore_cur[(m_block_safe, None)].iterator,
+                                warp_local_tidx,
+                                0,  # flag_offset
+                                lock_value,
+                            )
+                        for warp_group_idx in cutlass.range_constexpr(num_dQ_chunks):
                             cute.arch.barrier(
                                 barrier_id=int(NamedBarrierBwd.dQFullWG0) + warp_group_idx,
                                 number_of_threads=self.num_threads_per_warp_group
                             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_safe].iterator,
                                     self.tma_copy_bytes["dQ"],
                                 )
                             cute.arch.cp_async_bulk_commit_group()
+                        # Semaphore release: signal that this n_block is done with this m_block
+                        if const_expr(self.deterministic):
+                            cute.arch.cp_async_bulk_wait_group(0, read=read_flag)
+                            barrier.arrive_inc(
+                                mdQ_semaphore_cur[(m_block_safe, None)].iterator,
+                                warp_local_tidx,
+                                0,  # flag_offset
+                                1,
+                            )
                 else:
+                    assert not self.deterministic, (
+                        "Deterministic not implemented for block-sparse backward"
+                    )
                     dQaccum_store_block_sparse_bwd_sm90(
                         blocksparse_tensors,
                         batch_idx,
                         n_block,
                         sdQaccum,
                         gdQaccum,
+                        q_subtile_factor=self.q_subtile_factor,
                         m_block_max=m_block_max,
+                        num_dQ_warp_groups=self.num_wg_dQ,
                         num_threads_per_warp_group=self.num_threads_per_warp_group,
                         tma_copy_bytes_dQ=self.tma_copy_bytes["dQ"],
                     )
+            # For local masking + deterministic (non-spt): signal remaining m_blocks
+            # that this n_block won't visit, so they don't deadlock waiting.
+            if const_expr(
+                self.deterministic and not self.spt and block_info.window_size_left is not None
+            ):
+                m_block_global_max = cute.ceil_div(seqlen.seqlen_q, self.tile_m)
+                for m_block in cutlass.range(m_block_max, m_block_global_max, unroll=1):
+                    barrier.arrive_inc(
+                        mdQ_semaphore_cur[(m_block, None)].iterator,
+                        warp_local_tidx,
+                        0,  # flag_offset
+                        1,
+                    )
             tile_scheduler.advance_to_next_work()
             work_tile = tile_scheduler.get_current_work()
+        if const_expr(not self.deterministic):
+            cute.arch.cp_async_bulk_wait_group(0, read=True)

build/torch-cuda/flash_fwd.py CHANGED Viewed

The diff for this file is too large to render. See raw diff

build/torch-cuda/flash_fwd_combine.py CHANGED Viewed

@@ -2,7 +2,7 @@
 # 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
-from typing import Type, Optional
 from functools import partial
 import cuda.bindings.driver as cuda
@@ -10,11 +10,18 @@ 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 .cute_dsl_utils import assume_tensor_aligned
 from .seqlen_info import SeqlenInfo
 from cutlass.cute import FastDivmodDivisor
@@ -24,7 +31,8 @@ class FlashAttentionForwardCombine:
         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,
@@ -36,7 +44,8 @@ class FlashAttentionForwardCombine:
         :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
@@ -46,7 +55,8 @@ class FlashAttentionForwardCombine:
         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
@@ -58,7 +68,7 @@ class FlashAttentionForwardCombine:
         dtype,
         dtype_partial,
         head_dim,
-        m_block_size,
         k_block_size,
         log_max_splits,
         num_threads,
@@ -72,12 +82,12 @@ class FlashAttentionForwardCombine:
             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
@@ -124,15 +134,11 @@ class FlashAttentionForwardCombine:
         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
@@ -183,12 +189,12 @@ class FlashAttentionForwardCombine:
             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
@@ -201,7 +207,9 @@ class FlashAttentionForwardCombine:
         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
@@ -269,7 +277,7 @@ class FlashAttentionForwardCombine:
             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
             ]
@@ -289,11 +297,29 @@ class FlashAttentionForwardCombine:
         seqlen_divmod = FastDivmodDivisor(seqlen)
         head_divmod = FastDivmodDivisor(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,
@@ -303,6 +329,7 @@ class FlashAttentionForwardCombine:
             cu_seqlens,
             seqused,
             num_splits_dynamic_ptr,
             semaphore_to_reset,
             SharedStorage,
             self.smem_layout_lse,
@@ -314,6 +341,8 @@ class FlashAttentionForwardCombine:
             seqlen_divmod,
             head_divmod,
             varlen,
         ).launch(
             grid=grid_dim,
             block=[self.num_threads, 1, 1],
@@ -331,6 +360,7 @@ class FlashAttentionForwardCombine:
         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,
@@ -342,10 +372,21 @@ class FlashAttentionForwardCombine:
         seqlen_divmod: FastDivmodDivisor,
         head_divmod: FastDivmodDivisor,
         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
@@ -353,310 +394,315 @@ class FlashAttentionForwardCombine:
         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]
-            else:
-                mLSE_partial_cur = cute.domain_offset((offset, 0, 0), mLSE_partial)
-            mLSE_partial_copy = cute.tiled_divide(mLSE_partial_cur, (1,))
-            gmem_thr_copy_LSE = gmem_tiled_copy_LSE.get_slice(tidx)
-            tLSEsLSE = gmem_thr_copy_LSE.partition_D(sLSE)
-            # Create identity tensor for coordinate tracking
-            cLSE = cute.make_identity_tensor((self.max_splits, self.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 FastDivmodDivisor
-                    if const_expr(not varlen):
-                        head_idx, m_idx = divmod(idx, seqlen_divmod)
-                    else:
-                        head_idx = idx // seqlen
-                        m_idx = idx - head_idx * seqlen
-                    mLSE_partial_cur_copy = mLSE_partial_copy[None, m_idx, None, head_idx]
-                    for s in cutlass.range(cute.size(tLSEcLSE, mode=[1]), unroll_full=True):
-                        si = tLSEcLSE[0, s, 0][0]  # Get split coordinate
-                        if si < num_splits:
-                            cute.copy(
-                                gmem_thr_copy_LSE,
-                                mLSE_partial_cur_copy[None, si],
-                                tLSEsLSE[None, s, m],
-                            )
-                        else:
-                            tLSEsLSE[None, s, m].fill(-Float32.inf)
-                # 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]
-            else:
-                mO_partial_cur = cute.domain_offset((offset, 0, 0, 0), mO_partial)
-            # Precompute these values to avoid recomputing them in the loop
-            num_rows = const_expr(cute.size(tOcO, mode=[1]))
-            tOmidx = cute.make_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] = divmod(idx, seqlen_divmod)
-                else:
-                    tOhidx[m] = idx // seqlen
-                    tOmidx[m] = idx - tOhidx[m] * seqlen
-                tOrOptr[m] = utils.elem_pointer(
-                    mO_partial_cur, (tOmidx[m], k_block * self.k_block_size, 0, tOhidx[m])
-                ).toint()
-                if idx >= max_idx:
-                    tOhidx[m] = -1
-            tOpO = 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 = cute.arch.warp_reduction_max(
-                    ts2rrLSE[None, None, m]
-                    .load()
-                    .reduce(cute.ReductionOp.MAX, init_val=-Float32.inf, reduction_profile=0),
-                    threads_in_group=threads_per_col,
-                )
-                # 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] = cute.arch.warp_reduction_max(
-                    max_valid_idx, threads_in_group=threads_per_col
-                )
-                # Compute exp scales and sum
-                lse_max_cur = (
-                    0.0 if lse_max == -Float32.inf else lse_max
-                )  # In case all local LSEs are -inf
-                LOG2_E = math.log2(math.e)
-                lse_sum_cur = 0.0
-                for s in cutlass.range(cute.size(ts2rrLSE, mode=[1]), unroll_full=True):
-                    scale = cute.math.exp2(
-                        ts2rrLSE[0, s, m] * LOG2_E - (lse_max_cur * LOG2_E), fastmath=True
-                    )
-                    lse_sum_cur += scale
-                    ts2rrLSE[0, s, m] = scale  # Store scale for later use
-                lse_sum_cur = cute.arch.warp_reduction_sum(
-                    lse_sum_cur, threads_in_group=threads_per_col
-                )
-                lse_sum[m] = cute.math.log(lse_sum_cur, fastmath=True) + lse_max
-                # Normalize scales
-                inv_sum = (
-                    0.0 if (lse_sum_cur == 0.0 or lse_sum_cur != lse_sum_cur) else 1.0 / lse_sum_cur
-                )
-                ts2rrLSE[None, None, m].store(ts2rrLSE[None, None, m].load() * inv_sum)
-            # Store the scales exp(lse - lse_logsum) back to smem
-            cute.copy(s2r_tiled_copy_LSE, ts2rrLSE, ts2rsLSE)
-            # Store max valid split to smem
-            for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True):
-                if ts2rcLSE[0, 0, m][0] == 0:  # Only thread responsible for s=0 writes
-                    mi = ts2rcLSE[0, 0, m][1]
-                    if mi < self.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]
-                else:
-                    mLSE_cur = cute.domain_offset((offset, 0), mLSE)
-                if k_block == 0:  # Only first k_block writes LSE when mLSE is provided
-                    for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True):
-                        if ts2rcLSE[0, 0, m][0] == 0:  # Only thread responsible for s=0 writes
-                            mi = ts2rcLSE[0, 0, m][1]
-                            idx = m_block * self.m_block_size + mi
-                            if idx < max_idx:
-                                if const_expr(not varlen):
-                                    head_idx, m_idx = divmod(idx, seqlen_divmod)
-                                else:
-                                    head_idx = idx // seqlen
-                                    m_idx = idx - head_idx * seqlen
-                                mLSE_cur[m_idx, head_idx] = lse_sum[m]
-            # ===============================
-            # Step 6: Read O_partial and accumulate final O
-            # ===============================
-            cute.arch.sync_threads()
-            # Get max valid split for this thread
-            thr_max_valid_split = sMaxValidSplit[tOcO[0, 0, 0][0]]
-            for m in cutlass.range(1, cute.size(tOcO, mode=[1])):
-                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]
-            else:
-                mO_cur = cute.domain_offset((offset, 0, 0), mO)
-            mO_cur = utils.domain_offset_aligned((0, k_block * self.k_block_size, 0), mO_cur)
-            elems_per_store = const_expr(cute.size(gmem_tiled_copy_O.layout_tv_tiled[1]))
-            # 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(
@@ -665,7 +711,7 @@ class FlashAttentionForwardCombine:
         tOrOptr: cute.Tensor,
         tOsO_partial: cute.Tensor,
         tOhidx: cute.Tensor,
-        tOpO: cute.Tensor,
         tOcO: cute.Tensor,
         mO_cur_partial_layout: cute.Layout,
         split: Int32,
@@ -684,7 +730,7 @@ class FlashAttentionForwardCombine:
                 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],

 # 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
+from typing import Callable, Type, Optional
 from functools import partial
 import cuda.bindings.driver as cuda
 import cutlass
 import cutlass.cute as cute
 from cutlass.cute.nvgpu import cpasync
+from cutlass import Float32, Int32, Boolean, const_expr
+from .quack.cute_dsl_utils import ParamsBase
 from . import utils
 from .cute_dsl_utils import assume_tensor_aligned
 from .seqlen_info import SeqlenInfo
+from .tile_scheduler import (
+    SingleTileScheduler,
+    SingleTileVarlenScheduler,
+    TileSchedulerArguments,
+)
 from cutlass.cute import FastDivmodDivisor
         dtype: Type[cutlass.Numeric],
         dtype_partial: Type[cutlass.Numeric],
         head_dim: int,
+        num_head: int,
+        tile_m: int = 8,
         k_block_size: int = 64,
         log_max_splits: int = 4,
         num_threads: int = 256,
         :param dtype: output data type
         :param dtype_partial: partial accumulation data type
         :param head_dim: head dimension
+        :param num_head: number of heads
+        :param tile_m: m block size
         :param k_block_size: k block size
         :param log_max_splits: log2 of maximum splits
         :param num_threads: number of threads
         self.dtype = dtype
         self.dtype_partial = dtype_partial
         self.head_dim = head_dim
+        self.num_head = num_head
+        self.tile_m = tile_m
         self.k_block_size = k_block_size
         self.max_splits = 1 << log_max_splits
         self.num_threads = num_threads
         dtype,
         dtype_partial,
         head_dim,
+        tile_m,
         k_block_size,
         log_max_splits,
         num_threads,
             return False
         if num_threads % 32 != 0:
             return False
+        if tile_m % 8 != 0:
             return False
         max_splits = 1 << log_max_splits
         if max_splits > 256:
             return False
+        if (tile_m * max_splits) % num_threads != 0:
             return False
         return True
         lse_copy_bits = Float32.width  # 1 element per copy, width is in bits
         m_block_smem = (
             128
+            if self.tile_m % 128 == 0
             else (
                 64
+                if self.tile_m % 64 == 0
+                else (32 if self.tile_m % 32 == 0 else (16 if self.tile_m % 16 == 0 else 8))
             )
         )
         gmem_threads_per_row_lse = m_block_smem
             smem_lse_swizzle, 0, cute.make_ordered_layout((8, m_block_smem), order=(1, 0))
         )
         self.smem_layout_lse = cute.tile_to_shape(
+            smem_layout_atom_lse, (self.max_splits, self.tile_m), (0, 1)
         )
         # O partial shared memory layout (simple layout for pipeline stages)
         self.smem_layout_o = cute.make_ordered_layout(
+            (self.tile_m, self.k_block_size, self.stages), order=(1, 0, 2)
         )
     @cute.jit
         cu_seqlens: Optional[cute.Tensor] = None,
         seqused: Optional[cute.Tensor] = None,
         num_splits_dynamic_ptr: Optional[cute.Tensor] = None,
+        virtual_batch_idx: Optional[cute.Tensor] = None,
         semaphore_to_reset: Optional[cute.Tensor] = None,
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
         stream: cuda.CUstream = None,
     ):
         # Type checking
             sLSE: cute.struct.Align[
                 cute.struct.MemRange[Float32, cute.cosize(self.smem_layout_lse)], 128
             ]
+            sMaxValidSplit: cute.struct.Align[cute.struct.MemRange[Int32, self.tile_m], 128]
             sO: cute.struct.Align[
                 cute.struct.MemRange[self.dtype_partial, cute.cosize(self.smem_layout_o)], 128
             ]
         seqlen_divmod = FastDivmodDivisor(seqlen)
         head_divmod = FastDivmodDivisor(num_head)
+        if const_expr(varlen):
+            TileScheduler = SingleTileVarlenScheduler
+        else:
+            TileScheduler = SingleTileScheduler
+        tile_sched_args = TileSchedulerArguments(
+            num_block=cute.ceil_div(seqlen * num_head, self.tile_m),
+            num_head=cute.ceil_div(self.head_dim, self.k_block_size),
+            num_batch=batch_size,
+            num_splits=1,
+            seqlen_k=1,
+            headdim=1,
+            headdim_v=1,
+            total_q=mO_partial.shape[0] * num_head
+            if const_expr(cu_seqlens is not None)
+            else seqlen * batch_size * num_head,
+            tile_shape_mn=(self.tile_m, self.tile_m),
+            mCuSeqlensQ=cu_seqlens,
+            mSeqUsedQ=seqused,
+            qhead_per_kvhead_packgqa=self.num_head,
+            virtual_batch_idx_ptr=virtual_batch_idx,
         )
+        tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
+        grid_dim = TileScheduler.get_grid_shape(tile_sched_params)
         self.kernel(
             mO_partial,
             cu_seqlens,
             seqused,
             num_splits_dynamic_ptr,
+            virtual_batch_idx,
             semaphore_to_reset,
             SharedStorage,
             self.smem_layout_lse,
             seqlen_divmod,
             head_divmod,
             varlen,
+            tile_sched_params,
+            TileScheduler,
         ).launch(
             grid=grid_dim,
             block=[self.num_threads, 1, 1],
         cu_seqlens: Optional[cute.Tensor],
         seqused: Optional[cute.Tensor],
         num_splits_dynamic_ptr: Optional[cute.Tensor],
+        virtual_batch_idx: Optional[cute.Tensor],
         semaphore_to_reset: Optional[cute.Tensor],
         SharedStorage: cutlass.Constexpr,
         smem_layout_lse: cute.Layout | cute.ComposedLayout,
         seqlen_divmod: FastDivmodDivisor,
         head_divmod: FastDivmodDivisor,
         varlen: cutlass.Constexpr[bool],
+        tile_sched_params: ParamsBase,
+        TileScheduler: cutlass.Constexpr[Callable],
     ):
         # Thread and block indices
         tidx, _, _ = cute.arch.thread_idx()
+        tile_scheduler = TileScheduler.create(tile_sched_params)
+        work_tile = tile_scheduler.initial_work_tile_info()
+        m_block, k_block, maybe_virtual_batch, _ = work_tile.tile_idx
+        # Map virtual batch index to real batch index (for persistent tile schedulers)
+        batch_idx = (
+            virtual_batch_idx[maybe_virtual_batch]
+            if const_expr(virtual_batch_idx is not None and not varlen)
+            else maybe_virtual_batch
+        )
         # ///////////////////////////////////////////////////////////////////////////////
         # Get shared memory buffer
         smem = cutlass.utils.SmemAllocator()
         storage = smem.allocate(SharedStorage)
         sLSE = storage.sLSE.get_tensor(smem_layout_lse)
+        sMaxValidSplit = storage.sMaxValidSplit.get_tensor((self.tile_m,))
         sO = storage.sO.get_tensor(smem_layout_o)
+        # Handle semaphore reset — wait for dependent grids first
         if const_expr(semaphore_to_reset is not None):
+            bidx, bidy, bidz = cute.arch.block_idx()
             if (
                 tidx == 0
+                and bidx == cute.arch.grid_dim()[0] - 1
+                and bidy == cute.arch.grid_dim()[1] - 1
+                and bidz == cute.arch.grid_dim()[2] - 1
             ):
+                cute.arch.griddepcontrol_wait()
                 semaphore_to_reset[0] = 0
+        if work_tile.is_valid_tile:
+            # Get number of splits (use maybe_virtual_batch for per-batch-slot splits)
+            num_splits = (
+                num_splits_dynamic_ptr[maybe_virtual_batch]
+                if const_expr(num_splits_dynamic_ptr is not None)
+                else mLSE_partial.shape[1]
+            )
+            # 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,
+                # Don't need to pass in tile size since we won't use offset_padded
             )
+            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
+            # TODO: early exit for single split if dynamic — for now always merge so the
+            # num_splits_dynamic == 1 case still writes mO from mO_partial[0].
+            if (const_expr(num_splits_dynamic_ptr is None) or num_splits > 0) and (
+                const_expr(not varlen) or m_block * self.tile_m < max_idx
+            ):
+                # Wait for dependent grids (e.g., the main attention kernel that produces O_partial/LSE_partial)
+                cute.arch.griddepcontrol_wait()
+                # ===============================
+                # Step 1: Load LSE_partial from gmem to shared memory
+                # ===============================
+                mLSE_partial_cur = seqlen_info.offset_batch(mLSE_partial, batch_idx, dim=3)
+                mLSE_partial_copy = cute.tiled_divide(mLSE_partial_cur, (1,))
+                gmem_thr_copy_LSE = gmem_tiled_copy_LSE.get_slice(tidx)
+                tLSEsLSE = gmem_thr_copy_LSE.partition_D(sLSE)
+                # Create identity tensor for coordinate tracking
+                cLSE = cute.make_identity_tensor((self.max_splits, self.tile_m))
+                tLSEcLSE = gmem_thr_copy_LSE.partition_S(cLSE)
+                # Load LSE partial values
+                for m in cutlass.range(cute.size(tLSEcLSE, mode=[2]), unroll_full=True):
+                    mi = tLSEcLSE[0, 0, m][1]  # Get m coordinate
+                    idx = m_block * self.tile_m + mi
+                    if idx < max_idx:
+                        # Calculate actual sequence position and head using FastDivmodDivisor
+                        if const_expr(not varlen):
+                            head_idx, m_idx = divmod(idx, seqlen_divmod)
+                        else:
+                            head_idx = idx // seqlen
+                            m_idx = idx - head_idx * seqlen
+                        mLSE_partial_cur_copy = mLSE_partial_copy[None, m_idx, None, head_idx]
+                        for s in cutlass.range(cute.size(tLSEcLSE, mode=[1]), unroll_full=True):
+                            si = tLSEcLSE[0, s, 0][0]  # Get split coordinate
+                            if si < num_splits:
+                                cute.copy(
+                                    gmem_thr_copy_LSE,
+                                    mLSE_partial_cur_copy[None, si],
+                                    tLSEsLSE[None, s, m],
+                                )
+                            else:
+                                tLSEsLSE[None, s, m].fill(-Float32.inf)
+                    # 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.tile_m, self.k_block_size))
+                tOcO = gmem_thr_copy_O_partial.partition_D(cO)
+                tOsO_partial = gmem_thr_copy_O_partial.partition_D(sO)
+                mO_partial_cur = seqlen_info.offset_batch(mO_partial, batch_idx, dim=4)
+                # Precompute these values to avoid recomputing them in the loop
+                num_rows = const_expr(cute.size(tOcO, mode=[1]))
+                tOmidx = cute.make_rmem_tensor(num_rows, cutlass.Int32)
+                tOhidx = cute.make_rmem_tensor(num_rows, cutlass.Int32)
+                tOrOptr = cute.make_rmem_tensor(num_rows, cutlass.Int64)
                 for m in cutlass.range(num_rows, unroll_full=True):
+                    mi = tOcO[0, m, 0][0]  # m coordinate
+                    idx = m_block * self.tile_m + mi
+                    if const_expr(not varlen):
+                        tOhidx[m], tOmidx[m] = divmod(idx, seqlen_divmod)
+                    else:
+                        tOhidx[m] = idx // seqlen
+                        tOmidx[m] = idx - tOhidx[m] * seqlen
+                    tOrOptr[m] = utils.elem_pointer(
+                        mO_partial_cur, (tOmidx[m], k_block * self.k_block_size, 0, tOhidx[m])
+                    ).toint()
+                    if idx >= max_idx:
+                        tOhidx[m] = -1
+                tOpO = None
+                if const_expr(not self.is_even_k):
+                    tOpO = cute.make_rmem_tensor(cute.size(tOcO, mode=[2]), Boolean)
+                    for k in cutlass.range(cute.size(tOpO), unroll_full=True):
+                        tOpO[k] = (
+                            tOcO[0, 0, k][1] < mO_partial.shape[1] - k_block * self.k_block_size
+                        )
+                    # 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_rmem_tensor_like(ts2rsLSE)
+                cute.copy(s2r_tiled_copy_LSE, ts2rsLSE, ts2rrLSE)
+                # ===============================
+                # Step 4: Compute final LSE along split dimension
+                # ===============================
+                lse_sum = cute.make_rmem_tensor(cute.size(ts2rrLSE, mode=[2]), Float32)
+                ts2rcLSE = s2r_thr_copy_LSE.partition_D(cLSE)
+                # We compute the max valid split for each row to short-circuit the computation later
+                max_valid_split = cute.make_rmem_tensor(cute.size(ts2rrLSE, mode=[2]), Int32)
+                assert cute.size(ts2rrLSE, mode=[0]) == 1
+                # Compute max, scales, and final LSE for each row
+                for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True):
+                    # Find max LSE value across splits
+                    threads_per_col = const_expr(self.smem_threads_per_col_lse)
+                    lse_max = cute.arch.warp_reduction_max(
+                        ts2rrLSE[None, None, m]
+                        .load()
+                        .reduce(cute.ReductionOp.MAX, init_val=-Float32.inf, reduction_profile=0),
+                        threads_in_group=threads_per_col,
+                    )
+                    # 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] = cute.arch.warp_reduction_max(
+                        max_valid_idx, threads_in_group=threads_per_col
+                    )
+                    # Compute exp scales and sum
+                    lse_max_cur = (
+                        0.0 if lse_max == -Float32.inf else lse_max
+                    )  # In case all local LSEs are -inf
+                    LOG2_E = math.log2(math.e)
+                    lse_sum_cur = 0.0
+                    for s in cutlass.range(cute.size(ts2rrLSE, mode=[1]), unroll_full=True):
+                        scale = cute.math.exp2(
+                            ts2rrLSE[0, s, m] * LOG2_E - (lse_max_cur * LOG2_E), fastmath=True
+                        )
+                        lse_sum_cur += scale
+                        ts2rrLSE[0, s, m] = scale  # Store scale for later use
+                    lse_sum_cur = cute.arch.warp_reduction_sum(
+                        lse_sum_cur, threads_in_group=threads_per_col
+                    )
+                    lse_sum[m] = cute.math.log(lse_sum_cur, fastmath=True) + lse_max
+                    # Normalize scales
+                    inv_sum = (
+                        0.0
+                        if (lse_sum_cur == 0.0 or lse_sum_cur != lse_sum_cur)
+                        else 1.0 / lse_sum_cur
+                    )
+                    ts2rrLSE[None, None, m].store(ts2rrLSE[None, None, m].load() * inv_sum)
+                # Store the scales exp(lse - lse_logsum) back to smem
+                cute.copy(s2r_tiled_copy_LSE, ts2rrLSE, ts2rsLSE)
+                # Store max valid split to smem
+                for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True):
+                    if ts2rcLSE[0, 0, m][0] == 0:  # Only thread responsible for s=0 writes
+                        mi = ts2rcLSE[0, 0, m][1]
+                        if mi < self.tile_m:
+                            sMaxValidSplit[mi] = max_valid_split[m]
+                # ===============================
+                # Step 5: Store final LSE to gmem
+                # ===============================
+                if const_expr(mLSE is not None):
+                    if const_expr(cu_seqlens is None):
+                        mLSE_cur = mLSE[None, None, batch_idx]
+                    else:
+                        mLSE_cur = cute.domain_offset((offset, 0), mLSE)
+                    if k_block == 0:  # Only first k_block writes LSE when mLSE is provided
+                        for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True):
+                            if ts2rcLSE[0, 0, m][0] == 0:  # Only thread responsible for s=0 writes
+                                mi = ts2rcLSE[0, 0, m][1]
+                                idx = m_block * self.tile_m + mi
+                                if idx < max_idx:
+                                    if const_expr(not varlen):
+                                        head_idx, m_idx = divmod(idx, seqlen_divmod)
+                                    else:
+                                        head_idx = idx // seqlen
+                                        m_idx = idx - head_idx * seqlen
+                                    mLSE_cur[m_idx, head_idx] = lse_sum[m]
+                # ===============================
+                # Step 6: Read O_partial and accumulate final O
+                # ===============================
+                cute.arch.sync_threads()
+                # Get max valid split for this thread
+                thr_max_valid_split = sMaxValidSplit[tOcO[0, 0, 0][0]]
+                for m in cutlass.range(1, cute.size(tOcO, mode=[1]), unroll_full=True):
+                    thr_max_valid_split = max(thr_max_valid_split, sMaxValidSplit[tOcO[0, m, 0][0]])
+                tOrO_partial = cute.make_rmem_tensor_like(tOsO_partial[None, None, None, 0])
+                tOrO = cute.make_rmem_tensor_like(tOrO_partial, Float32)
+                tOrO.fill(0.0)
+                stage_load = self.stages - 1
+                stage_compute = 0
+                # Main accumulation loop
+                for s in cutlass.range(thr_max_valid_split + 1, unroll=4):
+                    # Get scales for this split
+                    scale = cute.make_rmem_tensor(num_rows, Float32)
+                    for m in cutlass.range(num_rows, unroll_full=True):
+                        scale[m] = sLSE[s, tOcO[0, m, 0][0]]  # Get scale from smem
+                    # Load next stage if needed
+                    split_to_load = s + self.stages - 1
+                    if split_to_load <= thr_max_valid_split:
+                        load_O_partial(split_to_load, stage_load)
+                    cute.arch.cp_async_commit_group()
+                    stage_load = 0 if stage_load == self.stages - 1 else stage_load + 1
+                    # Wait for the current stage to be ready
+                    cute.arch.cp_async_wait_group(self.stages - 1)
+                    # We don't need __syncthreads() because each thread is just reading its own data from smem
+                    # Copy from smem to registers
+                    cute.autovec_copy(tOsO_partial[None, None, None, stage_compute], tOrO_partial)
+                    stage_compute = 0 if stage_compute == self.stages - 1 else stage_compute + 1
+                    # Accumulate scaled partial results
+                    for m in cutlass.range(num_rows, unroll_full=True):
+                        if tOhidx[m] >= 0 and scale[m] > 0.0:
+                            tOrO[None, m, None].store(
+                                tOrO[None, m, None].load()
+                                + scale[m] * tOrO_partial[None, m, None].load().to(Float32)
+                            )
+                # ===============================
+                # Step 7: Write final O to gmem
+                # ===============================
+                rO = cute.make_rmem_tensor_like(tOrO, self.dtype)
+                rO.store(tOrO.load().to(self.dtype))
+                if const_expr(cu_seqlens is None):
+                    mO_cur = mO[None, None, None, batch_idx]
+                else:
+                    mO_cur = cute.domain_offset((offset, 0, 0), mO)
+                mO_cur = utils.domain_offset_aligned((0, k_block * self.k_block_size, 0), mO_cur)
+                elems_per_store = const_expr(cute.size(gmem_tiled_copy_O.layout_tv_tiled[1]))
+                # 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(
         tOrOptr: cute.Tensor,
         tOsO_partial: cute.Tensor,
         tOhidx: cute.Tensor,
+        tOpO: Optional[cute.Tensor],
         tOcO: cute.Tensor,
         mO_cur_partial_layout: cute.Layout,
         split: Int32,
                 mO_partial_cur_copy = cute.tiled_divide(mO_partial_cur, (elems_per_load,))
                 for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True):
                     k_idx = tOcO[0, 0, k][1] // elems_per_load
+                    if const_expr(tOpO is None) or tOpO[k]:
                         cute.copy(
                             gmem_tiled_copy_O_partial,
                             mO_partial_cur_copy[None, k_idx, split],

build/torch-cuda/flash_fwd_mla_sm100.py ADDED Viewed

The diff for this file is too large to render. See raw diff

build/torch-cuda/flash_fwd_sm100.py CHANGED Viewed

The diff for this file is too large to render. See raw diff

build/torch-cuda/flash_fwd_sm120.py ADDED Viewed

	@@ -0,0 +1,61 @@

+# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+# SM120 (Blackwell GeForce / DGX Spark) forward pass.
+#
+# SM120 uses the same SM80-era MMA instructions (mma.sync.aligned.m16n8k16) but has
+# a smaller shared memory capacity (99 KB vs 163 KB on SM80). This module subclasses
+# FlashAttentionForwardSm80 and overrides the SMEM capacity check accordingly.
+import cutlass
+import cutlass.utils as utils_basic
+from cutlass.base_dsl.arch import Arch
+from .flash_fwd import FlashAttentionForwardSm80
+class FlashAttentionForwardSm120(FlashAttentionForwardSm80):
+    def __init__(self, *args, **kwargs):
+        """Force SM80 code paths while the DSL still targets the resident SM120 GPU."""
+        super().__init__(*args, **kwargs)
+        self.arch = Arch.sm_80
+    @staticmethod
+    def can_implement(
+        dtype,
+        head_dim,
+        head_dim_v,
+        tile_m,
+        tile_n,
+        num_stages,
+        num_threads,
+        is_causal,
+        Q_in_regs=False,
+    ) -> bool:
+        """Check if the kernel can be implemented on SM120.
+        Same logic as SM80 but uses SM120's shared memory capacity (99 KB).
+        """
+        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
+        # Shared memory usage: Q tile + (K tile + V tile)
+        smem_usage_Q = tile_m * head_dim * 2
+        smem_usage_K = tile_n * head_dim * num_stages * 2
+        smem_usage_V = tile_n * head_dim_v * num_stages * 2
+        smem_usage_QV = (
+            (smem_usage_Q + smem_usage_V) if not Q_in_regs else max(smem_usage_Q, smem_usage_V)
+        )
+        smem_usage = smem_usage_QV + smem_usage_K
+        # SM120 has 99 KB shared memory (vs 163 KB on SM80)
+        smem_capacity = utils_basic.get_smem_capacity_in_bytes("sm_120")
+        if smem_usage > smem_capacity:
+            return False
+        if (tile_m * 2) % num_threads != 0:
+            return False
+        return True

build/torch-cuda/flash_fwd_sm90.py ADDED Viewed

	@@ -0,0 +1,1550 @@

+# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+# SM90 (Hopper) forward pass for flash attention, extracted from flash_fwd.py.
+from types import SimpleNamespace
+from typing import Callable, Literal, Optional
+from functools import partial
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Int32, const_expr
+from cutlass.cute.nvgpu import cpasync, warpgroup
+from cutlass.utils import LayoutEnum
+import cutlass.utils.hopper_helpers as sm90_utils_basic
+from cutlass import pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
+from cutlass.base_dsl.arch import Arch
+from .quack import copy_utils
+from .quack import layout_utils
+from .quack import sm90_utils
+from .cute_dsl_utils import assume_tensor_aligned
+from . import utils
+from .mask import AttentionMask
+from .softmax import Softmax, apply_score_mod_inner
+from .seqlen_info import SeqlenInfoQK
+from .block_info import BlockInfo
+from .block_sparsity import BlockSparseTensors
+from .block_sparse_utils import (
+    produce_block_sparse_loads,
+    consume_block_sparse_loads,
+)
+from . import pipeline as pipeline_custom
+from .pack_gqa import PackGQA, pack_gqa_layout, make_packgqa_tiled_tma_atom
+from .paged_kv import PagedKVManager
+from .named_barrier import NamedBarrierFwd
+from .quack.cute_dsl_utils import ParamsBase
+from .tile_scheduler import (
+    TileSchedulerArguments,
+    SingleTileScheduler,
+    SingleTileLPTScheduler,
+    SingleTileVarlenScheduler,
+)
+from cutlass.cute import FastDivmodDivisor
+from .flash_fwd import FlashAttentionForwardBase
+from .utils import AuxData
+class FlashAttentionForwardSm90(FlashAttentionForwardBase):
+    def __init__(
+        self,
+        *args,
+        intra_wg_overlap: bool = True,
+        mma_pv_is_rs: bool = True,
+        paged_kv_non_tma: bool = False,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.intra_wg_overlap = intra_wg_overlap
+        self.mma_pv_is_rs = mma_pv_is_rs
+        self.buffer_align_bytes = 1024
+        self.use_tma_KV = not paged_kv_non_tma
+        assert self.use_tma_KV or not (self.check_hdim_oob or self.check_hdim_v_oob), (
+            "Paged KV does not support irregular head dim"
+        )
+        self.cluster_shape_mn = (1, 1)
+        assert self.arch.is_family_of(Arch.sm_90a), "Only SM 9.x is supported"
+    def _get_smem_layout_atom(self):
+        sQ_layout_atom = warpgroup.make_smem_layout_atom(
+            sm90_utils_basic.get_smem_layout_atom(LayoutEnum.ROW_MAJOR, self.dtype, self.tile_hdim),
+            self.dtype,
+        )
+        sK_layout_atom = sQ_layout_atom
+        sV_layout_atom = warpgroup.make_smem_layout_atom(
+            sm90_utils_basic.get_smem_layout_atom(
+                LayoutEnum.ROW_MAJOR, self.dtype, self.tile_hdimv
+            ),
+            self.dtype,
+        )
+        sO_layout_atom = sV_layout_atom
+        if not self.mma_pv_is_rs:
+            sP_layout_atom = warpgroup.make_smem_layout_atom(
+                sm90_utils_basic.get_smem_layout_atom(
+                    LayoutEnum.ROW_MAJOR, self.dtype, self.tile_n
+                ),
+                self.dtype,
+            )
+        else:
+            sP_layout_atom = None
+        return sQ_layout_atom, sK_layout_atom, sV_layout_atom, sO_layout_atom, sP_layout_atom
+    def _get_tiled_mma(self):
+        tiled_mma_qk = sm90_utils_basic.make_trivial_tiled_mma(
+            self.dtype,
+            self.dtype,
+            warpgroup.OperandMajorMode.K,
+            warpgroup.OperandMajorMode.K,
+            Float32,
+            atom_layout_mnk=(self.tile_m // 64, 1, 1),
+            tiler_mn=(64, self.tile_n),
+        )
+        tiled_mma_pv = sm90_utils_basic.make_trivial_tiled_mma(
+            self.dtype,
+            self.dtype,
+            warpgroup.OperandMajorMode.K,
+            warpgroup.OperandMajorMode.MN,
+            Float32,
+            atom_layout_mnk=(self.tile_m // 64, 1, 1),  # Might need (1, 2, 1) for hdim 512
+            tiler_mn=(64, self.tile_hdimv),
+            a_source=warpgroup.OperandSource.RMEM
+            if self.mma_pv_is_rs
+            else warpgroup.OperandSource.SMEM,
+        )
+        return tiled_mma_qk, tiled_mma_pv
+    def _get_shared_storage_cls(self):
+        sQ_struct, sK_struct, sV_struct = [
+            cute.struct.Align[
+                cute.struct.MemRange[self.dtype, cute.cosize(layout)], self.buffer_align_bytes
+            ]
+            for layout in (self.sQ_layout, self.sK_layout, self.sV_layout)
+        ]
+        cosize_sQV = max(cute.cosize(self.sQ_layout), cute.cosize(self.sV_layout))
+        sQV_struct = cute.struct.Align[cute.struct.MemRange[self.dtype, cosize_sQV], 1024]
+        cosize_sP = cute.cosize(self.sP_layout) if const_expr(self.sP_layout is not None) else 0
+        sP_struct = cute.struct.Align[cute.struct.MemRange[self.dtype, cosize_sP], 1024]
+        # 1 stage * 2 for Q pipeline (full + empty), self.num_stages*2 for K, self.num_stages*2 for V,
+        mbar_ptr_Q_struct = cute.struct.MemRange[cutlass.Int64, 1 * 2]
+        mbar_ptr_K_struct = cute.struct.MemRange[cutlass.Int64, self.num_stages * 2]
+        mbar_ptr_V_struct = cute.struct.MemRange[cutlass.Int64, self.num_stages * 2]
+        @cute.struct
+        class SharedStorageQKV:
+            mbar_ptr_Q: mbar_ptr_Q_struct
+            mbar_ptr_K: mbar_ptr_K_struct
+            mbar_ptr_V: mbar_ptr_V_struct
+            sV: sV_struct
+            sQ: sQ_struct
+            sK: sK_struct
+            sP: sP_struct
+        @cute.struct
+        class SharedStorageSharedQV:
+            mbar_ptr_Q: mbar_ptr_Q_struct
+            mbar_ptr_K: mbar_ptr_K_struct
+            mbar_ptr_V: mbar_ptr_V_struct
+            sQ: sQV_struct
+            sK: sK_struct
+            sP: sP_struct
+        return SharedStorageQKV if const_expr(not self.Q_in_regs) else SharedStorageSharedQV
+    @cute.jit
+    def __call__(
+        self,
+        mQ: cute.Tensor,  # (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
+        mK: cute.Tensor,  # (b_k, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k or (num_pages, page_size, h_k, d) if there is page_table
+        mV: cute.Tensor,  # (b_k, s_k, h_k, dv) or (total_k, h_k, dv) if there is cu_seqlens_k or (num_pages, page_size, h_k, dv) if there is page_table
+        mO: cute.Tensor,  # (b, s_q, h, dv) or (total_q, h, dv) if there is cu_seqlens_q
+        mLSE: Optional[cute.Tensor],
+        softmax_scale: Float32,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        mSeqUsedQ: Optional[cute.Tensor] = None,
+        mSeqUsedK: Optional[cute.Tensor] = None,
+        mPageTable: Optional[cute.Tensor] = None,  # (b_k, max_num_pages_per_seq)
+        window_size_left: Int32 | int | None = None,
+        window_size_right: Int32 | int | None = None,
+        learnable_sink: Optional[cute.Tensor] = None,
+        blocksparse_tensors: Optional[BlockSparseTensors] = None,
+        aux_data: AuxData = AuxData(),
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
+        mCuTotalSplitsMBlocks: Optional[cute.Tensor] = None,
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
+        stream: cuda.CUstream = None,
+    ):
+        """Configures and launches the flash attention kernel.
+        mQ/mK/mV/mO has same data types(supports fp16 and bf16) and same layout:
+        (batch_size, seqlen_q, num_head, head_dim):(_, _, _, 1)
+        """
+        self._check_type(
+            *(
+                t.element_type if t is not None else None
+                for t in (mQ, mK, mV, mO, mLSE, mCuSeqlensQ, mCuSeqlensK, mSeqUsedQ, mSeqUsedK)
+            )
+        )
+        self.varlen_q = mCuSeqlensQ is not None or mSeqUsedQ is not None
+        mQ, mK, mV, mO = [assume_tensor_aligned(t) for t in (mQ, mK, mV, mO)]
+        QO_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1]
+        mQ, mO = [layout_utils.select(t, QO_layout_transpose) for t in (mQ, mO)]
+        KV_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensK is None) else [0, 2, 1]
+        mK, mV = [layout_utils.select(t, KV_layout_transpose) for t in (mK, mV)]
+        LSE_layout_transpose = [2, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 0]
+        mLSE = (
+            layout_utils.select(mLSE, LSE_layout_transpose)
+            if const_expr(mLSE is not None)
+            else None
+        )
+        tiled_mma_qk, tiled_mma_pv = self._get_tiled_mma()
+        self.num_mma_threads = tiled_mma_qk.size
+        self.num_threads_per_warp_group = 128
+        self.num_wg_mma = self.num_mma_threads // self.num_threads_per_warp_group
+        assert self.num_wg_mma in [1, 2, 3]
+        self.num_threads = self.num_threads_per_warp_group * (self.num_wg_mma + 1)
+        self.num_producer_threads = 32
+        self.num_Q_load_threads = self.num_threads_per_warp_group  # If not TMA_Q
+        self.num_epilogue_threads = self.num_mma_threads
+        self.num_mma_regs, self.num_producer_regs = {1: (256, 56), 2: (240, 24), 3: (160, 32)}[
+            self.num_wg_mma
+        ]
+        self.use_block_sparsity = cutlass.const_expr(blocksparse_tensors is not None)
+        self.use_scheduler_barrier = (
+            (self.num_wg_mma >= 2 and self.tile_hdim <= 128)
+            if const_expr(self.intra_wg_overlap)
+            else (self.num_wg_mma == 2)
+        )
+        self.use_tma_Q = self.arch >= Arch.sm_90 and not (
+            self.pack_gqa and self.tile_m % self.qhead_per_kvhead != 0
+        )
+        self.use_tma_O = self.use_tma_Q
+        # Producer needs more registers when doing cp.async Q or KV loads
+        if const_expr(self.num_wg_mma == 2 and (not self.use_tma_Q or not self.use_tma_KV)):
+            self.num_mma_regs, self.num_producer_regs = 224, 40
+        self.rescale_O_before_gemm = self.tile_hdimv > 128 and self.intra_wg_overlap
+        self._setup_attributes()
+        # TODO: we prob don't need most of what's in _setup_attributes
+        self.sQ_layout, self.sK_layout, self.sV_layout, self.sO_layout = [
+            sm90_utils.make_smem_layout(mX.element_type, LayoutEnum.ROW_MAJOR, shape, stage)
+            for mX, shape, stage in [
+                (mQ, (self.tile_m, self.tile_hdim), None),
+                (mK, (self.tile_n, self.tile_hdim), self.num_stages),
+                (mV, (self.tile_n, self.tile_hdimv), self.num_stages),
+                (mO, (self.tile_m, self.tile_hdimv), None),
+            ]
+        ]
+        self.sP_layout = None
+        if const_expr(not self.mma_pv_is_rs):
+            self.sP_layout = sm90_utils.make_smem_layout(
+                mV.element_type, LayoutEnum.ROW_MAJOR, (self.tile_m, self.tile_n)
+            )
+        SharedStorage = self._get_shared_storage_cls()
+        mQ_og, mO_og = mQ, mO
+        if const_expr(self.pack_gqa):
+            nheads_kv = mK.shape[2]
+            mQ = pack_gqa_layout(mQ, self.qhead_per_kvhead, nheads_kv, head_idx=2)
+            mO = pack_gqa_layout(mO, self.qhead_per_kvhead, nheads_kv, head_idx=2)
+            if const_expr(mLSE is not None):
+                mLSE = pack_gqa_layout(mLSE, self.qhead_per_kvhead, nheads_kv, head_idx=1)
+        # TMA
+        gmem_tiled_copy_Q = cpasync.CopyBulkTensorTileG2SOp()
+        gmem_tiled_copy_KV = cpasync.CopyBulkTensorTileG2SOp()  # Might multicast
+        gmem_tiled_copy_O = cpasync.CopyBulkTensorTileS2GOp()
+        self.tma_copy_bytes = {
+            name: cute.size_in_bytes(mX.element_type, cute.select(layout, mode=[0, 1]))
+            for name, mX, layout in [
+                ("Q", mQ, self.sQ_layout),
+                ("K", mK, self.sK_layout),
+                ("V", mV, self.sV_layout),
+            ]
+        }
+        make_tiled_tma_atom_fn = (
+            partial(make_packgqa_tiled_tma_atom, qhead_per_kvhead=self.qhead_per_kvhead, head_idx=2)
+            if const_expr(self.pack_gqa)
+            else cpasync.make_tiled_tma_atom
+        )
+        tma_atom_Q, tma_tensor_Q = None, None
+        if const_expr(self.use_tma_Q):
+            tma_atom_Q, tma_tensor_Q = make_tiled_tma_atom_fn(
+                gmem_tiled_copy_Q,
+                mQ_og if const_expr(self.pack_gqa) else mQ,
+                self.sQ_layout,
+                (self.tile_m, self.tile_hdim),  # No mcast
+            )
+        tma_atom_K, tma_tensor_K = None, None
+        tma_atom_V, tma_tensor_V = None, None
+        if const_expr(self.use_tma_KV):
+            tma_atom_K, tma_tensor_K = cpasync.make_tiled_tma_atom(
+                gmem_tiled_copy_KV,
+                mK,
+                cute.select(self.sK_layout, mode=[0, 1]),
+                (self.tile_n, self.tile_hdim),
+                1,  # No mcast for now
+            )
+            tma_atom_V, tma_tensor_V = cpasync.make_tiled_tma_atom(
+                gmem_tiled_copy_KV,
+                mV,
+                cute.select(self.sV_layout, mode=[0, 1]),
+                (self.tile_n, self.tile_hdimv),
+                1,  # No mcast for now
+            )
+        tma_atom_O, tma_tensor_O = None, None
+        if const_expr(self.use_tma_O):
+            mO_tma = mO_og if const_expr(self.pack_gqa) else mO
+            if const_expr(self.varlen_q):
+                mO_tma = copy_utils.create_ragged_tensor_for_tma(
+                    mO_tma, ragged_dim=0, ptr_shift=True
+                )
+            tma_atom_O, tma_tensor_O = make_tiled_tma_atom_fn(
+                gmem_tiled_copy_O,
+                mO_tma,
+                self.sO_layout,
+                (self.tile_m, self.tile_hdimv),  # No mcast
+            )
+        if const_expr(mCuSeqlensQ is not None or mSeqUsedQ is not None):
+            # TODO: dispatch to DynamicPersistentVarlenScheduler when appropriate
+            TileScheduler = SingleTileVarlenScheduler
+        else:
+            TileScheduler = (
+                SingleTileScheduler
+                if const_expr(not self.is_causal or self.is_local)
+                else SingleTileLPTScheduler
+            )
+        tile_sched_args = TileSchedulerArguments(
+            cute.ceil_div(cute.size(mQ.shape[0]), self.tile_m),
+            cute.size(mQ.shape[2]),
+            cute.size(mQ.shape[3])
+            if const_expr(mCuSeqlensQ is None)
+            else cute.size(mCuSeqlensQ.shape[0] - 1),
+            1,  # num_splits
+            cute.size(mK.shape[0])
+            if const_expr(mPageTable is None)
+            else mK.shape[0] * mPageTable.shape[1],
+            mQ.shape[1],
+            mV.shape[1],
+            total_q=cute.size(mQ.shape[0])
+            if const_expr(mCuSeqlensQ is not None)
+            else cute.size(mQ.shape[0]) * cute.size(mQ.shape[3]),
+            tile_shape_mn=(self.tile_m, self.tile_n),
+            mCuSeqlensQ=mCuSeqlensQ,
+            mSeqUsedQ=mSeqUsedQ,
+            qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
+            element_size=self.dtype.width // 8,
+            is_persistent=False,
+            lpt=self.is_causal or self.is_local,
+            cu_total_m_blocks_ptr=mCuTotalMBlocks,
+            cu_total_splits_m_blocks_ptr=mCuTotalSplitsMBlocks,
+        )
+        tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
+        grid_dim = TileScheduler.get_grid_shape(tile_sched_params)
+        softmax_scale_log2, softmax_scale = utils.compute_softmax_scale_log2(
+            softmax_scale, self.score_mod
+        )
+        window_size_left = Int32(window_size_left) if window_size_left is not None else None
+        window_size_right = Int32(window_size_right) if window_size_right is not None else None
+        fastdiv_mods = utils.compute_fastdiv_mods(
+            mQ, mK, self.qhead_per_kvhead, self.pack_gqa, aux_data.tensors, mPageTable
+        )
+        self.kernel(
+            tma_tensor_Q if const_expr(self.use_tma_Q) else mQ,
+            tma_tensor_K if const_expr(self.use_tma_KV) else mK,
+            tma_tensor_V if const_expr(self.use_tma_KV) else mV,
+            tma_tensor_O if const_expr(self.use_tma_O) else mO,
+            mLSE,
+            mCuSeqlensQ,
+            mCuSeqlensK,
+            mSeqUsedQ,
+            mSeqUsedK,
+            mPageTable,
+            tma_atom_Q,
+            tma_atom_K,
+            tma_atom_V,
+            tma_atom_O,
+            softmax_scale_log2,
+            softmax_scale,
+            window_size_left,
+            window_size_right,
+            learnable_sink,
+            blocksparse_tensors,
+            self.sQ_layout,
+            self.sK_layout,
+            self.sV_layout,
+            self.sO_layout,
+            self.sP_layout,
+            self.gmem_tiled_copy_Q,
+            self.gmem_tiled_copy_K,
+            self.gmem_tiled_copy_V,
+            self.gmem_tiled_copy_O,
+            tiled_mma_qk,
+            tiled_mma_pv,
+            tile_sched_params,
+            TileScheduler,
+            SharedStorage,
+            aux_data,
+            fastdiv_mods,
+        ).launch(
+            grid=grid_dim,
+            block=[self.num_threads, 1, 1],
+            stream=stream,
+            min_blocks_per_mp=1,
+        )
+    @cute.kernel
+    def kernel(
+        self,
+        mQ: cute.Tensor,
+        mK: cute.Tensor,
+        mV: cute.Tensor,
+        mO: cute.Tensor,
+        mLSE: Optional[cute.Tensor],
+        mCuSeqlensQ: Optional[cute.Tensor],
+        mCuSeqlensK: Optional[cute.Tensor],
+        mSeqUsedQ: Optional[cute.Tensor],
+        mSeqUsedK: Optional[cute.Tensor],
+        mPageTable: Optional[cute.Tensor],
+        tma_atom_Q: Optional[cute.CopyAtom],
+        tma_atom_K: Optional[cute.CopyAtom],
+        tma_atom_V: Optional[cute.CopyAtom],
+        tma_atom_O: Optional[cute.CopyAtom],
+        softmax_scale_log2: Float32,
+        softmax_scale: Optional[Float32],
+        window_size_left: Optional[Int32],
+        window_size_right: Optional[Int32],
+        learnable_sink: Optional[cute.Tensor],
+        blocksparse_tensors: Optional[BlockSparseTensors],
+        sQ_layout: cute.ComposedLayout,
+        sK_layout: cute.ComposedLayout,
+        sV_layout: cute.ComposedLayout,
+        sO_layout: cute.ComposedLayout,
+        sP_layout: cute.ComposedLayout | None,
+        gmem_tiled_copy_Q: cute.TiledCopy,
+        gmem_tiled_copy_K: cute.TiledCopy,
+        gmem_tiled_copy_V: cute.TiledCopy,
+        gmem_tiled_copy_O: cute.TiledCopy,
+        tiled_mma_qk: cute.TiledMma,
+        tiled_mma_pv: cute.TiledMma,
+        tile_sched_params: ParamsBase,
+        TileScheduler: cutlass.Constexpr[Callable],
+        SharedStorage: cutlass.Constexpr[Callable],
+        aux_data: AuxData = AuxData(),
+        fastdiv_mods=None,
+    ):
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+        # Prefetch tma descriptor
+        if warp_idx == 0:
+            for tma_atom in (tma_atom_Q, tma_atom_K, tma_atom_V, tma_atom_O):
+                if const_expr(tma_atom is not None):
+                    cpasync.prefetch_descriptor(tma_atom)
+        smem = cutlass.utils.SmemAllocator()
+        storage = smem.allocate(SharedStorage)
+        # Mbarrier / pipeline init
+        mbar_ptr_Q = storage.mbar_ptr_Q.data_ptr()
+        ThreadCooperativeGroup = partial(pipeline.CooperativeGroup, pipeline.Agent.Thread)
+        tma_warp = ThreadCooperativeGroup(1)
+        load_threads = ThreadCooperativeGroup(self.num_threads_per_warp_group)
+        mma_warps = ThreadCooperativeGroup(self.num_mma_threads // cute.arch.WARP_SIZE)
+        if const_expr(self.use_tma_Q):
+            pipeline_q = pipeline_custom.PipelineTmaAsync.create(
+                barrier_storage=mbar_ptr_Q,
+                num_stages=1,
+                producer_group=tma_warp,
+                consumer_group=mma_warps,
+                tx_count=self.tma_copy_bytes["Q"],
+                defer_sync=True,
+            )
+        else:
+            pipeline_q = pipeline_custom.PipelineCpAsync.create(
+                barrier_storage=mbar_ptr_Q,
+                num_stages=1,
+                producer_group=load_threads,
+                consumer_group=mma_warps,
+                defer_sync=True,
+                elect_one_release=True,
+                syncwarp_before_release=False,
+            )
+        if const_expr(self.use_tma_KV):
+            pipeline_k = pipeline_custom.PipelineTmaAsync.create(
+                barrier_storage=storage.mbar_ptr_K.data_ptr(),
+                num_stages=self.num_stages,
+                producer_group=tma_warp,
+                consumer_group=mma_warps,
+                tx_count=self.tma_copy_bytes["K"],
+                defer_sync=True,
+            )
+            pipeline_v = pipeline_custom.PipelineTmaAsync.create(
+                barrier_storage=storage.mbar_ptr_V.data_ptr(),
+                num_stages=self.num_stages,
+                producer_group=tma_warp,
+                consumer_group=mma_warps,
+                tx_count=self.tma_copy_bytes["V"],
+                defer_sync=True,
+            )
+        else:
+            pipeline_k = pipeline_custom.PipelineCpAsync.create(
+                barrier_storage=storage.mbar_ptr_K.data_ptr(),
+                num_stages=self.num_stages,
+                producer_group=load_threads,
+                consumer_group=mma_warps,
+                defer_sync=True,
+                elect_one_release=True,
+                syncwarp_before_release=False,
+            )
+            pipeline_v = pipeline_custom.PipelineCpAsync.create(
+                barrier_storage=storage.mbar_ptr_V.data_ptr(),
+                num_stages=self.num_stages,
+                producer_group=load_threads,
+                consumer_group=mma_warps,
+                defer_sync=True,
+                elect_one_release=True,
+                syncwarp_before_release=False,
+            )
+        # Cluster arrive after barrier init
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
+        # ///////////////////////////////////////////////////////////////////////////////
+        # Get shared memory buffer
+        # ///////////////////////////////////////////////////////////////////////////////
+        sQ = storage.sQ.get_tensor(sQ_layout.outer, swizzle=sQ_layout.inner)
+        sK = storage.sK.get_tensor(sK_layout.outer, swizzle=sK_layout.inner)
+        if const_expr(not self.Q_in_regs):
+            sV = storage.sV.get_tensor(sV_layout.outer, swizzle=sV_layout.inner)
+        else:
+            sV = storage.sQ.get_tensor(
+                sV_layout.outer, swizzle=sV_layout.inner, dtype=mV.element_type
+            )
+        # Transpose view of V to tensor with layout (head_dim_v, tile_n) for tiled mma
+        sVt = layout_utils.transpose_view(sV)
+        sP = None
+        if const_expr(sP_layout is not None):
+            sP = storage.sP.get_tensor(sP_layout.outer, swizzle=sP_layout.inner)
+        # reuse sQ's data iterator
+        sO = storage.sQ.get_tensor(sO_layout.outer, swizzle=sO_layout.inner, dtype=self.dtype)
+        block_info = BlockInfo(
+            self.tile_m,
+            self.tile_n,
+            self.is_causal,
+            self.is_local,
+            False,  # is_split_kv
+            window_size_left,
+            window_size_right,
+            qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
+        )
+        SeqlenInfoCls = partial(
+            SeqlenInfoQK.create,
+            seqlen_q_static=mQ.shape[0] if const_expr(not self.pack_gqa) else mQ.shape[0][1],
+            seqlen_k_static=mK.shape[0]
+            if const_expr(mPageTable is None)
+            else mK.shape[0] * mPageTable.shape[1],
+            mCuSeqlensQ=mCuSeqlensQ,
+            mCuSeqlensK=mCuSeqlensK,
+            mSeqUsedQ=mSeqUsedQ,
+            mSeqUsedK=mSeqUsedK,
+            mCuTotalMBlocks=(
+                blocksparse_tensors.cu_total_m_blocks if blocksparse_tensors is not None else None
+            ),
+            mCuBlockIdxOffsets=(
+                blocksparse_tensors.cu_block_idx_offsets
+                if blocksparse_tensors is not None
+                else None
+            ),
+            # Don't need to pass in tile_mn because we won't access offset_padded
+        )
+        AttentionMaskCls = partial(
+            AttentionMask,
+            self.tile_m,
+            self.tile_n,
+            window_size_left=window_size_left,
+            window_size_right=window_size_right,
+            qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
+        )
+        TileSchedulerCls = partial(TileScheduler.create, tile_sched_params)
+        # Cluster wait before starting
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
+        if warp_idx < 4:  # Producer
+            cute.arch.setmaxregister_decrease(self.num_producer_regs)
+            self.load(
+                mQ,
+                mK,
+                mV,
+                sQ,
+                sK,
+                sV,
+                tma_atom_Q,
+                tma_atom_K,
+                tma_atom_V,
+                pipeline_k,
+                pipeline_v,
+                pipeline_q,
+                gmem_tiled_copy_Q,
+                mPageTable,
+                blocksparse_tensors,
+                block_info,
+                SeqlenInfoCls,
+                TileSchedulerCls,
+            )
+        else:  # Consumer
+            cute.arch.setmaxregister_increase(self.num_mma_regs)
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Tile MMA compute thread partitions and allocate accumulators
+            # ///////////////////////////////////////////////////////////////////////////////
+            tidx, _, _ = cute.arch.thread_idx()
+            tidx = tidx - 128
+            self.mma(
+                tiled_mma_qk,
+                tiled_mma_pv,
+                mO,
+                mLSE,
+                sQ,
+                sK,
+                sVt,
+                sP,
+                sO,
+                learnable_sink,
+                pipeline_k,
+                pipeline_v,
+                pipeline_q,
+                gmem_tiled_copy_O,
+                tma_atom_O,
+                tidx,
+                softmax_scale_log2,
+                softmax_scale,
+                block_info,
+                SeqlenInfoCls,
+                AttentionMaskCls,
+                TileSchedulerCls,
+                blocksparse_tensors,
+                aux_data,
+                fastdiv_mods,
+            )
+    @cute.jit
+    def load(
+        self,
+        mQ: cute.Tensor,
+        mK: cute.Tensor,
+        mV: cute.Tensor,
+        sQ: cute.Tensor,
+        sK: cute.Tensor,
+        sV: cute.Tensor,
+        tma_atom_Q: Optional[cute.CopyAtom],
+        tma_atom_K: Optional[cute.CopyAtom],
+        tma_atom_V: Optional[cute.CopyAtom],
+        pipeline_k: pipeline.PipelineAsync,
+        pipeline_v: pipeline.PipelineAsync,
+        pipeline_q: pipeline.PipelineAsync,
+        gmem_tiled_copy_Q: cute.TiledCopy,
+        mPageTable: Optional[cute.Tensor],
+        blocksparse_tensors: Optional[BlockSparseTensors],
+        block_info: BlockInfo,
+        SeqlenInfoCls: Callable,
+        TileSchedulerCls: Callable,
+    ):
+        warp_idx_in_wg = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4
+        tidx, _, _ = cute.arch.thread_idx()
+        # TMA: only warp 0 loads. cp_async: all warps load.
+        # When not use_tma_Q, all 128 producer threads participate in Q loading.
+        is_load_warp = warp_idx_in_wg == 0 or const_expr(not self.use_tma_KV or not self.use_tma_Q)
+        # KV loading restricted to warp 0 for TMA, all warps for non-TMA KV
+        is_kv_load_warp = warp_idx_in_wg == 0 or const_expr(not self.use_tma_KV)
+        if is_load_warp:
+            q_producer_phase = Int32(1)
+            kv_producer_state = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Producer, self.num_stages
+            )
+            tile_scheduler = TileSchedulerCls()
+            work_tile = tile_scheduler.initial_work_tile_info()
+            while work_tile.is_valid_tile:
+                # if work_tile.is_valid_tile:
+                m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+                seqlen = SeqlenInfoCls(batch_idx)
+                mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx]
+                head_idx_kv = (
+                    head_idx // self.qhead_per_kvhead if const_expr(not self.pack_gqa) else head_idx
+                )
+                load_Q = None
+                if const_expr(self.use_tma_Q):
+                    gQ = cute.local_tile(mQ_cur, (self.tile_m, self.tile_hdim), (m_block, 0))
+                    load_Q, _, _ = copy_utils.tma_get_copy_fn(
+                        tma_atom_Q, 0, cute.make_layout(1), gQ, sQ, single_stage=True
+                    )
+                paged_kv_manager = None
+                tma_load_K_fn = None
+                tma_load_V_fn = None
+                if const_expr(self.use_tma_KV):
+                    # === TMA path (non-paged and paged with page_size == n_block_size) ===
+                    if const_expr(mPageTable is not None):
+                        # Paged TMA: keep page dimension indexable
+                        mK_cur = mK[None, None, head_idx_kv, None]
+                        mV_cur = mV[None, None, head_idx_kv, None]
+                        gK = cute.local_tile(mK_cur, (self.tile_n, self.tile_hdim), (0, 0, None))
+                        gV = cute.local_tile(mV_cur, (self.tile_n, self.tile_hdimv), (0, 0, None))
+                    else:
+                        # Non-paged TMA
+                        mK_cur = seqlen.offset_batch_K(mK, batch_idx, dim=3)[
+                            None, None, head_idx_kv
+                        ]
+                        mV_cur = seqlen.offset_batch_K(mV, batch_idx, dim=3)[
+                            None, None, head_idx_kv
+                        ]
+                        gK = cute.local_tile(mK_cur, (self.tile_n, self.tile_hdim), (None, 0))
+                        gV = cute.local_tile(mV_cur, (self.tile_n, self.tile_hdimv), (None, 0))
+                    # TODO: mcast
+                    tma_load_K_fn, _, _ = copy_utils.tma_get_copy_fn(
+                        tma_atom_K, 0, cute.make_layout(1), gK, sK
+                    )
+                    tma_load_K_fn = copy_utils.tma_producer_copy_fn(tma_load_K_fn, pipeline_k)
+                    tma_load_V_fn, _, _ = copy_utils.tma_get_copy_fn(
+                        tma_atom_V, 0, cute.make_layout(1), gV, sV
+                    )
+                    tma_load_V_fn = copy_utils.tma_producer_copy_fn(tma_load_V_fn, pipeline_v)
+                else:
+                    # === cp_async path (paged KV with page_size != n_block_size) ===
+                    paged_kv_manager = PagedKVManager.create(
+                        mPageTable,
+                        mK,
+                        mV,
+                        FastDivmodDivisor(mK.shape[0]),
+                        batch_idx,
+                        head_idx_kv,
+                        tidx,
+                        seqlen.seqlen_k,
+                        0,  # leftpad_k
+                        self.tile_n,
+                        self.tile_hdim,
+                        self.tile_hdimv,
+                        self.num_threads_per_warp_group,
+                        mK.element_type,
+                        arch=self.arch.major * 10 + self.arch.minor,
+                    )
+                load_K = partial(
+                    self.load_KV,
+                    tma_load_K_fn,
+                    paged_kv_manager,
+                    sK,
+                    pipeline_kv=pipeline_k,
+                    K_or_V="K",
+                )
+                load_V = partial(
+                    self.load_KV,
+                    tma_load_V_fn,
+                    paged_kv_manager,
+                    sV,
+                    pipeline_kv=pipeline_v,
+                    K_or_V="V",
+                )
+                pack_gqa = None
+                if const_expr(not self.use_tma_Q):
+                    pack_gqa = PackGQA(
+                        self.tile_m, self.tile_hdim, self.check_hdim_oob, self.qhead_per_kvhead
+                    )
+                if const_expr(not self.use_block_sparsity):
+                    n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block)
+                    # if cute.arch.thread_idx()[0] == 0:
+                    #     cute.printf("m_block = %d, n_block_min: %d, n_block_max: %d", m_block, n_block_min, n_block_max)
+                    # Clamp n_block to 0 when n_block_max == 0 (can happen with causal
+                    # + pack_gqa when seqlen_k < tile_n). TMA handles n_block=-1
+                    # gracefully (fills zeros), but cp.async would crash on
+                    # out-of-bounds page table access.
+                    n_block = (
+                        n_block_max - 1
+                        if const_expr(self.use_tma_KV)
+                        else cutlass.max(n_block_max - 1, 0)
+                    )
+                    page_idx = (
+                        mPageTable[batch_idx, n_block]
+                        if const_expr(mPageTable is not None and self.use_tma_KV)
+                        else None
+                    )
+                    # First iteration: load K on pipeline_k, Q on pipeline_q
+                    if is_kv_load_warp:
+                        pipeline_k.producer_acquire(kv_producer_state)
+                        if const_expr(not self.use_tma_KV):
+                            paged_kv_manager.load_page_table(n_block)
+                        load_K(block=n_block, producer_state=kv_producer_state, page_idx=page_idx)
+                    if const_expr(self.use_tma_Q):
+                        if warp_idx_in_wg == 0:
+                            pipeline_q.producer_acquire_w_index_phase(0, q_producer_phase)
+                            load_Q(tma_bar_ptr=pipeline_q.sync_object_full.get_barrier(0))
+                            q_producer_phase ^= 1
+                    else:
+                        pipeline_q.producer_acquire_w_index_phase(0, q_producer_phase)
+                        pack_gqa.load_Q(
+                            mQ_cur, sQ, gmem_tiled_copy_Q, tidx, m_block, seqlen.seqlen_q
+                        )
+                        cute.arch.cp_async_commit_group()
+                        pipeline_q.producer_commit_w_index(0)
+                        q_producer_phase ^= 1
+                    if is_kv_load_warp:
+                        if const_expr(not self.intra_wg_overlap or not self.use_tma_KV):
+                            pipeline_v.producer_acquire(kv_producer_state)
+                            load_V(
+                                block=n_block, producer_state=kv_producer_state, page_idx=page_idx
+                            )
+                            kv_producer_state.advance()
+                            for i in cutlass.range(n_block_max - 1 - n_block_min, unroll=1):
+                                n_block = n_block_max - 1 - i - 1
+                                page_idx = (
+                                    mPageTable[batch_idx, n_block]
+                                    if const_expr(mPageTable is not None and self.use_tma_KV)
+                                    else None
+                                )
+                                if const_expr(not self.use_tma_KV):
+                                    paged_kv_manager.load_page_table(n_block)
+                                pipeline_k.producer_acquire(kv_producer_state)
+                                load_K(
+                                    block=n_block,
+                                    producer_state=kv_producer_state,
+                                    page_idx=page_idx,
+                                )
+                                pipeline_v.producer_acquire(kv_producer_state)
+                                load_V(
+                                    block=n_block,
+                                    producer_state=kv_producer_state,
+                                    page_idx=page_idx,
+                                )
+                                kv_producer_state.advance()
+                        else:
+                            for i in cutlass.range(n_block_max - 1 - n_block_min, unroll=1):
+                                n_block_prev = n_block_max - i - 1
+                                n_block = n_block_prev - 1
+                                page_idx = (
+                                    mPageTable[batch_idx, n_block]
+                                    if const_expr(mPageTable is not None)
+                                    else None
+                                )
+                                page_idx_prev = (
+                                    mPageTable[batch_idx, n_block_prev]
+                                    if const_expr(mPageTable is not None)
+                                    else None
+                                )
+                                kv_producer_state_prev = kv_producer_state.clone()
+                                kv_producer_state.advance()
+                                pipeline_k.producer_acquire(kv_producer_state)
+                                load_K(
+                                    block=n_block,
+                                    producer_state=kv_producer_state,
+                                    page_idx=page_idx,
+                                )
+                                pipeline_v.producer_acquire(kv_producer_state_prev)
+                                load_V(
+                                    block=n_block_prev,
+                                    producer_state=kv_producer_state_prev,
+                                    page_idx=page_idx_prev,
+                                )
+                            n_block = n_block_min
+                            page_idx = (
+                                mPageTable[batch_idx, n_block]
+                                if const_expr(mPageTable is not None)
+                                else None
+                            )
+                            pipeline_v.producer_acquire(kv_producer_state)
+                            load_V(
+                                block=n_block, producer_state=kv_producer_state, page_idx=page_idx
+                            )
+                            kv_producer_state.advance()
+                else:
+                    # Block sparsity: use TMA closures directly (not paged)
+                    # Load Q on pipeline_q, separate from K/V pipeline
+                    if const_expr(self.use_tma_Q):
+                        if warp_idx_in_wg == 0:
+                            pipeline_q.producer_acquire_w_index_phase(0, q_producer_phase)
+                            load_Q(tma_bar_ptr=pipeline_q.sync_object_full.get_barrier(0))
+                            q_producer_phase ^= 1
+                    else:
+                        pipeline_q.producer_acquire_w_index_phase(0, q_producer_phase)
+                        pack_gqa.load_Q(
+                            mQ_cur, sQ, gmem_tiled_copy_Q, tidx, m_block, seqlen.seqlen_q
+                        )
+                        cute.arch.cp_async_commit_group()
+                        pipeline_q.producer_commit_w_index(0)
+                        q_producer_phase ^= 1
+                    if is_kv_load_warp:
+                        kv_producer_state = produce_block_sparse_loads(
+                            blocksparse_tensors,
+                            batch_idx,
+                            head_idx,
+                            m_block,
+                            seqlen,
+                            kv_producer_state,
+                            tma_load_K_fn,
+                            tma_load_V_fn,
+                            pipeline_k,
+                            pipeline_v,
+                            self.intra_wg_overlap,
+                            self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
+                            self.q_subtile_factor,
+                        )
+                tile_scheduler.prefetch_next_work()
+                tile_scheduler.advance_to_next_work()
+                work_tile = tile_scheduler.get_current_work()
+                # End of persistent scheduler loop
+            # Producer tail is only useful for cluster to avoid early exit of blocks.
+            # We only need producer_tail on V since that's the last that's loaded, we don't
+            # need it for Q (no cluster) and K.
+            if is_kv_load_warp:
+                pipeline_v.producer_tail(kv_producer_state)
+    @cute.jit
+    def load_KV(
+        self,
+        tma_load_fn: Optional[Callable],
+        paged_kv_manager: Optional[PagedKVManager],
+        sX: cute.Tensor,
+        block: Int32,
+        pipeline_kv: pipeline.PipelineAsync,
+        producer_state: pipeline.PipelineState,
+        K_or_V: Literal["K", "V"],
+        page_idx: Optional[Int32] = None,
+    ):
+        if const_expr(self.use_tma_KV):
+            src_idx = block if const_expr(page_idx is None) else page_idx
+            tma_load_fn(src_idx=src_idx, producer_state=producer_state)
+        else:
+            paged_kv_manager.load_KV(block, sX[None, None, producer_state.index], K_or_V)
+            cute.arch.cp_async_commit_group()
+        pipeline_kv.producer_commit(producer_state)
+    @cute.jit
+    def mma(
+        self,
+        tiled_mma_qk: cute.TiledMma,
+        tiled_mma_pv: cute.TiledMma,
+        mO: cute.Tensor,
+        mLSE: Optional[cute.Tensor],
+        sQ: cute.Tensor,
+        sK: cute.Tensor,
+        sVt: cute.Tensor,
+        sP: Optional[cute.Tensor],
+        sO: cute.Tensor,
+        learnable_sink: Optional[cute.Tensor],
+        pipeline_k: pipeline.PipelineAsync,
+        pipeline_v: pipeline.PipelineAsync,
+        pipeline_q: pipeline.PipelineAsync,
+        gmem_tiled_copy_O: cute.TiledCopy,
+        tma_atom_O: Optional[cute.CopyAtom],
+        tidx: Int32,
+        softmax_scale_log2: Float32,
+        softmax_scale: Optional[Float32],
+        block_info: BlockInfo,
+        SeqlenInfoCls: Callable,
+        AttentionMaskCls: Callable,
+        TileSchedulerCls: Callable,
+        blocksparse_tensors: Optional[BlockSparseTensors],
+        aux_data: AuxData = AuxData(),
+        fastdiv_mods=None,
+    ):
+        aux_tensors = aux_data.tensors
+        warp_group_idx = cute.arch.make_warp_uniform(tidx // self.num_threads_per_warp_group)
+        warp_group_thread_layout = cute.make_layout(
+            self.num_wg_mma, stride=self.num_threads_per_warp_group
+        )
+        thr_mma_qk = tiled_mma_qk.get_slice(tidx)
+        wg_mma_qk = tiled_mma_qk.get_slice(warp_group_thread_layout(warp_group_idx))
+        wg_mma_pv = tiled_mma_pv.get_slice(warp_group_thread_layout(warp_group_idx))
+        _, tSrQ, tSrK = sm90_utils.partition_fragment_ABC(
+            wg_mma_qk, (self.tile_m, self.tile_n, self.tile_hdim), sQ, sK
+        )
+        mma_qk_fn = partial(
+            sm90_utils.gemm_zero_init, tiled_mma_qk, (self.tile_m, self.tile_n), tSrQ, tSrK
+        )
+        acc_O, tOrP, tOrVt = sm90_utils.partition_fragment_ABC(
+            wg_mma_pv, (self.tile_m, self.tile_hdimv, self.tile_n), sP, sVt
+        )
+        mma_pv_fn = partial(sm90_utils.gemm_w_idx, tiled_mma_pv, acc_O, tOrP, tOrVt)
+        # ///////////////////////////////////////////////////////////////////////////////
+        # Smem copy atom tiling
+        # ///////////////////////////////////////////////////////////////////////////////
+        smem_copy_atom_P = utils.get_smem_store_atom(
+            self.arch.major * 10 + self.arch.minor, self.dtype
+        )
+        smem_thr_copy_P = cute.make_tiled_copy_C(smem_copy_atom_P, tiled_mma_qk).get_slice(tidx)
+        tPsP = smem_thr_copy_P.partition_D(sP) if const_expr(sP is not None) else None
+        smem_copy_params = SimpleNamespace(smem_thr_copy_P=smem_thr_copy_P, tPsP=tPsP)
+        self.mma_init()
+        q_consumer_phase = Int32(0)
+        kv_consumer_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, self.num_stages
+        )
+        tile_scheduler = TileSchedulerCls()
+        work_tile = tile_scheduler.initial_work_tile_info()
+        softmax = Softmax.create(
+            softmax_scale_log2,
+            num_rows=acc_O.shape[0][0] * acc_O.shape[1],
+            softmax_scale=softmax_scale,
+        )
+        # For RescaleOBeforeGemm: persistent scores_scale across iterations
+        scores_scale = None
+        if const_expr(self.rescale_O_before_gemm):
+            scores_scale = cute.make_rmem_tensor_like(softmax.row_max, Float32)
+        mma_one_n_block_all = partial(
+            self.mma_one_n_block_intrawg_overlap
+            if const_expr(self.intra_wg_overlap)
+            else self.mma_one_n_block,
+            mma_qk_fn=mma_qk_fn,
+            pipeline_k=pipeline_k,
+            pipeline_v=pipeline_v,
+            acc_O=acc_O,
+            tOrP=tOrP,
+            smem_copy_params=smem_copy_params,
+            check_inf=True,
+            scores_scale=scores_scale,
+        )
+        process_first_half_block = partial(
+            self.first_half_block_overlap,
+            mma_qk_fn=mma_qk_fn,
+            pipeline_k=pipeline_k,
+            tOrP=tOrP,
+            smem_copy_params=smem_copy_params,
+            scores_scale=scores_scale,
+            softmax=softmax,
+            acc_O=acc_O,
+        )
+        process_last_half_block = partial(
+            self.last_half_block_overlap,
+            pipeline_v=pipeline_v,
+            mma_pv_fn=mma_pv_fn,
+            scores_scale=scores_scale,
+            softmax=softmax,
+            acc_O=acc_O,
+        )
+        while work_tile.is_valid_tile:
+            # if work_tile.is_valid_tile:
+            # shape: (atom_v_m * rest_m)
+            m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+            seqlen = SeqlenInfoCls(batch_idx)
+            # Recompute fastdiv_mods if necessary for varlen with aux_tensors
+            recompute_fastdiv_mods_q = cutlass.const_expr(
+                aux_tensors is not None and (seqlen.has_cu_seqlens_q or seqlen.has_seqused_q)
+            )
+            recompute_fastdiv_mods_k = cutlass.const_expr(
+                aux_tensors is not None and (seqlen.has_cu_seqlens_k or seqlen.has_seqused_k)
+            )
+            if cutlass.const_expr(fastdiv_mods is not None):
+                seqlen_q_divmod, seqlen_k_divmod = fastdiv_mods
+                fastdiv_mods = (
+                    seqlen_q_divmod
+                    if not recompute_fastdiv_mods_q
+                    else FastDivmodDivisor(seqlen.seqlen_q),
+                    seqlen_k_divmod
+                    if not recompute_fastdiv_mods_k
+                    else FastDivmodDivisor(seqlen.seqlen_k),
+                )
+            mask = AttentionMaskCls(seqlen)
+            mask_fn = partial(
+                mask.apply_mask,
+                batch_idx=batch_idx,
+                head_idx=head_idx,
+                m_block=m_block,
+                thr_mma=thr_mma_qk,
+                mask_causal=self.is_causal,
+                mask_local=self.is_local,
+                aux_data=aux_data,
+                fastdiv_mods=fastdiv_mods,
+            )
+            score_mod_fn = None
+            if const_expr(self.score_mod is not None):
+                score_mod_fn = partial(
+                    self.apply_score_mod,
+                    thr_mma_qk,
+                    batch_idx,
+                    head_idx,
+                    m_block,
+                    softmax_scale=softmax_scale,
+                    aux_data=aux_data,
+                    fastdiv_mods=fastdiv_mods,
+                )
+            mma_one_n_block = partial(
+                mma_one_n_block_all, seqlen=seqlen, softmax=softmax, score_mod_fn=score_mod_fn
+            )
+            n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block)
+            pipeline_q.consumer_wait_w_index_phase(0, q_consumer_phase)
+            # For performance reason, we separate out two kinds of iterations:
+            # those that need masking on S, and those that don't.
+            # We need masking on S for the very last block when K and V has length not multiple of tile_n.
+            # We also need masking on S if it's causal, for the last several blocks.
+            # softmax.reset()  # Don't need reset as we explicitly call softmax w is_first=True
+            O_should_accumulate = False
+            # ==========================================
+            # MAINLOOP
+            # ==========================================
+            if const_expr(not self.use_block_sparsity):
+                # ==========================================
+                # No block-sparsity (original path)
+                # ==========================================
+                # First iteration with seqlen masking
+                if const_expr(self.intra_wg_overlap):
+                    kv_consumer_state = process_first_half_block(
+                        n_block=n_block_max - 1,
+                        seqlen=seqlen,
+                        kv_consumer_state=kv_consumer_state,
+                        mask_fn=partial(mask_fn, mask_mod=self.mask_mod),
+                        score_mod_fn=score_mod_fn,
+                        is_first_block=True,
+                    )
+                else:
+                    self.warp_scheduler_barrier_sync()
+                    kv_consumer_state = mma_one_n_block(
+                        kv_consumer_state,
+                        n_block=n_block_max - 1,
+                        seqlen=seqlen,
+                        mma_pv_fn=partial(mma_pv_fn, zero_init=True),
+                        is_first_n_block=True,
+                        mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=True),
+                    )
+                    O_should_accumulate = True
+                # if cute.arch.thread_idx()[0] == 128: cute.printf("m_block = {}, n_block_max = {}, n_block_min = {}", m_block, n_block_max, n_block_min)
+                n_block_max -= 1
+                # Next couple of iterations with causal masking
+                if const_expr(self.is_causal or self.is_local):
+                    n_block_min_causal_local_mask = block_info.get_n_block_min_causal_local_mask(
+                        seqlen, m_block, n_block_min
+                    )
+                    # if cute.arch.thread_idx()[0] == 128: cute.printf("n_block_min_causal_local_mask = {}", n_block_min_causal_local_mask)
+                    for n_tile in cutlass.range(
+                        n_block_max - n_block_min_causal_local_mask, unroll=1
+                    ):
+                        kv_consumer_state = mma_one_n_block(
+                            kv_consumer_state,
+                            n_block=n_block_max - 1 - n_tile,
+                            seqlen=seqlen,
+                            mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate),
+                            mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=False),
+                        )
+                        O_should_accumulate = True
+                    n_block_max = cutlass.min(n_block_max, n_block_min_causal_local_mask)
+                # The remaining iterations have no masking
+                n_block_min_before_local_mask = block_info.get_n_block_min_before_local_mask(
+                    seqlen, m_block, n_block_min
+                )
+                # if cute.arch.thread_idx()[0] == 128: cute.printf("n_block_min_before_local_mask = {}, n_block_min = {}", n_block_min_before_local_mask, n_block_min)
+                for n_tile in cutlass.range(n_block_max - n_block_min_before_local_mask, unroll=1):
+                    kv_consumer_state = mma_one_n_block(
+                        kv_consumer_state,
+                        n_block=n_block_max - 1 - n_tile,
+                        seqlen=seqlen,
+                        mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate),
+                        mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=False),
+                    )
+                    O_should_accumulate = True
+                # Separate iterations with local masking on the left
+                if const_expr(self.is_local and block_info.window_size_left is not None):
+                    n_block_max = cutlass.min(n_block_max, n_block_min_before_local_mask)
+                    for n_tile in cutlass.range(n_block_max - n_block_min, unroll=1):
+                        kv_consumer_state = mma_one_n_block(
+                            kv_consumer_state,
+                            n_block=n_block_max - 1 - n_tile,
+                            seqlen=seqlen,
+                            mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate),
+                            mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=False),
+                        )
+                        O_should_accumulate = True
+                # Release Q pipeline so the producer can load the next tile's Q
+                pipeline_q.consumer_release_w_index(0)
+                # Last "half" iteration
+                if const_expr(self.intra_wg_overlap):
+                    kv_consumer_state = process_last_half_block(
+                        kv_consumer_state=kv_consumer_state,
+                        zero_init=not O_should_accumulate,
+                    )
+                    O_should_accumulate = True
+                else:
+                    self.warp_scheduler_barrier_arrive()
+            else:
+                # ==========================================
+                # Block sparsity
+                # ==========================================
+                kv_consumer_state, O_should_accumulate, processed_any = consume_block_sparse_loads(
+                    blocksparse_tensors,
+                    batch_idx,
+                    head_idx,
+                    m_block,
+                    seqlen,
+                    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,
+                    self.mask_mod,
+                    fastdiv_mods,
+                    self.intra_wg_overlap,
+                    self.warp_scheduler_barrier_sync,
+                    self.warp_scheduler_barrier_arrive,
+                    self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
+                    self.q_subtile_factor,
+                )
+                # Release Q pipeline so the producer can load the next tile's Q
+                pipeline_q.consumer_release_w_index(0)
+                # Handle empty case (when no blocks to process)
+                if not processed_any:
+                    softmax.reset()
+                    acc_O.fill(0.0)
+            q_consumer_phase ^= 1
+            sink_val = None
+            if const_expr(learnable_sink is not None):
+                if const_expr(not self.pack_gqa):
+                    sink_val = Float32(learnable_sink[head_idx])
+                else:  # Each thread might have a different sink value due to different q_head
+                    sink_val = cute.make_rmem_tensor_like(softmax.row_max, Float32)
+                    cS = cute.make_identity_tensor((self.tile_m, self.tile_n))
+                    tScS_mn = layout_utils.reshape_acc_to_mn(thr_mma_qk.partition_C(cS))
+                    for r in cutlass.range(cute.size(sink_val), unroll_full=True):
+                        row = m_block * self.tile_m + tScS_mn[r][0]
+                        q_head_idx = row % self.qhead_per_kvhead + head_idx * self.qhead_per_kvhead
+                        sink_val[r] = Float32(learnable_sink[q_head_idx])
+            # normalize acc_O by row_sum and calculate the lse
+            row_scale = softmax.finalize(sink_val=sink_val)
+            softmax.rescale_O(acc_O, row_scale)
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Epilogue
+            # ///////////////////////////////////////////////////////////////////////////////
+            self.epilogue(
+                acc_O,
+                softmax.row_sum,
+                mO,
+                mLSE,
+                sO,
+                seqlen,
+                gmem_tiled_copy_O,
+                tma_atom_O,
+                tiled_mma_pv,
+                tidx,
+                m_block,
+                head_idx,
+                batch_idx,
+            )
+            tile_scheduler.advance_to_next_work()
+            work_tile = tile_scheduler.get_current_work()
+    @cute.jit
+    def first_half_block_overlap(
+        self,
+        n_block: Int32,
+        mma_qk_fn: Callable,
+        kv_consumer_state,
+        pipeline_k,
+        tOrP: cute.Tensor,
+        smem_copy_params: SimpleNamespace,
+        softmax: Softmax,
+        seqlen: SeqlenInfoQK,
+        scores_scale: Optional[cute.Tensor] = None,
+        acc_O: Optional[cute.Tensor] = None,
+        mask_fn: Callable = None,
+        score_mod_fn: Optional[Callable] = None,
+        is_first_block: bool = False,
+    ):
+        """Processes the first half block when using intra-warpgroup-overlap"""
+        pipeline_k.consumer_wait(kv_consumer_state, pipeline_k.consumer_try_wait(kv_consumer_state))
+        acc_S = mma_qk_fn(B_idx=kv_consumer_state.index, wg_wait=0)
+        pipeline_k.consumer_release(kv_consumer_state)
+        # Apply score modification if present
+        if const_expr(score_mod_fn is not None):
+            score_mod_fn(acc_S, n_block=n_block, seqlen=seqlen)
+        # Apply mask; mask_seqlen always True for first block
+        # Caveat: if full block further right than mask block, seqlen masking is redundant;
+        # however, masking is being applied anyway, so essentially no perf hit
+        mask_fn(acc_S, n_block=n_block, mask_seqlen=True)
+        row_scale = softmax.online_softmax(acc_S, is_first=is_first_block)
+        tOrP_acc = layout_utils.reshape_acc_to_frgA(acc_S)
+        tOrP_cur = (
+            tOrP
+            if const_expr(self.mma_pv_is_rs)
+            else cute.make_rmem_tensor_like(tOrP_acc, self.dtype)
+        )
+        tOrP_cur.store(tOrP_acc.load().to(self.dtype))
+        if const_expr(not self.mma_pv_is_rs):
+            tPrP = smem_copy_params.smem_thr_copy_P.retile(tOrP_cur)
+            cute.copy(smem_copy_params.smem_thr_copy_P, tPrP, smem_copy_params.tPsP)
+            # Fence and barrier to make smem store visible to WGMMA
+            cute.arch.fence_view_async_shared()
+            cute.arch.sync_warp()
+        # For RescaleOBeforeGemm: initialize acc_O
+        if const_expr(self.rescale_O_before_gemm):
+            acc_O.fill(0.0)
+            scores_scale.store(row_scale.load())
+        return kv_consumer_state
+    @cute.jit
+    def last_half_block_overlap(
+        self,
+        kv_consumer_state,
+        pipeline_v,
+        mma_pv_fn: Callable,
+        zero_init: bool,
+        scores_scale: Optional[cute.Tensor] = None,
+        softmax: Optional[Softmax] = None,
+        acc_O: Optional[cute.Tensor] = None,
+    ):
+        """Processes the final PV GEMM when using intra-warpgroup-overlap"""
+        # For RescaleOBeforeGemm: rescale O before the final PV GEMM
+        if const_expr(self.rescale_O_before_gemm):
+            softmax.rescale_O(acc_O, scores_scale)
+        pipeline_v.consumer_wait(kv_consumer_state, pipeline_v.consumer_try_wait(kv_consumer_state))
+        mma_pv_fn(B_idx=kv_consumer_state.index, zero_init=zero_init, wg_wait=0)
+        pipeline_v.consumer_release(kv_consumer_state)
+        kv_consumer_state.advance()
+        return kv_consumer_state
+    @cute.jit
+    def mma_one_n_block(
+        self,
+        smem_pipe_read: pipeline.PipelineState | pipeline_custom.PipelineStateSimple,
+        n_block: Int32,
+        mma_qk_fn: Callable,
+        mma_pv_fn: Callable,
+        pipeline_k: pipeline.PipelineAsync,
+        pipeline_v: pipeline.PipelineAsync,
+        acc_O: cute.Tensor,
+        tOrP: cute.Tensor,
+        smem_copy_params: SimpleNamespace,
+        softmax: Softmax,
+        seqlen: SeqlenInfoQK,
+        scores_scale: Optional[cute.Tensor] = None,  # not used
+        score_mod_fn: Optional[Callable] = None,
+        mask_fn: Optional[Callable] = None,
+        is_first_n_block: cutlass.Constexpr = False,
+        check_inf: cutlass.Constexpr = True,
+    ):
+        pipeline_k.consumer_wait(smem_pipe_read, pipeline_k.consumer_try_wait(smem_pipe_read))
+        # S = Q @ K.T
+        acc_S = mma_qk_fn(B_idx=smem_pipe_read.index, wg_wait=-1)
+        self.warp_scheduler_barrier_arrive()
+        warpgroup.wait_group(0)
+        pipeline_k.consumer_release(smem_pipe_read)
+        # handle score mods and masking
+        if const_expr(score_mod_fn is not None):
+            score_mod_fn(acc_S, n_block=n_block, seqlen=seqlen)
+        if const_expr(mask_fn is not None):
+            mask_fn(acc_S=acc_S, n_block=n_block)
+        row_scale = softmax.online_softmax(acc_S, is_first=is_first_n_block, check_inf=check_inf)
+        # if cute.arch.thread_idx()[0] == 0: cute.print_tensor(layout_utils.reshape_acc_to_mn(acc_S))
+        tOrP_acc = layout_utils.reshape_acc_to_frgA(acc_S)
+        tOrP_cur = (
+            tOrP
+            if const_expr(self.mma_pv_is_rs)
+            else cute.make_rmem_tensor_like(tOrP_acc, self.dtype)
+        )
+        # tOrP.store(tOrP_acc.load().to(self.dtype))
+        # the "to(self.dtype)" conversion fails to vectorize for block sizes other
+        # than 128 x 128, i.e. it calls convert on 1 fp32 element at a time instead of
+        # 2 elements. So we just call ptx directly.
+        utils.cvt_f16(tOrP_acc, tOrP_cur)
+        if const_expr(not self.mma_pv_is_rs):
+            tPrP = smem_copy_params.smem_thr_copy_P.retile(tOrP_cur)
+            cute.copy(smem_copy_params.smem_thr_copy_P, tPrP, smem_copy_params.tPsP)
+        softmax.rescale_O(acc_O, row_scale)
+        if const_expr(not self.mma_pv_is_rs):
+            # Fence and barrier to make sure smem store is visible to WGMMA
+            cute.arch.fence_view_async_shared()
+            cute.arch.sync_warp()  # Only need syncwarp since each warp is using its own P values for MmaPV
+        pipeline_v.consumer_wait(smem_pipe_read, pipeline_v.consumer_try_wait(smem_pipe_read))
+        self.warp_scheduler_barrier_sync()
+        # O += P @ V
+        mma_pv_fn(B_idx=smem_pipe_read.index, wg_wait=0)
+        pipeline_v.consumer_release(smem_pipe_read)
+        smem_pipe_read.advance()
+        return smem_pipe_read
+    @cute.jit
+    def mma_one_n_block_intrawg_overlap(
+        self,
+        smem_pipe_read: pipeline.PipelineState | pipeline_custom.PipelineStateSimple,
+        n_block: Int32,
+        mma_qk_fn: Callable,
+        mma_pv_fn: Callable,
+        pipeline_k: pipeline.PipelineAsync,
+        pipeline_v: pipeline.PipelineAsync,
+        acc_O: cute.Tensor,
+        tOrP: cute.Tensor,
+        smem_copy_params: SimpleNamespace,
+        softmax: Softmax,
+        seqlen: SeqlenInfoQK,
+        scores_scale: Optional[cute.Tensor] = None,
+        score_mod_fn: Optional[Callable] = None,
+        mask_fn: Optional[Callable] = None,
+        check_inf: cutlass.Constexpr = True,
+    ):
+        smem_pipe_read_v = smem_pipe_read.clone()
+        smem_pipe_read.advance()
+        pipeline_k.consumer_wait(smem_pipe_read, pipeline_k.consumer_try_wait(smem_pipe_read))
+        self.warp_scheduler_barrier_sync()
+        # S = Q @ K.T
+        acc_S = mma_qk_fn(B_idx=smem_pipe_read.index, wg_wait=-1)
+        # RescaleOBeforeGemm: rescale O while QK GEMM is in flight, before PV GEMM
+        if const_expr(self.rescale_O_before_gemm):
+            softmax.rescale_O(acc_O, scores_scale)
+        pipeline_v.consumer_wait(smem_pipe_read_v, pipeline_v.consumer_try_wait(smem_pipe_read_v))
+        # O += P @ V
+        mma_pv_fn(B_idx=smem_pipe_read_v.index, wg_wait=-1)
+        self.warp_scheduler_barrier_arrive()
+        warpgroup.wait_group(1)
+        pipeline_k.consumer_release(smem_pipe_read)
+        # handle score mods and masking
+        if const_expr(score_mod_fn is not None):
+            score_mod_fn(acc_S, n_block=n_block, seqlen=seqlen)
+        if const_expr(mask_fn is not None):
+            mask_fn(acc_S=acc_S, n_block=n_block)
+        # if cute.arch.thread_idx()[0] == 128: cute.print_tensor(layout_utils.reshape_acc_to_mn(acc_S))
+        row_scale = softmax.online_softmax(acc_S, check_inf=check_inf)
+        warpgroup.wait_group(0)
+        pipeline_v.consumer_release(smem_pipe_read_v)
+        tOrP_acc = layout_utils.reshape_acc_to_frgA(acc_S)
+        tOrP_cur = (
+            tOrP
+            if const_expr(self.mma_pv_is_rs)
+            else cute.make_rmem_tensor_like(tOrP_acc, self.dtype)
+        )
+        # tOrP_cur.store(tOrP_acc.load().to(self.dtype))
+        # the "to(self.dtype)" conversion fails to vectorize for block sizes other
+        # than 128 x 128, i.e. it calls convert on 1 fp32 element at a time instead of
+        # 2 elements. So we just call ptx directly.
+        utils.cvt_f16(tOrP_acc, tOrP_cur)
+        if const_expr(not self.mma_pv_is_rs):
+            tPrP = smem_copy_params.smem_thr_copy_P.retile(tOrP_cur)
+            cute.copy(smem_copy_params.smem_thr_copy_P, tPrP, smem_copy_params.tPsP)
+        if const_expr(not self.rescale_O_before_gemm):
+            softmax.rescale_O(acc_O, row_scale)
+        if const_expr(self.rescale_O_before_gemm):
+            scores_scale.store(row_scale.load())
+        if const_expr(not self.mma_pv_is_rs):
+            # Fence and barrier to make sure smem store is visible to WGMMA
+            cute.arch.fence_view_async_shared()
+            cute.arch.sync_warp()  # Only need syncwarp since each warp is using its own P values for MmaPV
+        return smem_pipe_read
+    @cute.jit
+    def mma_init(self):
+        warp_group_idx = utils.canonical_warp_group_idx(sync=False)
+        if const_expr(self.use_scheduler_barrier):
+            if warp_group_idx == 1:
+                cute.arch.barrier_arrive(
+                    barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1),
+                    number_of_threads=2 * self.num_threads_per_warp_group,
+                )
+    @cute.jit
+    def apply_score_mod(
+        self,
+        thr_mma_qk,
+        batch_idx,
+        head_idx,
+        m_block,
+        acc_S,
+        n_block,
+        softmax_scale,
+        seqlen,
+        aux_data: AuxData = AuxData(),
+        fastdiv_mods=None,
+    ):
+        # Prepare index tensor
+        cS = cute.make_identity_tensor((self.tile_m, self.tile_n))
+        cS = cute.domain_offset((m_block * self.tile_m, n_block * self.tile_n), cS)
+        tScS = thr_mma_qk.partition_C(cS)
+        apply_score_mod_inner(
+            acc_S,
+            tScS,
+            self.score_mod,
+            batch_idx,
+            head_idx,
+            softmax_scale,
+            self.score_vec_size,
+            self.qk_acc_dtype,
+            aux_data,
+            fastdiv_mods,
+            seqlen_info=seqlen,
+            constant_q_idx=None,
+            qhead_per_kvhead=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
+        )
+    def warp_scheduler_barrier_sync(self):
+        if const_expr(self.use_scheduler_barrier):
+            cute.arch.barrier(
+                barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1)
+                - 1
+                + utils.canonical_warp_group_idx(sync=False),
+                number_of_threads=2 * self.num_threads_per_warp_group,
+            )
+    def warp_scheduler_barrier_arrive(self):
+        if const_expr(self.use_scheduler_barrier):
+            assert self.num_wg_mma in [2, 3]
+            cur_wg = utils.canonical_warp_group_idx(sync=False) - 1
+            if const_expr(self.num_wg_mma == 2):
+                next_wg = 1 - cur_wg
+            else:
+                t = cur_wg + 1
+                next_wg = t % self.num_wg_mma
+            cute.arch.barrier_arrive(
+                barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1) + next_wg,
+                number_of_threads=2 * self.num_threads_per_warp_group,
+            )

build/torch-cuda/interface.py CHANGED Viewed

The diff for this file is too large to render. See raw diff

build/torch-cuda/mask.py CHANGED Viewed

@@ -1,108 +1,157 @@
 # 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 .quack import layout_utils
-from . import utils
 from .seqlen_info import SeqlenInfoQK
 @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)
 @cute.jit
-def mask_r2p_dual_bound(
-    X: cute.Tensor,
-    col_limit_left: Int32,  # Inclusive lower bound
-    col_limit_right: Int32,  # Exclusive upper bound
-) -> None:
-    """
-    Dual-bound masking using two bitmasks for SM100, following mask_r2p.
-    Masks elements where: NOT (col_limit_left <= col < col_limit_right)
-    Uses bit manipulation to create a range mask:
-        mask_right = (1 << right) - 1  -> bits (right-1)..0 are 1
-        mask_left  = (1 << left) - 1   -> bits (left-1)..0 are 1
-        mask_range = mask_range = mask_right & ~ mask_left -> bits (right-1)..left are 1
-    """
-    ncol = const_expr(cute.size(X.shape))
-    for s in cutlass.range_constexpr(cute.ceil_div(ncol, 24)):
-        right_s = max(col_limit_right - s * 24, 0)
-        left_s = max(col_limit_left - s * 24, 0)
-        # otherwise cute dsl complains about python int too large to convert into c long
-        right_s = min(right_s, 24)
-        left_s = min(left_s, 24)
-        # bits (right-1)..left are 1
-        mask_right = (1 << right_s) - 1
-        mask_left = (1 << left_s) - 1
-        mask_range = mask_right & ~mask_left
-        # 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_range & (1 << i))
-            c = s * 24 + i
             X[c] = X[c] if in_bound else -Float32.inf
@@ -137,7 +186,7 @@ class AttentionMask:
         mask_causal: cutlass.Constexpr[bool],
         mask_local: cutlass.Constexpr[bool] = False,
         mask_mod: cutlass.Constexpr[Optional[Callable]] = None,
-        aux_tensors: Optional[list] = None,
         fastdiv_mods=(None, None),
     ) -> None:
         assert not (mask_causal and mask_local), "mask_causal and mask_local cannot be both True"
@@ -161,8 +210,7 @@ class AttentionMask:
         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):
@@ -170,7 +218,8 @@ class AttentionMask:
                         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
@@ -183,7 +232,7 @@ class AttentionMask:
                 and fastdiv_mods[1] is not None
             )
             wrap_aux_indices = const_expr(
-                has_fastdiv and mask_seqlen and const_expr(aux_tensors is not None)
             )
             for r in cutlass.range_constexpr(nrow):
@@ -212,13 +261,14 @@ class AttentionMask:
                     head_idx_ssa = utils.scalar_to_ssa(head_idx_for_mod, cutlass.Int32)
                     q_idx_ssa = utils.scalar_to_ssa(row_for_mod, cutlass.Int32)
                     kv_idx_ssa = utils.scalar_to_ssa(col_for_mod, cutlass.Int32)
-                    mask_value = mask_mod(
                         batch_idx_ssa,
                         head_idx_ssa,
                         q_idx_ssa,
                         kv_idx_ssa,
                         self.seqlen_info,
-                        aux_tensors,
                     )
                     cond = cutlass.Boolean(utils.ssa_to_scalar(mask_value))
                     if const_expr(mask_seqlen):
@@ -272,7 +322,12 @@ class AttentionMask:
                                     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
@@ -284,6 +339,7 @@ class AttentionMask:
                         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
@@ -302,13 +358,22 @@ class AttentionMask:
                             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]
@@ -338,11 +403,18 @@ class AttentionMask:
                         # 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] = (
@@ -351,6 +423,194 @@ class AttentionMask:
                                 else acc_S_mn[r, c]
                             )
     @cute.jit
     def apply_mask_sm100(
         self,
@@ -365,10 +625,13 @@ class AttentionMask:
         mask_mod: cutlass.Constexpr[Optional[Callable]] = None,
         batch_idx: Int32 = None,
         head_idx: Int32 = None,
-        aux_tensors: Optional[list] = None,
         fastdiv_mods=(None, None),
         head_divmod=None,
         check_q_boundary: bool = False,
     ) -> 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)
@@ -382,8 +645,18 @@ class AttentionMask:
         if n_block < 0:
             n_block = 0
         seqlenk_col_limit = self.seqlen_k - n_block * self.tile_n
-        r2p = True
-        if const_expr(not mask_causal and not mask_local and mask_mod is None):
             if const_expr(mask_seqlen):
                 if const_expr(not r2p):
                     for i in cutlass.range(cute.size(tScS_t2r.shape), unroll_full=True):
@@ -392,65 +665,58 @@ class AttentionMask:
                         # 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)
         elif const_expr(not mask_causal and not mask_local and mask_mod is not None):
-            # Block sparse case w/ mask_mod
-            has_fastdiv = const_expr(
-                fastdiv_mods is not None
-                and fastdiv_mods[0] is not None
-                and fastdiv_mods[1] is not None
             )
-            batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32)
-            ncol = const_expr(cute.size(tScS_t2r.shape))
-            for i in cutlass.range_constexpr(ncol):
-                row_coord = tScS_t2r[i][0] if not self.swap_AB else tScS_t2r[i][1]
-                col_coord = tScS_t2r[i][1] if not self.swap_AB else tScS_t2r[i][0]
-                global_row = row_coord + m_block * self.tile_m
-                global_col = col_coord + n_block * self.tile_n
-                if const_expr(self.qhead_per_kvhead_packgqa != 1):
-                    assert head_divmod is not None
-                    mask_row, head_offset = divmod(global_row, head_divmod)
-                    head_idx_for_mod = head_idx * self.qhead_per_kvhead_packgqa + head_offset
-                else:
-                    head_idx_for_mod = head_idx
-                    mask_row = global_row
-                mask_row_for_mod = mask_row
-                if const_expr(has_fastdiv and aux_tensors is not None):
-                    if check_q_boundary:
-                        _, mask_row_for_mod = divmod(mask_row, fastdiv_mods[0])
-                global_col_for_mod = global_col
-                if const_expr(has_fastdiv and mask_seqlen and aux_tensors is not None):
-                    _, global_col_for_mod = divmod(global_col, fastdiv_mods[1])
-                head_idx_ssa = utils.scalar_to_ssa(head_idx_for_mod, cutlass.Int32)
-                mask_row_ssa = utils.scalar_to_ssa(mask_row_for_mod, cutlass.Int32)
-                kv_idx_ssa = utils.scalar_to_ssa(global_col_for_mod, cutlass.Int32)
-                mask_value = mask_mod(
-                    batch_idx_ssa,
-                    head_idx_ssa,
-                    mask_row_ssa,
-                    kv_idx_ssa,
-                    self.seqlen_info,
-                    aux_tensors,
                 )
-                cond = cutlass.Boolean(utils.ssa_to_scalar(mask_value))
-                acc_S[i] = acc_S[i] if cond else -Float32.inf
-                if const_expr(mask_seqlen):
-                    acc_S[i] = -Float32.inf if global_col >= self.seqlen_k else acc_S[i]
-                if check_q_boundary:
-                    acc_S[i] = -Float32.inf if mask_row >= self.seqlen_q else acc_S[i]
         else:  # Causal or local
-            causal_row_offset = 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:
@@ -460,15 +726,19 @@ class AttentionMask:
                     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
                 )
@@ -493,8 +763,15 @@ class AttentionMask:
                             else acc_S[i]
                         )
                 else:
-                    # XOR-based R2P dual bound masking
-                    mask_r2p_dual_bound(acc_S, col_limit_left, col_limit_right)
     @cute.jit
     def apply_mask_sm100_transposed(
@@ -510,7 +787,7 @@ class AttentionMask:
         mask_mod: cutlass.Constexpr[Optional[Callable]] = None,
         batch_idx: Int32 = None,
         head_idx: Int32 = None,
-        aux_tensors: Optional[list] = None,
         fastdiv_mods=(None, None),
         is_full_block: bool = False,
         check_m_boundary: bool = True,
@@ -518,7 +795,7 @@ class AttentionMask:
         """
         Backward pass: mask S = K @ Q.T where n_block tiles seqlen_k and m_block tiles seqlen_q.
-        Coordinate conventio:
         - ROW corresponds to Q (m_block)
         - COL corresponds to KV (n_block)
@@ -568,7 +845,7 @@ class AttentionMask:
                     and fastdiv_mods[1] is not None
                 )
                 wrap_aux_indices = const_expr(
-                    has_fastdiv and mask_seqlen and const_expr(aux_tensors is not None)
                 )
                 batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32)
                 head_idx_ssa = utils.scalar_to_ssa(head_idx, cutlass.Int32)
@@ -589,13 +866,14 @@ class AttentionMask:
                     q_idx_ssa = utils.scalar_to_ssa(q_idx_for_mod, cutlass.Int32)
                     kv_idx_ssa = utils.scalar_to_ssa(kv_idx_for_mod, cutlass.Int32)
-                    mask_value = mask_mod(
                         batch_idx_ssa,
                         head_idx_ssa,
                         q_idx_ssa,
                         kv_idx_ssa,
                         self.seqlen_info,
-                        aux_tensors,
                     )
                     cond = cutlass.Boolean(utils.ssa_to_scalar(mask_value))
                     acc_S[i] = acc_S[i] if cond else -cutlass.Float32.inf
@@ -634,7 +912,13 @@ class AttentionMask:
                         )
                 else:
                     num_rep = cute.size(tScS_t2r, mode=[0])  # 16 or 32
-                    mask_r2p_transposed(acc_S, row_limit_top, num_rep)
             else:
                 if const_expr(self.window_size_right is not None):
                     row_limit_top = causal_offset - self.window_size_right
@@ -645,9 +929,783 @@ class AttentionMask:
                 if const_expr(mask_seqlen):
                     if seqlenk_col_limit <= 0:
                         row_limit_top = self.tile_m
-                for i in cutlass.range(cute.size(acc_S.shape), unroll_full=True):
-                    row_idx = t0ScS_t2r[i][ROW]
-                    local_mask = row_idx < row_limit_top
-                    if const_expr(self.window_size_left is not None):
-                        local_mask |= row_idx > row_limit_bot
-                    acc_S[i] = -cutlass.Float32.inf if local_mask else acc_S[i]

 # Copyright (c) 2025, Tri Dao.
+from typing import Optional, Callable, TypeAlias, Tuple
 from dataclasses import dataclass
+import enum
 import cutlass
 import cutlass.cute as cute
+from cutlass import Float32, Int32, Uint32, const_expr
+from cutlass.cutlass_dsl import min as dsl_min
 from .quack import layout_utils
+from . import utils as utils
+from .block_info import BlockInfo
 from .seqlen_info import SeqlenInfoQK
+from .utils import AuxData
+MaskGenFn: TypeAlias = Callable[[int], Uint32]
+MASK_R2P_CHUNK_SIZE: int = 32
+@cute.jit
+def call_mask_mod(
+    mask_mod: cutlass.Constexpr,
+    batch_idx,
+    head_idx,
+    q_idx,
+    kv_idx,
+    seqlen_info,
+    aux_data: AuxData,
+):
+    # Compatibility shim for pre-aux_scalars mask_mod callables.
+    if const_expr(aux_data.scalars is not None):
+        return mask_mod(
+            batch_idx,
+            head_idx,
+            q_idx,
+            kv_idx,
+            seqlen_info,
+            aux_data.tensors,
+            aux_data.scalars,
+        )
+    return mask_mod(
+        batch_idx,
+        head_idx,
+        q_idx,
+        kv_idx,
+        seqlen_info,
+        aux_data.tensors,
+    )
+@cute.jit
+def r2p_bitmask_below(limit: Int32, s: int) -> Uint32:
+    """32-bit R2P bitmask keeping positions < limit (exclusive upper bound).
+    Positions 0..limit-1 in chunk `s` get bit=1 (keep), the rest bit=0 (mask).
+    Uses inline PTX to avoid shift-by-type-width UB.
+    """
+    m = max((s + 1) * MASK_R2P_CHUNK_SIZE - limit, 0)
+    return utils.shr_u32(Uint32(0xFFFFFFFF), Uint32(m))
+@cute.jit
+def r2p_bitmask_above(limit: Int32, s: int) -> Uint32:
+    """32-bit R2P bitmask keeping positions >= limit (inclusive lower bound).
+    Positions limit..31 in chunk `s` get bit=1 (keep), the rest bit=0 (mask).
+    Uses inline PTX to avoid shift-by-type-width UB.
+    """
+    n = max(limit - s * MASK_R2P_CHUNK_SIZE, 0)
+    return utils.shl_u32(Uint32(0xFFFFFFFF), Uint32(n))
 @cute.jit
+def mask_r2p_lambda(
+    X: cute.Tensor,
+    mask_gen_fn: cutlass.Constexpr[MaskGenFn],
+    rank1: bool = False,
+) -> None:
+    """Apply R2P masking with a custom bitmask generator.
+    mask_gen_fn(chunk_idx: constexpr int) -> Uint32:
+        Returns a 32-bit bitmask for the chunk. Bit i set means column
+        chunk_idx * chunk_size + i is KEPT; bit i clear means masked to -inf.
+    """
     ncol = const_expr(cute.size(X.shape[cute.rank(X) - 1]) if not rank1 else cute.size(X.shape))
+    # 32-column chunks. The mask_gen_fn returns a Uint32 bitmask (1=keep).
+    CHUNK_SIZE = MASK_R2P_CHUNK_SIZE
+    for s in cutlass.range_constexpr(cute.ceil_div(ncol, CHUNK_SIZE)):
+        mask = mask_gen_fn(s)
         # This needs to be range_constexpr, o/w the compiler can't generate the R2P instruction
+        for i in cutlass.range_constexpr(min(CHUNK_SIZE, ncol - s * CHUNK_SIZE)):
+            in_bound = cutlass.Boolean(mask & (Uint32(1) << i))
+            c = s * CHUNK_SIZE + i
             if const_expr(rank1):
                 X[c] = X[c] if in_bound else -Float32.inf
             else:
                 for r in cutlass.range_constexpr(cute.size(X.shape[0])):
                     X[r, c] = X[r, c] if in_bound else -Float32.inf
 @cute.jit
+def sm90_col_to_r2p_idx(col_limit: Int32) -> Int32:
+    """Transform SM90 MMA column coordinate to R2P element index.
+    SM90 MMA accumulator column indices are non-contiguous: 0, 1, 8, 9, 16, 17, ...
+    Element indices are contiguous: 0, 1, 2, 3, 4, 5, ...
+    This converts a column-space threshold to element-space for r2p_bitmask_below/above.
+    """
+    return col_limit // 8 * 2 + min(col_limit % 8, 2)
 @cute.jit
+def row_to_r2p_idx(x: Int32, num_rep: int, num_wg: int) -> Int32:
+    """Convert a row coordinate to an R2P element index in the warp-group interleaved layout.
+    In the SM100 backward pass, 2 warp groups share TMEM. The TMEM load atom
+    distributes rows in an interleaved pattern: elements 0..num_rep-1 map to
+    rows 0..num_rep-1 (warp group 0), elements num_rep..2*num_rep-1 map to
+    rows num_rep*num_wg..num_rep*num_wg+num_rep-1 (warp group 1), and so on.
+    Row-coordinate thresholds (causal limits, window bounds, uih_len) must be
+    converted to element indices before use with r2p_bitmask_above/below.
+    Rows not owned by this thread (in the gap between warp groups) are clamped
+    to the boundary element index, which is safe because R2P thresholds are
+    monotonic.
+    Example with num_rep=16, num_wg=2:
+        row  0 -> elem  0,  row 15 -> elem 15,
+        row 16 -> elem 16 (clamped), row 31 -> elem 16 (clamped),
+        row 32 -> elem 16, row 33 -> elem 17, row 47 -> elem 31.
+    """
+    return x // (num_rep * num_wg) * num_rep + min(x % (num_rep * num_wg), num_rep)
+@cute.jit
+def apply_packed_mask_chunk(
+    X: cute.Tensor,
+    chunk_idx: cutlass.Constexpr[int],
+    mask: Uint32,
+) -> None:
+    """Apply one 32-bit keep mask to one 32-column chunk.
+    The one-iteration chunk loop keeps the same lowering pattern as mask_r2p_lambda.
+    """
+    ncol = const_expr(cute.size(X.shape))
+    col_base = chunk_idx * MASK_R2P_CHUNK_SIZE
+    for s in cutlass.range_constexpr(1):
+        for i in cutlass.range_constexpr(
+            min(MASK_R2P_CHUNK_SIZE, ncol - col_base - s * MASK_R2P_CHUNK_SIZE)
+        ):
+            in_bound = cutlass.Boolean(mask & (Uint32(1) << i))
+            c = col_base + s * MASK_R2P_CHUNK_SIZE + i
             X[c] = X[c] if in_bound else -Float32.inf
         mask_causal: cutlass.Constexpr[bool],
         mask_local: cutlass.Constexpr[bool] = False,
         mask_mod: cutlass.Constexpr[Optional[Callable]] = None,
+        aux_data: AuxData = AuxData(),
         fastdiv_mods=(None, None),
     ) -> None:
         assert not (mask_causal and mask_local), "mask_causal and mask_local cannot be both True"
         seqlenk_col_limit = self.seqlen_k - n_block * self.tile_n - thr_col_offset
         if const_expr(not mask_causal and not mask_local and mask_mod is None):
             if const_expr(mask_seqlen):
+                r2p = const_expr(not self.swap_AB)
                 if const_expr(not r2p):
                     # traverse column index.
                     for c in cutlass.range(cute.size(tScS_mn.shape[1]), unroll_full=True):
                         for r in cutlass.range(cute.size(tScS_mn.shape[0]), unroll_full=True):
                             acc_S_mn[r, c] = -Float32.inf if oob else acc_S_mn[r, c]
                 else:
+                    seqlenk_col_limit_r2p = sm90_col_to_r2p_idx(seqlenk_col_limit)
+                    mask_r2p_lambda(acc_S_mn, lambda s: r2p_bitmask_below(seqlenk_col_limit_r2p, s))
         elif const_expr(
             not mask_causal and not mask_local and mask_mod is not None
                 and fastdiv_mods[1] is not None
             )
             wrap_aux_indices = const_expr(
+                has_fastdiv and mask_seqlen and const_expr(aux_data.tensors is not None)
             )
             for r in cutlass.range_constexpr(nrow):
                     head_idx_ssa = utils.scalar_to_ssa(head_idx_for_mod, cutlass.Int32)
                     q_idx_ssa = utils.scalar_to_ssa(row_for_mod, cutlass.Int32)
                     kv_idx_ssa = utils.scalar_to_ssa(col_for_mod, cutlass.Int32)
+                    mask_value = call_mask_mod(
+                        mask_mod,
                         batch_idx_ssa,
                         head_idx_ssa,
                         q_idx_ssa,
                         kv_idx_ssa,
                         self.seqlen_info,
+                        aux_data,
                     )
                     cond = cutlass.Boolean(utils.ssa_to_scalar(mask_value))
                     if const_expr(mask_seqlen):
                                     else acc_S_mn[r, c]
                                 )
                         else:
+                            col_limit_r2p = sm90_col_to_r2p_idx(col_limit_right)
+                            mask_r2p_lambda(
+                                acc_S_mn[r, None],
+                                lambda s: r2p_bitmask_below(col_limit_r2p, s),
+                                rank1=True,
+                            )
                 else:  # Local
                     local_row_offset_right = (
                         causal_row_offset + self.window_size_right
                         if const_expr(self.window_size_left is not None)
                         else None
                     )
+                    r2p_local = const_expr(not self.swap_AB)
                     for r in cutlass.range(cute.size(tScS_mn.shape[0]), unroll_full=True):
                         if const_expr(self.qhead_per_kvhead_packgqa == 1):
                             row_idx = tScS_mn[r, 0][0] + m_block * self.tile_m
                             if const_expr(self.window_size_left is not None)
                             else 0
                         )
+                        if const_expr(not r2p_local):
+                            # traverse column index.
+                            for c in cutlass.range(cute.size(tScS_mn.shape[1]), unroll_full=True):
+                                col_idx = t0ScS_mn[0, c][1]
+                                if col_idx >= col_limit_right or col_idx < col_limit_left:
+                                    acc_S_mn[r, c] = -Float32.inf
+                        else:
+                            col_limit_right_r2p = sm90_col_to_r2p_idx(col_limit_right)
+                            col_limit_left_r2p = sm90_col_to_r2p_idx(col_limit_left)
+                            def mask_gen_fn(s: int) -> Uint32:
+                                return r2p_bitmask_below(
+                                    col_limit_right_r2p, s
+                                ) & r2p_bitmask_above(col_limit_left_r2p, s)
+                            mask_r2p_lambda(acc_S_mn[r, None], mask_gen_fn, rank1=True)
             else:  # swap_AB
                 assert self.qhead_per_kvhead_packgqa == 1
                 thr_row_offset = tScS_mn[0][ROW]
                         # column, by setting row limit to be self.tile_m.
                         row_limit_top = (
                             self.tile_m
+                            if col0 >= seqlenk_col_limit and mask_seqlen
+                            else (
+                                col0 - causal_row_offset - self.window_size_right
+                                if const_expr(self.window_size_right is not None)
+                                else 0
+                            )
+                        )
+                        row_limit_bot = (
+                            col0 - causal_row_offset + self.window_size_left
+                            if const_expr(self.window_size_left is not None)
+                            else self.tile_m
                         )
                         for r in cutlass.range(cute.size(tScS_mn.shape[0]), unroll_full=True):
                             row_idx = t0ScS_mn[r, 0][ROW]
                             acc_S_mn[r, c] = (
                                 else acc_S_mn[r, c]
                             )
+    @cute.jit
+    def apply_mask_mod_sm100_scalar(
+        self,
+        acc_S: cute.Tensor,
+        tScS_t2r: cute.Tensor,
+        m_block: Int32,
+        n_block: Int32,
+        mask_seqlen: cutlass.Constexpr[bool],
+        mask_mod: cutlass.Constexpr[Callable],
+        batch_idx: Int32,
+        head_idx: Int32,
+        aux_data: AuxData = AuxData(),
+        fastdiv_mods=(None, None),
+        head_divmod=None,
+        check_q_boundary: bool = False,
+    ) -> None:
+        """Apply a scalar FlexAttention mask_mod to an SM100 accumulator fragment.
+        Each accumulator lane calls mask_mod once with logical (batch, head, q, kv)
+        indices. Pack-GQA rows are converted back to logical q/head indices before
+        the call. When aux tensors are present, indices are wrapped with fastdiv so
+        mask_mod never reads outside the per-example auxiliary storage.
+        """
+        has_fastdiv = const_expr(
+            fastdiv_mods is not None and fastdiv_mods[0] is not None and fastdiv_mods[1] is not None
+        )
+        batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32)
+        ncol = const_expr(cute.size(tScS_t2r.shape))
+        for i in cutlass.range(ncol, unroll_full=True):
+            row_coord = tScS_t2r[i][0] if not self.swap_AB else tScS_t2r[i][1]
+            col_coord = tScS_t2r[i][1] if not self.swap_AB else tScS_t2r[i][0]
+            global_row = row_coord + m_block * self.tile_m
+            global_col = col_coord + n_block * self.tile_n
+            if const_expr(self.qhead_per_kvhead_packgqa != 1):
+                assert head_divmod is not None
+                mask_row, head_offset = divmod(global_row, head_divmod)
+                head_idx_for_mod = head_idx * self.qhead_per_kvhead_packgqa + head_offset
+            else:
+                head_idx_for_mod = head_idx
+                mask_row = global_row
+            mask_row_for_mod = mask_row
+            if const_expr(has_fastdiv and aux_data.tensors is not None):
+                if check_q_boundary:
+                    _, mask_row_for_mod = divmod(mask_row, fastdiv_mods[0])
+            global_col_for_mod = global_col
+            if const_expr(has_fastdiv and mask_seqlen and aux_data.tensors is not None):
+                _, global_col_for_mod = divmod(global_col, fastdiv_mods[1])
+            head_idx_ssa = utils.scalar_to_ssa(head_idx_for_mod, cutlass.Int32)
+            mask_row_ssa = utils.scalar_to_ssa(mask_row_for_mod, cutlass.Int32)
+            kv_idx_ssa = utils.scalar_to_ssa(global_col_for_mod, cutlass.Int32)
+            mask_value = call_mask_mod(
+                mask_mod,
+                batch_idx_ssa,
+                head_idx_ssa,
+                mask_row_ssa,
+                kv_idx_ssa,
+                self.seqlen_info,
+                aux_data,
+            )
+            cond = cutlass.Boolean(utils.ssa_to_scalar(mask_value))
+            acc_S[i] = acc_S[i] if cond else -Float32.inf
+            if const_expr(mask_seqlen):
+                acc_S[i] = -Float32.inf if global_col >= self.seqlen_k else acc_S[i]
+            if check_q_boundary:
+                acc_S[i] = -Float32.inf if mask_row >= self.seqlen_q else acc_S[i]
+    @cute.jit
+    def apply_mask_mod_sm100_vector(
+        self,
+        acc_S: cute.Tensor,
+        tScS_t2r: cute.Tensor,
+        m_block: Int32,
+        n_block: Int32,
+        mask_seqlen: cutlass.Constexpr[bool],
+        mask_mod: cutlass.Constexpr[Callable],
+        batch_idx: Int32,
+        head_idx: Int32,
+        vec_size: cutlass.Constexpr[int],
+        aux_data: AuxData = AuxData(),
+        fastdiv_mods=(None, None),
+        head_divmod=None,
+        check_q_boundary: bool = False,
+    ) -> None:
+        """Apply a vectorized FlexAttention mask_mod to an SM100 fragment.
+        mask_mod receives vec_size adjacent KV indices for one logical q row and
+        returns bit-packed Uint32 keep masks. Low bits correspond to lower KV
+        indices. The packed masks are combined with sequence-boundary checks, then
+        applied in 32-column chunks so the final masking lowers to R2P.
+        """
+        has_fastdiv = const_expr(
+            fastdiv_mods is not None and fastdiv_mods[0] is not None and fastdiv_mods[1] is not None
+        )
+        batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32)
+        ncol = const_expr(cute.size(tScS_t2r.shape))
+        mask_vals_per_apply = const_expr(max(1, vec_size // 32))
+        calls_per_apply = const_expr(max(1, 32 // vec_size))
+        n_calls = const_expr(cute.ceil_div(ncol, vec_size))
+        mask_vals = cute.make_rmem_tensor(mask_vals_per_apply, dtype=cutlass.Uint32)
+        # Accumulate enough vector mask_mod calls to produce 32-bit chunks that
+        # apply_packed_mask_chunk can lower to R2P.
+        for s in cutlass.range_constexpr(n_calls):
+            if const_expr(s % calls_per_apply == 0):
+                for c in cutlass.range_constexpr(mask_vals_per_apply):
+                    mask_vals[c] = cutlass.Uint32(0)
+            i = s * vec_size
+            row_coord = tScS_t2r[i][0] if not self.swap_AB else tScS_t2r[i][1]
+            col_coord = tScS_t2r[i][1] if not self.swap_AB else tScS_t2r[i][0]
+            global_row = row_coord + m_block * self.tile_m
+            global_col = col_coord + n_block * self.tile_n
+            if const_expr(self.qhead_per_kvhead_packgqa != 1):
+                assert head_divmod is not None
+                mask_row, head_offset = divmod(global_row, head_divmod)
+                head_idx_for_mod = head_idx * self.qhead_per_kvhead_packgqa + head_offset
+            else:
+                head_idx_for_mod = head_idx
+                mask_row = global_row
+            mask_row_for_mod = mask_row
+            if const_expr(has_fastdiv and aux_data.tensors is not None):
+                if check_q_boundary:
+                    _, mask_row_for_mod = divmod(mask_row, fastdiv_mods[0])
+            head_idx_ssa = utils.scalar_to_ssa(head_idx_for_mod, cutlass.Int32).broadcast_to(
+                (vec_size,)
+            )
+            mask_row_ssa = utils.scalar_to_ssa(mask_row_for_mod, cutlass.Int32).broadcast_to(
+                (vec_size,)
+            )
+            batch_idx_ssa_call = batch_idx_ssa.broadcast_to((vec_size,))
+            kv_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32)
+            # Build the per-lane KV indices for this vectorized mask_mod call.
+            for j in cutlass.range_constexpr(min(vec_size, ncol - i)):
+                col_j_coord = tScS_t2r[i + j][1] if not self.swap_AB else tScS_t2r[i + j][0]
+                col_j_global = col_j_coord + n_block * self.tile_n
+                col_j_for_mod = col_j_global
+                if const_expr(has_fastdiv and mask_seqlen and aux_data.tensors is not None):
+                    _, col_j_for_mod = divmod(col_j_global, fastdiv_mods[1])
+                kv_idx_vec[j] = col_j_for_mod
+            kv_idx_ssa = kv_idx_vec.load()
+            # mask_value is already bit-packed by the vectorized mask_mod.
+            mask_value = call_mask_mod(
+                mask_mod,
+                batch_idx_ssa_call,
+                head_idx_ssa,
+                mask_row_ssa,
+                kv_idx_ssa,
+                self.seqlen_info,
+                aux_data,
+            )
+            # For vec_size < 32, multiple mask_mod calls fill one R2P chunk.
+            bit_offset = const_expr((s % calls_per_apply) * vec_size)
+            seqlen_thresh_call = (
+                self.seqlen_k - global_col if const_expr(mask_seqlen) else cutlass.Int32(0)
+            )
+            q_in_bounds = mask_row < self.seqlen_q if check_q_boundary else cutlass.Boolean(True)
+            for c in cutlass.range_constexpr(mask_vals_per_apply):
+                mask_val = mask_value[c]
+                if const_expr(vec_size < 32):
+                    lane_keep = utils.shr_u32(
+                        cutlass.Uint32(0xFFFFFFFF),
+                        cutlass.Uint32(32 - vec_size),
+                    )
+                    mask_val = mask_val & lane_keep
+                if const_expr(mask_seqlen):
+                    mask_val = mask_val & r2p_bitmask_below(seqlen_thresh_call, c)
+                if check_q_boundary:
+                    mask_val = mask_val if q_in_bounds else cutlass.Uint32(0)
+                mask_vals[c] = mask_vals[c] | (mask_val << bit_offset)
+            # Apply only when the 32-bit chunk is complete, or at the tile tail.
+            is_last_in_apply = const_expr(s % calls_per_apply == calls_per_apply - 1)
+            is_last_overall = const_expr(s == n_calls - 1)
+            if const_expr(is_last_in_apply or is_last_overall):
+                apply_idx = s // calls_per_apply
+                for c in cutlass.range_constexpr(mask_vals_per_apply):
+                    chunk_idx = apply_idx * mask_vals_per_apply + c
+                    # Skip packed chunks that start past the accumulator fragment.
+                    if const_expr(chunk_idx * 32 < ncol):
+                        apply_packed_mask_chunk(acc_S, chunk_idx, mask_vals[c])
     @cute.jit
     def apply_mask_sm100(
         self,
         mask_mod: cutlass.Constexpr[Optional[Callable]] = None,
         batch_idx: Int32 = None,
         head_idx: Int32 = None,
+        aux_data: AuxData = AuxData(),
         fastdiv_mods=(None, None),
         head_divmod=None,
+        vec_size: cutlass.Constexpr[int] = 1,
         check_q_boundary: bool = False,
+        r2p: bool = True,
+        rBitmask: Optional[cute.Tensor] = None,
     ) -> 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)
         if n_block < 0:
             n_block = 0
         seqlenk_col_limit = self.seqlen_k - n_block * self.tile_n
+        if const_expr(rBitmask is not None):
+            ncol_packed = const_expr(cute.size(rBitmask.shape[0]))
+            for i in cutlass.range_constexpr(ncol_packed):
+                col_start = 32 * i  # mask is bit-packed into uint32
+                curr_mask_val = rBitmask[i]
+                for j in cutlass.range_constexpr(32):
+                    curr_col = col_start + j
+                    mask = (curr_mask_val >> j) & 1
+                    acc_S[curr_col] = acc_S[curr_col] if cutlass.Boolean(mask) else -Float32.inf
+        elif const_expr(not mask_causal and not mask_local and mask_mod is None):
             if const_expr(mask_seqlen):
                 if const_expr(not r2p):
                     for i in cutlass.range(cute.size(tScS_t2r.shape), unroll_full=True):
                         # For some reason the 2 lines above generate really bad SASS
                         acc_S[i] = -Float32.inf if tScS_t2r[i][1] >= seqlenk_col_limit else acc_S[i]
                 else:
+                    mask_r2p_lambda(
+                        acc_S,
+                        lambda s: r2p_bitmask_below(seqlenk_col_limit, s),
+                        rank1=True,
+                    )
         elif const_expr(not mask_causal and not mask_local and mask_mod is not None):
+            # FlexAttention mask_mod vectorization is gated on `mask_mod.__vec_size__`.
+            # vec_size == 1 returns a scalar Boolean. vec_size > 1 returns packed
+            # Uint32 mask fragments: one word per 32 evaluated columns.
+            assert vec_size % 32 == 0 or 32 % vec_size == 0, (
+                "vec_size must divide 32 or be a multiple of 32"
             )
+            if const_expr(vec_size == 1):
+                self.apply_mask_mod_sm100_scalar(
+                    acc_S,
+                    tScS_t2r,
+                    m_block,
+                    n_block,
+                    mask_seqlen,
+                    mask_mod,
+                    batch_idx,
+                    head_idx,
+                    aux_data,
+                    fastdiv_mods,
+                    head_divmod,
+                    check_q_boundary,
+                )
+            else:
+                self.apply_mask_mod_sm100_vector(
+                    acc_S,
+                    tScS_t2r,
+                    m_block,
+                    n_block,
+                    mask_seqlen,
+                    mask_mod,
+                    batch_idx,
+                    head_idx,
+                    vec_size,
+                    aux_data,
+                    fastdiv_mods,
+                    head_divmod,
+                    check_q_boundary,
                 )
         else:  # Causal or local
+            causal_row_offset = self.seqlen_k - n_block * self.tile_n - self.seqlen_q
             row_idx = tScS_t2r[0][0] + m_block * self.tile_m
             if const_expr(self.qhead_per_kvhead_packgqa != 1):
                 row_idx = row_idx // self.qhead_per_kvhead_packgqa
             if const_expr(mask_causal):
+                col_limit_right = row_idx + causal_row_offset + 1
                 if const_expr(mask_seqlen):
                     col_limit_right = cutlass.min(col_limit_right, seqlenk_col_limit)
                 # if cute.arch.thread_idx()[0] % 32 == 0:
                     for i in cutlass.range(ncol, unroll_full=True):
                         acc_S[i] = -Float32.inf if tScS_t2r[i][1] >= col_limit_right else acc_S[i]
                 else:
+                    mask_r2p_lambda(
+                        acc_S,
+                        lambda s: r2p_bitmask_below(col_limit_right, s),
+                        rank1=True,
+                    )
             else:
                 local_row_offset_right = (
+                    causal_row_offset + 1 + self.window_size_right
                     if const_expr(self.window_size_right is not None)
                     else None
                 )
                 local_row_offset_left = (
+                    causal_row_offset - self.window_size_left
                     if const_expr(self.window_size_left is not None)
                     else None
                 )
                             else acc_S[i]
                         )
                 else:
+                    # Dual-bound R2P masking for SM100.
+                    # Masks elements where: NOT (col_limit_left <= col < col_limit_right)
+                    def mask_gen_fn(s: int) -> Uint32:
+                        return r2p_bitmask_below(col_limit_right, s) & r2p_bitmask_above(
+                            col_limit_left, s
+                        )
+                    mask_r2p_lambda(acc_S, mask_gen_fn, rank1=True)
     @cute.jit
     def apply_mask_sm100_transposed(
         mask_mod: cutlass.Constexpr[Optional[Callable]] = None,
         batch_idx: Int32 = None,
         head_idx: Int32 = None,
+        aux_data: AuxData = AuxData(),
         fastdiv_mods=(None, None),
         is_full_block: bool = False,
         check_m_boundary: bool = True,
         """
         Backward pass: mask S = K @ Q.T where n_block tiles seqlen_k and m_block tiles seqlen_q.
+        Coordinate convention:
         - ROW corresponds to Q (m_block)
         - COL corresponds to KV (n_block)
                     and fastdiv_mods[1] is not None
                 )
                 wrap_aux_indices = const_expr(
+                    has_fastdiv and mask_seqlen and const_expr(aux_data.tensors is not None)
                 )
                 batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32)
                 head_idx_ssa = utils.scalar_to_ssa(head_idx, cutlass.Int32)
                     q_idx_ssa = utils.scalar_to_ssa(q_idx_for_mod, cutlass.Int32)
                     kv_idx_ssa = utils.scalar_to_ssa(kv_idx_for_mod, cutlass.Int32)
+                    mask_value = call_mask_mod(
+                        mask_mod,
                         batch_idx_ssa,
                         head_idx_ssa,
                         q_idx_ssa,
                         kv_idx_ssa,
                         self.seqlen_info,
+                        aux_data,
                     )
                     cond = cutlass.Boolean(utils.ssa_to_scalar(mask_value))
                     acc_S[i] = acc_S[i] if cond else -cutlass.Float32.inf
                         )
                 else:
                     num_rep = cute.size(tScS_t2r, mode=[0])  # 16 or 32
+                    num_wg = 2
+                    row_limit = row_to_r2p_idx(row_limit_top, num_rep, num_wg)
+                    mask_r2p_lambda(
+                        acc_S,
+                        lambda s: r2p_bitmask_above(row_limit, s),
+                        rank1=True,
+                    )
             else:
                 if const_expr(self.window_size_right is not None):
                     row_limit_top = causal_offset - self.window_size_right
                 if const_expr(mask_seqlen):
                     if seqlenk_col_limit <= 0:
                         row_limit_top = self.tile_m
+                r2p = True
+                if const_expr(not r2p):
+                    for i in cutlass.range(cute.size(acc_S.shape), unroll_full=True):
+                        row_idx = t0ScS_t2r[i][ROW]
+                        local_mask = row_idx < row_limit_top
+                        if const_expr(self.window_size_left is not None):
+                            local_mask |= row_idx > row_limit_bot
+                        acc_S[i] = -cutlass.Float32.inf if local_mask else acc_S[i]
+                else:
+                    def mask_gen_fn(s: int) -> Uint32:
+                        num_rep = cute.size(tScS_t2r, mode=[0])
+                        num_wg = 2
+                        row_limit = row_to_r2p_idx(row_limit_top, num_rep, num_wg)
+                        mask = r2p_bitmask_above(row_limit, s)
+                        if const_expr(self.window_size_left is not None):
+                            row_limit_bottom = row_to_r2p_idx(row_limit_bot + 1, num_rep, num_wg)
+                            mask = mask & r2p_bitmask_below(row_limit_bottom, s)
+                        return mask
+                    mask_r2p_lambda(
+                        acc_S,
+                        mask_gen_fn,
+                        rank1=True,
+                    )
+# -----------------------------------------------------------------------------
+# SM100 FMHA fused-mask policy layer (separate from generic mask primitives).
+# -----------------------------------------------------------------------------
+class Sm100MaskEnum(enum.Enum):
+    """Enumeration of mask types for FMHA operations.
+    - RESIDUAL_MASK: Residual mask for handling variable sequence lengths
+    - WINDOW_MASK: Window mask for attention which also includes causal and no mask
+    - WINDOW_MASK_INFERENCE: Same as the window mask, but has the limitation that the end of q is aligned with the end of k
+    - WINDOW_MASK_BWD: Window mask for backward pass
+    - WINDOW_MASK_BWD_INFERENCE: Same as the window mask for backward pass, but has the limitation that the end of q is aligned with the end of k
+    """
+    NO_MASK = enum.auto()
+    RESIDUAL_MASK = enum.auto()
+    CAUSAL_MASK = enum.auto()
+    WINDOW_MASK = enum.auto()
+    WINDOW_MASK_INFERENCE = enum.auto()
+    # Deprecated the following types
+    WINDOW_MASK_BWD = enum.auto()
+    WINDOW_MASK_BWD_INFERENCE = enum.auto()
+    RESIDUAL_MASK_BWD = enum.auto()
+class Sm100FusedMask:
+    """A fused mask implementation for FMHA operations.
+    This class handles different types of attention masks including no mask,
+    residual mask for variable sequence lengths, and causal mask for
+    autoregressive attention patterns.
+    The class provides methods to:
+    - Calculate trip counts for different mask types
+    - Apply masks to attention scores
+    - Handle masked and unmasked trip calculations
+    """
+    def get_trip_count(
+        mask_type: Sm100MaskEnum,
+        blk_coord: cute.Coord,
+        tile_shape: cute.Shape,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        window_size_left: Optional[Int32] = None,
+        window_size_right: Optional[Int32] = None,
+    ) -> Int32:
+        """
+        Calculate the number of trips needed for the current block.
+        The trip count depends on the mask type and the block coordinates.
+        For causal masks, it considers the autoregressive constraint.
+        :param mask_type: Type of mask to use
+        :type mask_type: utils.Sm100MaskEnum
+        :param blk_coord: Block coordinates.
+        :type blk_coord: cute.Coord
+        :param tile_shape: Shape of the tile.
+        :type tile_shape: cute.Shape
+        :param seqlen_q: Query sequence length for attention computation.
+        :type seqlen_q: Int32
+        :param seqlen_k: Key sequence length for attention computation.
+        :type seqlen_k: Int32
+        :param window_size_left: Left-side sliding window size for attention masking.
+        :type window_size_left: Optional[Int32]
+        :param window_size_right: Right-side sliding window size for attention masking.
+        :type window_size_right: Optional[Int32]
+        :return: Number of trips needed.
+        :rtype: Int32
+        """
+        result = 0
+        offset = 0
+        if cutlass.const_expr(mask_type is Sm100MaskEnum.WINDOW_MASK_INFERENCE):
+            offset = seqlen_k - seqlen_q
+        if cutlass.const_expr(mask_type is Sm100MaskEnum.WINDOW_MASK_BWD_INFERENCE):
+            offset = seqlen_q - seqlen_k
+        if cutlass.const_expr(mask_type == Sm100MaskEnum.RESIDUAL_MASK):
+            result = cute.ceil_div(seqlen_k, tile_shape[1])
+        if cutlass.const_expr(mask_type is Sm100MaskEnum.RESIDUAL_MASK_BWD):
+            result = cute.ceil_div(seqlen_q, tile_shape[0])
+        if cutlass.const_expr(
+            mask_type == Sm100MaskEnum.WINDOW_MASK
+            or mask_type == Sm100MaskEnum.WINDOW_MASK_INFERENCE
+        ):
+            if cutlass.const_expr(window_size_right is None):
+                result = cute.ceil_div(seqlen_k, tile_shape[1])
+            else:
+                max_idx_q = (blk_coord[0] + 1) * tile_shape[0]
+                idx_k = max_idx_q + offset + window_size_right
+                tmp_blocks_k = cute.ceil_div(idx_k, tile_shape[1])
+                max_blocks_k = cute.ceil_div(seqlen_k, tile_shape[1])
+                result = dsl_min(max_blocks_k, tmp_blocks_k)
+        if cutlass.const_expr(
+            mask_type == Sm100MaskEnum.WINDOW_MASK_BWD
+            or mask_type == Sm100MaskEnum.WINDOW_MASK_BWD_INFERENCE
+        ):
+            if cutlass.const_expr(window_size_left is None):
+                result = cute.ceil_div(seqlen_q, tile_shape[0])
+            else:
+                max_idx_k = (blk_coord[1] + 1) * tile_shape[1]
+                idx_k = max_idx_k + offset + window_size_left
+                tmp_blocks_q = cute.ceil_div(idx_k, tile_shape[0])
+                max_blocks_q = cute.ceil_div(seqlen_q, tile_shape[0])
+                result = dsl_min(max_blocks_q, tmp_blocks_q)
+        start_block = Sm100FusedMask.get_trip_start(
+            mask_type,
+            blk_coord,
+            tile_shape,
+            seqlen_q,
+            seqlen_k,
+            window_size_left,
+            window_size_right,
+        )
+        result = result - start_block
+        return result
+    @cute.jit
+    def get_trip_start_count_via_block_info(
+        blk_coord: cute.Coord,
+        tile_shape: cute.Shape,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        is_causal: cutlass.Constexpr[bool] = False,
+        is_local: cutlass.Constexpr[bool] = False,
+        window_size_left: Optional[Int32] = None,
+        window_size_right: Optional[Int32] = None,
+    ) -> Tuple[Int32, Int32]:
+        block_info = BlockInfo(
+            tile_m=tile_shape[0],
+            tile_n=tile_shape[1],
+            is_causal=is_causal,
+            is_local=is_local and not is_causal,
+            window_size_left=window_size_left,
+            window_size_right=window_size_right,
+        )
+        seqlen_info = SeqlenInfoQK(
+            offset_q=Int32(0),
+            offset_k=Int32(0),
+            padded_offset_q=Int32(0),
+            padded_offset_k=Int32(0),
+            seqlen_q=seqlen_q,
+            seqlen_k=seqlen_k,
+            m_block_offset=Int32(0),
+            block_idx_offset=Int32(0),
+            num_n_blocks=cute.ceil_div(seqlen_k, tile_shape[1]),
+            has_cu_seqlens_q=False,
+            has_cu_seqlens_k=False,
+            has_seqused_q=False,
+            has_seqused_k=False,
+            has_cu_block_idx_offsets=False,
+        )
+        n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen_info, blk_coord[0])
+        return n_block_min, n_block_max - n_block_min
+    @cute.jit
+    def get_trip_mask_bounds_via_block_info(
+        blk_coord: cute.Coord,
+        tile_shape: cute.Shape,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        is_causal: cutlass.Constexpr[bool] = False,
+        is_local: cutlass.Constexpr[bool] = False,
+        window_size_left: Optional[Int32] = None,
+        window_size_right: Optional[Int32] = None,
+    ) -> Tuple[Int32, Int32]:
+        """Return SM100-style mask boundaries for dense iteration.
+        Returns:
+          - n_block_min_causal_local_mask: right-side masked region start
+          - n_block_min_before_local_mask: start of fully unmasked middle region
+        """
+        block_info = BlockInfo(
+            tile_m=tile_shape[0],
+            tile_n=tile_shape[1],
+            is_causal=is_causal,
+            is_local=is_local and not is_causal,
+            window_size_left=window_size_left,
+            window_size_right=window_size_right,
+        )
+        seqlen_info = SeqlenInfoQK(
+            offset_q=Int32(0),
+            offset_k=Int32(0),
+            padded_offset_q=Int32(0),
+            padded_offset_k=Int32(0),
+            seqlen_q=seqlen_q,
+            seqlen_k=seqlen_k,
+            m_block_offset=Int32(0),
+            block_idx_offset=Int32(0),
+            num_n_blocks=cute.ceil_div(seqlen_k, tile_shape[1]),
+            has_cu_seqlens_q=False,
+            has_cu_seqlens_k=False,
+            has_seqused_q=False,
+            has_seqused_k=False,
+            has_cu_block_idx_offsets=False,
+        )
+        n_block_min, _ = block_info.get_n_block_min_max(seqlen_info, blk_coord[0])
+        n_block_min_causal_local_mask = block_info.get_n_block_min_causal_local_mask(
+            seqlen_info, blk_coord[0], n_block_min
+        )
+        n_block_min_before_local_mask = block_info.get_n_block_min_before_local_mask(
+            seqlen_info, blk_coord[0], n_block_min
+        )
+        return n_block_min_causal_local_mask, n_block_min_before_local_mask
+    @cute.jit
+    def get_trip_start(
+        mask_type: Sm100MaskEnum,
+        blk_coord: cute.Coord,
+        tile_shape: cute.Shape,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        window_size_left: Optional[Int32] = None,
+        window_size_right: Optional[Int32] = None,
+    ) -> Int32:
+        """
+        Get the start of the trip for the current block.
+        :param mask_type: Type of mask to use
+        :type mask_type: utils.Sm100MaskEnum
+        :param blk_coord: Block coordinates.
+        :type blk_coord: cute.Coord
+        :param tile_shape: Shape of the tile.
+        :type tile_shape: cute.Shape
+        :param seqlen_q: Query sequence length for attention computation.
+        :type seqlen_q: Int32
+        :param seqlen_k: Key sequence length for attention computation.
+        :type seqlen_k: Int32
+        :param window_size_left: Left-side sliding window size for attention masking.
+        :type window_size_left: Optional[Int32]
+        :param window_size_right: Right-side sliding window size for attention masking.
+        :type window_size_right: Optional[Int32]
+        """
+        result = 0
+        offset = 0
+        if cutlass.const_expr(mask_type is Sm100MaskEnum.WINDOW_MASK_INFERENCE):
+            offset = seqlen_k - seqlen_q
+        if cutlass.const_expr(mask_type is Sm100MaskEnum.WINDOW_MASK_BWD_INFERENCE):
+            offset = seqlen_q - seqlen_k
+        if cutlass.const_expr(
+            mask_type is Sm100MaskEnum.WINDOW_MASK
+            or mask_type is Sm100MaskEnum.WINDOW_MASK_INFERENCE
+        ):
+            if cutlass.const_expr(window_size_left is not None):
+                min_idx_q = blk_coord[0] * tile_shape[0]
+                idx_k = min_idx_q + offset - window_size_left
+                tmp_blocks_k = idx_k // tile_shape[1]
+                result = max(tmp_blocks_k, result)
+        if cutlass.const_expr(
+            mask_type is Sm100MaskEnum.WINDOW_MASK_BWD
+            or mask_type is Sm100MaskEnum.WINDOW_MASK_BWD_INFERENCE
+        ):
+            if cutlass.const_expr(window_size_right is not None):
+                min_idx_k = blk_coord[1] * tile_shape[1]
+                idx_q = min_idx_k + offset - window_size_right
+                tmp_blocks_q = idx_q // tile_shape[0]
+                result = max(tmp_blocks_q, result)
+        return result
+    @cute.jit
+    def get_leading_mask_id(
+        mask_type: Sm100MaskEnum,
+        blk_coord: cute.Coord,
+        tile_shape: cute.Shape,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        window_size_left: Optional[Int32] = None,
+        window_size_right: Optional[Int32] = None,
+    ) -> Tuple[Int32, Int32]:
+        """
+        Get the begin and end tile idx for the leading mask.
+        :param mask_type: Type of mask to use
+        :type mask_type: utils.Sm100MaskEnum
+        :param blk_coord: Block coordinates.
+        :type blk_coord: cute.Coord
+        :param tile_shape: Shape of the tile.
+        :type tile_shape: cute.Shape
+        :param seqlen_q: Query sequence length for attention computation.
+        :type seqlen_q: Int32
+        :param seqlen_k: Key sequence length for attention computation.
+        :type seqlen_k: Int32
+        :param window_size_left: Left-side sliding window size for attention masking.
+        :type window_size_left: Optional[Int32]
+        :param window_size_right: Right-side sliding window size for attention masking.
+        :type window_size_right: Optional[Int32]
+        :return: Tuple of (begin, end) tile idx for the leading mask.
+        :rtype: Tuple[Int32, Int32]
+        """
+        offset = 0
+        if cutlass.const_expr(mask_type is Sm100MaskEnum.WINDOW_MASK_INFERENCE):
+            offset = seqlen_k - seqlen_q
+        if cutlass.const_expr(mask_type is Sm100MaskEnum.WINDOW_MASK_BWD_INFERENCE):
+            offset = seqlen_q - seqlen_k
+        leading_mask_begin = Sm100FusedMask.get_trip_start(
+            mask_type,
+            blk_coord,
+            tile_shape,
+            seqlen_q,
+            seqlen_k,
+            window_size_left,
+            window_size_right,
+        )
+        trip_count = Sm100FusedMask.get_trip_count(
+            mask_type,
+            blk_coord,
+            tile_shape,
+            seqlen_q,
+            seqlen_k,
+            window_size_left,
+            window_size_right,
+        )
+        leading_mask_end = leading_mask_begin
+        if cutlass.const_expr(
+            mask_type is Sm100MaskEnum.WINDOW_MASK
+            or mask_type is Sm100MaskEnum.WINDOW_MASK_INFERENCE
+        ):
+            if cutlass.const_expr(window_size_left is not None):
+                min_idx_q = (blk_coord[0] + 1) * tile_shape[0] + offset - window_size_left
+                leading_mask_end = dsl_min(
+                    cute.ceil_div(min_idx_q, tile_shape[1]) - 1,
+                    trip_count + leading_mask_begin - 1,
+                )
+            else:
+                leading_mask_end = leading_mask_begin - 1
+        elif cutlass.const_expr(
+            mask_type is Sm100MaskEnum.WINDOW_MASK_BWD
+            or mask_type is Sm100MaskEnum.WINDOW_MASK_BWD_INFERENCE
+        ):
+            if cutlass.const_expr(window_size_right is not None):
+                min_idx_k = (blk_coord[1] + 1) * tile_shape[1] + offset - window_size_right
+                leading_mask_end = cute.ceil_div(min_idx_k, tile_shape[0]) - 1
+            else:
+                leading_mask_end = leading_mask_begin - 1
+        return leading_mask_begin, leading_mask_end
+    @cute.jit
+    def get_trailing_mask_id(
+        mask_type: Sm100MaskEnum,
+        blk_coord: cute.Coord,
+        tile_shape: cute.Shape,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        window_size_left: Optional[Int32] = None,
+        window_size_right: Optional[Int32] = None,
+    ) -> Tuple[Optional[Int32], Optional[Int32]]:
+        """
+        Get the begin and end tile idx for the trailing mask.
+        :param mask_type: Type of mask to use
+        :type mask_type: utils.Sm100MaskEnum
+        :param blk_coord: Block coordinates.
+        :type blk_coord: cute.Coord
+        :param tile_shape: Shape of the tile.
+        :type tile_shape: cute.Shape
+        :param seqlen_q: Query sequence length for attention computation.
+        :type seqlen_q: Int32
+        :param seqlen_k: Key sequence length for attention computation.
+        :type seqlen_k: Int32
+        :param window_size_left: Left-side sliding window size for attention masking.
+        :type window_size_left: Optional[Int32]
+        :param window_size_right: Right-side sliding window size for attention masking.
+        :type window_size_right: Optional[Int32]
+        :return: Tuple of (begin, end) tile idx for the trailing mask.
+        :rtype: Tuple[Int32, Int32]
+        """
+        offset = 0
+        if cutlass.const_expr(mask_type is Sm100MaskEnum.WINDOW_MASK_INFERENCE):
+            offset = seqlen_k - seqlen_q
+        if cutlass.const_expr(mask_type is Sm100MaskEnum.WINDOW_MASK_BWD_INFERENCE):
+            offset = seqlen_q - seqlen_k
+        trip_start = Sm100FusedMask.get_trip_start(
+            mask_type,
+            blk_coord,
+            tile_shape,
+            seqlen_q,
+            seqlen_k,
+            window_size_left,
+            window_size_right,
+        )
+        trip_count = Sm100FusedMask.get_trip_count(
+            mask_type,
+            blk_coord,
+            tile_shape,
+            seqlen_q,
+            seqlen_k,
+            window_size_left,
+            window_size_right,
+        )
+        trailing_mask_begin, trailing_mask_end = None, None
+        if cutlass.const_expr(
+            mask_type is Sm100MaskEnum.WINDOW_MASK
+            or mask_type is Sm100MaskEnum.WINDOW_MASK_INFERENCE
+        ):
+            if cutlass.const_expr(window_size_right is not None):
+                min_idx_q = blk_coord[0] * tile_shape[0] + offset + window_size_right
+                trailing_mask_begin = dsl_min(
+                    min_idx_q // tile_shape[1], trip_count + trip_start - 1
+                )
+                trailing_mask_end = trip_count + trip_start - 1
+            else:
+                # last tile, we always apply mask on it regardless whether it's a residual tile
+                trailing_mask_begin = trip_count + trip_start - 1
+                trailing_mask_end = trip_count + trip_start - 1
+        else:
+            if cutlass.const_expr(window_size_left is not None):
+                min_idx_k = blk_coord[1] * tile_shape[1] + offset + window_size_left + 1
+                max_idx_k = (blk_coord[1] + 1) * tile_shape[1] + offset + window_size_left
+                trailing_mask_begin = dsl_min(
+                    cute.ceil_div(min_idx_k, tile_shape[0]) - 1,
+                    trip_count + trip_start - 1,
+                )
+                trailing_mask_end = dsl_min(
+                    cute.ceil_div(max_idx_k, tile_shape[0]) - 1,
+                    trip_count + trip_start - 1,
+                )
+            else:
+                # last tile, we always apply mask on it regardless whether it's a residual tile
+                trailing_mask_begin = trip_count + trip_start - 1
+                trailing_mask_end = trip_count + trip_start - 1
+        return trailing_mask_begin, trailing_mask_end
+    @cute.jit
+    def get_masked_leading_count(
+        mask_type: Sm100MaskEnum,
+        blk_coord: cute.Coord,
+        tile_shape: cute.Shape,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        window_size_left: Optional[Int32] = None,
+        window_size_right: Optional[Int32] = None,
+    ) -> Int32:
+        """
+        Calculate the number of masked trips for the leading mask.
+        This is used for blocks that need special handling due to masking.
+        :param mask_type: Type of mask to use
+        :type mask_type: utils.Sm100MaskEnum
+        :param blk_coord: Block coordinates.
+        :type blk_coord: cute.Coord
+        :param tile_shape: Shape of the tile.
+        :type tile_shape: cute.Shape
+        :param seqlen_q: Query sequence length for attention computation.
+        :type seqlen_q: Int32
+        :param seqlen_k: Key sequence length for attention computation.
+        :type seqlen_k: Int32
+        :param window_size_left: Left-side sliding window size for attention masking.
+        :type window_size_left: Optional[Int32]
+        :param window_size_right: Right-side sliding window size for attention masking.
+        :type window_size_right: Optional[Int32]
+        :return: Number of masked trips.
+        :rtype: Int32
+        """
+        result = 0
+        if cutlass.const_expr(
+            mask_type is not Sm100MaskEnum.RESIDUAL_MASK
+            and mask_type is not Sm100MaskEnum.RESIDUAL_MASK_BWD
+        ):
+            if cutlass.const_expr(window_size_left is not None or window_size_right is not None):
+                leading_mask_begin, leading_mask_end = Sm100FusedMask.get_leading_mask_id(
+                    mask_type,
+                    blk_coord,
+                    tile_shape,
+                    seqlen_q,
+                    seqlen_k,
+                    window_size_left,
+                    window_size_right,
+                )
+                result = max(leading_mask_end - leading_mask_begin + 1, 0)
+        return result
+    @cute.jit
+    def get_masked_trailing_count(
+        mask_type: Sm100MaskEnum,
+        blk_coord: cute.Coord,
+        tile_shape: cute.Shape,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        window_size_left: Optional[Int32] = None,
+        window_size_right: Optional[Int32] = None,
+        rem_count: Optional[Int32] = 0,
+    ) -> Int32:
+        """
+        Calculate the number of masked trips for the trailing mask.
+        This is used for blocks that need special handling due to masking.
+        :param mask_type: Type of mask to use
+        :type mask_type: utils.Sm100MaskEnum
+        :param blk_coord: Block coordinates.
+        :type blk_coord: cute.Coord
+        :param tile_shape: Shape of the tile.
+        :type tile_shape: cute.Shape
+        :param seqlen_q: Query sequence length for attention computation.
+        :type seqlen_q: Int32
+        :param seqlen_k: Key sequence length for attention computation.
+        :type seqlen_k: Int32
+        :param window_size_left: Left-side sliding window size for attention masking.
+        :type window_size_left: Optional[Int32]
+        :param window_size_right: Right-side sliding window size for attention masking.
+        :type window_size_right: Optional[Int32]
+        :param rem_count: Remaining count from previous calculations.
+        :type rem_count: Int32
+        :return: Number of masked trips.
+        :rtype: Int32
+        """
+        result = 0
+        if cutlass.const_expr(
+            mask_type is not Sm100MaskEnum.RESIDUAL_MASK
+            and mask_type is not Sm100MaskEnum.RESIDUAL_MASK_BWD
+        ):
+            if cutlass.const_expr(window_size_left is not None or window_size_right is not None):
+                trailing_mask_begin, trailing_mask_end = Sm100FusedMask.get_trailing_mask_id(
+                    mask_type,
+                    blk_coord,
+                    tile_shape,
+                    seqlen_q,
+                    seqlen_k,
+                    window_size_left,
+                    window_size_right,
+                )
+                leading_mask_begin, leading_mask_end = Sm100FusedMask.get_leading_mask_id(
+                    mask_type,
+                    blk_coord,
+                    tile_shape,
+                    seqlen_q,
+                    seqlen_k,
+                    window_size_left,
+                    window_size_right,
+                )
+                if cutlass.const_expr(
+                    trailing_mask_begin is not None and trailing_mask_end is not None
+                ):
+                    if trailing_mask_begin <= leading_mask_end:
+                        result = max(trailing_mask_end - leading_mask_end, 0)
+                    else:
+                        result = max(trailing_mask_end - trailing_mask_begin + 1, 0)
+        else:
+            if seqlen_k % tile_shape[1] != 0:
+                result = 1
+            else:
+                result = 0
+        return result + rem_count
+    @cute.jit
+    def get_unmasked_trip_count(
+        mask_type: Sm100MaskEnum,
+        blk_coord: cute.Coord,
+        tile_shape: cute.Shape,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        window_size_left: Optional[Int32] = None,
+        window_size_right: Optional[Int32] = None,
+    ) -> Int32:
+        """
+        Calculate the number of unmasked trips for the current block.
+        This represents the number of trips that don't require special
+        masking treatment.
+        :param mask_type: Type of mask to use
+        :type mask_type: utils.Sm100MaskEnum
+        :param blk_coord: Block coordinates.
+        :type blk_coord: cute.Coord
+        :param tile_shape: Shape of the tile.
+        :type tile_shape: cute.Shape
+        :param seqlen_q: Query sequence length for attention computation.
+        :type seqlen_q: Int32
+        :param seqlen_k: Key sequence length for attention computation.
+        :type seqlen_k: Int32
+        :param window_size_left: Left-side sliding window size for attention masking.
+        :type window_size_left: Optional[Int32]
+        :param window_size_right: Right-side sliding window size for attention masking.
+        :type window_size_right: Optional[Int32]
+        :return: Number of unmasked trips.
+        :rtype: Int32
+        """
+        result = (
+            Sm100FusedMask.get_trip_count(
+                mask_type,
+                blk_coord,
+                tile_shape,
+                seqlen_q,
+                seqlen_k,
+                window_size_left,
+                window_size_right,
+            )
+            - Sm100FusedMask.get_masked_leading_count(
+                mask_type,
+                blk_coord,
+                tile_shape,
+                seqlen_q,
+                seqlen_k,
+                window_size_left,
+                window_size_right,
+            )
+            - Sm100FusedMask.get_masked_trailing_count(
+                mask_type,
+                blk_coord,
+                tile_shape,
+                seqlen_q,
+                seqlen_k,
+                window_size_left,
+                window_size_right,
+                0,
+            )
+        )
+        return result
+    @cute.jit
+    def apply_mask(
+        mask_type: Sm100MaskEnum,
+        acc_qk: cute.Tensor,
+        index_qk: cute.Tensor,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        window_size_left: Optional[int] = None,
+        window_size_right: Optional[int] = None,
+        index_transform: cutlass.Constexpr = lambda index_q, index_k: (
+            index_q,
+            index_k,
+        ),
+    ):
+        """
+        Apply the appropriate mask to the attention scores.
+        This method modifies the attention scores (acc_qk) based on the mask type
+        and the positions in the index tensor.
+        :param mask_type: Type of mask to use
+        :type mask_type: utils.Sm100MaskEnum
+        :param acc_qk: Accumulated QK attention scores tensor.
+        :type acc_qk: cute.Tensor
+        :param index_qk: Index tensor containing position information.
+        :type index_qk: cute.Tensor
+        :param seqlen_k: Key sequence length for attention computation.
+        :type seqlen_k: Int32
+        :param seqlen_q: Query sequence length for attention computation.
+        :type seqlen_q: Optional[int]
+        :param window_size_left: Left-side sliding window size for attention masking.
+        :type window_size_left: Optional[int]
+        :param window_size_right: Right-side sliding window size for attention masking.
+        :type window_size_right: Optional[int]
+        """
+        offset = 0
+        # NOTE: causal masking in this repo aligns the *end* of Q with the *end* of K
+        # when seqlen_k != seqlen_q (same as the test/reference implementation):
+        #   k_index <= q_index + (seqlen_k - seqlen_q) + window_right
+        # In our kernels, causal is represented by (window_left is None, window_right is not None).
+        if cutlass.const_expr(window_size_left is None and window_size_right is not None):
+            offset = seqlen_k - seqlen_q
+        elif cutlass.const_expr(
+            mask_type is Sm100MaskEnum.WINDOW_MASK_INFERENCE
+            or mask_type is Sm100MaskEnum.WINDOW_MASK_BWD_INFERENCE
+        ):
+            offset = seqlen_k - seqlen_q
+        for i in cutlass.range_constexpr(cute.size(acc_qk), unroll_full=True):
+            index_q, index_k = index_transform(*index_qk[i])
+            if cutlass.const_expr(window_size_left is not None or window_size_right is not None):
+                if cutlass.const_expr(window_size_left is None):
+                    if index_q + offset + window_size_right < index_k:
+                        acc_qk[i] = -Float32.inf
+                    if index_k >= seqlen_k or index_q >= seqlen_q:  # residual mask
+                        acc_qk[i] = -Float32.inf
+                elif cutlass.const_expr(window_size_right is None):
+                    if index_q + offset - window_size_left > index_k:
+                        acc_qk[i] = -Float32.inf
+                    if index_k >= seqlen_k or index_q >= seqlen_q:  # residual mask
+                        acc_qk[i] = -Float32.inf
+                else:
+                    max_K_index = dsl_min(index_q + offset + window_size_right, seqlen_k)
+                    min_K_index = max(0, index_q + offset - window_size_left)
+                    if index_k > max_K_index or index_k < min_K_index:
+                        acc_qk[i] = -Float32.inf
+                    if index_k >= seqlen_k or index_q >= seqlen_q:  # residual mask
+                        acc_qk[i] = -Float32.inf
+            if cutlass.const_expr(
+                mask_type == Sm100MaskEnum.RESIDUAL_MASK
+                or mask_type == Sm100MaskEnum.RESIDUAL_MASK_BWD
+            ):
+                if index_k >= seqlen_k or index_q >= seqlen_q:
+                    acc_qk[i] = -Float32.inf
+    @cute.jit
+    def apply_mask_via_causal_local(
+        acc_qk: cute.Tensor,
+        index_qk: cute.Tensor,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        apply_semantic_window: cutlass.Constexpr[bool] = True,
+        is_causal: cutlass.Constexpr[bool] = False,
+        is_local: cutlass.Constexpr[bool] = False,
+        window_size_left: Optional[int] = None,
+        window_size_right: Optional[int] = None,
+        index_transform: cutlass.Constexpr = lambda index_q, index_k: (
+            index_q,
+            index_k,
+        ),
+    ):
+        """Apply forward mask without mask_type.
+        - If apply_semantic_window=True, apply causal/local window constraints.
+        - Always apply residual OOB masking (index_k>=seqlen_k or index_q>=seqlen_q).
+        """
+        offset = 0
+        if cutlass.const_expr(apply_semantic_window):
+            # Match WINDOW_MASK_INFERENCE semantics: end-align Q/K when lengths differ.
+            offset = seqlen_k - seqlen_q
+        for i in cutlass.range_constexpr(cute.size(acc_qk), unroll_full=True):
+            index_q, index_k = index_transform(*index_qk[i])
+            if cutlass.const_expr(apply_semantic_window):
+                if cutlass.const_expr(is_causal and not is_local):
+                    # Pure causal; tolerate both external forms:
+                    # - (None, None) from interface
+                    # - (None, 0) from fused-mask-style callers
+                    right = 0 if const_expr(window_size_right is None) else window_size_right
+                    if index_q + offset + right < index_k:
+                        acc_qk[i] = -Float32.inf
+                elif cutlass.const_expr(
+                    is_local or window_size_left is not None or window_size_right is not None
+                ):
+                    if cutlass.const_expr(window_size_left is None):
+                        if index_q + offset + window_size_right < index_k:
+                            acc_qk[i] = -Float32.inf
+                    elif cutlass.const_expr(window_size_right is None):
+                        if index_q + offset - window_size_left > index_k:
+                            acc_qk[i] = -Float32.inf
+                    else:
+                        max_K_index = dsl_min(index_q + offset + window_size_right, seqlen_k)
+                        min_K_index = max(0, index_q + offset - window_size_left)
+                        if index_k > max_K_index or index_k < min_K_index:
+                            acc_qk[i] = -Float32.inf
+            # Residual mask is always needed for boundary protection.
+            if index_k >= seqlen_k or index_q >= seqlen_q:
+                acc_qk[i] = -Float32.inf

build/torch-cuda/metadata.json CHANGED Viewed

@@ -1,9 +1,93 @@
 {
-  "version": 0,
   "license": "BSD-3-Clause",
   "python-depends": [
     "einops",
     "tvm-ffi",
     "nvidia-cutlass-dsl"
-  ]
 }

 {
+  "name": "flash-attn4",
+  "id": "_flash_attn4_cuda_d8eaf3c",
+  "version": 1,
+  "kernels-minver": "0.14.0",
   "license": "BSD-3-Clause",
+  "upstream": "https://github.com/Dao-AILab/flash-attention.git",
   "python-depends": [
     "einops",
     "tvm-ffi",
     "nvidia-cutlass-dsl"
+  ],
+  "kernel-depends": [],
+  "backend": {
+    "type": "cuda"
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "9nhmToMqhkzqb88qdOhfpDICvN3LHSZhEM8HXa7VlcU=",
+      "_ops.py": "PVG5218SwxBt8nM52gD8xBjus4sHeK7/H9pZkvPRxCc=",
+      "_private_for_testing.py": "QXVDjtjb5NEXpgkZ93UyWJgHbrhi757Jlc8lkkTKETw=",
+      "ampere_helpers.py": "ytrBcKfaN54nw20EQ2BIDXEM+ojQgtrSud7NWOSEcRE=",
+      "barrier.py": "dnr4B6p1IbjRXubJU09HRaetIiSw2iuBbMScy56mwRk=",
+      "blackwell_helpers.py": "1G7e3qNjDRq24hiED6qbAJaN+9Lt+VfKRdbBSQBlP5U=",
+      "block_info.py": "c0GOkiVrwE1jkpNJFT+RPO1xdRD5ZM6RImIsvyLWnLc=",
+      "block_sparse_utils.py": "Y4Puv8qFy3MmuN1wY7/SvPHiOkKdZC5ItN3JJhwGqvo=",
+      "block_sparsity.py": "0PG41ePOKP/KS1j6eIo9pkMp7a7vrdCAysgmU/Eul+I=",
+      "cache_utils.py": "LbHAKc04S2NK2yEvmR6aB8yZuCeT2MCOrVv5lvEzQhA=",
+      "compute_block_sparsity.py": "p2+8YK5TnSM+G7g02pkrQxllSjOg/Lb4hbgNsKsttrM=",
+      "copy_utils.py": "D1Y+7rZlRgLRVB9yAre4g4bE29vSStawglg01l5QjAk=",
+      "cu_blocks_kernel.py": "oZ/Lf41VMM5o6vzUuDvd/dwDyNm7hqFfHkeC/gnSKW8=",
+      "cute_dsl_ptxas.py": "IfBnTJ9amdfDOQkuSdWCLTh7CkZziIvs/xrAc8taxhk=",
+      "cute_dsl_utils.py": "6ko8k13DHTgblp8bhmMeqkTJ4izAlr0S0TDIGiofkrs=",
+      "fa_logging.py": "vHHt+ksct69kt43vCeeQ/0xv5kq0/Pjj2WmfROjTwKA=",
+      "fast_math.py": "fl+CK/TQujaWdVjzrpoUCLQQy3JsMSIs1hcDVp0b5zg=",
+      "flash_bwd.py": "TRhiH4s/hQpkpQAKPt8lRMpg/92QqYUIMCok60L9gUo=",
+      "flash_bwd_mla_dk_sm100.py": "nPKlnayjEyxu0T5JqrXjojn5wzBQmwW7O0RmekOu6Q8=",
+      "flash_bwd_mla_dq_dqv_sm100.py": "WM/DT6DLexdrSs4EWskTENaRwN9VdGbmbEOVqY1ITgU=",
+      "flash_bwd_mla_sm100.py": "k6mUzSMehuxfkOxnhOHQEh0VoWWu0Wqczqs0BoCPNyY=",
+      "flash_bwd_postprocess.py": "N3dE8x0gau28oJMZDaRnOtQ+Eu8dJdG32lo5c4RrDRw=",
+      "flash_bwd_preprocess.py": "0lZNAuKMo4hZb5EF1IjtoMvR5FSR0Iq1MTDKWOD25Gk=",
+      "flash_bwd_sm100.py": "32wNaSsj3PLh7FqTTgKan3Mwe9iXrKFZOXwjXFkiEFw=",
+      "flash_bwd_sm120.py": "pYu0Z1Yt1TGkLrGUDQbHj7e7plUvtCW9lE5kJhM+TeM=",
+      "flash_bwd_sm90.py": "DvV+aDvDjC50LnEXdXUizK6+maRzsjFmcGO+h8bLOGc=",
+      "flash_fwd.py": "CJCCeosdamts+KoMuX2wYzCTK8d9veU4W9xDd8Qrw3Q=",
+      "flash_fwd_combine.py": "BGXmjQjrNksA0+TpgF3Trh6KlmejTiSQD5TMbgTWujM=",
+      "flash_fwd_mla_sm100.py": "wsXfpdNvwO4bw4Ol38H5t9n57N+E5Ln8gpkNiazXA4E=",
+      "flash_fwd_sm100.py": "8kqCPNeGuI56BDiSIXdfFZalVLz8SWmGpkHlfvYwmXg=",
+      "flash_fwd_sm120.py": "RGmi1PkGYo48xqAWYimrIb7XrokTOKZYYQx5XkxDKco=",
+      "flash_fwd_sm90.py": "zmqsNatHicpxW+dSRnPn2gNMt+TaI8HX+VEOr2hlyns=",
+      "interface.py": "WSarPGybE0wvH3K6tYnGrldr/z2A7TQojBcd7Dz7iWc=",
+      "mask.py": "45XBbPI1GDqaoT0k3Nq83gKu61ZOGrK+R+mWzWXWw+k=",
+      "mma_sm100_desc.py": "rG7Ciy8aS3l8ORp4B7iPg7QwRS1BgmLD36ew81FuOfk=",
+      "named_barrier.py": "H4D5a2kKKIjuEwHTd2ZdtunddLQC+A6qhn049ylNgTs=",
+      "pack_gqa.py": "Oqx6zBC5MrUYnj5E/tTBMRuDOUV5kCu0slYHy6jDvDg=",
+      "paged_kv.py": "aCfJSKrg3NnhitFjG9FbShvP+PuAt+gQNz01AjgCtao=",
+      "pipeline.py": "hJvOkccr4KXHRYl54QyCkPReBSU6g2AOMPuVQMYe1Wo=",
+      "prepare_scheduler.py": "83KHAGQBk7uz7WdmCwV3/k7gEbhiUYEc3VVMkx+LwiU=",
+      "quack/__init__.py": "m24hqJucTK14gZS18kAZhF29vDzEWaplx6wp6XN4T70=",
+      "quack/compile_utils.py": "qJ3oTsDlbAiddrJHtEO7LPYVqn/s+neNfiw+/KvfXZU=",
+      "quack/copy_utils.py": "48qGZz3ql2hwxAH+rr8g7hJoiy2drx4YRnXjb3TcEJo=",
+      "quack/cute_dsl_utils.py": "4jBC37wrqgJU/2Y3HFVbi46rQBkade3+ymPcHIxAYQk=",
+      "quack/dsl/__init__.py": "7mF8m94hdEj3jrzX0sk0LBszrP7WDhdcUJd4lNDOpGs=",
+      "quack/dsl/cute_tensor_indexing.py": "1QnPrnekkECXkgHC5NWY6JxyDmTTf5kgClCIYwxMfG0=",
+      "quack/layout_utils.py": "MZncP1XzRhg+PShPbamPQ5Tq8U8ot2FtFH5uSeyJYZQ=",
+      "quack/rounding.py": "jI6fP/qCjSMh9s8PtjPDYfTPs200C3xx9zKll1hlodU=",
+      "quack/sm90_utils.py": "QaLZ3MLlp9dKOD+VwdhxmqEstPdq8CVGaNwffbwspNw=",
+      "quack/utils.py": "/FsIscjaxmOSj/g6frh4gLgB6iuIZ6e3EeGNunKWBuw=",
+      "seqlen_info.py": "ljfpUQcu9fILvJfyYYGkM85kiXjmNP1p748eivq0dwk=",
+      "sm100_hd256_2cta_fmha_backward.py": "BFD5O/rIJi5X//5ecE/h+zQHhb9HLI8TmpZHHhcVmns=",
+      "sm100_hd256_2cta_fmha_backward_dkdvkernel.py": "EJiDHIFGQ8A63CHSzFrDF5kigtWNZMqsvXk/28zVWfI=",
+      "sm100_hd256_2cta_fmha_backward_dqkernel.py": "SKQ8lUEl1P4ZcsHiMaGBqQJFL2o0oZTJhF8oOx6iXDs=",
+      "sm100_hd256_2cta_fmha_forward.py": "aJ0icWUbC3D2iLlj8SDXb+eaM0/U4mljFlIaH8iEZd4=",
+      "softmax.py": "eCMvLFCXjlrgm0s24D5rrdbFFRMyah1ESWjB3BEgLu4=",
+      "testing.py": "M0K0R1EmZm5b/wUDFs0qWGhS98gusj/SEXJHDMQKFf8=",
+      "tile_scheduler.py": "ELRFy6rI+J+8Z2fMDLxYrCiuUMSaBLa4h10dIe27bmA=",
+      "topk_gather_kv.py": "WchMfXog4eIIIOhTKu0vE5k3dfn41jVqrl9wtafJvRw=",
+      "utils.py": "j+SX9mJpWWA998oBdZgS28MoIV69XwTRN0mVA6Hcha8="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "commit": "a7f0afdb29a6a3372b1d47180cc0c182454c5e3b",
+      "dirty": false
+    },
+    "kernel": {
+      "commit": "d8eaf3cbe010f4795d0b0c4b0103862387ccdc97",
+      "dirty": false
+    }
+  }
 }

build/torch-cuda/metadata.json.sigstore ADDED Viewed

	@@ -0,0 +1 @@

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHLTCCBrKgAwIBAgIUC8FD02MzDpPjLuXMiJ3QOc8FBi4wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDA5MDU0NzQyWhcNMjYxMDA5MDU1NzQyWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEvU7iFTJAhfZ8mmvV3pYwcuyy5onUd8y8YkK3oDJrIYFKZZeP51FT0kntoP0n4AGSPRV4Ro9wrRR7jPlxS/0VyKOCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChkOGVhZjNjYmUwMTBmNDc5NWQwYjBjNGIwMTAzODYyMzg3Y2NkYzk3MBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChkOGVhZjNjYmUwMTBmNDc5NWQwYjBjNGIwMTAzODYyMzg3Y2NkYzk3MBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChkOGVhZjNjYmUwMTBmNDc5NWQwYjBjNGIwMTAzODYyMzg3Y2NkYzk3MB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZDhlYWYzY2JlMDEwZjQ3OTVkMGIwYzRiMDEwMzg2MjM4N2NjZGM5NzAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zNzg5MDA3MjU2Mi9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhHzP5igAABAMARzBFAiEAq9xS/SQssYBHs+y3wB9piTHTBtP5LMlf5JId4nIfA3sCIH5m7oZGU8V7XUlOnv+dvMJ3KQo46r1zQBPQkI9jcu5gMAoGCCqGSM49BAMDA2kAMGYCMQCrjTEdjE3f7UqDCOh3sVTk2MU9H7IJ2OkZwrbUwupS/xvK0ni6p5gl2AdH0pLBfkkCMQDkblT245iYWckRzNjXOOIUi4lNnrNhu5fdRtd9Z58/K5sQp21Z61tVdGfkOPF8UF4="},"tlogEntries":[{"logIndex":"3158296362","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1791524864","inclusionPromise":{"signedEntryTimestamp":"MEUCIHdcmCseFb7Vr5yOOIYHwnu8qbD12aVWYFJAvrIFWQlQAiEAiI/g9IvQbhN7/BHfD7fKpMlkheGsalXj3sjCz0fV9Eg="},"inclusionProof":{"logIndex":"3036392100","rootHash":"teeCGBfEA5uGpfYT9SeWxOIWLaZFezcGvaKdKW6ZXHk=","treeSize":"3036392181","hashes":["KyrJTiRKZ75+pT8Sx+ulXcJDqdovWH9dj1tSGpuz7bI=","GHJ44yU03czeAoYe92xG9t5dcLSj8DC6P93R9ir+YGw=","+MpZb3IeTXUQ1VXFXZgfVFzCvOX0MDxGanUQMFngYWw=","R62P9HI0DerDEF/J2NPUG+pRooZZjr87tb/6a02VAZs=","witl6lVGRqWEukI7TPfEwEterD0x0U/N+vLFrEgM3d0=","//9Wrtt/km05V/Uz1gBzHmw1++3e5qzCpZycY0id0hM=","5qY+JootET5QYakLD1VSLaWLYsBV5CNsny3SMmLgDck=","gnsD5AdzyM8v81+WoK8S73uW+oNypt2c3v4gw07kL6c=","wDpbIRy5hWWa+M5aB2/ZontGTidhd0Opkg9pOLRfJLc=","TDNBgrWNhMz/ZLmNeiUQjXJE6NabHjQEsHCekqo5DS8=","78RhvRg8R4WwHxVvzWwFGNaJwKZGwJSN7GYO7kn5vP8=","F5tQx3q6DM2bQD3L/G2KJ9rAjZ0VMzlAmuw1Urrltpg=","KbyZDOTON24dtLPkz0tlGLB0OsswC4djFUNHxHjMipI=","VPmQZaiIPP+jGVmSWjYHBv/V9OCTZbwibPFCdg+AHeQ=","GSGvSu0XsdO6cBgqJwZWqKj8191ujOCFNrbeVOnHXfU=","bDbLklSZE3cNPdbgkYjsldgt1tp4oyGy0Ey74MEDzVs=","kbqJ2Mh+r6zBJA5FYtUICq/XJ22GIh0GC+s04H3SZOs=","QsIovxEueg218cKeugrO+eGZOkloq5eRks0NcVdhADs=","QJUdtMG5fqTebI7bWTn0N8FLr9gz+BNgQr/7OsNQa1w=","hafZfJdzUX/yyCr3Anv7Qq1DWVaWxotaclNRcl/O9SQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3036392181\nteeCGBfEA5uGpfYT9SeWxOIWLaZFezcGvaKdKW6ZXHk=\n\n— rekor.sigstore.dev wNI9ajBFAiEApHSDrn9feEHa8pVsxt21mKS1GsE0GwW/yWXtEDWTr8UCIAt1IfpaBYiybTSBozWZJk38aPVtTQZe2Uk7H1KNyATZ\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJjODljNjJjNDZlMmNlNmY2MzNkMjM3NWIwMGY5ZDU2ZmUxZDM3MTY5MGI0YjEyZDIxODQ3ODliODU4ZTZhYjk2In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNVM1VGTmxzVkxDSVloNVFRanlicHhranp2YU1VM0J6MTk1T2hTZTk1ZTh3SWdVK0dVdDZGK1c4Wll4ZjBQUXFqU001bVhWeXlrSDQrZmZLZnVIS1F5cmRFPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1WRU5EUW5KTFowRjNTVUpCWjBsVlF6aEdSREF5VFhwRWNGQnFUSFZZVFdsS00xRlBZemhHUW1rMGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUVRWTlJGVXdUbnBSZVZkb1kwNU5hbGw0VFVSQk5VMUVWVEZPZWxGNVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVYyVlRkcFJsUktRV2htV2podGJYWldNM0JaZDJOMWVYazFiMjVWWkRoNU9GbHJTek1LYjBSS2NrbFpSa3RhV21WUU5URkdWREJyYm5SdlVEQnVORUZIVTFCU1ZqUlNiemwzY2xKU04ycFFiSGhUTHpCV2VVdFBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmhyVDBkV2FBcGFhazVxV1cxVmQwMVVRbTFPUkdNMVRsZFJkMWxxUW1wT1IwbDNUVlJCZWs5RVdYbE5lbWN6V1RKT2ExbDZhek5OUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUd0UFIxWm9XbXBPYWxsdFZYZE5WRUp0Q2s1RVl6Vk9WMUYzV1dwQ2FrNUhTWGROVkVGNlQwUlplVTE2WnpOWk1rNXJXWHByTTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUd0UFIxWm9DbHBxVG1wWmJWVjNUVlJDYlU1RVl6Vk9WMUYzV1dwQ2FrNUhTWGROVkVGNlQwUlplVTE2WnpOWk1rNXJXWHByTTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhUkdoc1dWZFplbGt5U214TlJFVjNXbXBSTTA5VVZtdE5SMGwzV1hwU2FVMUVSWGROZW1jeVRXcE5ORTR5VG1vS1drZE5OVTU2UVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlRucG5OVTFFUVROTmFsVXlUV2s1YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvU0hwUU5XbG5RVUZDUVUxQlVucENSa0ZwUlVGeE9YaFRMMU5SYzNOWlFraHpLM2t6ZDBJNWNHbFVTRlJDZEZBMVRFMXNaZ28xU2tsa05HNUpaa0V6YzBOSlNEVnROMjlhUjFVNFZqZFlWV3hQYm5ZclpIWk5Tak5MVVc4ME5uSXhlbEZDVUZGclNUbHFZM1UxWjAxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKclFVMUhXVU5OVVVOeWFsUkZaR3BGTTJZM1ZYRkVRMDlvTTNOV1ZHc3lUVlU1U0RkSlNqSlBhMXAzY21KVmQzVndVeTk0ZGtzS01HNXBObkExWjJ3eVFXUklNSEJNUW1acmEwTk5VVVJyWW14VU1qUTFhVmxYWTJ0U2VrNXFXRTlQU1ZWcE5HeE9ibkpPYUhVMVptUlNkR1E1V2pVNEx3cExOWE5SY0RJeFdqWXhkRlprUjJaclQxQkdPRlZHTkQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg4k983K1coZlkayD036+cJpsVAAkEZEsAMpxiUQTrwEACFDkWZsJr2b7zVR04BWQQhoVadpRtGA8yMDI2MTAwOTA1NDc0MlowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDA5MDU0NzQyWjAvBgkqhkiG9w0BCQQxIgQgLObLq7NGK87U9ONi2A7Ln6dRup/atfMcOdqLVbkL/G4wgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIxAIcZfruFVmPc3PiaAgWG2V/Lf7wzIFR8Zzf04nbAW4f9QWNBAschM6N1+f5egCBMbgIwO4DVCem+xvxWkNErxMg+VRttorrtEnVgXG9eqCuTc0xyNiLSbTkH2iLHCFVwFz7Y"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"yJxixG4s5vYz0jdbAPnVb+HTcWkLSxLSGEeJuFjmq5Y="},"signature":"MEUCIQCU3UFNlsVLCIYh5QQjybpxkjzvaMU3Bz195OhSe95e8wIgU+GUt6F+W8ZYxf0PQqjSM5mXVyykH4+ffKfuHKQyrdE="}}

build/torch-cuda/mma_sm100_desc.py CHANGED Viewed

@@ -83,9 +83,9 @@ def to_UMMA_format(cutlass_type) -> int:
     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}")

     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.Float8E4M3FN:
         return MXF8F6F4Format.E4M3
+    if cutlass_type is cutlass.Float8E5M2:
         return MXF8F6F4Format.E5M2
     raise TypeError(f"Unsupported CUTLASS scalar type for A/B: {cutlass_type!r}")

build/torch-cuda/named_barrier.py CHANGED Viewed

@@ -12,6 +12,19 @@ class NamedBarrierFwd(enum.IntEnum):
     PEmpty = enum.auto()
 class NamedBarrierBwd(enum.IntEnum):
     Epilogue = enum.auto()
     WarpSchedulerWG1 = enum.auto()
@@ -20,8 +33,10 @@ class NamedBarrierBwd(enum.IntEnum):
     PdS = enum.auto()
     dQFullWG0 = enum.auto()
     dQFullWG1 = enum.auto()
     dQEmptyWG0 = enum.auto()
     dQEmptyWG1 = enum.auto()
 class NamedBarrierBwdSm100(enum.IntEnum):
@@ -30,3 +45,19 @@ class NamedBarrierBwdSm100(enum.IntEnum):
     Compute = enum.auto()
     dQaccReduce = enum.auto()
     TmemPtr = enum.auto()

     PEmpty = enum.auto()
+class NamedBarrierFwdSm100(enum.IntEnum):
+    Epilogue = enum.auto()  # starts from 1 as barrier 0 is reserved for sync_threads()
+    TmemPtr = enum.auto()
+    SoftmaxStatsW0 = enum.auto()
+    SoftmaxStatsW1 = enum.auto()
+    SoftmaxStatsW2 = enum.auto()
+    SoftmaxStatsW3 = enum.auto()
+    SoftmaxStatsW4 = enum.auto()
+    SoftmaxStatsW5 = enum.auto()
+    SoftmaxStatsW6 = enum.auto()
+    SoftmaxStatsW7 = enum.auto()
 class NamedBarrierBwd(enum.IntEnum):
     Epilogue = enum.auto()
     WarpSchedulerWG1 = enum.auto()
     PdS = enum.auto()
     dQFullWG0 = enum.auto()
     dQFullWG1 = enum.auto()
+    dQFullWG2 = enum.auto()
     dQEmptyWG0 = enum.auto()
     dQEmptyWG1 = enum.auto()
+    dQEmptyWG2 = enum.auto()
 class NamedBarrierBwdSm100(enum.IntEnum):
     Compute = enum.auto()
     dQaccReduce = enum.auto()
     TmemPtr = enum.auto()
+class NamedBarrierFwdSm100_MLA2CTA(enum.IntEnum):
+    Epilogue = enum.auto()
+    TmemPtr = enum.auto()
+    Cpasync = enum.auto()
+    Softmax = enum.auto()
+    SoftmaxStatsFull = enum.auto()
+    SoftmaxStatsEmpty = enum.auto()
+class NamedBarrierBwdSm100_MLA2CTA(enum.IntEnum):
+    Epilogue = enum.auto()
+    TmemPtr = enum.auto()
+    Cpasync = enum.auto()
+    Softmax = enum.auto()

build/torch-cuda/pack_gqa.py CHANGED Viewed

@@ -1,25 +1,123 @@
 # Copyright (c) 2025, Tri Dao.
 import cutlass
 import cutlass.cute as cute
 from .quack import layout_utils
-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(
@@ -32,7 +130,7 @@ class PackGQA:
         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

 # Copyright (c) 2025, Tri Dao.
+from dataclasses import dataclass
+from typing import Union, Tuple
 import cutlass
 import cutlass.cute as cute
+from cutlass.cute.nvgpu import cpasync
 from .quack import layout_utils
+from . import utils as utils
+def pack_gqa_layout(T, qhead_per_kvhead, nheads_kv, head_idx):
+    """Reshape a tensor to fold qhead_per_kvhead into the seqlen dimension (mode 0).
+    The head dimension is at mode ``head_idx``.  Modes before it (1..head_idx-1)
+    are kept as-is (e.g. headdim for Q/O tensors), and modes after it are kept
+    as-is (e.g. batch).
+    For Q/O tensors (head_idx=2):
+        (seqlen_q, headdim, nheads, batch, ...) -> ((qhead_per_kvhead, seqlen_q), headdim, nheads_kv, batch, ...)
+    For LSE tensors (head_idx=1):
+        (seqlen_q, nheads, batch, ...) -> ((qhead_per_kvhead, seqlen_q), nheads_kv, batch, ...)
+    """
+    head_stride = T.stride[head_idx]
+    shape_packed = (
+        (qhead_per_kvhead, T.shape[0]),
+        *[T.shape[i] for i in range(1, head_idx)],
+        nheads_kv,
+        *[T.shape[i] for i in range(head_idx + 1, len(T.shape))],
+    )
+    stride_packed = (
+        (head_stride, T.stride[0]),
+        *[T.stride[i] for i in range(1, head_idx)],
+        head_stride * qhead_per_kvhead,
+        *[T.stride[i] for i in range(head_idx + 1, len(T.shape))],
+    )
+    return cute.make_tensor(T.iterator, cute.make_layout(shape_packed, stride=stride_packed))
+def make_packgqa_tiled_tma_atom(
+    op: cute.atom.CopyOp,
+    gmem_tensor: cute.Tensor,
+    smem_layout: Union[cute.Layout, cute.ComposedLayout],
+    cta_tiler: Tuple[int, int],
+    qhead_per_kvhead: int,
+    head_idx: int,
+):
+    # This packing and unpacking of the layout is so that we keep the same TMA dimension as usual.
+    # e.g. for (seqlen, d, nheads, b) layout, we still have 4D TMA after packing to
+    # ((nheads, seqlen), d, b).
+    # If we instead pack directly to ((qhead_per_kvhead, seqlen), d, nheads_kv, b) we'd have 5D TMA.
+    # Pack headdim and seqlen dim into 1: (seqlen, d, nheads, b) -> ((nheads, seqlen), d, b)
+    gmem_tensor = layout_utils.select(
+        gmem_tensor, [head_idx, *range(head_idx), *range(head_idx + 1, cute.rank(gmem_tensor))]
+    )
+    gmem_tensor = cute.group_modes(gmem_tensor, 0, 2)
+    assert cta_tiler[0] % qhead_per_kvhead == 0, (
+        "CTA tile size in the seqlen dimension must be divisible by qhead_per_kvhead"
+    )
+    tma_atom, tma_tensor = cpasync.make_tiled_tma_atom(
+        op,
+        gmem_tensor,
+        smem_layout,
+        ((qhead_per_kvhead, cta_tiler[0] // qhead_per_kvhead), cta_tiler[1]),  # No mcast
+    )
+    # Unpack from ((nheads, seqlen), d, b) -> ((qhead_per_kvhead, seqlen), d, nheads_kv, b)
+    T = tma_tensor
+    shape_packed = (
+        (qhead_per_kvhead, T.shape[0][1]),
+        *[T.shape[i] for i in range(1, head_idx)],
+        T.shape[0][0] // qhead_per_kvhead,
+        *[T.shape[i] for i in range(head_idx, len(T.shape))],
+    )
+    stride_packed = (
+        *[T.stride[i] for i in range(head_idx)],
+        T.stride[0][0] * qhead_per_kvhead,
+        *[T.stride[i] for i in range(head_idx, len(T.shape))],
+    )
+    tma_tensor = cute.make_tensor(T.iterator, cute.make_layout(shape_packed, stride=stride_packed))
+    return tma_atom, tma_tensor
+def unpack_gqa_layout(T, qhead_per_kvhead, head_idx):
+    """Reverse of pack_gqa_layout: unfold qhead_per_kvhead from the seqlen dimension (mode 0).
+    The head dimension is at mode ``head_idx``.  Modes before it (1..head_idx-1)
+    are kept as-is (e.g. headdim for Q/O tensors), and modes after it are kept
+    as-is (e.g. batch).
+    For Q/O tensors (head_idx=2):
+        ((qhead_per_kvhead, seqlen_q), headdim, nheads_kv, batch, ...) -> (seqlen_q, headdim, nheads, batch, ...)
+    For LSE tensors (head_idx=1):
+        ((qhead_per_kvhead, seqlen_q), nheads_kv, batch, ...) -> (seqlen_q, nheads, batch, ...)
+    """
+    seqlen_stride = T.stride[0][1]
+    head_stride = T.stride[0][0]
+    shape_unpacked = (
+        T.shape[0][1],
+        *[T.shape[i] for i in range(1, head_idx)],
+        T.shape[head_idx] * qhead_per_kvhead,
+        *[T.shape[i] for i in range(head_idx + 1, len(T.shape))],
+    )
+    stride_unpacked = (
+        seqlen_stride,
+        *[T.stride[i] for i in range(1, head_idx)],
+        head_stride,
+        *[T.stride[i] for i in range(head_idx + 1, len(T.shape))],
+    )
+    return cute.make_tensor(T.iterator, cute.make_layout(shape_unpacked, stride=stride_unpacked))
+@dataclass
 class PackGQA:
+    m_block_size: cutlass.Constexpr[int]
+    head_dim_padded: cutlass.Constexpr[int]
+    check_hdim_oob: cutlass.Constexpr[bool]
+    qhead_per_kvhead: cutlass.Constexpr[bool]
     @cute.jit
     def compute_ptr(
         num_threads: cutlass.Constexpr[int],
     ):
         num_ptr_per_thread = cute.ceil_div(cute.size(cRows), threads_per_row)
+        tPrPtr = cute.make_rmem_tensor(num_ptr_per_thread, cutlass.Int64)
         for i in cutlass.range_constexpr(num_ptr_per_thread):
             row = i * num_threads + cRows[tidx % threads_per_row][0]
             idx = block * self.m_block_size + row

build/torch-cuda/paged_kv.py CHANGED Viewed

@@ -28,6 +28,9 @@ class PagedKVManager(ParamsBase):
     head_dim_padded: cutlass.Constexpr[Int32]
     head_dim_v_padded: cutlass.Constexpr[Int32]
     gmem_threads_per_row: cutlass.Constexpr[Int32]
     page_entry_per_thread: Int32
     async_copy_elems: Int32
@@ -55,7 +58,11 @@ class PagedKVManager(ParamsBase):
         head_dim_v_padded: cutlass.Constexpr[Int32],
         num_threads: cutlass.Constexpr[Int32],
         dtype: Type[cutlass.Numeric],
     ):
         universal_copy_bits = 128
         async_copy_elems = universal_copy_bits // dtype.width
         dtype_bytes = dtype.width // 8
@@ -79,7 +86,8 @@ class PagedKVManager(ParamsBase):
         val_layout = cute.make_layout((1, async_copy_elems))
         gmem_tiled_copy_KV = cute.make_tiled_copy_tv(atom_async_copy, thr_layout, val_layout)
         gmem_thr_copy_KV = gmem_tiled_copy_KV.get_slice(thread_idx)
-        page_entry_per_thread = n_block_size // num_threads
         tPrPage = cute.make_rmem_tensor((page_entry_per_thread,), Int32)
         tPrPageOffset = cute.make_rmem_tensor((page_entry_per_thread,), Int32)
@@ -97,7 +105,9 @@ class PagedKVManager(ParamsBase):
         else:
             cV = cute.make_identity_tensor((n_block_size, head_dim_v_padded))
             tVcV = gmem_thr_copy_KV.partition_S(cV)
-            tVpV = utils.predicate_k(tVcV, limit=mV_paged.shape[0])
         return PagedKVManager(
             mPageTable,
@@ -111,6 +121,8 @@ class PagedKVManager(ParamsBase):
             num_threads,
             head_dim_padded,
             head_dim_v_padded,
             gmem_threads_per_row,
             page_entry_per_thread,
             async_copy_elems,
@@ -144,24 +156,24 @@ class PagedKVManager(ParamsBase):
             self.tPrPageOffset[i] = page_offset
     @cute.jit
-    def compute_X_ptr(self, K_or_V: str):
         tPrXPtr = cute.make_rmem_tensor((self.page_entry_per_thread,), cutlass.Int64)
         for i in cutlass.range(self.page_entry_per_thread, unroll=1):
             page = self.tPrPage[i]
             page_offset = self.tPrPageOffset[i]
-            if const_expr(K_or_V == "K"):
-                tPrXPtr[i] = utils.elem_pointer(self.mK_paged, (page_offset, 0, page)).toint()
             else:
-                tPrXPtr[i] = utils.elem_pointer(self.mV_paged, (0, page_offset, page)).toint()
         return tPrXPtr
     @cute.jit
-    def load_KV(self, n_block: Int32, sX: cute.Tensor, K_or_V: str):
-        assert K_or_V in ("K", "V")
-        tPrXPtr = self.compute_X_ptr(K_or_V)
-        # Finesse sX layout to be (M, N).
         sX_pi = cute.make_tensor(
             sX.iterator,
             cute.make_layout(
@@ -169,10 +181,45 @@ class PagedKVManager(ParamsBase):
                 stride=(sX.stride[0][0], (sX.stride[0][1], sX.stride[2])),
             ),
         )
         if const_expr(K_or_V == "V"):
-            # Need to transpose V
             sX_pi = cute.make_tensor(sX_pi.iterator, cute.select(sX_pi.layout, mode=[1, 0]))
         head_dim = self.head_dim_v_padded if const_expr(K_or_V == "V") else self.head_dim_padded
         cX = cute.make_identity_tensor((self.n_block_size, head_dim))
@@ -198,17 +245,4 @@ class PagedKVManager(ParamsBase):
             )
             mX_paged_cur = cute.make_tensor(x_gmem_ptr, cute.make_layout((head_dim,)))
             mX_paged_cur_copy = cute.tiled_divide(mX_paged_cur, (self.async_copy_elems,))
-            for k in cutlass.range_constexpr(cute.size(tXsX, mode=[2])):
-                ki = tXcX[0, 0, k][1] // self.async_copy_elems
-                mX_paged_cur_copy_ki = mX_paged_cur_copy[None, ki]
-                tXsX_k = tXsX[None, m, k]
-                mX_paged_cur_copy_ki = cute.make_tensor(
-                    mX_paged_cur_copy_ki.iterator, tXsX_k.layout
-                )
-                cute.copy(
-                    self.gmem_tiled_copy_KV,
-                    mX_paged_cur_copy_ki,
-                    tXsX_k,
-                    pred=should_load,
-                )

     head_dim_padded: cutlass.Constexpr[Int32]
     head_dim_v_padded: cutlass.Constexpr[Int32]
+    arch: cutlass.Constexpr[Int32]
+    v_gmem_transposed: cutlass.Constexpr[bool]
     gmem_threads_per_row: cutlass.Constexpr[Int32]
     page_entry_per_thread: Int32
     async_copy_elems: Int32
         head_dim_v_padded: cutlass.Constexpr[Int32],
         num_threads: cutlass.Constexpr[Int32],
         dtype: Type[cutlass.Numeric],
+        arch: cutlass.Constexpr[int] = 100,
     ):
+        # SM100 transposes V in gmem to (dv, page_size, num_pages);
+        # SM90 keeps V as (page_size, dv, num_pages), same layout as K.
+        v_gmem_transposed = arch != 90
         universal_copy_bits = 128
         async_copy_elems = universal_copy_bits // dtype.width
         dtype_bytes = dtype.width // 8
         val_layout = cute.make_layout((1, async_copy_elems))
         gmem_tiled_copy_KV = cute.make_tiled_copy_tv(atom_async_copy, thr_layout, val_layout)
         gmem_thr_copy_KV = gmem_tiled_copy_KV.get_slice(thread_idx)
+        # Include the final partially populated wave of rows.
+        page_entry_per_thread = (n_block_size + num_threads - 1) // num_threads
         tPrPage = cute.make_rmem_tensor((page_entry_per_thread,), Int32)
         tPrPageOffset = cute.make_rmem_tensor((page_entry_per_thread,), Int32)
         else:
             cV = cute.make_identity_tensor((n_block_size, head_dim_v_padded))
             tVcV = gmem_thr_copy_KV.partition_S(cV)
+            # When V is transposed in gmem, dv is shape[0]; otherwise dv is shape[1] (same as K)
+            V_limit = cute.size(mV_paged.shape[0 if v_gmem_transposed else 1])
+            tVpV = utils.predicate_k(tVcV, limit=V_limit)
         return PagedKVManager(
             mPageTable,
             num_threads,
             head_dim_padded,
             head_dim_v_padded,
+            arch,
+            v_gmem_transposed,
             gmem_threads_per_row,
             page_entry_per_thread,
             async_copy_elems,
             self.tPrPageOffset[i] = page_offset
     @cute.jit
+    def compute_X_ptr(self, K_or_V: str, d_offset: int = 0):
         tPrXPtr = cute.make_rmem_tensor((self.page_entry_per_thread,), cutlass.Int64)
+        mX = self.mK_paged if const_expr(K_or_V == "K") else self.mV_paged
+        # K is always (page_size, d, num_pages). V matches K when not transposed,
+        # but is (dv, page_size, num_pages) when transposed (SM100).
+        transposed = const_expr(K_or_V == "V" and self.v_gmem_transposed)
         for i in cutlass.range(self.page_entry_per_thread, unroll=1):
             page = self.tPrPage[i]
             page_offset = self.tPrPageOffset[i]
+            if const_expr(transposed):
+                tPrXPtr[i] = utils.elem_pointer(mX, (d_offset, page_offset, page)).toint()
             else:
+                tPrXPtr[i] = utils.elem_pointer(mX, (page_offset, d_offset, page)).toint()
         return tPrXPtr
     @cute.jit
+    def _flatten_smem_sm100(self, sX: cute.Tensor, K_or_V: str):
+        """Flatten SM100 smem ((a,b), cta_split, k) to (a,(b,k)); transpose V to (d,page_size)."""
         sX_pi = cute.make_tensor(
             sX.iterator,
             cute.make_layout(
                 stride=(sX.stride[0][0], (sX.stride[0][1], sX.stride[2])),
             ),
         )
         if const_expr(K_or_V == "V"):
             sX_pi = cute.make_tensor(sX_pi.iterator, cute.select(sX_pi.layout, mode=[1, 0]))
+        return sX_pi
+    @cute.jit
+    def _copy_row_async(
+        self,
+        tXsX: cute.Tensor,
+        tXcX: cute.Tensor,
+        mX_paged_cur_copy: cute.Tensor,
+        m: Int32,
+        should_load: cute.Tensor,
+    ):
+        """Issue cp.async copies for one row across all k-tiles."""
+        for k in cutlass.range_constexpr(cute.size(tXsX, mode=[2])):
+            ki = tXcX[0, 0, k][1] // self.async_copy_elems
+            mX_paged_cur_copy_ki = mX_paged_cur_copy[None, ki]
+            tXsX_k = tXsX[None, m, k]
+            mX_paged_cur_copy_ki = cute.make_tensor(mX_paged_cur_copy_ki.iterator, tXsX_k.layout)
+            cute.copy(
+                self.gmem_tiled_copy_KV,
+                mX_paged_cur_copy_ki,
+                tXsX_k,
+                pred=should_load,
+            )
+    @cute.jit
+    def load_KV(self, n_block: Int32, sX: cute.Tensor, K_or_V: str):
+        assert K_or_V in ("K", "V")
+        tPrXPtr = self.compute_X_ptr(K_or_V)
+        if const_expr(self.arch == 90):
+            # SM90: sX is already stage-sliced by caller (sK[None, None, stage]).
+            # Flatten hierarchical modes to get (n_block_size, head_dim).
+            sX_pi = cute.group_modes(sX, 0, 1)
+            # SM90 does NOT transpose V here (it's transposed via utils.transpose_view before MMA)
+        else:
+            sX_pi = self._flatten_smem_sm100(sX, K_or_V)
         head_dim = self.head_dim_v_padded if const_expr(K_or_V == "V") else self.head_dim_padded
         cX = cute.make_identity_tensor((self.n_block_size, head_dim))
             )
             mX_paged_cur = cute.make_tensor(x_gmem_ptr, cute.make_layout((head_dim,)))
             mX_paged_cur_copy = cute.tiled_divide(mX_paged_cur, (self.async_copy_elems,))
+            self._copy_row_async(tXsX, tXcX, mX_paged_cur_copy, m, should_load)

build/torch-cuda/pipeline.py CHANGED Viewed

@@ -1,6 +1,5 @@
 # Copyright (c) 2025, Tri Dao.
-# import math
 from typing import Optional
 from dataclasses import dataclass
@@ -11,12 +10,31 @@ from cutlass.pipeline import PipelineState
 from cutlass.pipeline import PipelineUserType
 from cutlass.pipeline import NamedBarrier as NamedBarrierOg
 from cutlass.pipeline import PipelineAsync as PipelineAsyncOg
 from cutlass.pipeline import PipelineTmaAsync as PipelineTmaAsyncOg
 from cutlass.pipeline import PipelineTmaUmma as PipelineTmaUmmaOg
 from cutlass.pipeline import PipelineUmmaAsync as PipelineUmmaAsyncOg
 from cutlass.pipeline import PipelineAsyncUmma as PipelineAsyncUmmaOg
 class PipelineStateSimple:
     """
     Pipeline state contains an index and phase bit corresponding to the current position in the circular buffer.
@@ -25,9 +43,6 @@ class PipelineStateSimple:
     """
     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
@@ -36,13 +51,10 @@ class PipelineStateSimple:
     @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:
@@ -50,11 +62,8 @@ class PipelineStateSimple:
     @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:
@@ -66,21 +75,6 @@ class PipelineStateSimple:
         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()]
@@ -94,7 +88,6 @@ 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))
@@ -102,14 +95,73 @@ def make_pipeline_state(type: PipelineUserType, stages: int):
         assert False, "Error: invalid PipelineUserType specified for make_pipeline_state."
 @dataclass(frozen=True)
 class NamedBarrier(NamedBarrierOg):
-    @staticmethod
-    def create(*args, **kwargs):
-        obj = NamedBarrierOg.create(*args, **kwargs)
-        # Can't assign to __class__ directly since the dataclass is frozen
-        object.__setattr__(obj, "__class__", NamedBarrier)
-        return obj
     @dsl_user_op
     def arrive_w_index(self, index: Int32, *, loc=None, ip=None) -> None:
@@ -134,72 +186,121 @@ class NamedBarrier(NamedBarrierOg):
         )
 @dataclass(frozen=True)
-class PipelineAsync(PipelineAsyncOg):
     @staticmethod
-    def create(*args, **kwargs):
         obj = PipelineAsyncOg.create(*args, **kwargs)
-        # Can't assign to __class__ directly since the dataclass is frozen
-        # obj.__class__ = PipelineAsync
         object.__setattr__(obj, "__class__", PipelineAsync)
         return obj
     @dsl_user_op
-    def producer_acquire_w_index_phase(
-        self,
-        index: Int32,
-        phase: Int32,
-        try_acquire_token: Optional[Boolean] = None,
-        *,
-        loc=None,
-        ip=None,
-    ):
-        if_generate(
-            try_acquire_token is None or try_acquire_token == 0,
-            lambda: self.sync_object_empty.wait(index, phase, loc=loc, ip=ip),
-            loc=loc,
-            ip=ip,
         )
     @dsl_user_op
-    def producer_commit_w_index(self, index: Int32, *, loc=None, ip=None):
-        self.sync_object_full.arrive(index, self.producer_mask, loc=loc, ip=ip)
-    @dsl_user_op
-    def consumer_wait_w_index_phase(
-        self,
-        index: Int32,
-        phase: Int32,
-        try_wait_token: Optional[Boolean] = None,
-        *,
-        loc=None,
-        ip=None,
-    ):
-        if_generate(
-            try_wait_token is None or try_wait_token == 0,
-            lambda: self.sync_object_full.wait(index, phase, loc=loc, ip=ip),
-            loc=loc,
-            ip=ip,
         )
-    @dsl_user_op
-    def consumer_release_w_index(self, index: Int32, *, loc=None, ip=None):
-        self.sync_object_empty.arrive(index, self.consumer_mask, loc=loc, ip=ip)
 @dataclass(frozen=True)
-class PipelineTmaAsync(PipelineTmaAsyncOg):
-    """
-    Override producer_acquire to take in extra_tx_count parameter.
-    """
     @staticmethod
-    def create(*args, **kwargs):
-        obj = PipelineTmaAsyncOg.create(*args, **kwargs)
-        # Can't assign to __class__ directly since the dataclass is frozen
-        object.__setattr__(obj, "__class__", PipelineTmaAsync)
         return obj
     @dsl_user_op
     def producer_acquire(
         self,
@@ -226,19 +327,15 @@ class PipelineTmaAsync(PipelineTmaAsyncOg):
             self.sync_object_full.arrive_and_expect_tx(state.index, tx_count, loc=loc, ip=ip)
-@dataclass(frozen=True)
-class PipelineTmaUmma(PipelineTmaUmmaOg):
-    """
-    Override producer_acquire to take in extra_tx_count parameter.
-    """
-    @staticmethod
-    def create(*args, **kwargs):
-        obj = PipelineTmaUmmaOg.create(*args, **kwargs)
-        # Can't assign to __class__ directly since the dataclass is frozen
-        # obj.__class__ = PipelineTmaUmma
-        object.__setattr__(obj, "__class__", PipelineTmaUmma)
-        return obj
     @dsl_user_op
     def producer_acquire(
@@ -279,162 +376,27 @@ class PipelineTmaUmma(PipelineTmaUmmaOg):
                 ip=ip,
             )
-    @dsl_user_op
-    def producer_acquire_w_index_phase(
-        self,
-        index: Int32,
-        phase: Int32,
-        try_acquire_token: Optional[Boolean] = None,
-        *,
-        loc=None,
-        ip=None,
-    ):
-        """
-        TMA producer commit conditionally waits on buffer empty and sets the transaction barrier for leader threadblocks.
-        """
-        if_generate(
-            try_acquire_token is None or try_acquire_token == 0,
-            lambda: self.sync_object_empty.wait(index, phase, loc=loc, ip=ip),
-            loc=loc,
-            ip=ip,
-        )
-        if_generate(
-            self.is_leader_cta,
-            lambda: self.sync_object_full.arrive(index, self.producer_mask, loc=loc, ip=ip),
-            loc=loc,
-            ip=ip,
-        )
-    @dsl_user_op
-    def consumer_wait_w_index_phase(
-        self,
-        index: Int32,
-        phase: Int32,
-        try_wait_token: Optional[Boolean] = None,
-        *,
-        loc=None,
-        ip=None,
-    ):
-        if_generate(
-            try_wait_token is None or try_wait_token == 0,
-            lambda: self.sync_object_full.wait(index, phase, loc=loc, ip=ip),
-            loc=loc,
-            ip=ip,
-        )
-    @dsl_user_op
-    def consumer_release_w_index(self, index: Int32, *, loc=None, ip=None):
-        """
-        UMMA consumer release buffer empty, cta_group needs to be provided.
-        """
-        self.sync_object_empty.arrive(index, self.consumer_mask, self.cta_group, loc=loc, ip=ip)
 @dataclass(frozen=True)
-class PipelineUmmaAsync(PipelineUmmaAsyncOg):
-    @staticmethod
-    def create(*args, **kwargs):
-        obj = PipelineUmmaAsyncOg.create(*args, **kwargs)
-        # Can't assign to __class__ directly since the dataclass is frozen
-        object.__setattr__(obj, "__class__", PipelineUmmaAsync)
-        return obj
-    @dsl_user_op
-    def producer_acquire_w_index_phase(
-        self,
-        index: Int32,
-        phase: Int32,
-        try_acquire_token: Optional[Boolean] = None,
-        *,
-        loc=None,
-        ip=None,
-    ):
-        if_generate(
-            try_acquire_token is None or try_acquire_token == 0,
-            lambda: self.sync_object_empty.wait(index, phase, loc=loc, ip=ip),
-            loc=loc,
-            ip=ip,
-        )
-    @dsl_user_op
-    def producer_commit_w_index(self, index: Int32, *, loc=None, ip=None):
-        """
-        UMMA producer commit buffer full, cta_group needs to be provided.
-        """
-        self.sync_object_full.arrive(index, self.producer_mask, self.cta_group, loc=loc, ip=ip)
-    @dsl_user_op
-    def consumer_wait_w_index_phase(
-        self,
-        index: Int32,
-        phase: Int32,
-        try_wait_token: Optional[Boolean] = None,
-        *,
-        loc=None,
-        ip=None,
-    ):
-        if_generate(
-            try_wait_token is None or try_wait_token == 0,
-            lambda: self.sync_object_full.wait(index, phase, loc=loc, ip=ip),
-            loc=loc,
-            ip=ip,
-        )
-    @dsl_user_op
-    def consumer_release_w_index(self, index: Int32, *, loc=None, ip=None):
-        self.sync_object_empty.arrive(index, self.consumer_mask, loc=loc, ip=ip)
 @dataclass(frozen=True)
-class PipelineAsyncUmma(PipelineAsyncUmmaOg):
-    @staticmethod
-    def create(*args, **kwargs):
-        obj = PipelineAsyncUmmaOg.create(*args, **kwargs)
-        # Can't assign to __class__ directly since the dataclass is frozen
-        object.__setattr__(obj, "__class__", PipelineAsyncUmma)
-        return obj
-    @dsl_user_op
-    def producer_acquire_w_index_phase(
-        self,
-        index: Int32,
-        phase: Int32,
-        try_acquire_token: Optional[Boolean] = None,
-        *,
-        loc=None,
-        ip=None,
-    ):
-        if_generate(
-            try_acquire_token is None or try_acquire_token == 0,
-            lambda: self.sync_object_empty.wait(index, phase, loc=loc, ip=ip),
-            loc=loc,
-            ip=ip,
-        )
-    @dsl_user_op
-    def producer_commit_w_index(self, index: Int32, *, loc=None, ip=None):
-        self.sync_object_full.arrive(index, self.producer_mask, loc=loc, ip=ip)
-    @dsl_user_op
-    def consumer_wait_w_index_phase(
-        self,
-        index: Int32,
-        phase: Int32,
-        try_wait_token: Optional[Boolean] = None,
-        *,
-        loc=None,
-        ip=None,
-    ):
-        if_generate(
-            try_wait_token is None or try_wait_token == 0,
-            lambda: self.sync_object_full.wait(index, phase, loc=loc, ip=ip),
-            loc=loc,
-            ip=ip,
-        )
-    @dsl_user_op
-    def consumer_release_w_index(self, index: Int32, *, loc=None, ip=None):
-        """
-        UMMA consumer release buffer empty, cta_group needs to be provided.
-        """
-        self.sync_object_empty.arrive(index, self.consumer_mask, self.cta_group, loc=loc, ip=ip)

 # Copyright (c) 2025, Tri Dao.
 from typing import Optional
 from dataclasses import dataclass
 from cutlass.pipeline import PipelineUserType
 from cutlass.pipeline import NamedBarrier as NamedBarrierOg
 from cutlass.pipeline import PipelineAsync as PipelineAsyncOg
+from cutlass.pipeline import PipelineCpAsync as PipelineCpAsyncOg
 from cutlass.pipeline import PipelineTmaAsync as PipelineTmaAsyncOg
 from cutlass.pipeline import PipelineTmaUmma as PipelineTmaUmmaOg
 from cutlass.pipeline import PipelineUmmaAsync as PipelineUmmaAsyncOg
 from cutlass.pipeline import PipelineAsyncUmma as PipelineAsyncUmmaOg
+def _override_create(parent_cls, child_cls):
+    """Create a static factory that constructs parent_cls then re-classes to child_cls."""
+    @staticmethod
+    def create(*args, **kwargs):
+        obj = parent_cls.create(*args, **kwargs)
+        # Can't assign to __class__ directly since the dataclass is frozen
+        object.__setattr__(obj, "__class__", child_cls)
+        return obj
+    return create
+def _make_state(index: Int32, phase: Int32) -> PipelineState:
+    """Construct a PipelineState from index and phase (count/stages unused by callers)."""
+    return PipelineState(stages=0, count=Int32(0), index=index, phase=phase)
 class PipelineStateSimple:
     """
     Pipeline state contains an index and phase bit corresponding to the current position in the circular buffer.
     """
     def __init__(self, stages: int, phase_index: Int32):
         self._stages = stages
         self._phase_index = phase_index
     @property
     def stages(self) -> int:
         return self._stages
     @property
     def index(self) -> Int32:
         if const_expr(self._stages == 1):
             return Int32(0)
         else:
     @property
     def phase(self) -> Int32:
         # PTX docs say that the phase parity needs to be 0 or 1, so by right we need to
         # take modulo 2. But in practice just passing the phase in without modulo works fine.
         if const_expr(self._stages == 1):
             return self._phase_index
         else:
         else:
             self._phase_index += 1
     def __extract_mlir_values__(self):
         phase_index = self._phase_index
         return [phase_index.ir_value()]
     Creates a pipeline state. Producers are assumed to start with an empty buffer and have a flipped phase bit of 1.
     """
     if type is PipelineUserType.Producer:
         return PipelineStateSimple(stages, Int32(stages))
     elif type is PipelineUserType.Consumer:
         return PipelineStateSimple(stages, Int32(0))
         assert False, "Error: invalid PipelineUserType specified for make_pipeline_state."
+# ── Shared helpers ───────────────────────────────────────────────────────────
+def _call_with_elect_one(parent_method, self, state, elect_one, syncwarp, loc, ip):
+    """Optionally wrap a parent pipeline method call in sync_warp + elect_one."""
+    if const_expr(elect_one):
+        if const_expr(syncwarp):
+            cute.arch.sync_warp()
+        with cute.arch.elect_one():
+            parent_method(self, state, loc=loc, ip=ip)
+    else:
+        parent_method(self, state, loc=loc, ip=ip)
+# ── Mixin: _w_index / _w_index_phase variants that delegate to parent ───────
+# Each parent class has PipelineState-based methods (producer_acquire, producer_commit,
+# consumer_wait, consumer_release). The _w_index_phase variants just construct a
+# PipelineState from (index, phase) and delegate.
+class _PipelineIndexPhaseMixin:
+    """Mixin providing _w_index_phase / _w_index methods that delegate to PipelineState-based parents."""
+    @dsl_user_op
+    def producer_acquire_w_index_phase(
+        self,
+        index: Int32,
+        phase: Int32,
+        try_acquire_token: Optional[Boolean] = None,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        state = _make_state(index, phase)
+        # Call the parent's producer_acquire (which takes PipelineState)
+        self.producer_acquire(state, try_acquire_token, loc=loc, ip=ip)
+    @dsl_user_op
+    def producer_commit_w_index(self, index: Int32, *, loc=None, ip=None):
+        state = _make_state(index, Int32(0))
+        self.producer_commit(state, loc=loc, ip=ip)
+    @dsl_user_op
+    def consumer_wait_w_index_phase(
+        self,
+        index: Int32,
+        phase: Int32,
+        try_wait_token: Optional[Boolean] = None,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        state = _make_state(index, phase)
+        self.consumer_wait(state, try_wait_token, loc=loc, ip=ip)
+    @dsl_user_op
+    def consumer_release_w_index(self, index: Int32, *, loc=None, ip=None):
+        state = _make_state(index, Int32(0))
+        self.consumer_release(state, loc=loc, ip=ip)
+# ── NamedBarrier ─────────────────────────────────────────────────────────────
 @dataclass(frozen=True)
 class NamedBarrier(NamedBarrierOg):
+    create = _override_create(NamedBarrierOg, None)  # patched below
     @dsl_user_op
     def arrive_w_index(self, index: Int32, *, loc=None, ip=None) -> None:
         )
+NamedBarrier.create = _override_create(NamedBarrierOg, NamedBarrier)
+# ── PipelineAsync ────────────────────────────────────────────────────────────
 @dataclass(frozen=True)
+class PipelineAsync(_PipelineIndexPhaseMixin, PipelineAsyncOg):
+    """
+    PipelineAsync with optional elect_one for producer_commit and consumer_release.
+    When elect_one_*=True (set at create time), only one elected thread per warp
+    signals the barrier arrive. This is useful when the mask count is set to 1 per warp.
+    Args (to create):
+        elect_one_commit: If True, only elected thread signals producer_commit.
+        syncwarp_before_commit: If True (default), issue syncwarp before elect_one.
+        elect_one_release: If True, only elected thread signals consumer_release.
+        syncwarp_before_release: If True (default), issue syncwarp before elect_one.
+            Set syncwarp to False when threads are already converged (e.g. after wgmma wait_group).
+    """
+    _elect_one_commit: bool = False
+    _syncwarp_before_commit: bool = True
+    _elect_one_release: bool = False
+    _syncwarp_before_release: bool = True
     @staticmethod
+    def create(
+        *args,
+        elect_one_commit: bool = False,
+        syncwarp_before_commit: bool = True,
+        elect_one_release: bool = False,
+        syncwarp_before_release: bool = True,
+        **kwargs,
+    ):
         obj = PipelineAsyncOg.create(*args, **kwargs)
         object.__setattr__(obj, "__class__", PipelineAsync)
+        object.__setattr__(obj, "_elect_one_commit", elect_one_commit)
+        object.__setattr__(obj, "_syncwarp_before_commit", syncwarp_before_commit)
+        object.__setattr__(obj, "_elect_one_release", elect_one_release)
+        object.__setattr__(obj, "_syncwarp_before_release", syncwarp_before_release)
         return obj
     @dsl_user_op
+    def producer_commit(self, state: PipelineState, *, loc=None, ip=None):
+        _call_with_elect_one(
+            PipelineAsyncOg.producer_commit,
+            self,
+            state,
+            self._elect_one_commit,
+            self._syncwarp_before_commit,
+            loc,
+            ip,
         )
     @dsl_user_op
+    def consumer_release(self, state: PipelineState, *, loc=None, ip=None):
+        _call_with_elect_one(
+            PipelineAsyncOg.consumer_release,
+            self,
+            state,
+            self._elect_one_release,
+            self._syncwarp_before_release,
+            loc,
+            ip,
         )
+    # _w_index variants inherited from _PipelineIndexPhaseMixin, which delegate
+    # to producer_commit / consumer_release above.
+# ── PipelineCpAsync ──────────────────────────────────────────────────────────
 @dataclass(frozen=True)
+class PipelineCpAsync(_PipelineIndexPhaseMixin, PipelineCpAsyncOg):
+    _elect_one_release: bool = False
+    _syncwarp_before_release: bool = True
     @staticmethod
+    def create(
+        *args,
+        elect_one_release: bool = False,
+        syncwarp_before_release: bool = True,
+        **kwargs,
+    ):
+        obj = PipelineCpAsyncOg.create(*args, **kwargs)
+        object.__setattr__(obj, "__class__", PipelineCpAsync)
+        object.__setattr__(obj, "_elect_one_release", elect_one_release)
+        object.__setattr__(obj, "_syncwarp_before_release", syncwarp_before_release)
         return obj
+    @dsl_user_op
+    def consumer_release(self, state: PipelineState, *, loc=None, ip=None):
+        _call_with_elect_one(
+            PipelineCpAsyncOg.consumer_release,
+            self,
+            state,
+            self._elect_one_release,
+            self._syncwarp_before_release,
+            loc,
+            ip,
+        )
+    # _w_index variants inherited from _PipelineIndexPhaseMixin.
+# ── PipelineTmaAsync ────────────────────────────────────────────────────────
+@dataclass(frozen=True)
+class PipelineTmaAsync(_PipelineIndexPhaseMixin, PipelineTmaAsyncOg):
+    """Override producer_acquire to take in extra_tx_count parameter."""
     @dsl_user_op
     def producer_acquire(
         self,
             self.sync_object_full.arrive_and_expect_tx(state.index, tx_count, loc=loc, ip=ip)
+PipelineTmaAsync.create = _override_create(PipelineTmaAsyncOg, PipelineTmaAsync)
+# ── PipelineTmaUmma ─────────────────────────────────────────────────────────
+@dataclass(frozen=True)
+class PipelineTmaUmma(_PipelineIndexPhaseMixin, PipelineTmaUmmaOg):
+    """Override producer_acquire to take in extra_tx_count parameter."""
     @dsl_user_op
     def producer_acquire(
                 ip=ip,
             )
+PipelineTmaUmma.create = _override_create(PipelineTmaUmmaOg, PipelineTmaUmma)
+# ── PipelineUmmaAsync ───────────────────────────────────────────────────────
 @dataclass(frozen=True)
+class PipelineUmmaAsync(_PipelineIndexPhaseMixin, PipelineUmmaAsyncOg):
+    pass
+PipelineUmmaAsync.create = _override_create(PipelineUmmaAsyncOg, PipelineUmmaAsync)
+# ── PipelineAsyncUmma ───────────────────────────────────────────────────────
 @dataclass(frozen=True)
+class PipelineAsyncUmma(_PipelineIndexPhaseMixin, PipelineAsyncUmmaOg):
+    pass
+PipelineAsyncUmma.create = _override_create(PipelineAsyncUmmaOg, PipelineAsyncUmma)

build/torch-cuda/prepare_scheduler.py ADDED Viewed

	@@ -0,0 +1,387 @@

+# A reimplementation of https://github.com/Dao-AILab/flash-attention/blob/main/hopper/flash_prepare_scheduler.cu
+# from CUTLASS C++ to Cute-DSL.
+from typing import Tuple, Optional, NamedTuple
+import operator
+import torch
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32, const_expr, Float32
+from cutlass.cute import FastDivmodDivisor
+from . import utils as utils
+class SchedulerMetadataTensorsTorch(NamedTuple):
+    """Class to store scheduler metadata for varlen"""
+    # tensors of shape (batch)
+    num_m_blocks_ptr: Optional[torch.Tensor]
+    num_splits_dynamic_ptr: Optional[torch.Tensor]
+    virtual_batch_idx_ptr: Optional[torch.Tensor]
+    num_nheads_in_l2_ptr: Optional[torch.Tensor]
+    # tensor of shape (1)
+    tile_count_semaphore: Optional[torch.Tensor]
+    # tensors of shape (batch + 1)
+    # cu_total_m_blocks[b+1] = sum_{i<=b} num_m_blocks[i]
+    # cu_total_splits_m_blocks[b+1] = sum_{i<=b} num_m_blocks[i] * num_splits_dynamic[i]
+    cu_total_m_blocks: Optional[torch.Tensor] = None
+    cu_total_splits_m_blocks: Optional[torch.Tensor] = None
+    blocks_to_batch_idx: Optional[torch.Tensor] = None
+class FlashPrepareScheduler:
+    def __init__(
+        self,
+        num_warps: int,
+        tile_m: int,
+        tile_n: int,
+        nheads: int,
+        nheads_kv: int,
+        headdim: int,
+        headdim_v: Optional[int] = None,
+        is_causal: bool = False,
+        packgqa: bool = False,
+        sort: bool = False,
+        zfill_padded_output: bool = False,
+    ):
+        self.num_warps = num_warps
+        self.is_causal = is_causal
+        self.packgqa = packgqa
+        assert not sort, "LPT batch sort not yet implemented"
+        self.sort = sort
+        self.num_threads_per_warp = 32
+        self.tile_m = tile_m
+        self.tile_n = tile_n
+        self.d = headdim
+        self.dv = headdim_v if headdim_v is not None else headdim
+        self.k_num_batch_per_warp = 31
+        self.k_smem_size = 1
+        self.zfill_padded_output = zfill_padded_output
+        # for pack gqa, query heads per kv head is combined with seqlen_q
+        self.nheads_computed = nheads if not self.packgqa else nheads_kv
+        # L2 cache calculations
+        self.qhead_per_khead = nheads // nheads_kv
+        self.size_l2_divisor = (
+            1
+            if self.qhead_per_khead == 1
+            else (
+                2
+                if self.qhead_per_khead <= 2
+                else (4 if self.qhead_per_khead <= 4 else (8 if self.qhead_per_khead <= 8 else 16))
+            )
+        )
+        self.size_l2 = (32 * 1024 * 1024) // self.size_l2_divisor
+        element_size = 2
+        self.size_one_kvblock = self.tile_n * (self.d + self.dv) * element_size
+        self.max_kvblocks_in_l2 = (
+            self.size_l2 + self.size_one_kvblock - 1
+        ) // self.size_one_kvblock
+    @staticmethod
+    def get_grid_shape(num_batch: int) -> Tuple[int, int, int]:
+        num_ctas = (num_batch + (31 * 32 - 1)) // (31 * 32)
+        return (num_ctas, 1, 1)
+    @cute.jit
+    def __call__(
+        self,
+        seqlen_q_static: int,
+        seqlen_k_static: int,
+        seqlen_k_new_static: int,
+        mCuSeqlensQ: Optional[cute.Tensor],
+        mCuSeqlensK: Optional[cute.Tensor],
+        mCuSeqlensKNew: Optional[cute.Tensor],
+        mSeqUsedQ: Optional[cute.Tensor],
+        mSeqUsedK: Optional[cute.Tensor],
+        mLeftPadK: Optional[cute.Tensor],
+        num_batch: int,
+        num_splits_static: int,
+        tile_count_semaphore: Optional[cute.Tensor],
+        num_m_blocks_ptr: Optional[cute.Tensor],
+        num_splits_dynamic_ptr: Optional[cute.Tensor],
+        virtual_batch_idx_ptr: Optional[cute.Tensor],
+        num_nheads_in_l2_ptr: Optional[cute.Tensor],
+        n_blocks_per_split: Optional[int],  # overrides heuristic
+        stream: cuda.CUstream,
+    ):
+        tile_m_divmod = FastDivmodDivisor(self.tile_m)
+        tile_n_divmod = FastDivmodDivisor(self.tile_n)
+        @cute.struct
+        class SharedStorage:
+            total_blocks_smem: cute.struct.MemRange[Int32, self.k_smem_size]
+        self.shared_storage = SharedStorage
+        block = (32 * self.num_warps, 1, 1)
+        grid = self.get_grid_shape(num_batch)
+        hardware_info = cutlass.utils.HardwareInfo()
+        num_sm = hardware_info.get_device_multiprocessor_count()
+        self.kernel(
+            seqlen_q_static,
+            seqlen_k_static,
+            seqlen_k_new_static,
+            mCuSeqlensQ,
+            mCuSeqlensK,
+            mCuSeqlensKNew,
+            mSeqUsedQ,
+            mSeqUsedK,
+            mLeftPadK,
+            num_batch,
+            num_sm,
+            num_splits_static,
+            tile_m_divmod,
+            tile_n_divmod,
+            tile_count_semaphore,
+            num_m_blocks_ptr,
+            num_splits_dynamic_ptr,
+            virtual_batch_idx_ptr,
+            num_nheads_in_l2_ptr,
+            n_blocks_per_split,
+        ).launch(
+            grid=grid,
+            block=block,
+            stream=stream,
+            smem=self.shared_storage.size_in_bytes(),
+        )
+    @cute.kernel
+    def kernel(
+        self,
+        seqlen_q_static: Int32,
+        seqlen_k_static: Int32,
+        seqlen_k_new_static: Int32,
+        mCuSeqlensQ: Optional[cute.Tensor],
+        mCuSeqlensK: Optional[cute.Tensor],
+        mCuSeqlensKNew: Optional[cute.Tensor],
+        mSeqUsedQ: Optional[cute.Tensor],
+        mSeqUsedK: Optional[cute.Tensor],
+        mLeftPadK: Optional[cute.Tensor],
+        num_batch: Int32,
+        num_sm: Int32,
+        num_splits_static: Int32,
+        tile_m_divmod: FastDivmodDivisor,
+        tile_n_divmod: FastDivmodDivisor,
+        tile_count_semaphore: Optional[cute.Tensor],
+        num_m_blocks_ptr: Optional[cute.Tensor],
+        num_splits_dynamic_ptr: Optional[cute.Tensor],
+        virtual_batch_idx_ptr: Optional[cute.Tensor],
+        num_nheads_in_l2_ptr: Optional[cute.Tensor],
+        n_blocks_per_split: Optional[Int32],
+    ):
+        bidx, _, _ = cute.arch.block_idx()
+        tidx, _, _ = cute.arch.thread_idx()
+        grid_dimx, _, _ = cute.arch.grid_dim()
+        warp_idx = cute.arch.warp_idx()
+        lane_idx = cute.arch.lane_idx()
+        smem = cutlass.utils.SmemAllocator()
+        storage = smem.allocate(self.shared_storage)
+        total_blocks_smem = storage.total_blocks_smem.get_tensor((1,))
+        if tidx == 0:
+            total_blocks_smem[0] = Int32(0)
+        cute.arch.sync_threads()
+        if const_expr(tile_count_semaphore is not None):
+            if tidx == 0:
+                tile_count_semaphore[0] = Int32(0)
+        batch_cta_idx_offset = bidx * self.k_num_batch_per_warp * self.num_warps
+        bidb_start = batch_cta_idx_offset + self.k_num_batch_per_warp * warp_idx
+        batch_idx = lane_idx + bidb_start
+        num_m_blocks, _ = self.get_num_m_blocks_and_seqlen(
+            lane_idx,
+            batch_idx,
+            mSeqUsedQ,
+            mCuSeqlensQ,
+            seqlen_q_static,
+            tile_m_divmod,
+            num_batch,
+        )
+        num_n_blocks = self.get_num_n_blocks(
+            lane_idx,
+            batch_idx,
+            mSeqUsedK,
+            mCuSeqlensK,
+            mCuSeqlensKNew,
+            seqlen_k_static,
+            seqlen_k_new_static,
+            mLeftPadK,
+            tile_n_divmod,
+            num_batch,
+        )
+        num_splits_dynamic = Int32(1)
+        if const_expr(n_blocks_per_split is not None):
+            num_splits_dynamic = cutlass.min(
+                cute.ceil_div(num_n_blocks, n_blocks_per_split), num_splits_static
+            )
+            if const_expr(self.zfill_padded_output):
+                num_splits_dynamic = cutlass.max(num_splits_dynamic, Int32(1))
+            if num_splits_dynamic > 0:
+                num_n_blocks = cute.ceil_div(num_n_blocks, num_splits_dynamic)
+        else:
+            if grid_dimx > 1 or num_splits_static == 1:
+                num_splits_dynamic = Int32(1)
+            else:
+                total_blocks = num_m_blocks * num_n_blocks
+                total_blocks = utils.warp_reduce(total_blocks, operator.add)
+                if lane_idx == 0:
+                    utils.atomic_add_i32(total_blocks, total_blocks_smem.iterator)
+                cute.arch.sync_threads()
+                total_blocks = total_blocks_smem[0]
+                sm_margin = max(Float32(num_sm) / 128 + 0.001, 1.1)  # e.g. 148/128 = 1.15625
+                blocks_per_sm = cutlass.max(
+                    Int32(
+                        (
+                            Float32(total_blocks)
+                            * sm_margin
+                            * Float32(self.nheads_computed)
+                            / Float32(num_sm)
+                        )
+                    ),
+                    Int32(1),
+                )
+                # blocks_per_sm = cute.ceil_div(total_blocks * self.nheads_computed, num_sm)
+                num_splits_dynamic = cutlass.min(
+                    cute.ceil_div(num_n_blocks, blocks_per_sm), num_splits_static
+                )
+                if const_expr(self.zfill_padded_output):
+                    num_splits_dynamic = cutlass.max(num_splits_dynamic, Int32(1))
+                if num_splits_dynamic > 0:
+                    num_n_blocks = cute.ceil_div(num_n_blocks, num_splits_dynamic)
+        if const_expr(self.sort):
+            # TODO: Implement sort logic
+            pass
+        if batch_idx < num_batch and lane_idx < self.k_num_batch_per_warp:
+            if const_expr(num_m_blocks_ptr is not None):
+                num_m_blocks_ptr[batch_idx] = num_m_blocks
+            if const_expr(num_splits_dynamic_ptr is not None):
+                num_splits_dynamic_ptr[batch_idx] = num_splits_dynamic
+            if const_expr(num_nheads_in_l2_ptr is not None):
+                nheads_in_l2 = self.get_num_nheads_in_l2(num_n_blocks)
+                num_nheads_in_l2_ptr[batch_idx] = nheads_in_l2
+    @cute.jit
+    def get_num_m_blocks_and_seqlen(
+        self,
+        lane_idx: Int32,
+        batch_idx: Int32,
+        mSeqUsedQ: Optional[cute.Tensor],
+        mCuSeqlensQ: Optional[cute.Tensor],
+        seqlen_q_static: Int32,
+        tile_m_divmod: FastDivmodDivisor,
+        num_batch: Int32,
+    ):
+        seqlen = Int32(0)
+        if const_expr(mSeqUsedQ is not None):
+            seqlen = mSeqUsedQ[batch_idx] if batch_idx < num_batch else Int32(0)
+        elif const_expr(mCuSeqlensQ is not None):
+            # Since k_num_batch_per_warp = 31, lane 31 never processes batches
+            # So shuffle_down is safe: lane 30 gets lane 31's value (which is 0)
+            # Only access cu_seqlens if batch_idx is valid (0 to num_batch inclusive)
+            cur_cu_seqlen = Int32(0)
+            if batch_idx <= num_batch:
+                cur_cu_seqlen = mCuSeqlensQ[batch_idx]
+            next_cu_seqlen = cute.arch.shuffle_sync_down(cur_cu_seqlen, offset=1)
+            seqlen = next_cu_seqlen - cur_cu_seqlen
+        else:
+            seqlen = seqlen_q_static
+        seqlen_for_blocks = seqlen
+        if const_expr(self.packgqa):
+            seqlen_for_blocks = seqlen * self.qhead_per_khead
+        num_m_blocks = (
+            (seqlen_for_blocks + self.tile_m - 1) // tile_m_divmod
+            if batch_idx < num_batch and lane_idx < self.k_num_batch_per_warp
+            else Int32(0)
+        )
+        return (num_m_blocks, seqlen)
+    @cute.jit
+    def get_num_n_blocks(
+        self,
+        lane_idx: Int32,
+        batch_idx: Int32,
+        mSeqUsedK: Optional[cute.Tensor],
+        mCuSeqlensK: Optional[cute.Tensor],
+        mCuSeqlensKNew: Optional[cute.Tensor],
+        seqlen_k_static: Int32,
+        seqlen_k_new_static: Int32,
+        mLeftPadK: Optional[cute.Tensor],
+        tile_n_divmod: FastDivmodDivisor,
+        num_batch: Int32,
+    ):
+        leftpad_k = (
+            mLeftPadK[batch_idx]
+            if const_expr(mLeftPadK is not None) and batch_idx < num_batch
+            else Int32(0)
+        )
+        seqlen = Int32(0)
+        if const_expr(mSeqUsedK is not None):
+            seqlen = mSeqUsedK[batch_idx] if batch_idx < num_batch else Int32(0)
+        elif const_expr(mCuSeqlensK is not None):
+            # Since k_num_batch_per_warp = 31, lane 31 never processes batches
+            # So shuffle_down is safe: lane 30 gets lane 31's value (which is 0)
+            # Only access cu_seqlens if batch_idx is valid (0 to num_batch inclusive)
+            cur_cu_seqlen = Int32(0)
+            if batch_idx <= num_batch:
+                cur_cu_seqlen = mCuSeqlensK[batch_idx]
+            next_cu_seqlen = cute.arch.shuffle_sync_down(cur_cu_seqlen, offset=1)
+            seqlen = next_cu_seqlen - cur_cu_seqlen
+        else:
+            seqlen = seqlen_k_static
+        seqlen_new = Int32(0)
+        if const_expr(mCuSeqlensKNew is not None):
+            # Since k_num_batch_per_warp = 31, lane 31 never processes batches
+            # So shuffle_down is safe: lane 30 gets lane 31's value (which is 0)
+            # Only access cu_seqlens if batch_idx is valid (0 to num_batch inclusive)
+            cur_cu_seqlen_new = Int32(0)
+            if batch_idx <= num_batch:
+                cur_cu_seqlen_new = mCuSeqlensKNew[batch_idx]
+            next_cu_seqlen_new = cute.arch.shuffle_sync_down(cur_cu_seqlen_new, offset=1)
+            seqlen_new = next_cu_seqlen_new - cur_cu_seqlen_new
+        else:
+            seqlen_new = seqlen_k_new_static
+        seqlen = seqlen - leftpad_k + seqlen_new
+        return (
+            (seqlen + self.tile_n - 1) // tile_n_divmod
+            if batch_idx < num_batch and lane_idx < self.k_num_batch_per_warp
+            else Int32(0)
+        )
+    @cute.jit
+    def get_num_nheads_in_l2(
+        self,
+        num_n_blocks: Int32,
+    ):
+        max_kvblocks_in_l2 = self.max_kvblocks_in_l2
+        qhead_per_khead = self.qhead_per_khead
+        nheads_in_l2 = Int32(16)
+        if num_n_blocks * Int32(16) <= max_kvblocks_in_l2:
+            nheads_in_l2 = Int32(16)
+        elif num_n_blocks * Int32(8) <= max_kvblocks_in_l2:
+            nheads_in_l2 = Int32(8)
+        elif num_n_blocks * Int32(4) <= max_kvblocks_in_l2:
+            nheads_in_l2 = Int32(4)
+        elif num_n_blocks * Int32(2) <= max_kvblocks_in_l2:
+            nheads_in_l2 = Int32(2)
+        else:
+            nheads_in_l2 = Int32(1)
+        if const_expr(not self.packgqa):
+            nheads_in_l2 *= qhead_per_khead
+        return cutlass.min(nheads_in_l2, self.nheads_computed)

build/torch-cuda/quack/__init__.py CHANGED Viewed

	@@ -0,0 +1,9 @@

+"""Vendored subset of quack (https://github.com/Dao-AILab/quack).
+Importing `.dsl` installs quack's CuTe tensor indexing patch. It must happen
+here, exactly as upstream's `quack/__init__.py` does it, because
+`quack.copy_utils` indexes tensors with `...` and CuTe does not support that
+without the patch.
+"""
+from . import dsl  # noqa: F401

build/torch-cuda/quack/activation.py DELETED Viewed

@@ -1,568 +0,0 @@
-# Copyright (c) 2025, Tri Dao.
-import math
-from typing import Tuple
-from functools import partial
-import cutlass.cute as cute
-from cutlass import Float32, Boolean, const_expr
-from cutlass.cutlass_dsl import T, dsl_user_op
-from cutlass._mlir.dialects import llvm, nvvm
-F32_or_F32x2 = Float32 | Tuple[Float32, Float32]
-sub_packed_f32x2 = partial(
-    cute.arch.calc_packed_f32x2_op,
-    src_c=None,
-    calc_func=nvvm.sub_packed_f32x2,
-)
-@dsl_user_op
-def tanh(a: float | Float32, *, loc=None, ip=None) -> Float32:
-    return Float32(
-        llvm.inline_asm(
-            T.f32(),
-            [Float32(a).ir_value(loc=loc, ip=ip)],
-            "tanh.approx.f32 $0, $1;",
-            "=f,f",
-            has_side_effects=False,
-            is_align_stack=False,
-            asm_dialect=llvm.AsmDialect.AD_ATT,
-        )
-    )
-@dsl_user_op
-def sigmoid(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
-    if const_expr(not isinstance(x, tuple)):
-        # return 0.5 + 0.5 * cute.math.tanh(0.5 * x, fastmath=True)
-        return 0.5 + 0.5 * tanh(0.5 * x)
-    else:
-        x_half = cute.arch.mul_packed_f32x2((0.5, 0.5), x)
-        tanh_x_half = (tanh(x_half[0]), tanh(x_half[1]))
-        return cute.arch.fma_packed_f32x2(tanh_x_half, (0.5, 0.5), (0.5, 0.5))
-@dsl_user_op
-def dsigmoid_from_output(out: Float32, dout: Float32, *, loc=None, ip=None) -> Float32:
-    # return dout * out * (1.0 - out)
-    return dout * (out - out * out)
-@dsl_user_op
-def relu(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
-    if const_expr(not isinstance(x, tuple)):
-        return cute.arch.fmax(x, Float32(0.0))
-    else:
-        return cute.arch.fmax(x[0], Float32(0.0)), cute.arch.fmax(x[1], Float32(0.0))
-@dsl_user_op
-@cute.jit
-def drelu(
-    x: F32_or_F32x2, dout: F32_or_F32x2, *, loc=None, ip=None
-) -> Tuple[F32_or_F32x2, F32_or_F32x2]:
-    if const_expr(not isinstance(x, tuple)):
-        x_pos = Boolean(x > 0)
-        return dout if x_pos else Float32(0.0), cute.arch.fmax(x, Float32(0.0))
-    else:
-        x0_pos = Boolean(x[0] > 0)
-        x1_pos = Boolean(x[1] > 0)
-        dx = (dout[0] if x0_pos else Float32(0.0), dout[1] if x1_pos else Float32(0.0))
-        return dx, relu(x)
-@dsl_user_op
-def relu_sq(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
-    if const_expr(not isinstance(x, tuple)):
-        return cute.arch.fmax(x, Float32(0.0)) * x
-    else:
-        relu_x = (cute.arch.fmax(x[0], Float32(0.0)), cute.arch.fmax(x[1], Float32(0.0)))
-        return cute.arch.mul_packed_f32x2(relu_x, x)
-@dsl_user_op
-@cute.jit
-def drelu_sq(
-    x: F32_or_F32x2, dout: F32_or_F32x2, *, loc=None, ip=None
-) -> Tuple[F32_or_F32x2, F32_or_F32x2]:
-    """
-    ReLU squared backward pass: computes gradient w.r.t. x and recomputes forward
-    Given: relu_sq_out = max(x, 0) * x, and dout = grad w.r.t. relu_sq_out
-    Returns: (dx, relu_sq_out) where:
-    - dx = dout * 2 * x if x > 0, else 0
-    - relu_sq_out = max(x, 0) * x
-    """
-    if const_expr(not isinstance(x, tuple)):
-        relu_x = relu(x)
-        relu_sq_out = relu_x * x
-        # Derivative: d/dx[max(x,0) * x] = 2*x if x > 0, else 0
-        dx = 2.0 * (dout * relu_x)
-        return dx, relu_sq_out
-    else:
-        relu_x = relu(x)
-        relu_sq_out = cute.arch.mul_packed_f32x2(relu_x, x)
-        dx = cute.arch.mul_packed_f32x2((2.0, 2.0), cute.arch.mul_packed_f32x2(dout, relu_x))
-        return dx, relu_sq_out
-@dsl_user_op
-def gelu_tanh_approx(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
-    """
-    gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
-            = 0.5 * x * (1 + tanh(x * (0.797885 + 0.0356774 * x * x)))
-    """
-    sqrt_2_over_pi = math.sqrt(2 / math.pi)  # ~0.797885
-    sqrt_2_over_pi_coeff = 0.044715 * sqrt_2_over_pi  # ~0.0356774
-    if const_expr(not isinstance(x, tuple)):
-        return 0.5 * (
-            x
-            # Currently cute.math.tanh(x, fastmath=True) generates very slow code
-            # * (1 + cute.math.tanh(x * (sqrt_2_over_pi + sqrt_2_over_pi_coeff * (x * x)), fastmath=True))
-            * (1.0 + tanh(x * (sqrt_2_over_pi + sqrt_2_over_pi_coeff * (x * x))))
-        )
-    else:
-        x_sq = cute.arch.mul_packed_f32x2(x, x)
-        x_sq_scaled = cute.arch.fma_packed_f32x2(
-            x_sq, (sqrt_2_over_pi_coeff, sqrt_2_over_pi_coeff), (sqrt_2_over_pi, sqrt_2_over_pi)
-        )
-        z = cute.arch.mul_packed_f32x2(x, x_sq_scaled)
-        tanh_z = (tanh(z[0]), tanh(z[1]))
-        x_tanh_z = cute.arch.fma_packed_f32x2(tanh_z, x, x)
-        return cute.arch.mul_packed_f32x2((0.5, 0.5), x_tanh_z)
-@dsl_user_op
-def dgelu_tanh_approx(
-    x: F32_or_F32x2, dout: F32_or_F32x2, *, loc=None, ip=None
-) -> Tuple[F32_or_F32x2, F32_or_F32x2]:
-    """
-    GELU tanh approximation backward pass: computes gradient w.r.t. x and recomputes forward
-    Given: gelu_out = 0.5 * x * (1 + tanh(x * (c1 + c2 * x^2))), and dout = grad w.r.t. gelu_out
-    Returns: (dx, gelu_out)
-    Derivative uses the chain rule:
-    d/dx[gelu(x)] = 0.5 * (1 + tanh(z)) + 0.5 * x * sech^2(z) * dz/dx
-    where z = x * (c1 + c2 * x^2), dz/dx = c1 + 3 * c2 * x^2
-    and sech^2(z) = 1 - tanh^2(z)
-    """
-    sqrt_2_over_pi = math.sqrt(2 / math.pi)  # c1 ~0.797885
-    sqrt_2_over_pi_coeff = 0.044715 * sqrt_2_over_pi  # c2 ~0.0356774
-    sqrt_2_over_pi_coeff_3 = 3.0 * sqrt_2_over_pi_coeff  # c3 ~0.01070322
-    if const_expr(not isinstance(x, tuple)):
-        # Compute z = x * (c1 + c2 * x^2)
-        x_sq = x * x
-        # tanh_z = cute.math.tanh(x * (sqrt_2_over_pi + sqrt_2_over_pi_coeff * x_sq), fastmath=True)
-        tanh_z = tanh(x * (sqrt_2_over_pi + sqrt_2_over_pi_coeff * x_sq))
-        half_tanh_z_plus_one = 0.5 + 0.5 * tanh_z
-        gelu_out = x * half_tanh_z_plus_one
-        # Compute gradient
-        # sech^2(z) = 1 - tanh^2(z)
-        sech2_z = 1 - tanh_z * tanh_z
-        # dz/dx = c1 + 3 * c2 * x^2
-        dz_dx = sqrt_2_over_pi + sqrt_2_over_pi_coeff_3 * x_sq
-        # d/dx[gelu(x)] = 0.5 * (1 + tanh(z)) + 0.5 * x * sech^2(z) * dz/dx
-        dgelu = half_tanh_z_plus_one + x * (0.5 * (sech2_z * dz_dx))
-        dx = dout * dgelu
-        return dx, gelu_out
-    else:
-        # Compute z = x * (c1 + c2 * x^2)
-        x_sq = cute.arch.mul_packed_f32x2(x, x)
-        x_sq_scaled = cute.arch.fma_packed_f32x2(
-            x_sq, (sqrt_2_over_pi_coeff, sqrt_2_over_pi_coeff), (sqrt_2_over_pi, sqrt_2_over_pi)
-        )
-        z = cute.arch.mul_packed_f32x2(x, x_sq_scaled)
-        tanh_z = (tanh(z[0]), tanh(z[1]))
-        half_tanh_z_plus_one = cute.arch.fma_packed_f32x2(tanh_z, (0.5, 0.5), (0.5, 0.5))
-        gelu_out = cute.arch.mul_packed_f32x2(x, half_tanh_z_plus_one)
-        # Compute gradient
-        # sech^2(z) = 1 - tanh^2(z)
-        sech2_z = cute.arch.fma_packed_f32x2(tanh_z, (-tanh_z[0], -tanh_z[1]), (1.0, 1.0))
-        # dz/dx = c1 + 3 * c2 * x^2
-        dz_dx = cute.arch.fma_packed_f32x2(
-            x_sq, (sqrt_2_over_pi_coeff_3, sqrt_2_over_pi_coeff_3), (sqrt_2_over_pi, sqrt_2_over_pi)
-        )
-        # d/dx[gelu(x)] = 0.5 * (1 + tanh(z)) + 0.5 * x * sech^2(z) * dz/dx
-        sech2_dz_dx = cute.arch.mul_packed_f32x2(sech2_z, dz_dx)
-        x_sech2_dz_dx = cute.arch.mul_packed_f32x2(x, sech2_dz_dx)
-        dgelu = cute.arch.fma_packed_f32x2(x_sech2_dz_dx, (0.5, 0.5), half_tanh_z_plus_one)
-        dx = cute.arch.mul_packed_f32x2(dout, dgelu)
-        return dx, gelu_out
-@dsl_user_op
-@cute.jit
-def softplus(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
-    if const_expr(not isinstance(x, tuple)):
-        use_linear = Boolean(x > 20.0)
-        return (
-            cute.math.log(Float32(cute.math.exp(x, fastmath=True)) + 1.0, fastmath=True)
-            if not use_linear
-            else x
-        )
-    else:
-        log2_e = math.log2(math.e)
-        x_log2e = cute.arch.mul_packed_f32x2(x, (log2_e, log2_e))
-        x_exp = (cute.math.exp(x_log2e[0], fastmath=True), cute.math.exp(x_log2e[1], fastmath=True))
-        x_exp_p1 = cute.arch.add_packed_f32x2(x_exp, (1.0, 1.0))
-        log_x_exp_p1 = (
-            cute.math.log2(x_exp_p1[0], fastmath=True),
-            cute.math.log2(x_exp_p1[1], fastmath=True),
-        )
-        ln2 = math.log(2.0)
-        softplus_x = cute.arch.mul_packed_f32x2(log_x_exp_p1, (ln2, ln2))
-        use_linear_0 = Boolean(x[0] > 20.0)
-        use_linear_1 = Boolean(x[1] > 20.0)
-        return (
-            softplus_x[0] if not use_linear_0 else x[0],
-            softplus_x[1] if not use_linear_1 else x[1],
-        )
-@dsl_user_op
-@cute.jit
-def dsoftplus_from_output(out: Float32, dout: Float32, *, loc=None, ip=None) -> Float32:
-    use_linear = Boolean(out > 20.0)
-    # dx = dout * (1.0 - cute.math.exp(-out, fastmath=True)) if not use_linear else dout
-    dx = dout - dout * cute.math.exp(-out, fastmath=True)
-    return dx if not use_linear else dout
-@dsl_user_op
-def silu(x: F32_or_F32x2, *, already_halved: bool = False, loc=None, ip=None) -> F32_or_F32x2:
-    """
-    silu(x) = x * sigmoid(x) = x * (1 + tanh(x / 2)) / 2 = (0.5 * x) * tanh(0.5 * x) + (0.5 * x)
-    This compiles down to 3 SASS instructions: FMUL to get 0.5 * x, MUFU.TANH, and FFMA.
-    """
-    if const_expr(not isinstance(x, tuple)):
-        x_half = 0.5 * x if const_expr(not already_halved) else x
-        # return x_half * cute.math.tanh(x_half, fastmath=True) + x_half
-        return x_half * tanh(x_half) + x_half
-    else:
-        x_half = cute.arch.mul_packed_f32x2((0.5, 0.5), x) if const_expr(not already_halved) else x
-        tanh_x_half = (tanh(x_half[0]), tanh(x_half[1]))
-        return cute.arch.fma_packed_f32x2(x_half, tanh_x_half, x_half)
-@dsl_user_op
-def swiglu(x: F32_or_F32x2, y: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
-    if const_expr(not isinstance(x, tuple)):
-        return silu(x) * y
-    else:
-        return cute.arch.mul_packed_f32x2(silu(x), y)
-@dsl_user_op
-def dswiglu(
-    x: F32_or_F32x2,
-    y: F32_or_F32x2,
-    dout: F32_or_F32x2,
-    *,
-    already_halved: bool = False,
-    loc=None,
-    ip=None,
-) -> Tuple[F32_or_F32x2, F32_or_F32x2, F32_or_F32x2]:
-    """
-    SwiGLU backward pass: computes gradients w.r.t. x (gate) and y (up projection)
-    Given: swiglu_out = silu(x) * y, and dout = grad w.r.t. swiglu_out
-    Returns: (dx, dy, swiglu_out) where dx = dout * y * d_silu(x), dy = dout * silu(x)
-    d_silu(x) = sigmoid(x) * (1 + x * (1 - sigmoid(x)))
-    This has been optimized to use fewer instructions (i.e. we expand things out
-    to use FFMA instead of FADD and FMUL).
-    """
-    if const_expr(not isinstance(x, tuple)):
-        # Compute sigmoid(x) using tanh: sigmoid(x) = 0.5 * (1 + tanh(0.5 * x))
-        # FMUL, MUFU.TANH, then FFMA
-        if const_expr(not already_halved):
-            sigmoid_x = sigmoid(x)
-            silu_x = x * sigmoid_x  # FMUL
-        else:
-            tanh_x = tanh(x)  # MUFU.TANH
-            sigmoid_x = 0.5 * tanh_x + 0.5  # FFMA
-            silu_x = x * tanh_x + x  # FFMA
-        silu_x_dout = silu_x * dout  # FMUL
-        #   d_silu(x) * dout
-        # = sigmoid_x * (1 + x * (1 - sigmoid_x)) * dout
-        # = (sigmoid_x + sigmoid_x * x * (1 - sigmoid_x)) * dout
-        # = (sigmoid_x + silu_x * (1 - sigmoid_x)) * dout
-        # = (sigmoid_x + silu_x - silu_x * sigmoid_x) * dout
-        # = (sigmoid_x - silu_x * sigmoid_x) * dout + silu_x * dout
-        d_silu_x_dout = (sigmoid_x - silu_x * sigmoid_x) * dout + silu_x_dout  # FFMA, FFMA
-        dx = d_silu_x_dout * y  # FMUL
-        dy = silu_x_dout
-        swiglu_out = silu_x * y  # FMUL
-        # Overall it's 1 MUFU.TANH, 5 FMUL, 3 FFMA
-        return dx, dy, swiglu_out
-    else:
-        # Compute sigmoid(x) and silu(x)
-        if const_expr(not already_halved):
-            sigmoid_x = sigmoid(x)
-            silu_x = cute.arch.mul_packed_f32x2(x, sigmoid_x)
-        else:
-            tanh_x = (tanh(x[0]), tanh(x[1]))
-            sigmoid_x = cute.arch.fma_packed_f32x2(tanh_x, (0.5, 0.5), (0.5, 0.5))
-            silu_x = cute.arch.fma_packed_f32x2(x, tanh_x, x)
-        silu_x_dout = cute.arch.mul_packed_f32x2(silu_x, dout)
-        # d_silu(x) * dout = (sigmoid_x - silu_x * sigmoid_x) * dout + silu_x * dout
-        sigmoid_x_minus_silu_x_sigmoid_x = cute.arch.fma_packed_f32x2(
-            sigmoid_x, (-silu_x[0], -silu_x[1]), sigmoid_x
-        )
-        d_silu_x_dout = cute.arch.fma_packed_f32x2(
-            sigmoid_x_minus_silu_x_sigmoid_x, dout, silu_x_dout
-        )
-        dx = cute.arch.mul_packed_f32x2(d_silu_x_dout, y)
-        dy = silu_x_dout
-        swiglu_out = cute.arch.mul_packed_f32x2(silu_x, y)
-        return dx, dy, swiglu_out
-@dsl_user_op
-def swiglu_oai(
-    x: F32_or_F32x2, y: F32_or_F32x2, alpha: float = 1.702, *, loc=None, ip=None
-) -> F32_or_F32x2:
-    """The swiglu variant used in gpt-oss, which has a scaling factor on x and bias of 1 to y.
-    https://github.com/openai/gpt-oss/blob/7be9334950053a888e24887a57dac797a17d6e00/gpt_oss/torch/model.py#L249
-    x * sigmoid(alpha * x) * (y + 1)
-    Compile down to FMUL, FMUL, TANH, FFMA, FFMA
-    """
-    # Compute sigmoid(alpha * x) using tanh: sigmoid(z) = 0.5 * (1 + tanh(z/2))
-    if const_expr(not isinstance(x, tuple)):
-        x_half = 0.5 * x
-        # silu_x = x_half * cute.math.tanh(alpha * x_half, fastmath=True) + x_half
-        silu_x = x_half * tanh(alpha * x_half) + x_half
-        return silu_x * y + silu_x
-    else:
-        x_half = cute.arch.mul_packed_f32x2((0.5, 0.5), x)
-        alpha_x_half = cute.arch.mul_packed_f32x2((alpha, alpha), x_half)
-        tanh_alpha_x_half = (tanh(alpha_x_half[0]), tanh(alpha_x_half[1]))
-        silu_x = cute.arch.fma_packed_f32x2(x_half, tanh_alpha_x_half, x_half)
-        return cute.arch.fma_packed_f32x2(silu_x, y, silu_x)
-@dsl_user_op
-def dswiglu_oai(
-    x: F32_or_F32x2, y: F32_or_F32x2, dout: F32_or_F32x2, alpha: float = 1.702, *, loc=None, ip=None
-) -> Tuple[F32_or_F32x2, F32_or_F32x2, F32_or_F32x2]:
-    """
-    Swiglu OAI backward pass: computes gradients w.r.t. x and y
-    Given: swiglu_oai_out = x * sigmoid(alpha * x) * (y + 1), and dout = grad w.r.t. swiglu_oai_out
-    Returns: (dx, dy, swiglu_oai_out)
-    Derivative of x * sigmoid(alpha * x) w.r.t. x:
-    d/dx[x * sigmoid(alpha * x)] = sigmoid(alpha * x) + alpha * x * sigmoid(alpha * x) * (1 - sigmoid(alpha * x))
-    """
-    if const_expr(not isinstance(x, tuple)):
-        # Compute sigmoid(alpha * x) using tanh: sigmoid(z) = 0.5 * (1 + tanh(z/2))
-        alpha_x_half = (0.5 * alpha) * x  # FMUL
-        # MUFU.TANH, then FFMA
-        # sigmoid_alpha_x = 0.5 + 0.5 * cute.math.tanh(alpha_x_half, fastmath=True)
-        sigmoid_alpha_x = 0.5 + 0.5 * tanh(alpha_x_half)
-        silu_x = x * sigmoid_alpha_x  # FMUL
-        silu_x_dout = silu_x * dout  # FMUL
-        # FFMA, FFMA, FMUL
-        d_silu_x_dout = (sigmoid_alpha_x + alpha * (silu_x - silu_x * sigmoid_alpha_x)) * dout
-        dx = d_silu_x_dout * y + d_silu_x_dout  # FFMA, instead of multiply by y + 1
-        dy = silu_x_dout
-        swiglu_out = silu_x * y + silu_x  # FFMA, instead of multiply by y + 1
-        # Overall it's 1 MUFU.TANH, 4 FMUL, 5 FFMA
-        return dx, dy, swiglu_out
-    else:
-        # Compute sigmoid(alpha * x)
-        alpha_x_half = cute.arch.mul_packed_f32x2(((0.5 * alpha), (0.5 * alpha)), x)
-        tanh_alpha_x_half = (tanh(alpha_x_half[0]), tanh(alpha_x_half[1]))
-        sigmoid_alpha_x = cute.arch.fma_packed_f32x2(tanh_alpha_x_half, (0.5, 0.5), (0.5, 0.5))
-        silu_x = cute.arch.mul_packed_f32x2(x, sigmoid_alpha_x)
-        silu_x_dout = cute.arch.mul_packed_f32x2(silu_x, dout)
-        # d_silu_x_dout = (sigmoid_alpha_x + alpha * (silu_x - silu_x * sigmoid_alpha_x)) * dout
-        silu_x_minus_product = cute.arch.fma_packed_f32x2(
-            silu_x, (-sigmoid_alpha_x[0], -sigmoid_alpha_x[1]), silu_x
-        )
-        sigmoid_plus_alpha_diff = cute.arch.fma_packed_f32x2(
-            (alpha, alpha), silu_x_minus_product, sigmoid_alpha_x
-        )
-        d_silu_x_dout = cute.arch.mul_packed_f32x2(sigmoid_plus_alpha_diff, dout)
-        dx = cute.arch.fma_packed_f32x2(d_silu_x_dout, y, d_silu_x_dout)
-        dy = silu_x_dout
-        swiglu_out = cute.arch.fma_packed_f32x2(silu_x, y, silu_x)
-        return dx, dy, swiglu_out
-@dsl_user_op
-def glu(x: F32_or_F32x2, y: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
-    """GLU: Gated Linear Unit
-    glu(x, y) = sigmoid(x) * y
-    Using tanh to compute sigmoid: sigmoid(x) = 0.5 * (1 + tanh(x/2))
-    """
-    if const_expr(not isinstance(x, tuple)):
-        sigmoid_x = sigmoid(x)  # FMUL, MUFU.TANH, then FFMA
-        return sigmoid_x * y  # FMUL
-    else:
-        sigmoid_x = sigmoid(x)
-        return cute.arch.mul_packed_f32x2(sigmoid_x, y)
-@dsl_user_op
-def dglu(
-    x: F32_or_F32x2, y: F32_or_F32x2, dout: F32_or_F32x2, *, loc=None, ip=None
-) -> Tuple[F32_or_F32x2, F32_or_F32x2, F32_or_F32x2]:
-    """
-    GLU backward pass: computes gradients w.r.t. x (gate) and y (up projection)
-    Given: glu_out = sigmoid(x) * y, and dout = grad w.r.t. glu_out
-    Returns: (dx, dy, glu_out) where:
-    - dx = dout * y * sigmoid(x) * (1 - sigmoid(x))
-    - dy = dout * sigmoid(x)
-    - glu_out = sigmoid(x) * y
-    """
-    if const_expr(not isinstance(x, tuple)):
-        # Compute sigmoid(x) using tanh: sigmoid(x) = 0.5 * (1 + tanh(x/2))
-        sigmoid_x = sigmoid(x)  # FMUL, MUFU.TANH, then FFMA
-        sigmoid_x_dout = sigmoid_x * dout  # FMUL
-        glu_out = sigmoid_x * y  # FMUL
-        # dx = y * sigmoid(x) * (1 - sigmoid(x)) * dout
-        #    = y * (1 - sigmoid(x)) * sigmoid_x_dout
-        #    = (y - y * sigmoid(x)) * sigmoid_x_dout
-        #    = (y - glu_out) * sigmoid_x_dout
-        dx = (y - glu_out) * sigmoid_x_dout  # FADD, FMUL
-        dy = sigmoid_x_dout
-        # Total: 1 MUFU.TANH, 4 FMUL, 1 FADD, 1 FFMA
-        return dx, dy, glu_out
-    else:
-        sigmoid_x = sigmoid(x)
-        sigmoid_x_dout = cute.arch.mul_packed_f32x2(sigmoid_x, dout)
-        glu_out = cute.arch.mul_packed_f32x2(sigmoid_x, y)
-        # dx = (y - glu_out) * sigmoid_x_dout
-        y_minus_glu_out = sub_packed_f32x2(y, glu_out)
-        dx = cute.arch.mul_packed_f32x2(y_minus_glu_out, sigmoid_x_dout)
-        dy = sigmoid_x_dout
-        return dx, dy, glu_out
-@dsl_user_op
-def reglu(x: F32_or_F32x2, y: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
-    """ReGLU: ReLU Gated Linear Unit
-    reglu(x, y) = relu(x) * y = max(x, 0) * y
-    """
-    if const_expr(not isinstance(x, tuple)):
-        return cute.arch.fmax(x, Float32(0.0)) * y
-    else:
-        relu_x = relu(x)
-        return cute.arch.mul_packed_f32x2(relu_x, y)
-@dsl_user_op
-@cute.jit
-def dreglu(
-    x: F32_or_F32x2, y: F32_or_F32x2, dout: F32_or_F32x2, *, loc=None, ip=None
-) -> Tuple[F32_or_F32x2, F32_or_F32x2, F32_or_F32x2]:
-    """
-    ReGLU backward pass: computes gradients w.r.t. x (gate) and y (up projection)
-    Given: reglu_out = relu(x) * y, and dout = grad w.r.t. reglu_out
-    Returns: (dx, dy, reglu_out) where:
-    - dx = dout * y if x > 0, else 0
-    - dy = dout * relu(x)
-    - reglu_out = relu(x) * y
-    """
-    if const_expr(not isinstance(x, tuple)):
-        x_pos = Boolean(x > 0)
-        relu_x = cute.arch.fmax(x, Float32(0.0))
-        dx = (dout * y) if x_pos else Float32(0.0)
-        dy = dout * relu_x
-        reglu_out = relu_x * y
-        return dx, dy, reglu_out
-    else:
-        x0_pos = Boolean(x[0] > 0)
-        x1_pos = Boolean(x[1] > 0)
-        relu_x = relu(x)
-        dout_y = cute.arch.mul_packed_f32x2(dout, y)
-        dx = ((dout_y[0] if x0_pos else Float32(0.0)), (dout_y[1] if x1_pos else Float32(0.0)))
-        dy = cute.arch.mul_packed_f32x2(dout, relu_x)
-        reglu_out = cute.arch.mul_packed_f32x2(relu_x, y)
-        return dx, dy, reglu_out
-@dsl_user_op
-def geglu(x: F32_or_F32x2, y: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
-    """GeGLU: GELU Gated Linear Unit
-    geglu(x, y) = gelu(x) * y
-    Uses the tanh approximation of GELU
-    """
-    if const_expr(not isinstance(x, tuple)):
-        return gelu_tanh_approx(x) * y
-    else:
-        return cute.arch.mul_packed_f32x2(gelu_tanh_approx(x), y)
-@dsl_user_op
-def dgeglu(
-    x: F32_or_F32x2, y: F32_or_F32x2, dout: F32_or_F32x2, *, loc=None, ip=None
-) -> Tuple[F32_or_F32x2, F32_or_F32x2, F32_or_F32x2]:
-    """
-    GeGLU backward pass: computes gradients w.r.t. x (gate) and y (up projection)
-    Given: geglu_out = gelu(x) * y, and dout = grad w.r.t. geglu_out
-    Returns: (dx, dy, geglu_out) where:
-    - dx = dout * y * d_gelu(x)
-    - dy = dout * gelu(x)
-    - geglu_out = gelu(x) * y
-    """
-    if const_expr(not isinstance(x, tuple)):
-        # Reuse dgelu_tanh_approx to compute d_gelu(x) * dout and gelu(x)
-        dgelu_x_dout, gelu_x = dgelu_tanh_approx(x, dout)
-        # Compute gradients for geglu
-        dx = dgelu_x_dout * y
-        dy = gelu_x * dout
-        geglu_out = gelu_x * y
-        return dx, dy, geglu_out
-    else:
-        # Reuse dgelu_tanh_approx to compute d_gelu(x) * dout and gelu(x)
-        dgelu_x_dout, gelu_x = dgelu_tanh_approx(x, dout)
-        # Compute gradients for geglu
-        dx = cute.arch.mul_packed_f32x2(dgelu_x_dout, y)
-        dy = cute.arch.mul_packed_f32x2(gelu_x, dout)
-        geglu_out = cute.arch.mul_packed_f32x2(gelu_x, y)
-        return dx, dy, geglu_out
-# ============================================================================
-# Activation name -> function maps
-# ============================================================================
-act_fn_map = {
-    None: None,
-    "silu": silu,
-    "relu": relu,
-    "relu_sq": relu_sq,
-    "gelu_tanh_approx": gelu_tanh_approx,
-}
-dact_fn_map = {
-    None: None,
-    "relu": drelu,
-    "relu_sq": drelu_sq,
-    "gelu_tanh_approx": dgelu_tanh_approx,
-}
-gate_fn_map = {
-    "swiglu": swiglu,
-    "swiglu_oai": swiglu_oai,
-    "reglu": reglu,
-    "geglu": geglu,
-    "glu": glu,
-}
-dgate_fn_map = {
-    "swiglu": dswiglu,
-    "swiglu_oai": dswiglu_oai,
-    "reglu": dreglu,
-    "geglu": dgeglu,
-    "glu": dglu,
-}

build/torch-cuda/quack/copy_utils.py CHANGED Viewed

@@ -1,22 +1,118 @@
-# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
-from typing import Optional, Type, Tuple, Callable, Sequence
 from functools import partial
 import cutlass
 import cutlass.cute as cute
 from cutlass import Int32, Int16, Boolean, const_expr
-from cutlass.cute.nvgpu import cpasync, warp, warpgroup
 from cutlass.cute.nvgpu.tcgen05.mma import CtaGroup  # noqa
 from cutlass.cutlass_dsl import dsl_user_op
 import cutlass.pipeline
-from cutlass._mlir.dialects import llvm
 from cutlass._mlir import ir
 from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir
 Sm100MmaPeerBitMask = 0xFEFFFFFF
 @dsl_user_op
@@ -41,6 +137,30 @@ def cvt_copy(
     cute.copy(tiled_copy, 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_rmem_tensor_like(src, src.element_type, loc=loc, ip=ip)
@@ -48,6 +168,13 @@ def load_s2r(src: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor:
     return dst
 @dsl_user_op
 def load_s2r_retile(
     tiled_copy: cute.TiledCopy,
@@ -68,11 +195,19 @@ def load_s2r_retile(
 @dsl_user_op
 def load_t2r(
-    thr_copy: cute.ThrCopy, shape: cute.Shape, src: cute.Tensor, *, loc=None, ip=None
 ) -> cute.Tensor:
-    cDst = cute.make_identity_tensor(shape)
-    dst = cute.make_rmem_tensor(thr_copy.partition_D(cDst).shape, src.element_type, loc=loc, ip=ip)
-    cute.copy(thr_copy, src, dst, loc=loc, ip=ip)
     return dst
@@ -294,18 +429,14 @@ def as_position_independent_swizzle_tensor(tensor: cute.Tensor) -> cute.Tensor:
     return cute.make_tensor(cute.recast_ptr(tensor.iterator, dtype=tensor.element_type), new_layout)
-def partition_D_position_independent(
-    thr_copy: cute.core.ThrCopy, tensor: cute.Tensor
-) -> cute.Tensor:
     return cute.make_tensor(
         swizzle_ptr(thr_copy.partition_D(tensor).iterator),
         thr_copy.partition_D(as_position_independent_swizzle_tensor(tensor)).layout,
     )
-def partition_S_position_independent(
-    thr_copy: cute.core.ThrCopy, tensor: cute.Tensor
-) -> cute.Tensor:
     return cute.make_tensor(
         swizzle_ptr(thr_copy.partition_S(tensor).iterator),
         thr_copy.partition_S(as_position_independent_swizzle_tensor(tensor)).layout,
@@ -346,12 +477,12 @@ def sm90_get_smem_load_op(
 def get_smem_store_atom(
-    arch: cutlass.Constexpr[int],
     element_type: Type[cute.Numeric],
     transpose: bool = False,
     major_mode_size: Optional[int] = None,
 ) -> cute.CopyAtom:
-    if const_expr(arch < 90 or element_type.width != 16):
         return cute.make_copy_atom(
             cute.nvgpu.CopyUniversalOp(),
             element_type,
@@ -370,12 +501,12 @@ def get_smem_store_atom(
 def get_smem_load_atom(
-    arch: cutlass.Constexpr[int],
     element_type: Type[cute.Numeric],
     transpose: bool = False,
     major_mode_size: Optional[int] = None,
 ) -> cute.CopyAtom:
-    if const_expr(arch < 90 or element_type.width != 16):
         return cute.make_copy_atom(
             cute.nvgpu.CopyUniversalOp(),
             element_type,
@@ -394,26 +525,35 @@ def get_smem_load_atom(
 def get_smem_store_C(
-    tiled_mma: cute.TiledMma,
     sC: cute.Tensor,
     tidx: Int32,
-    arch: int,
     transpose: bool = False,
     position_independent=False,
     major_mode_size: Optional[int] = None,
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]:
     dtype = sC.element_type
-    copy_atom = get_smem_store_atom(arch, dtype, transpose, major_mode_size=major_mode_size)
-    tiled_copy = cute.make_tiled_copy_C(copy_atom, tiled_mma)
     thr_copy = tiled_copy.get_slice(tidx)
     if const_expr(not position_independent):
         tRS_sC = thr_copy.partition_D(sC)
     else:
         tRS_sC = partition_D_position_independent(thr_copy, sC)
-    def copy_fn(src: cute.Tensor, dst_idx: Optional[Int32] = None, **new_kwargs):
-        dst_tensor = tRS_sC if const_expr(dst_idx is None) else tRS_sC[None, None, None, dst_idx]
         cvt_copy(tiled_copy, src, dst_tensor, retile=True, **new_kwargs)
     return copy_fn, thr_copy, tRS_sC
@@ -422,19 +562,18 @@ def get_smem_load_C(
     tiled_mma: cute.TiledMma,
     sC: cute.Tensor,
     tidx: Int32,
-    arch: int,
     transpose: bool = False,
     position_independent=False,
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]:
     dtype = sC.element_type
-    copy_atom = get_smem_load_atom(arch, dtype, transpose)
     tiled_copy = cute.make_tiled_copy_C(copy_atom, tiled_mma)
     thr_copy = tiled_copy.get_slice(tidx)
     if const_expr(not position_independent):
         tSR_sC = thr_copy.partition_S(sC)
     else:
         tSR_sC = partition_S_position_independent(thr_copy, sC)
-    copy_atom_RS = get_smem_store_atom(arch, dtype, transpose)
     thr_copy_RS = cute.make_tiled_copy_C(copy_atom_RS, tiled_mma).get_slice(tidx)
     tRS_shape = thr_copy_RS.partition_S(cute.make_identity_tensor(sC.shape[:2])).shape
@@ -448,10 +587,18 @@ def get_smem_load_C(
 def epilog_smem_copy_atom(
     tiled_mma: cute.TiledMma, epi_tile: cute.Shape, transpose: bool = False
 ) -> cute.TiledCopy:
-    copy_atom_C = cute.make_copy_atom(
-        warp.StMatrix8x8x16bOp(transpose, num_matrices=4 if epi_tile[1] % 16 == 0 else 2),
-        cutlass.Float16,  # this is just to get the right source layout
-    )
     tiled_copy_C_atom = cute.make_tiled_copy_C_atom(copy_atom_C, tiled_mma)
     return tiled_copy_C_atom
@@ -461,13 +608,12 @@ def get_smem_store_epi(
     epi_tile: cute.Shape,
     sC: Optional[cute.Tensor],
     tidx: Int32,
-    arch: int,
     transpose: bool = False,
     position_independent=False,
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor, cute.Tensor]:
     dtype = sC.element_type if const_expr(sC is not None) else cutlass.Float16
     tiled_copy_C_atom = epilog_smem_copy_atom(tiled_mma, epi_tile)
-    copy_atom = get_smem_store_atom(arch, dtype, transpose)
     tiled_copy = cute.make_tiled_copy_S(copy_atom, tiled_copy_C_atom)
     thr_copy = tiled_copy.get_slice(tidx)
     tRS_sC = None
@@ -488,11 +634,11 @@ def get_smem_store_epi(
 def get_smem_store_A(
-    tiled_mma: cute.TiledMma, sA: cute.Tensor, tidx: Int32, arch: int, position_independent=False
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]:
     dtype = sA.element_type
-    transpose = tiled_mma.op.a_major_mode == warpgroup.OperandMajorMode.MN
-    copy_atom = get_smem_store_atom(arch, dtype, transpose)
     tiled_copy = cute.make_tiled_copy_A(copy_atom, tiled_mma)
     thr_copy = tiled_copy.get_slice(tidx)
     if const_expr(not position_independent):
@@ -510,13 +656,12 @@ def get_smem_load_A(
     tiled_mma: cute.TiledMma,
     sA: cute.Tensor,
     tidx: Int32,
-    arch: int,
     with_dst_tensor: bool = False,
     position_independent=False,
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]:
     dtype = sA.element_type
-    transpose = tiled_mma.op.a_major_mode == warpgroup.OperandMajorMode.MN
-    copy_atom = get_smem_load_atom(arch, dtype, transpose)
     tiled_copy = cute.make_tiled_copy_A(copy_atom, tiled_mma)
     thr_copy = tiled_copy.get_slice(tidx)
     if const_expr(not position_independent):
@@ -545,19 +690,16 @@ def cpasync_reduce_bulk_add_f32(
     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,
     )
@@ -648,33 +790,34 @@ def tma_gather4_load(
     """
     if len(row_indices) != 4:
         raise ValueError(f"gather4 requires exactly 4 row indices, got {len(row_indices)}")
-    col_val = Int32(col_idx).ir_value()
-    row_vals = [Int32(row_idx).ir_value() for row_idx in row_indices]
     # Convert pointers to integer addresses
-    desc_addr = tma_desc_ptr.toint(loc=loc, ip=ip).ir_value()
-    dst_addr = dst_smem_ptr.toint(loc=loc, ip=ip).ir_value()
     mbar_addr = mbarrier_ptr.toint(loc=loc, ip=ip)
     if num_cta > 1:
         # Executed by both CTAs. Set peer bit to 0 so that the
         # transaction bytes will update CTA0's barrier.
         mbar_addr = mbar_addr & Sm100MmaPeerBitMask
-    mbar_addr = mbar_addr.ir_value()
     # Handle multicast_mask - may already be ir.Value or Python int
     multicast_mask_val = None
     if multicast_mask is not None:
-        multicast_mask_val = Int16(multicast_mask).ir_value()
     assert multicast_mask_val is None, "multicast is not supported yet"
     # Emit inline PTX for TMA gather4
     # PTX: cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4.mbarrier::complete_tx::bytes
     #      [dstMem], [tensorMap, {col, row0, row1, row2, row3}], [smem_bar];
     ptx = (
-        f"cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4.mbarrier::complete_tx::bytes.cta_group::{num_cta} "
-        "[$0], [$1, {$2, $3, $4, $5, $6}], [$7];"
     )
-    llvm.inline_asm(
-        None,
-        [
             dst_addr,
             desc_addr,
             col_val,
@@ -684,11 +827,6 @@ def tma_gather4_load(
             row_vals[3],
             mbar_addr,
         ],
-        ptx,
-        "r,l,r,r,r,r,r,r",  # constraints: register, long, 6x register
-        has_side_effects=True,
-        is_align_stack=False,
-        asm_dialect=llvm.AsmDialect.AD_ATT,
         loc=loc,
         ip=ip,
     )
@@ -775,6 +913,161 @@ def tma_get_copy_fn(
     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(
@@ -787,6 +1080,18 @@ def tma_producer_copy_fn(copy: Callable, pipeline: cutlass.pipeline.PipelineAsyn
     return copy_fn
 @cute.jit
 def gather_m_get_copy_fn(
     thr_copy_A: cute.ThrCopy,
@@ -796,17 +1101,17 @@ def gather_m_get_copy_fn(
     limit_m: Int32,
     limit_k: Int32,
 ) -> Callable:
-    tile_shape_mk = (cute.size(sA, mode=[0]), cute.size(sA, mode=[1]))
-    tAsA = thr_copy_A.partition_D(sA)
     # k-major
     assert tAsA.shape[2] == 1
     tAsA = cute.group_modes(cute.slice_(tAsA, (None, None, 0, None)), 0, 2)
-    is_even_m_smem = tile_shape_mk[0] % thr_copy_A.tiler_mn[0].shape == 0
     if const_expr(not is_even_m_smem):
-        limit_m = min(limit_m, tile_shape_mk[0])
     elems_per_load = cute.size(tAsA.shape[0][0])
-    cA = cute.make_identity_tensor(tile_shape_mk)
     tAcA = thr_copy_A.partition_S(cA)
     t0AcA = thr_copy_A.get_slice(0).partition_S(cA)
     # Instead of comparing tAcA to limit_m, we instead compare t0AcA to limit_m - tAcA[0][0]
@@ -828,13 +1133,13 @@ def gather_m_get_copy_fn(
         else:
             m_idx[m] = 0  # It's ok to load row 0 in the case of OOB
-    mA_k = cute.logical_divide(mA, (None, tile_shape_mk[1]))
-    def copy_fn(src_idx, dst_idx, pred: bool = False):
         tApA_k = None
         if const_expr(pred):
             tApA_k = cute.make_rmem_tensor(cols_per_thread, Boolean)
-            limit_k_cur = limit_k - src_idx * tile_shape_mk[1]
             for k in cutlass.range(cols_per_thread, unroll_full=True):
                 tApA_k[k] = t0AcA[0, 0, k][1] < limit_k_cur
         mA_cur = mA_k[None, (None, src_idx)]
@@ -997,11 +1302,162 @@ def gather_m_get_tma_copy_fn(
     tma_gather4_load_fn = partial(tma_gather4_load, tma_desc_ptr, num_cta=cta_group)
     def copy_fn(src_idx, dst_idx, tma_bar_ptr: cute.Pointer):
         col_idx = tile_K * src_idx
         for m in cutlass.range(cute.size(tSR_rAIdx, mode=[1]), unroll_full=True):
             row_indices = [tSR_rAIdx[v, m] for v in range(4)]
-            smem_ptr = tSR_sA[None, m, None, dst_idx].iterator
             with cute.arch.elect_one():
                 tma_gather4_load_fn(smem_ptr, tma_bar_ptr, col_idx, row_indices)
     return copy_fn

+# Copyright (c) 2025-2026, QuACK team.
+from typing import Any, Optional, Type, Tuple, Callable, Sequence
 from functools import partial
 import cutlass
 import cutlass.cute as cute
+import cutlass.utils.blackwell_helpers as sm100_utils
 from cutlass import Int32, Int16, Boolean, const_expr
+from cutlass.base_dsl.arch import Arch
+from cutlass.cute.nvgpu import cpasync, tcgen05, warp
 from cutlass.cute.nvgpu.tcgen05.mma import CtaGroup  # noqa
 from cutlass.cutlass_dsl import dsl_user_op
+from cutlass.utils import LayoutEnum
 import cutlass.pipeline
 from cutlass._mlir import ir
 from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir
+from . import layout_utils
+from .utils import make_vector
 Sm100MmaPeerBitMask = 0xFEFFFFFF
+_TCGEN05_TMEM_OPS = (
+    tcgen05.Ld16x128bOp,
+    tcgen05.Ld16x256bOp,
+    tcgen05.Ld16x32bx2Op,
+    tcgen05.Ld16x64bOp,
+    tcgen05.Ld32x32bOp,
+    tcgen05.LdRed16x32bx2Op,
+    tcgen05.LdRed32x32bOp,
+    tcgen05.St16x128bOp,
+    tcgen05.St16x256bOp,
+    tcgen05.St16x32bx2Op,
+    tcgen05.St16x64bOp,
+    tcgen05.St32x32bOp,
+)
+_TCGEN05_TMEM_STORE_OPS = (
+    tcgen05.St16x128bOp,
+    tcgen05.St16x256bOp,
+    tcgen05.St16x32bx2Op,
+    tcgen05.St16x64bOp,
+    tcgen05.St32x32bOp,
+)
+def tmem_store_atom_from_load_atom(
+    copy_atom_t2r: Any,
+    src_dtype: Type[cutlass.Numeric],
+    dst_dtype: Type[cutlass.Numeric],
+) -> cute.CopyAtom:
+    """Return the matching tcgen05 R2T store atom for a selected T2R load atom.
+    `src_dtype` is the register fragment dtype loaded by T2R; `dst_dtype` is
+    the TMEM element dtype to store. Ratio 1 uses CUTLASS's operation-family
+    mapping directly. Ratio 2 is intentionally narrow: we allow the current
+    Ld32x32b path by halving repeat, and the widest 16dp path by halving the
+    vector width. Narrower 16dp cross-family mappings are not mirrored by
+    CUTLASS's same-family helper, so they assert until validated.
+    C++ CuTe's operation-family mapping is `cute::TMEM::tmem_load_to_store`:
+    https://github.com/NVIDIA/cutlass/blob/main/include/cute/atom/copy_traits_sm100.hpp#L3274
+    """
+    load_op = copy_atom_t2r.op if const_expr(hasattr(copy_atom_t2r, "op")) else copy_atom_t2r
+    if const_expr(hasattr(load_op, "op")):
+        load_op = load_op.op
+    assert src_dtype.width >= dst_dtype.width, "TMEM R2T helper only supports narrowing stores"
+    assert src_dtype.width % dst_dtype.width == 0, "TMEM source/destination widths must divide"
+    ratio = src_dtype.width // dst_dtype.width
+    assert ratio in (1, 2), "TMEM R2T helper only supports src/dst width ratio 1 or 2"
+    repeat = load_op.repeat
+    unpack = tcgen05.Unpack.NONE
+    if const_expr(getattr(load_op, "pack", None) == tcgen05.Pack.PACK_16b_IN_32b):
+        unpack = tcgen05.Unpack.UNPACK_32b_IN_16b
+    if const_expr(isinstance(load_op, tcgen05.Ld16x64bOp)):
+        assert ratio == 1, "No validated ratio-2 store mapping for Ld16x64bOp"
+        store_op = tcgen05.St16x64bOp(repeat, unpack)
+    elif const_expr(isinstance(load_op, tcgen05.Ld16x128bOp)):
+        assert ratio == 1, "No validated ratio-2 store mapping for Ld16x128bOp"
+        store_op = tcgen05.St16x128bOp(repeat, unpack)
+    elif const_expr(isinstance(load_op, tcgen05.Ld16x256bOp)):
+        store_op = (
+            tcgen05.St16x256bOp(repeat, unpack)
+            if const_expr(ratio == 1)
+            else tcgen05.St16x128bOp(repeat, unpack)
+        )
+    elif const_expr(isinstance(load_op, tcgen05.Ld16x32bx2Op)):
+        assert ratio == 1, "No validated ratio-2 store mapping for Ld16x32bx2Op"
+        store_op = tcgen05.St16x32bx2Op(repeat, unpack)
+    elif const_expr(isinstance(load_op, tcgen05.Ld32x32bOp)):
+        if const_expr(ratio == 2):
+            assert repeat.value % 2 == 0, "Ld32x32b ratio-2 store needs even repeat"
+            repeat = tcgen05.Repetition(repeat.value // 2)
+        store_op = tcgen05.St32x32bOp(repeat, unpack)
+    else:
+        raise TypeError(f"Unsupported TMEM load op for store conversion: {type(load_op)}")
+    return cute.make_copy_atom(store_op, dst_dtype)
+def _tmem_copy_reg_tv_layout(tiled_copy: cute.TiledCopy):
+    """Return the register-side TV layout for a tcgen05 tmem copy."""
+    op = tiled_copy.op
+    if const_expr(hasattr(op, "op")):
+        op = op.op
+    if const_expr(isinstance(op, _TCGEN05_TMEM_OPS)):
+        # TMEM stores read from registers; all other TMEM copy ops here write
+        # registers, including LdRed* reductions that upstream is_tmem_load
+        # intentionally does not classify as plain loads.
+        return (
+            tiled_copy.layout_src_tv_tiled
+            if const_expr(isinstance(op, _TCGEN05_TMEM_STORE_OPS))
+            else tiled_copy.layout_dst_tv_tiled
+        )
+    raise TypeError(f"Cannot infer tmem copy direction from tiled_copy.op={op}")
 @dsl_user_op
     cute.copy(tiled_copy, src, dst, pred=pred, loc=loc, ip=ip, **kwargs)
+@dsl_user_op
+def sr_cvt_copy(
+    tiled_copy: cute.TiledCopy,
+    src: cute.Tensor,
+    dst: cute.Tensor,
+    seed: Int32,
+    tidx: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    """Like cvt_copy but uses stochastic rounding for FP32 -> BF16 conversion."""
+    assert isinstance(src.iterator, cute.Pointer) and src.memspace == cute.AddressSpace.rmem
+    from .rounding import convert_f32_to_bf16_sr
+    from cutlass.cute.tensor import TensorSSA
+    src_cvt = cute.make_rmem_tensor_like(src, dst.element_type)
+    src_vec = src.load()
+    raw_vec = convert_f32_to_bf16_sr(src_vec, seed, tidx, loc=loc, ip=ip)
+    src_cvt.store(TensorSSA(raw_vec, src_vec.shape, dst.element_type))
+    src = src_cvt
+    cute.copy(tiled_copy, src, dst, loc=loc, ip=ip)
 @dsl_user_op
 def load_s2r(src: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor:
     dst = cute.make_rmem_tensor_like(src, src.element_type, loc=loc, ip=ip)
     return dst
+@dsl_user_op
+def contiguous(src: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor:
+    dst = cute.make_rmem_tensor(src.shape, src.element_type, loc=loc, ip=ip)
+    cute.autovec_copy(src, dst, loc=loc, ip=ip)
+    return dst
 @dsl_user_op
 def load_s2r_retile(
     tiled_copy: cute.TiledCopy,
 @dsl_user_op
 def load_t2r(
+    tiled_copy: cute.TiledCopy,
+    src: cute.Tensor,
+    *,
+    loc=None,
+    ip=None,
 ) -> cute.Tensor:
+    """Load one tmem tile partition into rmem, deriving the rmem shape from `src`.
+    `src` should already be indexed to the tile being copied, with any
+    stage/subtile modes removed.
+    """
+    dst = tmem_reg_frag(tiled_copy, src, loc=loc, ip=ip)
+    cute.copy(tiled_copy, src, dst, loc=loc, ip=ip)
     return dst
     return cute.make_tensor(cute.recast_ptr(tensor.iterator, dtype=tensor.element_type), new_layout)
+def partition_D_position_independent(thr_copy: cute.ThrCopy, tensor: cute.Tensor) -> cute.Tensor:
     return cute.make_tensor(
         swizzle_ptr(thr_copy.partition_D(tensor).iterator),
         thr_copy.partition_D(as_position_independent_swizzle_tensor(tensor)).layout,
     )
+def partition_S_position_independent(thr_copy: cute.ThrCopy, tensor: cute.Tensor) -> cute.Tensor:
     return cute.make_tensor(
         swizzle_ptr(thr_copy.partition_S(tensor).iterator),
         thr_copy.partition_S(as_position_independent_swizzle_tensor(tensor)).layout,
 def get_smem_store_atom(
     element_type: Type[cute.Numeric],
     transpose: bool = False,
     major_mode_size: Optional[int] = None,
 ) -> cute.CopyAtom:
+    arch = cutlass.base_dsl.BaseDSL._get_dsl().get_arch_enum()
+    if const_expr(arch < Arch.sm_90 or element_type.width != 16):
         return cute.make_copy_atom(
             cute.nvgpu.CopyUniversalOp(),
             element_type,
 def get_smem_load_atom(
     element_type: Type[cute.Numeric],
     transpose: bool = False,
     major_mode_size: Optional[int] = None,
 ) -> cute.CopyAtom:
+    arch = cutlass.base_dsl.BaseDSL._get_dsl().get_arch_enum()
+    if const_expr(arch < Arch.sm_90 or element_type.width != 16):
         return cute.make_copy_atom(
             cute.nvgpu.CopyUniversalOp(),
             element_type,
 def get_smem_store_C(
+    tiled_mma: cute.TiledMma | cute.TiledCopy,
     sC: cute.Tensor,
     tidx: Int32,
     transpose: bool = False,
     position_independent=False,
     major_mode_size: Optional[int] = None,
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]:
     dtype = sC.element_type
+    if const_expr(isinstance(tiled_mma, cute.TiledCopy)):
+        tiled_copy_t2r = tiled_mma
+        layout = LayoutEnum.COL_MAJOR if const_expr(transpose) else LayoutEnum.ROW_MAJOR
+        copy_atom = sm100_utils.get_smem_store_op(
+            layout, dtype, tiled_copy_t2r.value_type, tiled_copy_t2r
+        )
+        tiled_copy = cute.make_tiled_copy_D(copy_atom, tiled_copy_t2r)
+    else:
+        copy_atom = get_smem_store_atom(dtype, transpose, major_mode_size=major_mode_size)
+        tiled_copy = cute.make_tiled_copy_C(copy_atom, tiled_mma)
     thr_copy = tiled_copy.get_slice(tidx)
     if const_expr(not position_independent):
         tRS_sC = thr_copy.partition_D(sC)
     else:
         tRS_sC = partition_D_position_independent(thr_copy, sC)
+    def copy_fn(src: cute.Tensor, dst_idx: Optional[Int32] = None, fence=False, **new_kwargs):
+        dst_tensor = tRS_sC if const_expr(dst_idx is None) else tRS_sC[..., dst_idx]
         cvt_copy(tiled_copy, src, dst_tensor, retile=True, **new_kwargs)
+        if const_expr(fence):
+            cute.arch.fence_view_async_shared()
     return copy_fn, thr_copy, tRS_sC
     tiled_mma: cute.TiledMma,
     sC: cute.Tensor,
     tidx: Int32,
     transpose: bool = False,
     position_independent=False,
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]:
     dtype = sC.element_type
+    copy_atom = get_smem_load_atom(dtype, transpose)
     tiled_copy = cute.make_tiled_copy_C(copy_atom, tiled_mma)
     thr_copy = tiled_copy.get_slice(tidx)
     if const_expr(not position_independent):
         tSR_sC = thr_copy.partition_S(sC)
     else:
         tSR_sC = partition_S_position_independent(thr_copy, sC)
+    copy_atom_RS = get_smem_store_atom(dtype, transpose)
     thr_copy_RS = cute.make_tiled_copy_C(copy_atom_RS, tiled_mma).get_slice(tidx)
     tRS_shape = thr_copy_RS.partition_S(cute.make_identity_tensor(sC.shape[:2])).shape
 def epilog_smem_copy_atom(
     tiled_mma: cute.TiledMma, epi_tile: cute.Shape, transpose: bool = False
 ) -> cute.TiledCopy:
+    arch = cutlass.base_dsl.BaseDSL._get_dsl().get_arch_enum()
+    if const_expr(arch < Arch.sm_90):
+        copy_atom_C = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            cutlass.Float16,  # this is just to get the right source layout
+            num_bits_per_copy=(2 if not transpose else 1) * cutlass.Float16.width,
+        )
+    else:
+        copy_atom_C = cute.make_copy_atom(
+            warp.StMatrix8x8x16bOp(transpose, num_matrices=4 if epi_tile[1] % 16 == 0 else 2),
+            cutlass.Float16,  # this is just to get the right source layout
+        )
     tiled_copy_C_atom = cute.make_tiled_copy_C_atom(copy_atom_C, tiled_mma)
     return tiled_copy_C_atom
     epi_tile: cute.Shape,
     sC: Optional[cute.Tensor],
     tidx: Int32,
     transpose: bool = False,
     position_independent=False,
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor, cute.Tensor]:
     dtype = sC.element_type if const_expr(sC is not None) else cutlass.Float16
+    copy_atom = get_smem_store_atom(dtype, transpose)
     tiled_copy_C_atom = epilog_smem_copy_atom(tiled_mma, epi_tile)
     tiled_copy = cute.make_tiled_copy_S(copy_atom, tiled_copy_C_atom)
     thr_copy = tiled_copy.get_slice(tidx)
     tRS_sC = None
 def get_smem_store_A(
+    tiled_mma: cute.TiledMma, sA: cute.Tensor, tidx: Int32, position_independent=False
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]:
     dtype = sA.element_type
+    transpose = tiled_mma.op.a_major_mode == cute.nvgpu.OperandMajorMode.MN
+    copy_atom = get_smem_store_atom(dtype, transpose)
     tiled_copy = cute.make_tiled_copy_A(copy_atom, tiled_mma)
     thr_copy = tiled_copy.get_slice(tidx)
     if const_expr(not position_independent):
     tiled_mma: cute.TiledMma,
     sA: cute.Tensor,
     tidx: Int32,
     with_dst_tensor: bool = False,
     position_independent=False,
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]:
     dtype = sA.element_type
+    transpose = tiled_mma.op.a_major_mode == cute.nvgpu.OperandMajorMode.MN
+    copy_atom = get_smem_load_atom(dtype, transpose)
     tiled_copy = cute.make_tiled_copy_A(copy_atom, tiled_mma)
     thr_copy = tiled_copy.get_slice(tidx)
     if const_expr(not position_independent):
     loc=None,
     ip=None,
 ):
+    smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip)
     # cache_hint = cutlass.Int64(0x14F0000000000000)  # EVICT_LAST
+    cute.arch.inline_ptx(
+        "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32 [{$r0}], [{$r1}], {$r2};",
+        read_only_args=[gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes)],
         # [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",
+        loc=loc,
+        ip=ip,
     )
     """
     if len(row_indices) != 4:
         raise ValueError(f"gather4 requires exactly 4 row indices, got {len(row_indices)}")
+    col_val = Int32(col_idx)
+    row_vals = [Int32(row_idx) for row_idx in row_indices]
     # Convert pointers to integer addresses
+    desc_addr = tma_desc_ptr.toint(loc=loc, ip=ip)
+    dst_addr = dst_smem_ptr.toint(loc=loc, ip=ip)
     mbar_addr = mbarrier_ptr.toint(loc=loc, ip=ip)
     if num_cta > 1:
         # Executed by both CTAs. Set peer bit to 0 so that the
         # transaction bytes will update CTA0's barrier.
         mbar_addr = mbar_addr & Sm100MmaPeerBitMask
+    mbar_addr = Int32(mbar_addr)
     # Handle multicast_mask - may already be ir.Value or Python int
     multicast_mask_val = None
     if multicast_mask is not None:
+        multicast_mask_val = Int16(multicast_mask)
     assert multicast_mask_val is None, "multicast is not supported yet"
     # Emit inline PTX for TMA gather4
     # PTX: cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4.mbarrier::complete_tx::bytes
     #      [dstMem], [tensorMap, {col, row0, row1, row2, row3}], [smem_bar];
     ptx = (
+        "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4.mbarrier::complete_tx::bytes."
+        f"cta_group::{num_cta} "
+        "[{$r0}], [{$r1}, {{$r2}, {$r3}, {$r4}, {$r5}, {$r6}}], [{$r7}];"
     )
+    cute.arch.inline_ptx(
+        ptx,
+        read_only_args=[
             dst_addr,
             desc_addr,
             col_val,
             row_vals[3],
             mbar_addr,
         ],
         loc=loc,
         ip=ip,
     )
     return (copy_tma if const_expr(not single_stage) else copy_tma_single_stage), s, g
+def s2t_get_copy_fn(
+    src_tensor: cute.Tensor,
+    dst_tensor: cute.Tensor,
+    cta_group: tcgen05.CtaGroup,
+) -> Callable:
+    """
+    Make tiledCopy for smem to tmem load, then return a copy function over stages.
+    :param src_tensor: The source tensor in smem
+    :param dst_tensor: The destination tensor in tmem
+    """
+    assert src_tensor.element_type == dst_tensor.element_type
+    # (MMA, MMA_MN, MMA_K, STAGE)
+    src_compact = cute.filter_zeros(src_tensor)
+    # (MMA, MMA_MN, MMA_K)
+    dst_compact = cute.filter_zeros(dst_tensor)
+    # Make S2T CopyAtom and tiledCopy.
+    copy_atom = cute.make_copy_atom(tcgen05.Cp4x32x128bOp(cta_group), dst_tensor.element_type)
+    tiled_copy = tcgen05.make_s2t_copy(copy_atom, dst_compact)
+    thr_copy = tiled_copy.get_slice(0)
+    # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
+    src_partition = tcgen05.get_s2t_smem_desc_tensor(tiled_copy, thr_copy.partition_S(src_compact))
+    # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
+    dst_partition = thr_copy.partition_D(dst_compact)
+    @dsl_user_op
+    def copy_s2t(stage_idx, *, loc=None, ip=None, **new_kwargs):
+        # Stage slice of partitioned source tensor: ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
+        stage_coord = (None, None, None, None, stage_idx)
+        cute.copy(
+            tiled_copy, src_partition[stage_coord], dst_partition, loc=loc, ip=ip, **new_kwargs
+        )
+    return copy_s2t
+# tcgen05 TMEM <-> RMEM helpers (t2r loads / r2t stores).
+#
+# The register-side fragment of a tmem copy is derivable by layout algebra,
+# with no reference to the original (pre-partition) tile tensor:
+# - the per-thread register VALUE shape is mode 1 of the tiled copy's
+#   register-side TV layout (`layout_dst_tv_tiled` for loads,
+#   `layout_src_tv_tiled` for stores). The tmem-side partition can't supply
+#   it: tmem partitioning is warp-collective, so its value mode counts tmem
+#   cells across the whole warp, not per-thread register elements.
+# - the tile-iteration modes are shared between partition_S and partition_D
+#   (same tiler over the same tile extent), so they can be read off whichever
+#   side was already partitioned.
+# This kills the make-a-fake/identity-tensor-and-partition_D dance previously
+# needed at every t2r site.
+def tmem_reg_frag(
+    tiled_copy: cute.TiledCopy,
+    partitioned: cute.Tensor,
+    num_extra_modes: int = 0,
+    dtype: Optional[Type[cutlass.Numeric]] = None,
+    *,
+    loc=None,
+    ip=None,
+) -> cute.Tensor:
+    """Allocate the per-thread register fragment for ONE tile of a tcgen05
+    tmem copy, given any partitioned view of it (tmem or otherwise).
+    `partitioned` is (V, iter..., extra...) as produced by partition_S/_D;
+    the trailing `num_extra_modes` modes (stage, epi-subtile, ...) are
+    excluded from the fragment and indexed at copy time instead. `dtype`
+    defaults to `partitioned.element_type`.
+    The register side is inferred from the tcgen05 load/store op: destination
+    for t2r loads, source for r2t stores."""
+    tv = _tmem_copy_reg_tv_layout(tiled_copy)
+    val_shape = tv.shape[1]
+    rank = cute.rank(partitioned.shape)
+    iters = tuple(partitioned.shape[i] for i in range(1, rank - num_extra_modes))
+    frag_dtype = partitioned.element_type if const_expr(dtype is None) else dtype
+    return cute.make_rmem_tensor((val_shape, *iters), frag_dtype, loc=loc, ip=ip)
+def r2s_partition_from_t2r(
+    tiled_copy_t2r: cute.TiledCopy,
+    s: cute.Tensor,
+    tidx: Int32,
+    transpose: bool = False,
+    position_independent=False,
+) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
+    """SMEM-store (r2s) side chained off a tmem-load tiled copy: the r2s copy
+    inherits the t2r copy's per-thread value ownership via make_tiled_copy_D,
+    so the loaded fragment can be stored (post-conversion) without a shuffle.
+    By default the store atom is selected like SM100 GEMM epilogues:
+    `get_smem_store_op(layout, dst_dtype, tiled_copy_t2r.value_type, tiled_copy_t2r)`,
+    so the stmatrix shape follows the tmem-load atom. `transpose=True` maps to
+    COL_MAJOR, otherwise ROW_MAJOR.
+    `s` is the staged SMEM tile; its trailing stage mode is excluded from the
+    register fragment. `position_independent=True` partitions through a
+    position-independent swizzle view, matching `get_smem_store_C`.
+    Returns `(tiled_copy, tRS_r, tRS_s)`; store via
+    `cute.copy(tiled_copy, tRS_r, tRS_s[..., idx])`."""
+    dtype = s.element_type
+    layout = LayoutEnum.COL_MAJOR if const_expr(transpose) else LayoutEnum.ROW_MAJOR
+    copy_atom = sm100_utils.get_smem_store_op(
+        layout, dtype, tiled_copy_t2r.value_type, tiled_copy_t2r
+    )
+    tiled_copy = cute.make_tiled_copy_D(copy_atom, tiled_copy_t2r)
+    thr_copy = tiled_copy.get_slice(tidx)
+    if const_expr(not position_independent):
+        tRS_s = thr_copy.partition_D(s)
+    else:
+        tRS_s = partition_D_position_independent(thr_copy, s)
+    rank = cute.rank(tRS_s.shape)
+    frag_shape = tuple(tRS_s.shape[i] for i in range(rank - 1))
+    tRS_r = cute.make_rmem_tensor(frag_shape, dtype)
+    return tiled_copy, tRS_r, tRS_s
+def s2r_partition_from_t2r(
+    tiled_copy_t2r: cute.TiledCopy,
+    s: cute.Tensor,
+    tidx: Int32,
+    r_layout: cute.Layout,
+    copy_atom: Optional[cute.CopyAtom] = None,
+    transpose: bool = False,
+    position_independent=False,
+) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor]:
+    """SMEM-load (s2r) counterpart of `r2s_partition_from_t2r` (ldmatrix vs
+    stmatrix), for reading an epilogue input that was TMA-staged into an
+    epi-tile SMEM buffer (e.g. C in gemm, z in ssd) into registers
+    element-aligned with the t2r fragments. The register fragment is
+    allocated with `r_layout` (pass the r2s fragment's layout) so its linear
+    element order matches the t2r/r2s fragments; `tSR_r` is its retiled view
+    for the collective copy. Per-warp SMEM footprints of this load and the
+    chained r2s store coincide, so reusing one buffer for input then output
+    is warp-local (no inter-warp hazard).
+    `position_independent=True` partitions through a position-independent
+    swizzle view, matching `get_smem_load_C`.
+    Returns `(tiled_copy, tRS_r, tSR_r, tSR_s)`; load via
+    `cute.copy(tiled_copy, tSR_s[..., idx], tSR_r)` then read `tRS_r`."""
+    dtype = s.element_type
+    if const_expr(copy_atom is None):
+        copy_atom = cute.make_copy_atom(
+            warp.LdMatrix8x8x16bOp(transpose=transpose, num_matrices=4), dtype
+        )
+    tiled_copy = cute.make_tiled_copy_D(copy_atom, tiled_copy_t2r)
+    thr_copy = tiled_copy.get_slice(tidx)
+    if const_expr(not position_independent):
+        tSR_s = thr_copy.partition_S(s)
+    else:
+        tSR_s = partition_S_position_independent(thr_copy, s)
+    tRS_r = cute.make_rmem_tensor(r_layout, dtype)
+    tSR_r = tiled_copy.retile(tRS_r)
+    return tiled_copy, tRS_r, tSR_r, tSR_s
 def tma_producer_copy_fn(copy: Callable, pipeline: cutlass.pipeline.PipelineAsync):
     def copy_fn(src_idx, producer_state: cutlass.pipeline.PipelineState, **new_kwargs):
         copy(
     return copy_fn
+def chain_tma_producer_copy_fns(copy_fns: Sequence[Optional[Callable]]):
+    if not any(fn is not None for fn in copy_fns):
+        return None
+    def copy_fn(src_idx, producer_state: cutlass.pipeline.PipelineState, **new_kwargs):
+        for fn in copy_fns:
+            if const_expr(fn is not None):
+                fn(src_idx=src_idx, producer_state=producer_state, **new_kwargs)
+    return copy_fn
 @cute.jit
 def gather_m_get_copy_fn(
     thr_copy_A: cute.ThrCopy,
     limit_m: Int32,
     limit_k: Int32,
 ) -> Callable:
+    tile_M, tile_K = cute.size(sA, mode=[0]), cute.size(sA, mode=[1])
+    tAsA = partition_D_position_independent(thr_copy_A, sA)
     # k-major
     assert tAsA.shape[2] == 1
     tAsA = cute.group_modes(cute.slice_(tAsA, (None, None, 0, None)), 0, 2)
+    is_even_m_smem = tile_M % thr_copy_A.tiler_mn[0].shape == 0
     if const_expr(not is_even_m_smem):
+        limit_m = min(limit_m, tile_M)
     elems_per_load = cute.size(tAsA.shape[0][0])
+    cA = cute.make_identity_tensor((tile_M, tile_K))
     tAcA = thr_copy_A.partition_S(cA)
     t0AcA = thr_copy_A.get_slice(0).partition_S(cA)
     # Instead of comparing tAcA to limit_m, we instead compare t0AcA to limit_m - tAcA[0][0]
         else:
             m_idx[m] = 0  # It's ok to load row 0 in the case of OOB
+    mA_k = cute.logical_divide(mA, (None, tile_K))
+    def copy_fn(src_idx, dst_idx, pred: cutlass.Constexpr[bool] = False):
         tApA_k = None
         if const_expr(pred):
             tApA_k = cute.make_rmem_tensor(cols_per_thread, Boolean)
+            limit_k_cur = limit_k - src_idx * tile_K
             for k in cutlass.range(cols_per_thread, unroll_full=True):
                 tApA_k[k] = t0AcA[0, 0, k][1] < limit_k_cur
         mA_cur = mA_k[None, (None, src_idx)]
     tma_gather4_load_fn = partial(tma_gather4_load, tma_desc_ptr, num_cta=cta_group)
     def copy_fn(src_idx, dst_idx, tma_bar_ptr: cute.Pointer):
+        tSR_sA_cur = tSR_sA[None, None, None, dst_idx]
         col_idx = tile_K * src_idx
         for m in cutlass.range(cute.size(tSR_rAIdx, mode=[1]), unroll_full=True):
             row_indices = [tSR_rAIdx[v, m] for v in range(4)]
+            smem_ptr = tSR_sA_cur[None, m, None].iterator
             with cute.arch.elect_one():
                 tma_gather4_load_fn(smem_ptr, tma_bar_ptr, col_idx, row_indices)
     return copy_fn
+@cute.jit
+def gather_k_get_tma_copy_fn(
+    tma_atom: cute.CopyAtom,
+    sA: cute.Tensor,  # ((4, tile_K/4), (tile_M,), STAGE) — K-grouped load layout
+    sAIdx: cute.Tensor,  # (tile_K, a_prefetch_stage) — K indices in smem
+    col_idx: Int32,  # M offset in global tensor (contiguous dim for M-major)
+    warp_idx: Int32,
+    num_warps: int,
+    num_cta: int = 1,
+) -> Tuple[Callable, Callable]:
+    """Build a copy function for TMA gather4 in K dimension (M-major A).
+    Each gather4 instruction loads 4 K-columns × tile_M contiguous M-elements.
+    col_idx is the absolute M position in the global tensor.
+    K indices come from sAIdx (prefetched to smem by the scheduler warp).
+    Returns copy_fn(src_idx, dst_idx, tma_bar_ptr) which:
+      Issues gather4 calls with those K indices as row_indices
+    """
+    tile_K = cute.size(sAIdx, mode=[0])
+    assert tile_K % 4 == 0
+    cta_group = num_cta
+    # Tiled copy for loading K indices from smem to registers (4 per vector, across warps)
+    copy_AIdx_s2r = cute.make_tiled_copy_tv(
+        cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Int32, num_bits_per_copy=128),
+        cute.make_layout(num_warps),  # thr_layout
+        cute.make_layout(4),  # val_layout — 4 K indices per gather4
+    )
+    warp_idx = cute.arch.make_warp_uniform(warp_idx)
+    warp_copy_AIdx_s2r = copy_AIdx_s2r.get_slice(warp_idx)
+    tSR_sAIdx = warp_copy_AIdx_s2r.partition_S(sAIdx)  # (((4,1),4,4))
+    # ((4,1),4,(64,2),(1,4)):((64,0),1024,(1,4096),(0,8192))
+    tSR_sA = warp_copy_AIdx_s2r.partition_S(layout_utils.transpose_view(sA))
+    tma_desc_ptr = get_tma_desc_addr(tma_atom)
+    tma_gather4_load_fn = partial(tma_gather4_load, tma_desc_ptr, num_cta=cta_group)
+    def prefetch_from_smem_fn(
+        a_prefetch_pipeline,
+        src_idx,
+        dst_idx,
+        a_prefetch_consumer_state,
+    ) -> cute.Tensor:
+        a_prefetch_pipeline.consumer_wait(a_prefetch_consumer_state)
+        tSR_rAIdx = load_s2r(tSR_sAIdx[None, None, dst_idx])
+        cute.arch.sync_warp()
+        with cute.arch.elect_one():
+            a_prefetch_pipeline.consumer_release(a_prefetch_consumer_state)
+        return tSR_rAIdx
+    def copy_fn(src_idx, dst_idx, tSR_rAIdx, tma_bar_ptr: cute.Pointer):
+        # Issue gather4: col_idx = M position, row_indices = 4 K positions
+        tSR_sA_cur = tSR_sA[None, None, None, dst_idx]
+        gather_dim = cute.size(tSR_sA_cur, mode=[2, 0])  # Typically 64
+        for k in cutlass.range(cute.size(tSR_rAIdx, mode=[1]), unroll_full=True):
+            row_indices = [tSR_rAIdx[v, k] for v in range(4)]
+            for m in cutlass.range(cute.size(tSR_sA_cur, mode=[2, 1]), unroll_full=True):
+                smem_ptr = tSR_sA_cur[None, k, (None, m)].iterator
+                with cute.arch.elect_one():
+                    tma_gather4_load_fn(
+                        smem_ptr, tma_bar_ptr, col_idx + m * gather_dim, row_indices
+                    )
+    return copy_fn, prefetch_from_smem_fn
+# ---------------------------------------------------------------------------
+# Store helpers
+# ---------------------------------------------------------------------------
+@dsl_user_op
+@cute.jit
+def store(
+    ptr: cute.Pointer,
+    val,
+    pred: Optional[Boolean] = None,
+    cop: cutlass.Constexpr = None,
+    *,
+    loc=None,
+    ip=None,
+):
+    """Store a scalar value via cute.arch.store.
+    ptr:  cute.Pointer (any address space).
+    val:  DSL Numeric value.
+    pred: None → unconditional.  DSL Boolean → skipped when pred == 0.
+    cop:  Cache operator — "wb" (default), "cg", "cs" (streaming), "wt".
+    """
+    if const_expr(pred is None):
+        cute.arch.store(ptr.llvm_ptr, type(val)(val), cop=cop, loc=loc, ip=ip)
+    else:
+        if pred:
+            cute.arch.store(ptr.llvm_ptr, type(val)(val), cop=cop, loc=loc, ip=ip)
+@dsl_user_op
+@cute.jit
+def store_v2(
+    ptr: cute.Pointer,
+    v0,
+    v1,
+    pred: Optional[Boolean] = None,
+    cop: cutlass.Constexpr = None,
+    *,
+    loc=None,
+    ip=None,
+):
+    """Vectorized store of 2 elements via cute.arch.store.
+    Packs v0, v1 into an MLIR <2 x T> vector.
+    ptr:  cute.Pointer (any address space, must be aligned for vector width).
+    cop:  Cache operator — "wb" (default), "cg", "cs" (streaming), "wt".
+    """
+    vec = make_vector(type(v0), v0, v1, loc=loc, ip=ip)
+    if const_expr(pred is None):
+        cute.arch.store(ptr.llvm_ptr, vec, cop=cop, loc=loc, ip=ip)
+    else:
+        if pred:
+            cute.arch.store(ptr.llvm_ptr, vec, cop=cop, loc=loc, ip=ip)
+@dsl_user_op
+@cute.jit
+def store_v4(
+    ptr: cute.Pointer,
+    v0,
+    v1,
+    v2,
+    v3,
+    pred: Optional[Boolean] = None,
+    cop: cutlass.Constexpr = None,
+    *,
+    loc=None,
+    ip=None,
+):
+    """Vectorized store of 4 elements via cute.arch.store.
+    Packs v0–v3 into an MLIR <4 x T> vector.
+    ptr:  cute.Pointer (any address space, must be aligned for vector width).
+    cop:  Cache operator — "wb" (default), "cg", "cs" (streaming), "wt".
+    """
+    vec = make_vector(type(v0), v0, v1, v2, v3, loc=loc, ip=ip)
+    if const_expr(pred is None):
+        cute.arch.store(ptr.llvm_ptr, vec, cop=cop, loc=loc, ip=ip)
+    else:
+        if pred:
+            cute.arch.store(ptr.llvm_ptr, vec, cop=cop, loc=loc, ip=ip)

build/torch-cuda/quack/cute_dsl_utils.py CHANGED Viewed

@@ -4,6 +4,9 @@ from typing import Tuple, get_origin
 from functools import lru_cache
 from dataclasses import dataclass, fields
 import torch
 try:
@@ -14,7 +17,6 @@ except ImportError:
 import cutlass
 import cutlass.cute as cute
 from cutlass import Int32, Int64, Float16, BFloat16, Float32
-from cutlass.base_dsl.typing import JitArgument
 from cutlass.base_dsl.tvm_ffi_builder import spec
 from cutlass.cutlass_dsl import NumericMeta
@@ -61,15 +63,64 @@ torch2cute_dtype_map = {
 @lru_cache
-def get_max_active_clusters(cluster_size):
-    return cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_size=cluster_size)
 @lru_cache
-def get_device_capacity(device: torch.device = None) -> Tuple[int, int]:
     return torch.cuda.get_device_capability(device)
 def _partition_fields(obj):
     """Split dataclass fields into (constexpr_dict, non_constexpr_dict) by type."""
     all_fields = {field.name: getattr(obj, field.name) for field in fields(obj)}
@@ -138,28 +189,3 @@ class ParamsBase:
         return values
     __new_from_mlir_values__ = _new_from_mlir_values
-@dataclass
-class ArgumentsBase(JitArgument):
-    def __c_pointers__(self):
-        _, non_constexpr_fields = _partition_fields(self)
-        c_ptrs = []
-        for obj in non_constexpr_fields.values():
-            if hasattr(obj, "__c_pointers__"):
-                c_ptrs.extend(obj.__c_pointers__())
-        return c_ptrs
-    def __get_mlir_types__(self):
-        _, non_constexpr_fields = _partition_fields(self)
-        types, self._values_pos = [], []
-        for obj in non_constexpr_fields.values():
-            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
-    __new_from_mlir_values__ = _new_from_mlir_values

 from functools import lru_cache
 from dataclasses import dataclass, fields
+import os
+import re
 import torch
 try:
 import cutlass
 import cutlass.cute as cute
 from cutlass import Int32, Int64, Float16, BFloat16, Float32
 from cutlass.base_dsl.tvm_ffi_builder import spec
 from cutlass.cutlass_dsl import NumericMeta
 @lru_cache
+def get_device_multiprocessor_count(device_id: int = 0) -> int:
+    return cutlass.utils.HardwareInfo(device_id).get_device_multiprocessor_count()
+@lru_cache
+def get_max_active_clusters(
+    cluster_size: int,
+    device_capacity: Tuple[int, int] | None = None,
+    device_id: int = 0,
+) -> int:
+    if device_capacity is None:
+        device_capacity = get_device_capacity()
+    if device_capacity[0] < 9:
+        if cluster_size != 1:
+            raise ValueError("SM8x kernels do not support CTA clusters; cluster_size must be 1")
+        return get_device_multiprocessor_count(device_id)
+    return cutlass.utils.HardwareInfo(device_id).get_max_active_clusters(cluster_size=cluster_size)
+def _parse_arch_str(arch_str: str) -> Tuple[int, int]:
+    """Parse arch string (e.g. 'sm_90', 'sm90', '90', 'sm_100a') to (major, minor) tuple."""
+    match = re.match(r"^(?:sm_?)?(\d+)(\d)([af]?)$", arch_str.strip(), re.IGNORECASE)
+    if not match:
+        raise ValueError(f"Invalid QUACK_ARCH format: {arch_str!r} (expected e.g. '90', 'sm_90')")
+    major, minor, _ = match.groups()
+    return int(major), int(minor)
 @lru_cache
+def _get_device_capacity_cached(device: torch.device = None) -> Tuple[int, int]:
+    """Return (major, minor) device capability.
+    Override with QUACK_ARCH (e.g. 'sm_90' or '90') for CPU-only compilation
+    without a GPU present.
+    """
+    arch_override = os.environ.get("QUACK_ARCH")
+    if arch_override is not None:
+        return _parse_arch_str(arch_override)
     return torch.cuda.get_device_capability(device)
+def get_device_capacity(
+    device: torch.device | torch.Tensor | None = None,
+) -> Tuple[int, int]:
+    """Return (major, minor) device capability.
+    Override with QUACK_ARCH (e.g. 'sm_90' or '90') for CPU-only compilation
+    without a GPU present.
+    Accepts either a ``torch.device`` or a tensor and canonicalizes to the
+    underlying device before consulting the cached helper. This avoids leaking
+    tensors through the LRU cache key.
+    """
+    if isinstance(device, torch.Tensor):
+        device = device.device
+    return _get_device_capacity_cached(device)
 def _partition_fields(obj):
     """Split dataclass fields into (constexpr_dict, non_constexpr_dict) by type."""
     all_fields = {field.name: getattr(obj, field.name) for field in fields(obj)}
         return values
     __new_from_mlir_values__ = _new_from_mlir_values

build/torch-cuda/quack/dsl/__init__.py ADDED Viewed

	@@ -0,0 +1,12 @@

+"""CuTe DSL helpers and integration hooks.
+Importing this package installs quack's Pythonic CuTe tensor indexing, which
+monkey-patches CuTe's tensor classes process-wide so that `...` and `:` work in
+`__getitem__`/`__setitem__`. `quack.copy_utils` relies on that sugar (e.g.
+`tRS_sC[..., dst_idx]`), and the sm90 backward kernel reaches it.
+Upstream's `quack.dsl` also exports `cute_op` from `torch_library_op`; nothing
+vendored here uses it, so it is not copied.
+"""
+from . import cute_tensor_indexing  # noqa: F401

build/torch-cuda/quack/dsl/cute_tensor_indexing.py ADDED Viewed

	@@ -0,0 +1,139 @@

+# Copyright (c) 2026, Tri Dao.
+"""Small compatibility layer for more Pythonic CuTe tensor indexing.
+CuTe uses ``None`` as its underscore/full-mode slice marker, e.g. ``A[i, None]``.
+This module teaches CuTe tensors the equivalent Python spelling ``:`` and expands
+``...`` to the right number of full-mode slices.
+Importing this module intentionally mutates CuTe's tensor classes process-wide.
+The original methods are retained as ``_quack_original_getitem`` and
+``_quack_original_setitem`` for debugging or manual rollback.
+"""
+from __future__ import annotations
+from typing import Any
+from cutlass.cutlass_dsl import dsl_user_op
+import cutlass.cute.tensor as _cute_tensor
+_ORIGINAL_GETITEM_ATTR = "_quack_original_getitem"
+_ORIGINAL_SETITEM_ATTR = "_quack_original_setitem"
+_PATCHED_ATTR = "_quack_extended_indexing"
+_PATCHED_SETITEM_ATTR = f"{_PATCHED_ATTR}_setitem"
+def _is_full_slice(idx: Any) -> bool:
+    return isinstance(idx, slice) and idx.start is None and idx.stop is None and idx.step is None
+def _index_uses_ellipsis(idx: Any) -> bool:
+    if idx is Ellipsis:
+        return True
+    if isinstance(idx, tuple):
+        return any(_index_uses_ellipsis(item) for item in idx)
+    return False
+def _shape_rank(shape: Any, idx: Any = None) -> int:
+    if shape is None:
+        suffix = f" in {idx!r}" if idx is not None else ""
+        raise ValueError(f"tensor shape is required to expand ellipsis{suffix}")
+    return _cute_tensor.rank(shape)
+def _shape_mode(shape: Any, mode: int) -> Any:
+    if isinstance(shape, tuple) and mode < len(shape):
+        return shape[mode]
+    return None
+def _canonicalize_cute_tensor_index(idx: Any, tensor_shape: Any = None) -> Any:
+    """Convert Python indexing sugar to CuTe's coordinate convention.
+    ``:`` becomes ``None`` (CuTe's full-mode/underscore marker) and ``...`` expands
+    within the current hierarchy level using ``tensor_shape``. Other slices like
+    ``1:4`` are intentionally rejected because CuTe tensor slicing only supports
+    keeping an entire mode or selecting a single coordinate.
+    """
+    if idx is Ellipsis:
+        return (None,) * _shape_rank(tensor_shape, idx)
+    if _is_full_slice(idx):
+        return None
+    if isinstance(idx, slice):
+        raise ValueError(f"CuTe Tensor indexing only supports full slices ':', got {idx!r}")
+    if not isinstance(idx, tuple):
+        return idx
+    ellipsis_count = sum(item is Ellipsis for item in idx)
+    if ellipsis_count > 1:
+        raise ValueError("CuTe Tensor indexing supports at most one ellipsis per tuple level")
+    explicit_modes = len(idx) - ellipsis_count
+    fill_modes = 0
+    if ellipsis_count:
+        tensor_rank = _shape_rank(tensor_shape, idx)
+        fill_modes = tensor_rank - explicit_modes
+        if fill_modes < 0:
+            raise ValueError(
+                f"ellipsis cannot expand index {idx!r} for rank-{tensor_rank} CuTe Tensor mode"
+            )
+    result: list[Any] = []
+    mode = 0
+    for item in idx:
+        if item is Ellipsis:
+            result.extend([None] * fill_modes)
+            mode += fill_modes
+        else:
+            result.append(_canonicalize_cute_tensor_index(item, _shape_mode(tensor_shape, mode)))
+            mode += 1
+    return tuple(result)
+def _make_getitem(original_getitem: Any) -> Any:
+    @dsl_user_op
+    def _getitem(self: Any, idx: Any, *, loc: Any = None, ip: Any = None) -> Any:
+        tensor_shape = self.shape if _index_uses_ellipsis(idx) else None
+        idx = _canonicalize_cute_tensor_index(idx, tensor_shape)
+        return original_getitem(self, idx, loc=loc, ip=ip)
+    return _getitem
+def _make_setitem(original_setitem: Any) -> Any:
+    @dsl_user_op
+    def _setitem(self: Any, idx: Any, data: Any, *, loc: Any = None, ip: Any = None) -> Any:
+        tensor_shape = self.shape if _index_uses_ellipsis(idx) else None
+        idx = _canonicalize_cute_tensor_index(idx, tensor_shape)
+        return original_setitem(self, idx, data, loc=loc, ip=ip)
+    return _setitem
+def patch_cute_tensor_indexing() -> None:
+    """Monkey patch CuTe Tensor indexing with ``:``, ``...`` sugar.
+    The patch is idempotent and keeps the original CuTe implementation for all
+    canonical coordinates, so existing ``A[i, j, None]`` code continues to behave
+    exactly as before. It is a process-wide mutation of CuTe's tensor classes.
+    """
+    for cls in (_cute_tensor._Tensor, _cute_tensor.TensorSSA):
+        if _PATCHED_ATTR not in cls.__dict__:
+            setattr(cls, _ORIGINAL_GETITEM_ATTR, cls.__getitem__)
+            cls.__getitem__ = _make_getitem(cls.__getitem__)  # type: ignore[method-assign]
+            setattr(cls, _PATCHED_ATTR, True)
+    # TensorSSA has no upstream __setitem__, so only _Tensor needs the store path patched.
+    tensor_cls = _cute_tensor._Tensor
+    if _PATCHED_SETITEM_ATTR not in tensor_cls.__dict__:
+        setattr(tensor_cls, _ORIGINAL_SETITEM_ATTR, tensor_cls.__setitem__)
+        tensor_cls.__setitem__ = _make_setitem(tensor_cls.__setitem__)  # type: ignore[method-assign]
+        setattr(tensor_cls, _PATCHED_SETITEM_ATTR, True)
+patch_cute_tensor_indexing()
+__all__ = ["patch_cute_tensor_indexing"]

build/torch-cuda/quack/layout_utils.py CHANGED Viewed

@@ -18,6 +18,19 @@ def select(a: cute.Tensor, mode: list[int]) -> cute.Tensor:
     return cute.make_tensor(a.iterator, cute.select(a.layout, mode))
 def expand(a: cute.Tensor, dim: int, size: Int32 | int) -> cute.Tensor:
     shape = (*a.shape[:dim], size, *a.shape[dim:])
     stride = (*a.layout.stride[:dim], 0, *a.layout.stride[dim:])
@@ -217,15 +230,16 @@ def convert_layout_acc_frgA(acc_layout: cute.Layout) -> cute.Layout:
         )
     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],
             ),
@@ -264,7 +278,7 @@ def convert_layout_zero_stride(
 def mma_partition_C_vec(
-    sVec: cute.Tensor, thr_mma: cute.core.ThrMma, expand_shape: int, is_colvec: bool
 ) -> cute.Tensor:
     assert cute.rank(sVec) == 2
     assert sVec.stride[0] == 1
@@ -281,7 +295,7 @@ def mma_partition_C_vec(
 def mma_partition_A_vec(
-    sVec: cute.Tensor, thr_mma: cute.core.ThrMma, expand_shape: int, is_colvec: bool
 ) -> cute.Tensor:
     assert cute.rank(sVec) == 2
     assert sVec.stride[0] == 1
@@ -295,3 +309,77 @@ def mma_partition_A_vec(
     sVec_mma = cute.make_tensor(sVec.iterator, cute.make_layout(shape, stride=stride))
     tC_sVec = make_acc_tensor_mn_view(thr_mma.partition_A(sVec_mma))
     return tC_sVec[None, 0, None] if const_expr(is_colvec) else tC_sVec[0, None, None]

     return cute.make_tensor(a.iterator, cute.select(a.layout, mode))
+def concat_to_interleave(a: cute.Tensor, dim: int) -> cute.Tensor:
+    """Reshape a concat [first_half; second_half] layout to interleaved along `dim`.
+    Splits dimension `dim` (size 2N) into hierarchical (2, N) so that elements
+    from the first half and second half alternate: [first_0, second_0, first_1, ...].
+    Used to convert gated MLP weight layout from concat [gate; up] to interleaved.
+    """
+    half = cute.size(a, mode=[dim]) // 2
+    shape = (*a.shape[:dim], (2, half), *a.shape[dim + 1 :])
+    stride = (*a.stride[:dim], (half * a.stride[dim], a.stride[dim]), *a.stride[dim + 1 :])
+    return cute.make_tensor(a.iterator, cute.make_layout(shape, stride=stride))
 def expand(a: cute.Tensor, dim: int, size: Int32 | int) -> cute.Tensor:
     shape = (*a.shape[:dim], size, *a.shape[dim:])
     stride = (*a.layout.stride[:dim], 0, *a.layout.stride[dim:])
         )
     else:  # Sm80
         # (4, MMA_M, MMA_N) -> (4, MMA_M, (2, MMA_N / 2))
+        assert acc_layout.shape[2] % 2 == 0
         l = cute.logical_divide(acc_layout, (None, None, 2))
         rA_mma_view = cute.make_layout(
             (
+                (l.shape[0][0], l.shape[0][1], l.shape[2][0]),
                 l.shape[1],
                 l.shape[2][1],
             ),
             stride=(
+                (l.stride[0][0], l.stride[0][1], l.stride[2][0]),
                 l.stride[1],
                 l.stride[2][1],
             ),
 def mma_partition_C_vec(
+    sVec: cute.Tensor, thr_mma: cute.ThrMma, expand_shape: int, is_colvec: bool
 ) -> cute.Tensor:
     assert cute.rank(sVec) == 2
     assert sVec.stride[0] == 1
 def mma_partition_A_vec(
+    sVec: cute.Tensor, thr_mma: cute.ThrMma, expand_shape: int, is_colvec: bool
 ) -> cute.Tensor:
     assert cute.rank(sVec) == 2
     assert sVec.stride[0] == 1
     sVec_mma = cute.make_tensor(sVec.iterator, cute.make_layout(shape, stride=stride))
     tC_sVec = make_acc_tensor_mn_view(thr_mma.partition_A(sVec_mma))
     return tC_sVec[None, 0, None] if const_expr(is_colvec) else tC_sVec[0, None, None]
+def copy_partition_S_vec(
+    sVec: cute.Tensor, thr_copy: cute.ThrCopy, expand_shape: int, is_colvec: bool
+) -> cute.Tensor:
+    assert cute.rank(sVec) == 2
+    assert sVec.stride[0] == 1
+    stage = sVec.shape[1]
+    shape = (
+        (sVec.shape[0], expand_shape, stage)
+        if const_expr(is_colvec)
+        else (expand_shape, sVec.shape[0], stage)
+    )
+    stride = (1, 0, sVec.stride[1]) if const_expr(is_colvec) else (0, 1, sVec.stride[1])
+    sVec_thr = cute.make_tensor(sVec.iterator, cute.make_layout(shape, stride=stride))
+    tC_sVec = reshape_acc_to_mn(thr_copy.partition_S(sVec_thr))
+    return tC_sVec[None, 0, None] if const_expr(is_colvec) else tC_sVec[0, None, None]
+def copy_partition_D_vec(
+    sVec: cute.Tensor, thr_copy: cute.ThrCopy, expand_shape: int, is_colvec: bool
+) -> cute.Tensor:
+    assert cute.rank(sVec) == 2
+    assert sVec.stride[0] == 1
+    stage = sVec.shape[1]
+    shape = (
+        (sVec.shape[0], expand_shape, stage)
+        if const_expr(is_colvec)
+        else (expand_shape, sVec.shape[0], stage)
+    )
+    stride = (1, 0, sVec.stride[1]) if const_expr(is_colvec) else (0, 1, sVec.stride[1])
+    sVec_thr = cute.make_tensor(sVec.iterator, cute.make_layout(shape, stride=stride))
+    tC_sVec = reshape_acc_to_mn(thr_copy.partition_D(sVec_thr))
+    return tC_sVec[None, 0, None] if const_expr(is_colvec) else tC_sVec[0, None, None]
+def tile_atom_to_shape_SF_strided(
+    shape: cute.Shape,
+    sf_vec_size: int,
+    sf_strides,
+) -> cute.Layout:
+    """Build an SFA/SFB layout matching `shape` (A or B operand shape) but
+    honoring the scale tensor's actual strides instead of hardcoded packed
+    ones.
+    Mirrors `cutlass.utils.blockscaled_layout.tile_atom_to_shape_SF(shape,
+    sf_vec_size)`, except outer-mode strides come from `sf_strides` (pass
+    `mSFA.stride` / `mSFB.stride` directly). The inner 512-B atom
+    `((32, 4), (sf_vec_size, 4)) : ((16, 4), (0, 1))` is hardware-fixed.
+    Implementation uses `cute.blocked_product(atom, outer)`; `blocked_product`
+    scales the outer layout's strides by `cosize(atom) == 512`, so we divide
+    the byte strides by 512 (one tile) before handing them in.
+    Args:
+        shape: A/B operand shape. Rank-3 `(m/n, k, l)` or rank-2
+            `(total_mn, k)` (varlen_m).
+        sf_vec_size: Scale factor vector size (16 or 32).
+        sf_strides: Strides of the scale tensor, which has logical shape
+            `(L, rmn, rk, 512)` (rank 4). Only `sf_strides[0..2]` are used:
+            `sf_strides[1]` as the rmn stride, `sf_strides[2]` as the rk
+            stride, and `sf_strides[0]` as the L stride (only for rank-3
+            `shape`).
+    """
+    from cutlass.utils.blockscaled_layout import BlockScaledBasicChunk
+    atom = BlockScaledBasicChunk(sf_vec_size).layout
+    rmn = cute.ceil_div(shape[0], 128)
+    rk = cute.ceil_div(shape[1], sf_vec_size * 4)
+    outer = cute.make_layout((rmn, rk), stride=(sf_strides[1] // 512, sf_strides[2] // 512))
+    sf_layout = cute.blocked_product(atom, outer)
+    if const_expr(len(shape) == 3):
+        sf_layout = cute.append(sf_layout, cute.make_layout(shape[2], stride=sf_strides[0]))
+    return sf_layout

build/torch-cuda/quack/rounding.py ADDED Viewed

	@@ -0,0 +1,231 @@

+# Copyright (c) 2025-2026, Vijay Thakkar, Tri Dao.
+"""Rounding mode control and stochastic rounding primitives for GEMM epilogues.
+Provides a RoundingMode enum for configuring how epilogues downconvert the
+accumulator dtype (typically FP32) to the output dtype before storing to gmem.
+Stochastic rounding (RS) uses the hardware cvt.rs.satfinite.bf16x2.f32 PTX
+instruction and is only supported on Blackwell (SM100+) GPUs.
+"""
+from enum import IntEnum
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Uint32
+from cutlass._mlir import ir
+from cutlass._mlir.dialects import llvm, vector
+from cutlass.cutlass_dsl import dsl_user_op, Int32, T
+class RoundingMode(IntEnum):
+    """Rounding modes for epilogue dtype downconversion.
+    RN — Round to nearest even (default hardware behavior)
+    RS — Stochastic rounding (SM100+ only, BF16 output only)
+    """
+    RN = 0
+    RS = 1
+# Odd strides used to derive distinct Philox counters across output tiles and subtiles.
+EPILOGUE_SR_SEED_M_STRIDE = 65537
+EPILOGUE_SR_SEED_N_STRIDE = 257
+EPILOGUE_SR_SEED_BATCH_STRIDE = 17
+EPILOGUE_SR_SEED_SUBTILE_STRIDE = 7
+EPILOGUE_SR_SEED_AUX_OUT_SALT = 0x9E3779B1
+PHILOX_N_ROUNDS_DEFAULT = 7
+PHILOX_ROUND_A = 0xD2511F53
+PHILOX_ROUND_B = 0xCD9E8D57
+PHILOX_KEY_A = 0x9E3779B9
+PHILOX_KEY_B = 0xBB67AE85
+@dsl_user_op
+def epilogue_sr_seed(
+    base_seed: Int32,
+    tile_coord_mnkl: cute.Coord,
+    subtile_idx,
+    *,
+    loc=None,
+    ip=None,
+) -> Int32:
+    return base_seed + (
+        tile_coord_mnkl[0] * EPILOGUE_SR_SEED_M_STRIDE
+        + tile_coord_mnkl[1] * EPILOGUE_SR_SEED_N_STRIDE
+        + tile_coord_mnkl[3] * EPILOGUE_SR_SEED_BATCH_STRIDE
+        + subtile_idx * EPILOGUE_SR_SEED_SUBTILE_STRIDE
+    )
+@dsl_user_op
+def epilogue_aux_out_sr_seed(
+    base_seed: Int32,
+    tile_coord_mnkl: cute.Coord,
+    subtile_idx,
+    *,
+    loc=None,
+    ip=None,
+) -> Int32:
+    return epilogue_sr_seed(
+        base_seed + EPILOGUE_SR_SEED_AUX_OUT_SALT,
+        tile_coord_mnkl,
+        subtile_idx,
+        loc=loc,
+        ip=ip,
+    )
+@dsl_user_op
+def mul_wide_u32(a: Uint32, b: Uint32, *, loc=None, ip=None) -> tuple:
+    """Unsigned 32b x 32b -> 64 wide multiply via PTX `mul.wide.u32`.
+    Returns (hi, lo) as a pair of Uint32 values.
+    """
+    prod = cute.arch.mul_wide(Uint32(a), Uint32(b), loc=loc, ip=ip)
+    # ptxas folds the shift and to(Uint32) into the same IMAD.WIDE.U32 register pair
+    # that the previous inline PTX mov.b64 split produced.
+    hi = (prod >> 32).to(Uint32, loc=loc, ip=ip)
+    lo = prod.to(Uint32, loc=loc, ip=ip)
+    return hi, lo
+@dsl_user_op
+def cvt_f32x2_bf16x2_rs(
+    a: Float32,
+    b: Float32,
+    rand_bits: Uint32,
+    *,
+    loc=None,
+    ip=None,
+) -> cutlass.Int32:
+    """Convert 2 FP32 values to packed BF16x2 using stochastic rounding.
+    Uses Blackwell PTX instruction: cvt.rs.satfinite.bf16x2.f32 dst, src_hi, src_lo, rand
+    """
+    return cutlass.Int32(
+        llvm.inline_asm(
+            T.i32(),
+            [
+                Float32(a).ir_value(loc=loc, ip=ip),
+                Float32(b).ir_value(loc=loc, ip=ip),
+                Uint32(rand_bits).ir_value(loc=loc, ip=ip),
+            ],
+            "cvt.rs.satfinite.bf16x2.f32 $0, $2, $1, $3;",
+            "=r,f,f,r",
+            has_side_effects=False,
+            is_align_stack=False,
+        )
+    )
+@dsl_user_op
+def philox(
+    counter: Uint32,
+    key: Uint32,
+    n_rounds: int = PHILOX_N_ROUNDS_DEFAULT,
+    *,
+    loc=None,
+    ip=None,
+) -> tuple:
+    """Philox 4x32b counter-based random number generator.
+    Given a 32b counter and a 32b key, returns four pseudo-random uint32 words
+    produced by running n_rounds of the Philox 4x32 bijection. Each round
+    performs two wide 32x32->64 multiplies with the Philox constants.
+    """
+    c0 = Uint32(counter)
+    c1 = Uint32(0)
+    c2 = Uint32(0)
+    c3 = Uint32(0)
+    k0 = Uint32(key)
+    k1 = Uint32(0)
+    round_a = Uint32(PHILOX_ROUND_A)
+    round_b = Uint32(PHILOX_ROUND_B)
+    key_a = Uint32(PHILOX_KEY_A)
+    key_b = Uint32(PHILOX_KEY_B)
+    for _ in range(n_rounds):
+        hi_b, lo_b = mul_wide_u32(c2, round_b, loc=loc, ip=ip)
+        hi_a, lo_a = mul_wide_u32(c0, round_a, loc=loc, ip=ip)
+        c0 = hi_b ^ c1 ^ k0
+        c2 = hi_a ^ c3 ^ k1
+        c1 = lo_b
+        c3 = lo_a
+        k0 = k0 + key_a
+        k1 = k1 + key_b
+    return c0, c1, c2, c3
+@dsl_user_op
+def convert_f32_to_bf16_sr(
+    src_vec,
+    seed: Int32,
+    tid: Int32,
+    *,
+    loc=None,
+    ip=None,
+):
+    """Convert an MLIR FP32 vector to BF16 with stochastic rounding.
+    Processes elements in pairs using Philox PRNG for entropy and the hardware
+    cvt.rs.satfinite.bf16x2.f32 instruction.
+    """
+    src_vec_type = ir.VectorType(src_vec.type)
+    num_elems = src_vec_type.shape[0]
+    assert num_elems % 2 == 0, f"requires even number of elements, got {num_elems}"
+    num_pairs = num_elems // 2
+    assert num_pairs % 4 == 0, (
+        f"num_pairs must be divisible by 4 for stochastic rounding, got {num_pairs}"
+    )
+    dst_mlir_type = cutlass.BFloat16.mlir_type
+    dst_vec_type = ir.VectorType.get([num_elems], dst_mlir_type, loc=loc)
+    i32_vec_type = ir.VectorType.get([num_pairs], Int32.mlir_type, loc=loc)
+    i32_vec = llvm.mlir_undef(i32_vec_type, loc=loc, ip=ip)
+    for pair_idx in range(num_pairs):
+        lo_idx = pair_idx * 2
+        hi_idx = pair_idx * 2 + 1
+        src_lo = vector.extract(
+            src_vec,
+            dynamic_position=[],
+            static_position=[lo_idx],
+            loc=loc,
+            ip=ip,
+        )
+        src_hi = vector.extract(
+            src_vec,
+            dynamic_position=[],
+            static_position=[hi_idx],
+            loc=loc,
+            ip=ip,
+        )
+        group_idx = pair_idx // 4
+        intra_idx = pair_idx % 4
+        if intra_idx == 0:
+            counter = cutlass.Uint32(group_idx << 16) | cutlass.Uint32(tid)
+            rand_batch = philox(counter, cutlass.Uint32(seed))
+        entropy = rand_batch[intra_idx]
+        packed_i32 = cvt_f32x2_bf16x2_rs(Float32(src_lo), Float32(src_hi), entropy, loc=loc, ip=ip)
+        packed_i32_val = cutlass.Int32(packed_i32).ir_value(loc=loc, ip=ip)
+        i32_vec = vector.insert(
+            packed_i32_val,
+            i32_vec,
+            dynamic_position=[],
+            static_position=[pair_idx],
+            loc=loc,
+            ip=ip,
+        )
+    dst_vec = llvm.bitcast(dst_vec_type, i32_vec, loc=loc, ip=ip)
+    return dst_vec

build/torch-cuda/quack/sm90_utils.py CHANGED Viewed

@@ -1,6 +1,6 @@
 # Copyright (c) 2025, Tri Dao.
-from typing import Type, Union, Optional
 import cutlass
 import cutlass.cute as cute
@@ -42,6 +42,38 @@ def make_smem_layout(
 make_smem_layout_epi = make_smem_layout
 @dsl_user_op
 def partition_for_epilogue(
     cT: cute.Tensor,

 # Copyright (c) 2025, Tri Dao.
+from typing import Literal, Type, Union, Optional
 import cutlass
 import cutlass.cute as cute
 make_smem_layout_epi = make_smem_layout
+def make_tiled_mma(
+    a_dtype: Type[Numeric],
+    a_major: Literal["K", "MN"],
+    b_major: Literal["K", "MN"],
+    tiler_n: int,
+    source: Literal["SS", "RS"] = "SS",
+    atom_layout_mnk: tuple = (1, 1, 1),
+    swap_AB: bool = False,
+    b_dtype: Optional[Type[Numeric]] = None,
+    acc_dtype: Type[Numeric] = Float32,
+) -> cute.TiledMma:
+    """`b_dtype` defaults to `a_dtype`; pass it for mixed-precision MMAs (e.g. fp8).
+    `acc_dtype` defaults to Float32."""
+    if b_dtype is None:
+        b_dtype = a_dtype
+    mode = {"K": cute.nvgpu.OperandMajorMode.K, "MN": cute.nvgpu.OperandMajorMode.MN}
+    a_mode, b_mode = mode[a_major], mode[b_major]
+    if swap_AB:
+        a_mode, b_mode = b_mode, a_mode
+    a_source = warpgroup.OperandSource.RMEM if source == "RS" else warpgroup.OperandSource.SMEM
+    return sm90_utils_og.make_trivial_tiled_mma(
+        a_dtype,
+        b_dtype,
+        a_mode,
+        b_mode,
+        acc_dtype,
+        atom_layout_mnk=atom_layout_mnk,
+        tiler_mn=(64, tiler_n),
+        a_source=a_source,
+    )
 @dsl_user_op
 def partition_for_epilogue(
     cT: cute.Tensor,

build/torch-cuda/quack/utils.py ADDED Viewed

	@@ -0,0 +1,201 @@

+# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
+import math
+from typing import Optional, Tuple
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Int32, const_expr
+from cutlass._mlir.dialects import llvm, vector
+from cutlass.cutlass_dsl import T, dsl_user_op
+@dsl_user_op
+def elem_pointer(x: cute.Tensor, coord: cute.Coord, *, loc=None, ip=None) -> cute.Pointer:
+    return x.iterator + cute.crd2idx(coord, x.layout, loc=loc, ip=ip)
+@cute.jit
+def load_scalar_or_pointer(x, dtype=Float32):
+    if const_expr(isinstance(x, cute.Pointer)):
+        return dtype(cute.make_tensor(x, cute.make_layout(1))[0])
+    else:
+        return x
+@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."""
+    dsmem_ptr = cute.arch.map_dsmem_ptr(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip)
+    return Int32(dsmem_ptr.toint(loc=loc, ip=ip))
+@dsl_user_op
+def store_shared_remote(
+    val: float | Float32 | Int32 | cutlass.Int64,
+    smem_ptr: cute.Pointer,
+    mbar_ptr: cute.Pointer,
+    peer_cta_rank_in_cluster: cute.typing.Int,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    remote_smem_ptr_i32 = set_block_rank(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip)
+    remote_mbar_ptr_i32 = set_block_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip)
+    if const_expr(isinstance(val, float)):
+        val = Float32(val)
+    assert isinstance(val, (Float32, Int32, cutlass.Int64)), "val must be Float32, Int32, or Int64"
+    suffix = {Float32: "f32", Int32: "s32", cutlass.Int64: "s64"}[type(val)]
+    cute.arch.inline_ptx(
+        f"st.async.shared::cluster.mbarrier::complete_tx::bytes.{suffix} "
+        "[{$r0}], {$r1}, [{$r2}];",
+        read_only_args=[remote_smem_ptr_i32, val, remote_mbar_ptr_i32],
+        loc=loc,
+        ip=ip,
+    )
+@dsl_user_op
+def store_shared_remote_x4(
+    val0: Float32 | Int32,
+    val1: Float32 | Int32,
+    val2: Float32 | Int32,
+    val3: Float32 | Int32,
+    smem_ptr: cute.Pointer,
+    mbar_ptr: cute.Pointer,
+    peer_cta_rank_in_cluster: cute.typing.Int,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    remote_smem_ptr_i32 = set_block_rank(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip)
+    remote_mbar_ptr_i32 = set_block_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip)
+    assert isinstance(val0, (Float32, Int32)), "val must be Float32, or Int32"
+    dtype = Float32 if isinstance(val0, Float32) else Int32
+    suffix = {Float32: "f32", Int32: "s32"}[dtype]
+    cute.arch.inline_ptx(
+        "{\n\t"
+        f".reg .v4 .{suffix} abcd;\n\t"
+        f"mov.{suffix} abcd.x, {{$r2}};\n\t"
+        f"mov.{suffix} abcd.y, {{$r3}};\n\t"
+        f"mov.{suffix} abcd.z, {{$r4}};\n\t"
+        f"mov.{suffix} abcd.w, {{$r5}};\n\t"
+        f"st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.{suffix} "
+        "[{$r0}], abcd, [{$r1}];\n\t"
+        "}\n",
+        read_only_args=[
+            remote_smem_ptr_i32,
+            remote_mbar_ptr_i32,
+            dtype(val0),
+            dtype(val1),
+            dtype(val2),
+            dtype(val3),
+        ],
+        loc=loc,
+        ip=ip,
+    )
+@cute.jit
+def fill_oob(tXsX: cute.Tensor, tXpX: Optional[cute.Tensor], fill_value: cute.Numeric) -> None:
+    """Fill out-of-bounds values in shared memory tensor.
+    Args:
+        tXsX: Shared memory tensor to fill
+        tXpX: Predicate tensor indicating valid elements
+        fill_value: Value to fill OOB locations with
+    """
+    tXrX_fill = cute.make_rmem_tensor_like(tXsX[(None, 0), None, 0])
+    tXrX_fill.fill(fill_value)
+    for rest_v in cutlass.range_constexpr(tXsX.shape[0][1]):
+        for rest_k in cutlass.range_constexpr(tXsX.shape[2]):
+            if const_expr(tXpX is not None):
+                if not tXpX[rest_v, 0, rest_k]:
+                    cute.autovec_copy(tXrX_fill, tXsX[(None, rest_v), None, rest_k])
+            else:
+                cute.autovec_copy(tXrX_fill, tXsX[(None, rest_v), None, rest_k])
+# ---------------------------------------------------------------------------
+# General-purpose DSL store / vector helpers
+# ---------------------------------------------------------------------------
+@dsl_user_op
+def make_vector(elem_type, *values, loc=None, ip=None):
+    """Build an MLIR vector <N x elem_type> from N scalar DSL values.
+    Example: make_vector(cutlass.Uint32, v0, v1) -> <2 x i32> MLIR vector
+    """
+    from cutlass._mlir import ir
+    n = len(values)
+    mlir_ty = elem_type.mlir_type
+    vec_ty = ir.VectorType.get([n], mlir_ty)
+    vec = llvm.mlir_undef(vec_ty, loc=loc, ip=ip)
+    for i, v in enumerate(values):
+        vec = vector.insert(
+            elem_type(v).ir_value(loc=loc, ip=ip),
+            vec,
+            dynamic_position=[],
+            static_position=[i],
+            loc=loc,
+            ip=ip,
+        )
+    return vec
+@dsl_user_op
+def f32x2_to_i64(a: Float32, b: Float32, *, loc=None, ip=None) -> cutlass.Int64:
+    vec_f32x2 = vector.from_elements(
+        T.vector(2, T.f32()), (a.ir_value(), b.ir_value()), loc=loc, ip=ip
+    )
+    vec_i64x1 = vector.bitcast(T.vector(1, T.i64()), vec_f32x2)
+    res = cutlass.Int64(
+        vector.extract(vec_i64x1, dynamic_position=[], static_position=[0], loc=loc, ip=ip)
+    )
+    return res
+@dsl_user_op
+def i64_to_f32x2(c: cutlass.Int64, *, loc=None, ip=None) -> Tuple[Float32, Float32]:
+    vec_i64x1 = vector.from_elements(T.vector(1, T.i64()), (c.ir_value(),), loc=loc, ip=ip)
+    vec_f32x2 = vector.bitcast(T.vector(2, T.f32()), vec_i64x1)
+    res0 = Float32(
+        vector.extract(vec_f32x2, dynamic_position=[], static_position=[0], loc=loc, ip=ip)
+    )
+    res1 = Float32(
+        vector.extract(vec_f32x2, dynamic_position=[], static_position=[1], loc=loc, ip=ip)
+    )
+    return res0, res1
+@cute.jit
+def warp_prefix_sum(val: Int32, lane: Optional[Int32] = None) -> Int32:
+    if const_expr(lane is None):
+        lane = cute.arch.lane_idx()
+    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
+    return val
+@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)