diff --git a/README.md b/README.md
index 8f679dd51b88bd69102391f3129669a5dc808517..a262fc28298fb339b675e191083adfb251724801 100644
--- a/README.md
+++ b/README.md
@@ -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
+kernel_module = get_kernel("kernels-community/flash-attn4", version=1)
+_private_for_testing = kernel_module._private_for_testing
 
-flash_attn_func(...)
+_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`.
diff --git a/build/torch-cuda/__init__.py b/build/torch-cuda/__init__.py
index 563d5a7752c613f31f9d9d3f78ecfa20d8ae4f70..b978bed15ee7cc47dd0a26dba0a9bd19ea95aa04 100644
--- a/build/torch-cuda/__init__.py
+++ b/build/torch-cuda/__init__.py
@@ -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
+__version__ = "4.0.0.beta30"
 
 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
-
+# 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",
 ]
diff --git a/build/torch-cuda/_ops.py b/build/torch-cuda/_ops.py
index 2845815a5f29f2755156b0be676b74d2e96d96e8..48b3e9b7adf32a9f69826149d7afc5e382f8db04 100644
--- a/build/torch-cuda/_ops.py
+++ b/build/torch-cuda/_ops.py
@@ -1,8 +1,44 @@
 import torch
-ops = torch.ops._flash_attn4_06b9ab8
 
-def add_op_namespace_prefix(op_name: str):
+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"_flash_attn4_06b9ab8::{op_name}"
\ No newline at end of file
+    return f"{_OPS_NAME}::{op_name}"
\ No newline at end of file
diff --git a/build/torch-cuda/_private_for_testing.py b/build/torch-cuda/_private_for_testing.py
new file mode 100644
index 0000000000000000000000000000000000000000..1b618eadf584939205205cd0e37d345e8feea2c2
--- /dev/null
+++ b/build/torch-cuda/_private_for_testing.py
@@ -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",
+]
diff --git a/build/torch-cuda/benchmark.py b/build/torch-cuda/benchmark.py
deleted file mode 100644
index 9a7820e7b0c7c70c3fea1c4239d8785561790de1..0000000000000000000000000000000000000000
--- a/build/torch-cuda/benchmark.py
+++ /dev/null
@@ -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
diff --git a/build/torch-cuda/blackwell_helpers.py b/build/torch-cuda/blackwell_helpers.py
index 3c2663936f12d4399788b874966cf05bb1941fce..2594dbbcd23b9b79fc2420bed065b57a567503e2 100644
--- a/build/torch-cuda/blackwell_helpers.py
+++ b/build/torch-cuda/blackwell_helpers.py
@@ -10,6 +10,26 @@ from cutlass._mlir.dialects import llvm
 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,
@@ -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))
+    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(
@@ -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"
+                    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,
@@ -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"
+                    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,
@@ -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))
+    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(
@@ -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"
+            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::f16 [$0], smem_desc_a, smem_desc_b, idesc, 1;\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]))
             )
@@ -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"
+            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::f16 [$0], [tmem_a + {hex(offset_a[k])}], 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]))
             )
@@ -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))
+    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(
@@ -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"
+            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"
@@ -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"
+                    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]))
             )
@@ -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"
+            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"
@@ -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"
+                    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,
@@ -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"
+                        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]))
                 )
@@ -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))
+    kind = _tcgen05_mma_kind(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"
+            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::f16 [tmem_acc], smem_desc_a, smem_desc_b, idesc, {{$5, $6, $7, $8}}, 1;\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]))
             )
@@ -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"
+            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::f16 [$0], [tmem_a], smem_desc_b, idesc, {{$4, $5, $6, $7}}, 1;\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]))
             )
@@ -783,6 +807,7 @@ def gemm_ptx_precomputed(
     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)
@@ -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"
+            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"
@@ -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"
+                    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)
             )
@@ -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"
+            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"
@@ -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"
+                    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,
@@ -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"
+                        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)
                 )
@@ -1019,6 +1044,7 @@ def gemm_ptx_precomputed_varname(
     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])
@@ -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"
+            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"
@@ -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"
+                    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)
             )
diff --git a/build/torch-cuda/block_info.py b/build/torch-cuda/block_info.py
index a5a2544a7883f3719953ec91ee246fdd41faa8b1..d2c860c71ab4612c94260490a36c4d07ba85bafd 100644
--- a/build/torch-cuda/block_info.py
+++ b/build/torch-cuda/block_info.py
@@ -6,7 +6,7 @@ import cutlass
 import cutlass.cute as cute
 from cutlass import Int32, const_expr
 
-from .seqlen_info import SeqlenInfoQK
+from .seqlen_info import SeqlenInfoQK, SeqlenInfoQKNewK
 
 
 @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
+    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: cutlass.Int32 = 0,
-        num_splits: cutlass.Int32 = 1,
+        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)):
@@ -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
-            )
+            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
@@ -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_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,
@@ -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))
+
+    @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
diff --git a/build/torch-cuda/block_sparse_utils.py b/build/torch-cuda/block_sparse_utils.py
index 10213d571f853a2f95fdfe9c241e9f7d916e1e92..5823090fb85850a27b5772c8aa5afbfcfc0da036 100644
--- a/build/torch-cuda/block_sparse_utils.py
+++ b/build/torch-cuda/block_sparse_utils.py
@@ -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 typing import Callable, Optional, Tuple
 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
+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]
@@ -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).
+# - 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,
-    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,
+    kv_subtile_factor: cutlass.Constexpr[int] = 1,
 ):
-    """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
+    """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.
 
@@ -98,44 +188,40 @@ def load_block_list(
 
     """
     if block_count > 0:
+        total_blocks = block_count * kv_subtile_factor
         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]
+            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 = block_indices[block_count - 1]
+            n_block_first = sparse_physical_n_block_forward(
+                block_indices,
+                total_blocks - 1,
+                kv_subtile_factor,
+            )
             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))
-
+                pipeline_k.producer_acquire(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]
+            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)
@@ -153,10 +239,11 @@ def finish_overlap_v_load(
     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]
+        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()
@@ -185,20 +272,21 @@ def produce_block_sparse_loads(
     batch_idx,
     head_idx,
     m_block,
+    seqlen_info: SeqlenInfoQK,
     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,
+    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.
@@ -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
-
+    # 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 = 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
+    (
+        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 Q+K load.
+        # 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,
-            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,
+            kv_subtile_factor=kv_subtile_factor,
         )
 
         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,
+                kv_subtile_factor=kv_subtile_factor,
             )
     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.
+        # 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,
-            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,
+            kv_subtile_factor=kv_subtile_factor,
         )
 
         if full_empty:
@@ -282,13 +369,18 @@ def produce_block_sparse_loads(
                     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]
-                n_block_full_first = curr_full_block_idx[curr_full_block_cnt - 1]
+                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)
@@ -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_subtile_factor=kv_subtile_factor,
                 )
 
                 kv_producer_state = finish_overlap_v_load(
@@ -318,24 +407,21 @@ def produce_block_sparse_loads(
                     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 (skipping the Q
-                # reload because the masked list already issued it).
+                # 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,
-                    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,
+                    kv_subtile_factor=kv_subtile_factor,
                 )
 
     return kv_producer_state
@@ -347,7 +433,7 @@ def consume_block_sparse_loads(
     batch_idx,
     head_idx,
     m_block,
-    seqlen,
+    seqlen_info,
     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]
+    (
+        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
 
@@ -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,
+                seqlen=seqlen_info,
                 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,
+                    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),
                 )
@@ -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,
+                    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,
@@ -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,
+                    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),
                 )
@@ -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,
+                    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),
                 )
@@ -525,10 +616,20 @@ def consume_block_sparse_loads(
     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_count,
+    block_begin,
+    block_end,
     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,
+    kv_subtile_factor: cutlass.Constexpr[int] = 1,
 ):
-    """SM100 version of load_block_list (no intra_wg_overlap, no extra_tx_count)."""
-    if block_count > 0:
+    """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 = block_indices[block_count - 1]
+        n_block_first = physical_n_block(block_end - 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]
+        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)
@@ -573,6 +679,9 @@ def produce_block_sparse_loads_sm100(
     batch_idx,
     head_idx,
     m_block,
+    seqlen_info: SeqlenInfoQK,
+    split_idx: Int32,
+    num_splits: Int32,
     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,
+    kv_subtile_factor: cutlass.Constexpr[int] = 1,
 ):
     """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
+    (
+        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
+    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
 
@@ -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,
+            full_begin,
+            full_end,
             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,
+            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,
-            curr_mask_block_cnt,
+            mask_begin,
+            mask_end,
             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,
+            kv_subtile_factor=kv_subtile_factor,
         )
         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,
+                full_begin,
+                full_end,
                 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,
+                kv_subtile_factor=kv_subtile_factor,
             )
 
     if q_phase_flipped:
@@ -666,19 +794,39 @@ def get_total_block_count(
     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_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]
+    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
@@ -691,7 +839,7 @@ def handle_block_sparse_empty_tile_correction_sm100(
     is_split_kv: cutlass.Constexpr,
     learnable_sink,
     mLSE,
-    seqlen,
+    seqlen_info,
     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,
+    thr_mma_pv: cute.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,
+    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,
@@ -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)
+                    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, fastmath=True
+                        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
@@ -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)
+
+        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.seqlen_q,
+            seqlen_info.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],
+            gO_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_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,
@@ -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,
+    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)
 
-    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
+    (
+        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,
+    )
 
-    total_block_cnt = curr_mask_block_cnt + curr_full_block_cnt
+    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:
-        # 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]
+        if split_mask_block_cnt > 0:
             (
                 mma_si_consumer_phase,
                 si_corr_producer_phase,
                 s0_s1_sequence_phase,
-            ) = softmax_step(
+            ) = 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,
-                mask_n_block,
-                is_first=True,
-                mask_fn=partial(mask_fn, mask_seqlen=True, check_q_boundary=check_m_boundary),
+                check_m_boundary,
+                True,
+                False,
             )
-            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
-                    ),
-                )
+        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,
@@ -940,24 +1161,24 @@ def get_total_q_block_count_bwd(
     batch_idx,
     head_idx,
     n_block,
-    subtile_factor: cutlass.Constexpr = 1,
+    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
+    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
+    return total * q_subtile_factor
 
 
 @cute.jit
-def produce_block_sparse_q_loads_bwd_sm100(
+def produce_block_sparse_q_loads_bwd_sm100_default(
     blocksparse_tensors: BlockSparseTensors,
     batch_idx,
     head_idx,
     n_block,
-    # Pipeline states (will be returned after advancing)
+    # Pipeline states 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,
+    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,
 ):
-    """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.
-    """
+    """Produce SM100 backward block-sparse Q/dO loads for 1CTA and non-hdim192 2CTA."""
     (
         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
+        blocksparse_tensors, batch_idx, head_idx, n_block, q_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,
+    # 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,
-            subtile_factor,
-            m_block_max,
+            q_subtile_factor=q_subtile_factor,
+            m_block_max=m_block_max,
         )
-        m_block_safe = m_block
+        # with q_subtile > 1 we need to guard against fully OOB regions
         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
+            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),
                 )
-                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)
+            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_safe, producer_state=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_safe],
+                        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)
-                load_dO(m_block_safe, producer_state=producer_state_dO_dPsum)
+                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_safe],
+                        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()
 
-    return producer_state_Q_LSE, producer_state_dO_dPsum
+            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
@@ -1081,14 +1552,14 @@ def get_block_sparse_iteration_info_bwd(
     batch_idx,
     head_idx,
     n_block,
-    subtile_factor: cutlass.Constexpr = 1,
+    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
+    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
+    total_count = sparse_block_count * q_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,
+    q_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_iter_idx = iter_idx // q_subtile_factor
+    subtile_offset = iter_idx % q_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
+    return sparse_m_block * q_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)
+    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 = (
@@ -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,
+    q_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
+    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
+    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(
@@ -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
+        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(
@@ -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,
+    q_subtile_factor: cutlass.Constexpr = 1,
     m_block_max: int = 0,
-    aux_tensors=None,
+    aux_data: AuxData = AuxData(),
     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
+    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,
+        aux_data=aux_data,
         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,
+        aux_data=aux_data,
         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
+    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(
@@ -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
+        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(
@@ -1390,18 +1868,18 @@ def _store_one_dQaccum_sm90(
     m_block,
     sdQaccum: cute.Tensor,
     gdQaccum: cute.Tensor,
-    num_mma_warp_groups: cutlass.Constexpr,
+    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_mma_warp_groups):
-        cute.arch.cp_async_bulk_wait_group(num_mma_warp_groups - 1 - warp_group_idx, read=True)
+    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_mma_warp_groups):
+    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,
@@ -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,
+                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,
+    q_subtile_factor: cutlass.Constexpr,
     m_block_max: int,
-    num_mma_warp_groups: cutlass.Constexpr,
+    num_dQ_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
+    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
+    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_mma_warp_groups,
+                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 * 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
+        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_mma_warp_groups,
+                    num_dQ_warp_groups,
                     num_threads_per_warp_group,
                     tma_copy_bytes_dQ,
                 )
diff --git a/build/torch-cuda/block_sparsity.py b/build/torch-cuda/block_sparsity.py
index 8ba66e95434f7b856a3019eb7ae0c5c377eef895..1c71f515fdac0f05f3e98c26fb2d699729acb883 100644
--- a/build/torch-cuda/block_sparsity.py
+++ b/build/torch-cuda/block_sparsity.py
@@ -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
+    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):
-        if len(values) == 2:
-            values = (*values, None, None)
-        return BlockSparseTensors(*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):
@@ -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
+    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(
@@ -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],
+    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]:
@@ -81,12 +268,37 @@ def _check_and_expand_block(
     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,
@@ -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_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.
@@ -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 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:
-        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:
+        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}."
             )
-        sparse_block_size_q = min_block_size
+        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(
@@ -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:
+    # [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}."
+            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,
+    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 subtile_factor * m_block_size.
+    by q_subtile_factor * m_block_size.
     """
-    sparse_block_size_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, n_block_size)
+    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
@@ -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],
+    expected_count_shape: Tuple[int, ...],
+    expected_index_shape: Tuple[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")
 
+    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,
     )
 
 
@@ -294,6 +536,8 @@ def get_block_sparse_broadcast_pattern(
         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))
@@ -302,6 +546,15 @@ def get_block_sparse_broadcast_pattern(
     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,
     *,
@@ -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]:
+    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 = q_stage * m_block_size, n_block_size
+        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:
+    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."
         )
-    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,
+    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 (
-        normalized_tensors,
-        get_block_sparse_broadcast_pattern(normalized_tensors),
-        q_subtile_factor,
+    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,
     )
 
 
@@ -356,21 +636,24 @@ def normalize_block_sparse_config_bwd(
     seqlen_q: int,
     seqlen_k: int,
     block_size: tuple[int, int],
-    subtile_factor: 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 = subtile_factor * m_block_size, n_block_size
+        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 != subtile_factor * m_block_size:
+    if sparse_block_size_q != 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}."
+            f"Block sparsity expects sparse_block_size_q={q_subtile_factor * m_block_size} "
+            f"for q_subtile_factor={q_subtile_factor}."
         )
-    if sparse_block_size_kv != n_block_size:
+    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]={n_block_size} to match tile_n."
+            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,
@@ -379,7 +662,8 @@ def normalize_block_sparse_config_bwd(
         seqlen_k,
         m_block_size,
         n_block_size,
-        subtile_factor,
+        q_subtile_factor,
+        kv_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})."
+            f"BLOCK_SIZE=({q_subtile_factor * m_block_size}, {expected_sparse_block_size_kv})."
         ),
     )
     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,
-    ) = [
+    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)
+        for t in (tensors.mask_block_cnt, tensors.mask_block_idx)
     ]
-    (
-        full_block_cnt_tensor,
-        full_block_idx_tensor,
-    ) = [
+    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 (full_block_cnt, full_block_idx)
+        for t in (tensors.dq_write_order, tensors.dq_write_order_full)
     ]
 
     return BlockSparseTensors(
@@ -431,6 +713,10 @@ def to_cute_block_sparse_tensors(
         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,
     )
 
 
diff --git a/build/torch-cuda/cache_utils.py b/build/torch-cuda/cache_utils.py
index 8606f04b62b74455e838cbb83bbb6e0aa5b85515..e494911d2ec3fb4a8cbac6584479c5c49141b555 100644
--- a/build/torch-cuda/cache_utils.py
+++ b/build/torch-cuda/cache_utils.py
@@ -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
+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.
@@ -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)
+                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:
-                logger.debug("Cache miss on disk for key hash %s", sha256_hex)
+                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:
@@ -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)
+                fa_log(1, f"Skipping export, already on disk: {obj_path}")
                 return
-            logger.debug("Exporting compiled function to disk: %s", obj_path)
+            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,
             )
-            logger.debug("Successfully exported compiled function to disk: %s", obj_path)
+            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()
@@ -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)
+        fa_log(1, f"Clearing persistent cache at {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)
+        fa_log(1, f"Creating persistent JIT cache at {path}")
         return JITPersistentCache(path)
     else:
-        logger.debug("Persistent cache disabled, using in-memory JIT cache")
+        fa_log(1, "Persistent cache disabled, using in-memory JIT cache")
         return JITCache()
diff --git a/build/torch-cuda/compute_block_sparsity.py b/build/torch-cuda/compute_block_sparsity.py
index 73e684fa6943765c5bbe1708396d9c426c0f2226..13c0272846beb593643fdb355170c55108568334 100644
--- a/build/torch-cuda/compute_block_sparsity.py
+++ b/build/torch-cuda/compute_block_sparsity.py
@@ -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 .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:
@@ -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,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        mSeqUsedQ: Optional[cute.Tensor] = None,
+        mSeqUsedK: Optional[cute.Tensor] = None,
+        aux_data: AuxData = AuxData(),
     ):
-        self.mask_cnt, self.mask_idx, self.full_cnt, self.full_idx = blocksparse_tensors
+        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 self.full_cnt is not None and self.full_idx is not None, (
+            assert full_cnt is not None and 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(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
@@ -72,46 +98,46 @@ class BlockSparsityKernel:
             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(
-            self.mask_cnt,
-            self.mask_idx,
-            self.full_cnt,
-            self.full_idx,
-            num_n_blocks,
+            blocksparse_tensors,
             seqlen_q,
             seqlen_k,
-            aux_tensors,
+            batch_size,
+            mCuSeqlensQ,
+            mCuSeqlensK,
+            mSeqUsedQ,
+            mSeqUsedK,
+            mCuTotalMBlocks,
+            mCuBlockIdxOffsets,
+            aux_data,
         ).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,
+        blocksparse_tensors: BlockSparseTensors,
         seqlen_q: Int32,
         seqlen_k: Int32,
-        aux_tensors: Optional[list] = None,
+        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()
-        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))
         )
+        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)
 
-        for n_block in cutlass.range(num_n_blocks, unroll_full=True):
-            m_base = m_block * self.tile_mn[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):
-                # 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)
+                # 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
 
-                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)
+                kv_idx = n_base
+                if tidx == 1 or tidx == 3:
+                    kv_idx = n_right
                 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)
+                    kv_idx = n_mid
 
-                # 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_result = Boolean(False)
+                thread_is_valid = Boolean(False)
+                if tidx < 5:
                     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(
+                        call_mask_mod(
+                            self.mask_mod,
                             ssa(batch_idx),
                             ssa(head_idx),
-                            q_idx_ssa,
-                            kv_idx_ssa,
+                            ssa(q_idx_sample),
+                            ssa(kv_idx),
                             seqlen,
-                            aux_tensors,
+                            aux_data,
                         )
                     )
-                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)
+                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
+                # 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)
-                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)
+                is_interior = (n_base + self.tile_mn[1]) <= seqlen_k
 
-                    # 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
+                if is_interior:
+                    if thread_in_bounds:
+                        for c in cutlass.range(self.tile_mn[1], unroll_full=True):
                             mask_val = ssa_to_scalar(
-                                self.mask_mod(
+                                call_mask_mod(
+                                    self.mask_mod,
                                     ssa(batch_idx),
                                     ssa(head_idx),
-                                    q_idx_ssa,
-                                    kv_idx_ssa,
+                                    ssa(q_idx_thread),
+                                    ssa(n_base + c),
                                     seqlen,
-                                    aux_tensors,
+                                    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)
 
-                            # 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
+                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:
-                    # 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)
-
+                    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()
 
-                # Thread 0 ORs all warp results together
+                # 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 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)
+                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:
@@ -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
+                    curr_mask_idx[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
+                    curr_full_idx[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
+            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(
@@ -282,11 +339,18 @@ def compute_block_sparsity(
     seqlen_q,
     seqlen_k,
     mask_mod: Callable,
-    aux_tensors: Optional[list],  # list[cute.Tensor]
+    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,
-) -> Tuple[BlockSparseTensors, BlockSparseTensorsTorch]:
+) -> 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.
+        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:
-        A tuple of `BlockSparseTensors` and `BlockSparseTensorsTorch`.
+        BlockSparseTensorsTorch
     """
-    # Check if mask_mod is marked as suitable for 5-point fast sampling
+    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
 
-    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(
+    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
         )
-        if compute_full_blocks
-        else None
-    )
+        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)
-    blocksparse_tensors = to_cute_block_sparse_tensors(
-        blocksparse_tensors_torch, enable_tvm_ffi=True
-    )
+    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),
@@ -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"
+            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",
         )
 
-    compute_block_sparsity.compile_cache[compile_key](
-        blocksparse_tensors_torch[:4],
-        seqlen_q,
-        seqlen_k,
-        aux_tensors,
-    )
+    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, blocksparse_tensors_torch
+    return blocksparse_tensors_torch
 
 
 compute_block_sparsity.compile_cache = {}
diff --git a/build/torch-cuda/cu_blocks_kernel.py b/build/torch-cuda/cu_blocks_kernel.py
new file mode 100644
index 0000000000000000000000000000000000000000..58946e7e7397bd76f23fa1e6a80941ceac80eb9c
--- /dev/null
+++ b/build/torch-cuda/cu_blocks_kernel.py
@@ -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)
diff --git a/build/torch-cuda/cute_dsl_utils.py b/build/torch-cuda/cute_dsl_utils.py
index 0cf0b605be35bc955d31d90ee86546641def4cf8..faa25808d5261ed0b7a7d8bbf6e10511dea3a5d4 100644
--- a/build/torch-cuda/cute_dsl_utils.py
+++ b/build/torch-cuda/cute_dsl_utils.py
@@ -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
+from functools import lru_cache
 
 import torch
+from torch._subclasses.fake_tensor import FakeTensor
 
 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,
+    torch.float8_e4m3fn: cutlass.Float8E4M3FN,
+    torch.float8_e5m2: cutlass.Float8E5M2,
 }
 
 
@@ -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
+@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)
         )
-    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
+    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):
@@ -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 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:
@@ -146,15 +173,37 @@ def to_cute_aux_tensor(t, enable_tvm_ffi=True):
     )
 
 
+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 tuple(
-        (
-            getattr(t, "__assumed_align__", 0),
-            getattr(t, "__leading_dim__", -1),
-            hasattr(t, "__leading_dim__"),
+    """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())
         )
-        for t in aux_tensors
-    )
+        metadata.append((tensor.dtype, getattr(tensor, "__assumed_align__", None), static_strides))
+    return tuple(metadata)
 
 
 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())
+    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 ---")
diff --git a/build/torch-cuda/fa_logging.py b/build/torch-cuda/fa_logging.py
new file mode 100644
index 0000000000000000000000000000000000000000..63189cd5d65adba3eb303d8db8757243ced1d0a3
--- /dev/null
+++ b/build/torch-cuda/fa_logging.py
@@ -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)
diff --git a/build/torch-cuda/flash_attn4/__init__.py b/build/torch-cuda/flash_attn4/__init__.py
deleted file mode 100644
index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000
--- a/build/torch-cuda/flash_attn4/__init__.py
+++ /dev/null
@@ -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")))
diff --git a/build/torch-cuda/flash_bwd.py b/build/torch-cuda/flash_bwd.py
index 824081219cff37f2ec005b374bde2f4e2a53d68e..b012c51ae2b2c32a790d5d29ed149676763a8c14 100644
--- a/build/torch-cuda/flash_bwd.py
+++ b/build/torch-cuda/flash_bwd.py
@@ -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
+from cutlass import 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 .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:
@@ -38,6 +42,7 @@ 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,
@@ -45,6 +50,8 @@ class FlashAttentionBackwardSm80:
         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.
 
@@ -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
+        # 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
@@ -77,6 +85,7 @@ class FlashAttentionBackwardSm80:
         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
@@ -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
+        self.score_mod = score_mod
+        self.score_mod_bwd = score_mod_bwd
 
     @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]):
+        if cutlass.const_expr(mQ_type not in [cutlass.Float16, cutlass.BFloat16]):
             raise TypeError("Only Float16 or BFloat16 is supported")
-        if cutlass.const_expr(not mLSE_type in [cutlass.Float32]):
+        if cutlass.const_expr(mLSE_type not in [cutlass.Float32]):
             raise TypeError("LSE tensor must be Float32")
-        if cutlass.const_expr(not mdPsum_type in [cutlass.Float32]):
+        if cutlass.const_expr(mdPsum_type not in [cutlass.Float32]):
             raise TypeError("dPsum tensor must be Float32")
-        if cutlass.const_expr(not mdQaccum_type in [cutlass.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")
@@ -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,
+        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, "semaphore not supported yet"
+        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)))
@@ -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,
+            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 = softmax_scale * math.log2(math.e)
+        softmax_scale_log2, _ = utils.compute_softmax_scale_log2(
+            softmax_scale, self.score_mod
+        )
         self.kernel(
             mQ,
             mK,
@@ -439,6 +460,8 @@ class FlashAttentionBackwardSm80:
             mSeqUsedK,
             softmax_scale,
             softmax_scale_log2,
+            window_size_left,
+            window_size_right,
             self.sQ_layout,
             self.sK_layout,
             self.sV_layout,
@@ -458,6 +481,7 @@ class FlashAttentionBackwardSm80:
             SharedStorage,
             tile_sched_params,
             TileScheduler,
+            aux_data,
         ).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,
+        window_size_left: Optional[Int32],
+        window_size_right: Optional[Int32],
         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],
+        aux_data: AuxData = AuxData(),
     ):
         # 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,
-                )
+            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
 
             # ///////////////////////////////////////////////////////////////////////////////
@@ -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
+                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])
@@ -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 = 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)
 
@@ -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=softmax_scale,
                 softmax_scale_log2=softmax_scale_log2,
+                aux_data=aux_data,
             )
 
-            # ///////////////////////////////////////////////////////////////////////////////
-            # 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)
+            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()
 
-            # ///////////////////////////////////////////////////////////////////////////////
-            # 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,
+                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 = 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)
+                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: cutlass.Float32,
         softmax_scale_log2: cutlass.Float32,
+        aux_data: AuxData = AuxData(),
         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 = 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()
@@ -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,
         )
+        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)
-        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 = 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()
@@ -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]))
+            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))
@@ -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 = 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,
@@ -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)
+            tdKrdS = layout_utils.reshape_acc_to_frgA(rdS)
         else:
             tdKrdS = mma_params.tdKrdS
         sm80_utils.gemm(
diff --git a/build/torch-cuda/flash_bwd_mla_dk_sm100.py b/build/torch-cuda/flash_bwd_mla_dk_sm100.py
new file mode 100644
index 0000000000000000000000000000000000000000..9b4071b75d5f837811095f60c2a60bb35bdb2f53
--- /dev/null
+++ b/build/torch-cuda/flash_bwd_mla_dk_sm100.py
@@ -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=}"
+            )
diff --git a/build/torch-cuda/flash_bwd_mla_dq_dqv_sm100.py b/build/torch-cuda/flash_bwd_mla_dq_dqv_sm100.py
new file mode 100644
index 0000000000000000000000000000000000000000..23fe6005cab0ebce155149e29244deef29bd585d
--- /dev/null
+++ b/build/torch-cuda/flash_bwd_mla_dq_dqv_sm100.py
@@ -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
diff --git a/build/torch-cuda/flash_bwd_mla_sm100.py b/build/torch-cuda/flash_bwd_mla_sm100.py
new file mode 100644
index 0000000000000000000000000000000000000000..fc061a6ed4bdc18dd40d80264f75471da9ba36f4
--- /dev/null
+++ b/build/torch-cuda/flash_bwd_mla_sm100.py
@@ -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()
diff --git a/build/torch-cuda/flash_bwd_postprocess.py b/build/torch-cuda/flash_bwd_postprocess.py
index 684045652ac4f6f374542ce52899f37a18a89e59..2aa0553b048b411d70d6fcb0314c91d02de350ef 100644
--- a/build/torch-cuda/flash_bwd_postprocess.py
+++ b/build/torch-cuda/flash_bwd_postprocess.py
@@ -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 operator
+from typing import Callable, NamedTuple, Optional, Type
 
 import cuda.bindings.driver as cuda
 
@@ -31,39 +32,49 @@ from .tile_scheduler import (
 )
 
 
+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: Literal[80, 90, 100],
+        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], (
-            "Only Ampere (8.x), Hopper (9.x), and Blackwell (10.x, 11.x) are supported"
+        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
-        # 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.use_2cta_instrs = use_2cta_instrs and arch // 10 in [10, 11] 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):
+        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)
@@ -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)
+        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,
@@ -125,7 +136,7 @@ class FlashAttentionBackwardPostprocess:
                 cta_group,
                 (self.tile_m, self.tile_hdim),
             )
-        if const_expr(self.arch in [80, 90]):
+        if const_expr(self.arch // 10 in [8, 9, 12]):
             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):
+        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 == 90):
+        elif const_expr(self.arch // 10 == 9):
             num_threads_per_warp_group = 128
-            num_mma_warp_groups = self.num_threads // 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_mma_warp_groups)),  # thr_layout
+                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_mma_warp_groups, num_mma_warp_groups)
+                (self.tile_m * self.tile_hdim // num_wg_mma, num_wg_mma)
             )
         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):
+        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 == 90):
+        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)
+                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
@@ -211,7 +226,10 @@ class FlashAttentionBackwardPostprocess:
         scale: cutlass.Float32,
         mCuSeqlensQ: Optional[cute.Tensor],
         mSeqUsedQ: Optional[cute.Tensor],
-        stream: cuda.CUstream,
+        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]):
@@ -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")
+        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)]
 
@@ -253,6 +284,7 @@ class FlashAttentionBackwardPostprocess:
             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)
@@ -264,6 +296,7 @@ class FlashAttentionBackwardPostprocess:
             mdQ,
             mCuSeqlensQ,
             mSeqUsedQ,
+            sink_tensors,
             scale,
             self.tiled_mma,
             self.dQ_swapAB,
@@ -288,6 +321,7 @@ class FlashAttentionBackwardPostprocess:
         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,
@@ -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]):
+        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
@@ -323,6 +357,73 @@ class FlashAttentionBackwardPostprocess:
 
         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.
@@ -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
+                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]
@@ -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):
+            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)
@@ -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_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(
@@ -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)
+                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
@@ -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]):
+                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_fragment(acc_shape, cutlass.Float32)
+                    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
@@ -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)
+                    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
@@ -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]):
+                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
                     )
@@ -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]):
+                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
diff --git a/build/torch-cuda/flash_bwd_preprocess.py b/build/torch-cuda/flash_bwd_preprocess.py
index 629888989082917b05152b1e15228d1ed27dc943..12726a38e488d14fa08737531bcdd45981334f9b 100644
--- a/build/torch-cuda/flash_bwd_preprocess.py
+++ b/build/torch-cuda/flash_bwd_preprocess.py
@@ -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.
+#
+# 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, Literal
+from typing import Callable, Type, Optional
 
 import cuda.bindings.driver as cuda
 
 import cutlass
 import cutlass.cute as cute
-from cutlass import Float32
+from cutlass import Float32, const_expr
+from cutlass.cutlass_dsl import Arch, BaseDSL
 
-from .quack import copy_utils
+from .quack import copy_utils, layout_utils
 
 from . import utils
-from .cute_dsl_utils import assume_tensor_aligned
-from .seqlen_info import SeqlenInfoQK
+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:
@@ -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,
+        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
@@ -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 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.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.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, m_block_size, num_threads) -> bool:
+    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 m_block_size: m block size
-        :type m_block_size: int
+        :param tile_m: m block size
+        :type tile_m: 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
+        if num_threads < tile_m:  # 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.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
@@ -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],
+        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],
-        mSeqUsedQ: Optional[cute.Tensor],
-        stream: cuda.CUstream,
+        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 cutlass.const_expr(not (mO.element_type == mdO.element_type)):
+        if 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]):
+        if 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]):
+        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 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]):
+        if const_expr(mLSE is not None):
+            if 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]):
+        if const_expr(mLSElog2 is not None):
+            if 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)]
+        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()
 
-        if cutlass.const_expr(mCuSeqlensQ is not None):
+        # (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[1]
+            num_head = mO.shape[2]
             num_batch = mCuSeqlensQ.shape[0] - 1
         else:
             TileScheduler = SingleTileScheduler
             num_head = mO.shape[2]
-            num_batch = mO.shape[0]
+            num_batch = mO.shape[3]
 
         tile_sched_args = TileSchedulerArguments(
-            num_block=cute.ceil_div(mO.shape[1], self.m_block_size),
+            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[2],
-            total_q=mO.shape[0],
-            tile_shape_mn=(self.m_block_size, 1),
+            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,
-            mdPsum,
+            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,
@@ -188,6 +283,7 @@ class FlashAttentionBackwardPreprocess:
             grid=grid_dim,
             block=[self.num_threads, 1, 1],
             stream=stream,
+            use_pdl=self.use_pdl,
         )
 
     @cute.kernel
@@ -195,12 +291,16 @@ class FlashAttentionBackwardPreprocess:
         self,
         mO: cute.Tensor,
         mdO: cute.Tensor,
-        mdPsum: 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,
@@ -213,149 +313,159 @@ class FlashAttentionBackwardPreprocess:
         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 = SeqlenInfoQK.create(
-                batch_idx,
-                mO.shape[1],
-                0,
-                mCuSeqlensQ=mCuSeqlensQ,
-                mCuSeqlensK=None,
-                mSeqUsedQ=mSeqUsedQ,
-                mSeqUsedK=None,
+            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
 
-            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))
+            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)
-
-            # ///////////////////////////////////////////////////////////////////////////////
-            # 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))
+            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 = 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])
+            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
-                # (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,
-                    )
+                # 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
-            dpsum = (tOrO.load().to(Float32) * tOrdO.load().to(Float32)).reduce(
+            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
-            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
+            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(dP_sum), unroll_full=True):
+                for m in cutlass.range(cute.size(PdP_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
+                    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 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,)
+            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_fragment_like(tdQgdQaccum)
+                zero = cute.make_rmem_tensor_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)
+            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.m_block_size:
-                    gLSElog2[tidx] = lse * LOG2_E if lse != -Float32.inf else 0.0
+                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
diff --git a/build/torch-cuda/flash_bwd_sm100.py b/build/torch-cuda/flash_bwd_sm100.py
index 1e99dbab83126b538d459edf746bd51d236735cb..fadb071a4ce48f552fb87c4148a9ec6fe282cea4 100644
--- a/build/torch-cuda/flash_bwd_sm100.py
+++ b/build/torch-cuda/flash_bwd_sm100.py
@@ -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 .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,
+    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
 
@@ -58,22 +78,28 @@ class FlashAttentionBackwardSm100:
         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,
-        subtile_factor: cutlass.Constexpr[int] = 1,
+        q_subtile_factor: cutlass.Constexpr[int] = 1,
+        kv_subtile_factor: cutlass.Constexpr[int] = 1,
     ):
-        # padding head_dim to a multiple of 16 as k_block_size
-        hdim_multiple_of = 16
+        # 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
@@ -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.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"
@@ -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)
+        # 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.subtile_factor = subtile_factor
+        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
@@ -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
 
+        # 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
@@ -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 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
@@ -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
+        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
@@ -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,
+        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
@@ -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)]
+        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]
@@ -706,7 +751,11 @@ class FlashAttentionBackwardSm100:
             TileScheduler = SingleTileLPTBwdScheduler
         else:
             TileScheduler = SingleTileScheduler
-        self.spt = (self.is_causal or self.is_local) and self.deterministic
+        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
@@ -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
+            cu_total_m_blocks_ptr=mCuTotalMBlocks,
             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
+                tmem_dealloc_mbar: 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]
+                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[
@@ -861,7 +912,7 @@ class FlashAttentionBackwardSm100:
                     cutlass.Int64, self.dQaccum_reduce_stage // 2
                 ]
                 tmem_holding_buf: Int32
-                tmem_dealloc_mbar_ptr: Int64
+                tmem_dealloc_mbar: 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):
+        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
             )
@@ -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."
+        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 or aux_tensors is not None):
+        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 or aux tensors in bwd"
+                "Variable sequence length is not supported yet for blocksparse in bwd"
             )
 
         self.kernel(
@@ -997,7 +1044,7 @@ class FlashAttentionBackwardSm100:
             window_size_left,
             window_size_right,
             tile_sched_params,
-            aux_tensors,
+            aux_data,
             fastdiv_mods,
             blocksparse_tensors,
         ).launch(
@@ -1009,6 +1056,16 @@ class FlashAttentionBackwardSm100:
             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,
@@ -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,
+        aux_data: AuxData = AuxData(),
         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,
+            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,
+            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
         )
+        # 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,
@@ -1209,6 +1267,16 @@ class FlashAttentionBackwardSm100:
             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(
@@ -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,
-        )
+        AttentionMaskCls = self._generate_attention_mask_cls(window_size_left, window_size_right)
         #  EMPTY
         # (15)
         if warp_idx == self.empty_warp_id:
@@ -1440,6 +1501,7 @@ class FlashAttentionBackwardSm100:
                     block_info,
                     SeqlenInfoCls,
                     TileSchedulerCls,
+                    blocksparse_tensors,
                 )
 
         #  LOAD
@@ -1536,6 +1598,7 @@ class FlashAttentionBackwardSm100:
                 pipeline_dKV,
                 pipeline_dP,
                 pipeline_dQ,
+                pipeline_P,
                 block_info,
                 SeqlenInfoCls,
                 TileSchedulerCls,
@@ -1574,6 +1637,7 @@ class FlashAttentionBackwardSm100:
                 pipeline_dS,
                 pipeline_dKV,
                 pipeline_dP,
+                pipeline_P,
                 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,
+                aux_data,
                 fastdiv_mods,
                 blocksparse_tensors,
             )
@@ -1631,6 +1695,7 @@ class FlashAttentionBackwardSm100:
         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)
@@ -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
             )
+            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:
-                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,
+        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],
@@ -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
+            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]
@@ -1835,6 +1912,7 @@ class FlashAttentionBackwardSm100:
                 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)
@@ -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)
+            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)
@@ -1880,6 +1959,7 @@ class FlashAttentionBackwardSm100:
                 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)
@@ -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:
+            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:
-                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:
+                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,
@@ -2060,6 +2099,10 @@ class FlashAttentionBackwardSm100:
                             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):
@@ -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)
+                                if const_expr(load_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):
+                            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)
 
@@ -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):
+                                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)
@@ -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)
+                                    if const_expr(load_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):
+                                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)
 
@@ -2171,27 +2214,104 @@ class FlashAttentionBackwardSm100:
 
                         #### Tail ####
                         if const_expr(should_load_Q):
-                            if const_expr(tma_atom_Qt is not None):
+                            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):
-                    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)
+                    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:
-                    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)
-
+                    (
+                        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()
@@ -2232,6 +2352,7 @@ class FlashAttentionBackwardSm100:
         pipeline_dKV: PipelineAsync,
         pipeline_dP: PipelineAsync,
         pipeline_dQ: PipelineAsync,
+        pipeline_P: Optional[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
+        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
         )
@@ -2366,8 +2491,8 @@ class FlashAttentionBackwardSm100:
                     blocksparse_tensors,
                     batch_idx,
                     head_idx,
-                    n_block,
-                    subtile_factor=self.subtile_factor,
+                    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)
@@ -2392,7 +2517,7 @@ class FlashAttentionBackwardSm100:
                     # 4. dV   = P.T  @ dO
                     # 5. dQ   = dS   @ K
 
-                    main_loop_iters = m_block_max - m_block_min
+                    main_loop_iters = block_iter_count
 
                     # 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
+                    accumulate_dV = False
                     # -----------------------------------------------------------
                     ###### Prologue
                     # -----------------------------------------------------------
-                    # 1. S  = Q0 @ K.T
-                    # 2. dP = V @ dOt.T
-                    # 3. dV = P @ dO
+                    # 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 = K @ Q
+                    # 1) S.T = K @ Q.T
                     handle_Q = pipeline_Q_consumer.wait_and_advance()
-                    pipeline_S_P.sync_object_empty.wait(0, producer_phase_acc)
+                    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 = V @ dOt.T
+                    # 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)
@@ -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()
+                    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
                     # -----------------------------------------------------------
-                    # 1. S  = K    @ Q.T
-                    # 2. dQ = dS   @ K
-                    # 3. dK = dS.T @ Q
-                    # 4. dP = V    @ dOt.T
-                    # 5. dV = P.T  @ dO
+                    # 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
@@ -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()
+                        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)
@@ -2647,22 +2796,34 @@ class FlashAttentionBackwardSm100:
                         pipeline_dS.consumer_release(consumer_state_dS)
                         consumer_state_dS.advance()
 
-                        # (4) dP = V @ dO.T
+                        # (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
+                        # (5) dV += P.T @ dO   (split_P_dS: done at (1b) of the next trip)
                         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()
+                        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
 
-                    pipeline_S_P.sync_object_full.arrive(0, pipeline_S_P.producer_mask, cta_group)
+                    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)
@@ -2753,13 +2914,17 @@ class FlashAttentionBackwardSm100:
         n_block,
         softmax_scale,
         seqlen_info,
-        aux_tensors=None,
+        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)
-        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)
+        # 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)
 
@@ -2772,7 +2937,7 @@ class FlashAttentionBackwardSm100:
             softmax_scale,
             self.vec_size,
             self.qk_acc_dtype,
-            aux_tensors,
+            aux_data,
             fastdiv_mods,
             seqlen_info,
             constant_q_idx=None,
@@ -2790,7 +2955,7 @@ class FlashAttentionBackwardSm100:
         head_idx,
         softmax_scale,
         seqlen_info,
-        aux_tensors=None,
+        aux_data: AuxData = AuxData(),
         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,
+            aux_data,
             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,
+        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,
@@ -2835,6 +3000,7 @@ class FlashAttentionBackwardSm100:
         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,
@@ -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,
+        aux_data: AuxData = AuxData(),
         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
+        # 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))
-        # tdS overlap with tdP
         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))
 
@@ -2953,6 +3123,12 @@ class FlashAttentionBackwardSm100:
             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
         )
@@ -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)
+
+            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)
-            n_block_for_cluster = n_block // self.cta_group_size
+            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=n_block_for_cluster,
+                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_tensors=aux_tensors,
+                aux_data=aux_data,
                 fastdiv_mods=fastdiv_mods,
             )
 
             # prefetch_LSE = not self.is_causal
             prefetch_LSE = False
-            # some tiles might be empty due to block sparsity
+
+            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,
@@ -3006,22 +3209,19 @@ class FlashAttentionBackwardSm100:
                     blocksparse_tensors,
                     batch_idx,
                     head_idx,
-                    n_block,
-                    subtile_factor=self.subtile_factor,
+                    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)
-            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):
+                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,
@@ -3029,27 +3229,34 @@ class FlashAttentionBackwardSm100:
                         curr_q_idx,
                         curr_full_cnt,
                         curr_full_idx,
-                        subtile_factor=self.subtile_factor,
+                        q_subtile_factor=self.q_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)
+                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)
 
-                pipeline_S_P.consumer_wait(consumer_state_S_P_dP)
+                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_fragment(tScS_t2r.shape, Float32)
+                tSrS_t2r = cute.make_rmem_tensor(tScS_t2r.shape, Float32)
                 cute.copy(thr_copy_t2r, tStS_t2r, tSrS_t2r)
 
-                if const_expr(self.tile_hdim == 192):
+                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()
@@ -3080,7 +3287,7 @@ class FlashAttentionBackwardSm100:
                         n_block,
                         softmax_scale,
                         seqlen,
-                        aux_tensors,
+                        aux_data,
                         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_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]
@@ -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()
+                        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],
@@ -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):
+                # 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()
                 # ---------------------------------------------
@@ -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)
+                    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()
-                    self.compute_sync_barrier.arrive_and_wait()
+                    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]
@@ -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.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))
+                        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(
-                            (n_block * self.tile_n, m_block * self.tile_m), cS_bwd
+                            (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)
@@ -3204,7 +3428,7 @@ class FlashAttentionBackwardSm100:
                             head_idx,
                             softmax_scale,
                             seqlen,
-                            aux_tensors,
+                            aux_data,
                             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()
+                # 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
@@ -3413,12 +3643,67 @@ class FlashAttentionBackwardSm100:
             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.core.ThrMma,
+        thr_mma_dQ: cute.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
+            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):
@@ -3485,13 +3771,43 @@ class FlashAttentionBackwardSm100:
             )
 
             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 = self.is_causal and not self.tile_hdim == 192
-            delay_semaphore_release = not self.tile_hdim == 192
+            delay_semaphore_release = not self.tile_hdim == 192 and not self.use_block_sparsity
 
-            # some tiles might be empty due to 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,
@@ -3502,22 +3818,30 @@ class FlashAttentionBackwardSm100:
                     blocksparse_tensors,
                     batch_idx,
                     head_idx,
-                    n_block,
-                    subtile_factor=self.subtile_factor,
+                    n_block_sparse,
+                    q_subtile_factor=self.q_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
+            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,
@@ -3525,16 +3849,13 @@ class FlashAttentionBackwardSm100:
                         curr_q_idx,
                         curr_full_cnt,
                         curr_full_idx,
-                        subtile_factor=self.subtile_factor,
+                        q_subtile_factor=self.q_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
+                    m_block_oob_upper = m_block >= m_block_max
                 pipeline_dQ.consumer_wait(dQ_consumer_state)
                 # TMEM -> RMEM
-                tdQrdQ_t2r = cute.make_fragment(tdQrdQ_t2r_shape, Float32)
+                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()
@@ -3542,6 +3863,8 @@ class FlashAttentionBackwardSm100:
                     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 = (
@@ -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
+                        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,
                             )
-                            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:
+                    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,
@@ -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)
+                    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()
-                    # 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:
+                        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()
-                    barrier.arrive_inc(
-                        mdQ_semaphore_cur[m_block, None].iterator, tidx, cta_rank_in_cluster, 1
-                    )
+                    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:
@@ -3639,7 +3963,10 @@ class FlashAttentionBackwardSm100:
                     )
 
             if const_expr(
-                self.deterministic and not self.spt and block_info.window_size_left is not None
+                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):
@@ -3650,6 +3977,9 @@ class FlashAttentionBackwardSm100:
             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,
@@ -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,
+        thr_mma_dV: cute.ThrMma,
+        thr_mma_dK: cute.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
+        # 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, tdVtdV)
+        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)
@@ -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)
+        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()
@@ -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)
+        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))
@@ -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)
 
+        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)
+            cute.copy(tiled_gmem_store_dV, tdVrdV_r2s, tdVgdV_r2g, pred=tdVpdV)
 
         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)
+        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)
@@ -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)
+        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()
@@ -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)
+        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
@@ -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)
 
+        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)
+            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,
@@ -3795,7 +4156,7 @@ class FlashAttentionBackwardSm100:
         head_idx: Int32,
         n_block: Int32,
         seqlen,
-        thr_mma: cute.core.ThrMma,
+        thr_mma: cute.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):
+            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):
@@ -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)
+            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"
@@ -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 = cute.make_rmem_tensor(tdKVrdKV_t2r.shape, dtype)  # (32 columns)
             tdKVrdKV.store(tdKVrdKV_t2r.load().to(dtype))
 
             # RMEM -> SMEM -- copy, fence and barrier
diff --git a/build/torch-cuda/flash_bwd_sm120.py b/build/torch-cuda/flash_bwd_sm120.py
new file mode 100644
index 0000000000000000000000000000000000000000..297504a2fadf5117f1dc8897743a6d3af38c9473
--- /dev/null
+++ b/build/torch-cuda/flash_bwd_sm120.py
@@ -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
diff --git a/build/torch-cuda/flash_bwd_sm90.py b/build/torch-cuda/flash_bwd_sm90.py
index be4affab559ef02744809a3b2981fe6222e3f040..9e75bfec6ed86af8aa12609c3611557d38de0494 100644
--- a/build/torch-cuda/flash_bwd_sm90.py
+++ b/build/torch-cuda/flash_bwd_sm90.py
@@ -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 .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,
@@ -46,6 +53,8 @@ class FlashAttentionBackwardSm90:
         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,
@@ -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,
+        q_subtile_factor: cutlass.Constexpr[int] = 1,
+        dQ_single_wg: bool = False,
     ):
         self.dtype = dtype
-        # padding head_dim to a multiple of 16 as k_block_size
-        hdim_multiple_of = 16
+        # 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
@@ -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.is_local = is_local
+        self.deterministic = deterministic
         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.num_wg_mma = (self.num_threads // 128) - 1
         self.mma_dkv_is_rs = (
             AtomLayoutMSdP == 1
-            and AtomLayoutNdKV == self.num_mma_warp_groups
+            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_mma_warp_groups == 2, "GQA backward assumes 2 warp groups"
+            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.
-        # 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
+        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(
@@ -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),
+        # 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_mma_warp_groups, self.num_mma_warp_groups)
+            (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_mma_warp_groups)),
+            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_mma_warp_groups // self.AtomLayoutMSdP)
+        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,
@@ -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],
+            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_mma_warp_groups // self.AtomLayoutNdKV)
+        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 = [
@@ -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],
+                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,
@@ -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)
+        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,
@@ -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],
+            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
 
@@ -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,
+        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 Sm90"
-        )
+        # 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(
             *(
@@ -330,23 +380,45 @@ class FlashAttentionBackwardSm90:
             )
         )
 
+        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)
         ]
 
-        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)]
+        # 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 = [layout_utils.select(t, layout_transpose) for t in (mdK, mdV)]
+            mdK, mdV = [_qkv_transpose(t) for t in (mdK, mdV)]
         else:
-            accum_transpose = [2, 1, 0]  # (b, n, s*h) -> (s*h, n, b)
+            # 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)]
-        LSE_dPsum_dQaccum_transpose = [2, 1, 0]  # (b, n, s) -> (s, n, b)
+        # 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
@@ -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
+        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()
@@ -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.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
@@ -404,38 +491,60 @@ class FlashAttentionBackwardSm90:
             (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,
+                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,
+                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
 
-        TileScheduler = SingleTileScheduler
+        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(mQ.shape[3]),
+            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(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,
+            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=False,
+            lpt=self.spt,
+            head_swizzle=self.deterministic,
+            cu_total_m_blocks_ptr=mCuTotalMBlocks,
         )
 
         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):
+        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)
@@ -461,6 +570,11 @@ class FlashAttentionBackwardSm90:
 
         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,
@@ -477,6 +591,10 @@ class FlashAttentionBackwardSm90:
             mLSE,
             mdPsum,
             mdQaccum,
+            mCuSeqlensQ,
+            mCuSeqlensK,
+            mSeqUsedQ,
+            mSeqUsedK,
             self.sQ_layout,
             self.sK_layout,
             self.sV_layout,
@@ -493,15 +611,21 @@ class FlashAttentionBackwardSm90:
             tile_sched_params,
             TileScheduler,
             SharedStorage,
-            aux_tensors,
+            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
@@ -522,6 +646,10 @@ class FlashAttentionBackwardSm90:
         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,
@@ -538,19 +666,23 @@ class FlashAttentionBackwardSm90:
         tile_sched_params: ParamsBase,
         TileScheduler: cutlass.Constexpr[Callable],
         SharedStorage: cutlass.Constexpr[Callable],
-        aux_tensors: Optional[list] = None,
+        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:
-            cpasync.prefetch_descriptor(tma_atom_Q)
-            cpasync.prefetch_descriptor(tma_atom_K)
-            cpasync.prefetch_descriptor(tma_atom_V)
-            cpasync.prefetch_descriptor(tma_atom_dO)
+            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)
@@ -604,25 +736,27 @@ class FlashAttentionBackwardSm90:
             self.is_causal,
             self.is_local,
             False,  # is_split_kv
-            None,
-            None,
+            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=None,
-            mCuSeqlensK=None,
-            mSeqUsedQ=None,
-            mSeqUsedK=None,
+            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=None,
-            window_size_right=None,
+            window_size_left=window_size_left,
+            window_size_right=window_size_right,
             swap_AB=self.SdP_swapAB,
         )
         TileSchedulerCls = partial(TileScheduler.create, tile_sched_params)
@@ -663,18 +797,20 @@ class FlashAttentionBackwardSm90:
                     TileSchedulerCls,
                     SeqlenInfoCls,
                     blocksparse_tensors,
+                    mdQ_semaphore,
                 )
         else:
-            cute.arch.setmaxregister_increase(self.num_mma_regs)
             tidx, _, _ = cute.arch.thread_idx()
             tidx = tidx - 128
-            self.mma(
+            mma_args = (
                 tiled_mma_SdP,
                 tiled_mma_dK,
                 tiled_mma_dV,
                 tiled_mma_dQ,
                 mdK,
                 mdV,
+                mdK_semaphore,
+                mdV_semaphore,
                 mdQaccum,
                 sQ,
                 sK,
@@ -697,11 +833,24 @@ class FlashAttentionBackwardSm90:
                 SeqlenInfoCls,
                 AttentionMaskCls,
                 TileSchedulerCls,
-                aux_tensors,
+                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(
@@ -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]
+                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))
-                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]
+                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))
-                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
+                    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,
-                        subtile_factor=self.subtile_factor,
+                        q_subtile_factor=self.q_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)
+                        # 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
@@ -836,6 +994,8 @@ class FlashAttentionBackwardSm90:
                             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,
@@ -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,
+                            q_subtile_factor=self.q_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,
+        thr_mma_SdP: cute.ThrMma,
         batch_idx,
         head_idx,
         m_block,
         n_block,
         softmax_scale,
         seqlen_info: SeqlenInfoQK,
-        aux_tensors=None,
+        aux_data: AuxData = AuxData(),
         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,
+            aux_data,
             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,
+        thr_mma_SdP: cute.ThrMma,
         batch_idx,
         head_idx,
         m_block,
         n_block,
         softmax_scale,
         seqlen_info: SeqlenInfoQK,
-        aux_tensors=None,
+        aux_data: AuxData = AuxData(),
         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,
+            aux_data,
             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,
+        mdK_semaphore: Optional[cute.Tensor],
+        mdV_semaphore: Optional[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,
+        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_mma_warp_groups, stride=self.num_threads_per_warp_group
+            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 = tiled_mma_dQ.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(
@@ -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
-        )
+        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
+        # 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, self.arch, transpose=self.SdP_swapAB
+                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, self.arch, transpose=self.SdP_swapAB
+            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(
@@ -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)
+        # 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
         )
-        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,
+            is_dQ_wg=is_dQ_wg,
         )
 
         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)
+
+            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,
@@ -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
+                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,
-                    subtile_factor=self.subtile_factor,
+                    q_subtile_factor=self.q_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,
+                        aux_data=aux_data,
                         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,
+                        q_subtile_factor=self.q_subtile_factor,
                         m_block_max=m_block_max,
-                        aux_tensors=aux_tensors,
+                        aux_data=aux_data,
                         fastdiv_mods=fastdiv_mods,
                     )
 
@@ -1216,10 +1435,12 @@ class FlashAttentionBackwardSm90:
                     head_idx,
                     batch_idx,
                     qhead_per_kvhead_divmod,
+                    mdK_semaphore,
+                    mdV_semaphore,
                 )
             else:
-                # Block sparsity: KV tile with zero Q blocks produces no dK/dV; write zeros.
-                if const_expr(self.use_block_sparsity):
+                # 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(
@@ -1239,6 +1460,8 @@ class FlashAttentionBackwardSm90:
                         head_idx,
                         batch_idx,
                         qhead_per_kvhead_divmod,
+                        mdK_semaphore,
+                        mdV_semaphore,
                     )
 
             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)
 
+    @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,
@@ -1266,16 +1505,17 @@ class FlashAttentionBackwardSm90:
         pipeline_dO: cutlass.pipeline.PipelineAsync,
         tLSEsLSE: cute.Tensor,
         tLSEsdPsum: cute.Tensor,
-        tdQsdQaccum: 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_dO if const_expr(self.Q_stage == self.dO_stage) else consumer_state_Q
+            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
@@ -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)
+        # 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(
@@ -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)
+        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 - tLSErLSE[r], fastmath=True
+                    acc_S_mn[r, c] * softmax_scale_log2 - lse_val, 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])):
+            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] - tLSErdPsum[r])
+                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)
@@ -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)
+        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
 
-        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,
-        )
+            # (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)
 
-        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)
+            # 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()
@@ -1408,6 +1666,8 @@ class FlashAttentionBackwardSm90:
         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
@@ -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]
+            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(
@@ -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
+                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, self.arch, transpose=self.dKV_swapAB
+                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()
@@ -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))
+            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
-            mdKaccum_cur = mdK[None, head_idx_kv, batch_idx]
+            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,))
-            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((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)
 
-            cute.arch.cp_async_bulk_wait_group(0, read=True)
+            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)
@@ -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):
+                    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_mma_warp_groups,
+                            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=True)
+            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)
@@ -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):
+                    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_mma_warp_groups,
+                            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(
@@ -1515,21 +1807,45 @@ class FlashAttentionBackwardSm90:
         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)
-            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,)
+            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) or m_block_min < m_block_max
+                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(
@@ -1537,7 +1853,7 @@ class FlashAttentionBackwardSm90:
                     batch_idx,
                     head_idx,
                     n_block,
-                    subtile_factor=self.subtile_factor,
+                    q_subtile_factor=self.q_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
-                            )
+                        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,
                             )
 
-                        for warp_group_idx in cutlass.range_constexpr(self.num_mma_warp_groups):
+                        # 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
@@ -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,
+                                    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,
@@ -1579,13 +1927,29 @@ class FlashAttentionBackwardSm90:
                         n_block,
                         sdQaccum,
                         gdQaccum,
-                        subtile_factor=self.subtile_factor,
+                        q_subtile_factor=self.q_subtile_factor,
                         m_block_max=m_block_max,
-                        num_mma_warp_groups=self.num_mma_warp_groups,
+                        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()
 
-        cute.arch.cp_async_bulk_wait_group(0, read=True)
+        if const_expr(not self.deterministic):
+            cute.arch.cp_async_bulk_wait_group(0, read=True)
diff --git a/build/torch-cuda/flash_fwd.py b/build/torch-cuda/flash_fwd.py
index 1824cf84782fd33cd214b02e65265b5dc3a88ece..484205249d3bae2e09a1c2640e58abc2c7caf88c 100644
--- a/build/torch-cuda/flash_fwd.py
+++ b/build/torch-cuda/flash_fwd.py
@@ -7,22 +7,21 @@
 
 import math
 from types import SimpleNamespace
-from typing import Type, Callable, Optional, List
+from typing import Type, Callable, Optional
 from functools import partial
 
 import cuda.bindings.driver as cuda
 
 import cutlass
 import cutlass.cute as cute
-from cutlass import Constexpr, Float32, Int32, const_expr, Boolean
-from cutlass.cute.nvgpu import cpasync, warp, warpgroup
+from cutlass import Float32, Int32, const_expr
+from cutlass.cute.nvgpu import cpasync, warp
 import cutlass.utils as utils_basic
-from cutlass.utils import LayoutEnum
-import cutlass.utils.hopper_helpers as sm90_utils_basic
+from cutlass.base_dsl.arch import Arch
+from cutlass.cutlass_dsl import BaseDSL
 
 from .quack import copy_utils
 from .quack import layout_utils
-from .quack import sm90_utils
 
 from . import ampere_helpers as sm80_utils
 from .cute_dsl_utils import assume_tensor_aligned
@@ -31,26 +30,14 @@ 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
-from .pack_gqa import PackGQA
+from .pack_gqa import PackGQA, pack_gqa_layout
 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 .block_sparsity import BlockSparseTensors
+from .tile_scheduler import SingleTileScheduler, SingleTileVarlenScheduler, TileSchedulerArguments
+from .utils import AuxData
 
 
 class FlashAttentionForwardBase:
-    arch: int = 80
 
     def __init__(
         self,
@@ -69,7 +56,7 @@ class FlashAttentionForwardBase:
         score_mod: Optional[cutlass.Constexpr] = None,
         mask_mod: Optional[cutlass.Constexpr] = None,
         has_aux_tensors: bool = False,
-        q_subtile_factor: int | None = None,
+        q_subtile_factor: int = 1,
     ):
         """Initializes the configuration for a flash attention kernel.
 
@@ -113,9 +100,21 @@ class FlashAttentionForwardBase:
         self.score_mod = score_mod
         self.mask_mod = mask_mod
         self.qk_acc_dtype = Float32
-        self.vec_size: cutlass.Constexpr = getattr(
+        self.score_vec_size: cutlass.Constexpr = getattr(
             score_mod, "__vec_size__", 1 if cutlass.const_expr(has_aux_tensors) else 2
         )
+        if self.score_vec_size > 2:
+            raise ValueError(
+                f"score_mod vec_size {self.score_vec_size} not supported on Sm80/90/120 "
+                "due to accumulator thread ownership pattern."
+            )
+        self.mask_vec_size: cutlass.Constexpr = getattr(mask_mod, "__vec_size__", 1)
+        if self.mask_vec_size > 1:
+            raise ValueError(
+                f"mask_mod vec_size {self.mask_vec_size} not supported on Sm80/90/120 "
+                "due to accumulator thread ownership pattern."
+            )
+        self.arch = BaseDSL._get_dsl().get_arch_enum()
 
     @staticmethod
     def can_implement(
@@ -318,7 +317,8 @@ class FlashAttentionForwardBase:
         mO: cute.Tensor,
         mLSE: Optional[cute.Tensor],
         softmax_scale: Float32,
-        stream: cuda.CUstream,
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
+        stream: cuda.CUstream = None,
     ):
         """Configures and launches the flash attention kernel.
 
@@ -351,7 +351,7 @@ class FlashAttentionForwardBase:
         cute.arch.barrier(
             barrier_id=int(NamedBarrierFwd.Epilogue), number_of_threads=self.num_epilogue_threads
         )
-        smem_copy_atom_O = utils.get_smem_store_atom(self.arch, self.dtype)
+        smem_copy_atom_O = utils.get_smem_store_atom(self.arch.major * 10 + self.arch.minor, self.dtype)
         smem_thr_copy_O = cute.make_tiled_copy_C(smem_copy_atom_O, tiled_mma).get_slice(tidx)
         taccOrO = smem_thr_copy_O.retile(rO)
         taccOsO = smem_thr_copy_O.partition_D(sO)
@@ -366,11 +366,7 @@ class FlashAttentionForwardBase:
 
         # Write LSE from rmem -> gmem
         if const_expr(mLSE is not None):
-            if const_expr(not seqlen.has_cu_seqlens_q):
-                mLSE_cur = mLSE[None, head_idx, batch_idx]
-            else:
-                offset = seqlen.offset_q if const_expr(not self.pack_gqa) else (0, seqlen.offset_q)
-                mLSE_cur = cute.domain_offset((offset,), mLSE[None, head_idx])
+            mLSE_cur = seqlen.offset_batch_Q(mLSE, batch_idx, dim=2)[None, head_idx]
             if const_expr(not self.pack_gqa):
                 gLSE = cute.local_tile(mLSE_cur, (self.tile_m,), (m_block,))
                 gLSE_expanded_layout = cute.append(
@@ -384,7 +380,7 @@ class FlashAttentionForwardBase:
                 t0accOcO = layout_utils.reshape_acc_to_mn(thr_mma.get_slice(0).partition_C(cO))
                 # Only the thread corresponding to column 0 writes out the lse to gmem
                 if taccOcO[0][1] == 0:
-                    for m in cutlass.range_constexpr(cute.size(taccOgLSE.shape[1])):
+                    for m in cutlass.range(cute.size(taccOgLSE.shape[1]), unroll_full=True):
                         if (
                             t0accOcO[m, 0][0]
                             < seqlen.seqlen_q - m_block * self.tile_m - taccOcO[0][0]
@@ -393,11 +389,8 @@ class FlashAttentionForwardBase:
             else:
                 pack_gqa.store_LSE(mLSE_cur, lse, tiled_mma, tidx, m_block, seqlen.seqlen_q)
 
-        if const_expr(not seqlen.has_cu_seqlens_q):
-            mO_cur = mO[None, None, head_idx, batch_idx]
-        else:
-            offset = seqlen.offset_q if const_expr(not self.pack_gqa) else (0, seqlen.offset_q)
-            mO_cur = cute.domain_offset((offset, 0), mO[None, None, head_idx])
+        ragged = self.use_tma_O and (seqlen.has_cu_seqlens_q or seqlen.has_seqused_q)
+        mO_cur = seqlen.offset_batch_Q(mO, batch_idx, dim=3, ragged=ragged)[None, None, head_idx]
         # thr_mma = tiled_mma.get_slice(tidx)
         # taccOgO = thr_mma.partition_C(gO)
         # cute.autovec_copy(rO, taccOgO)
@@ -634,12 +627,21 @@ class FlashAttentionForwardSm80(FlashAttentionForwardBase):
         mV: cute.Tensor,
         mO: cute.Tensor,
         mLSE: Optional[cute.Tensor],
-        stream: cuda.CUstream,
-        softmax_scale: Optional[Float32] = None,
-        window_size_left: Optional[Int32] = None,
-        window_size_right: Optional[Int32] = None,
+        softmax_scale: Float32,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        mSeqUsedQ: Optional[cute.Tensor] = None,
+        mSeqUsedK: Optional[cute.Tensor] = None,
+        mPageTable: Optional[cute.Tensor] = None,
+        window_size_left: Int32 | int | None = None,
+        window_size_right: Int32 | int | None = None,
         learnable_sink: Optional[cute.Tensor] = None,
-        aux_tensors=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.
 
@@ -648,49 +650,69 @@ class FlashAttentionForwardSm80(FlashAttentionForwardBase):
         """
         assert learnable_sink is None, "Learnable sink is not supported in this kernel"
         self._check_type(
-            *(t.element_type if t is not None else None for t in (mQ, mK, mV, mO, mLSE))
+            *(t.element_type if t is not None else None for t in (mQ, mK, mV, mO, mLSE, mCuSeqlensQ, mCuSeqlensK, mSeqUsedQ, mSeqUsedK))
         )
         tiled_mma_qk, tiled_mma_pv = self._get_tiled_mma()
         self.num_mma_threads = tiled_mma_pv.size
         self.num_producer_threads = self.num_threads
         self.num_Q_load_threads = self.num_threads
         self.num_epilogue_threads = self.num_threads
-        # self.use_tma_O = self.arch >= 90 and mCuSeqlensQ is None
-        self.use_tma_O = self.arch >= 90
+        self.use_tma_O = Arch.sm_90 <= self.arch < Arch.sm_120
         self._setup_attributes()
         SharedStorage = self._get_shared_storage_cls()
         mQ, mK, mV, mO = [assume_tensor_aligned(t) for t in (mQ, mK, mV, mO)]
-        mQ, mK, mV, mO = [
-            cute.make_tensor(t.iterator, cute.select(t.layout, mode=[1, 3, 2, 0]))
-            for t in (mQ, mK, mV, mO)
+        # Layout permutation: 4D non-varlen vs 3D varlen
+        QO_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1]
+        KV_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensK is None) else [0, 2, 1]
+        mQ, mO = [
+            cute.make_tensor(t.iterator, cute.select(t.layout, mode=QO_layout_transpose))
+            for t in (mQ, mO)
         ]
-        mLSE = cute.make_tensor(mLSE.iterator, cute.select(mLSE.layout, mode=[2, 1, 0]))
-        # grid_dim: (m_block, num_head, batch_size)
-        grid_dim = (
-            cute.ceil_div(mQ.shape[0], self.tile_m),
-            cute.size(mQ.shape[2]),
-            cute.size(mQ.shape[3]),
-        )
-        LOG2_E = math.log2(math.e)
-        if const_expr(self.score_mod is None):
-            softmax_scale_log2 = Float32(softmax_scale * LOG2_E)
-            softmax_scale = None
+        mK, mV = [
+            cute.make_tensor(t.iterator, cute.select(t.layout, mode=KV_layout_transpose))
+            for t in (mK, mV)
+        ]
+        if const_expr(mLSE is not None):
+            LSE_layout_transpose = [2, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 0]
+            mLSE = cute.make_tensor(mLSE.iterator, cute.select(mLSE.layout, mode=LSE_layout_transpose))
+        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)
+        # TileScheduler for varlen, simple grid for non-varlen
+        if const_expr(mCuSeqlensQ is not None or mSeqUsedQ is not None):
+            TileScheduler = SingleTileVarlenScheduler
         else:
-            # NB: If a user passes in a score mod, we want to apply the score-mod in the sm_scaled qk
-            # But in the original base 10. We hijack softmax_scale_log2 to just be the change of base
-            # and correctly apply the softmax_scale prior to score_mod in the softmax step
-            softmax_scale_log2 = Float32(LOG2_E)
-            softmax_scale = Float32(softmax_scale)
-
-        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
-            )
-            seqlen_k = cute.size(mK.shape[0])
-            seqlen_q_divmod = FastDivmodDivisor(seqlen_q)
-            seqlen_k_divmod = FastDivmodDivisor(seqlen_k)
-            fastdiv_mods = (seqlen_q_divmod, seqlen_k_divmod)
+            TileScheduler = SingleTileScheduler
+        num_batch = (
+            mCuSeqlensQ.shape[0] - 1
+            if const_expr(mCuSeqlensQ is not None)
+            else cute.size(mQ.shape[3])
+        )
+        tile_sched_args = TileSchedulerArguments(
+            num_block=cute.ceil_div(cute.size(mQ.shape[0]), self.tile_m),
+            num_head=cute.size(mQ.shape[2]),
+            num_batch=num_batch,
+            num_splits=1,
+            seqlen_k=0,
+            headdim=mQ.shape[1],
+            headdim_v=mV.shape[1],
+            total_q=cute.size(mQ.shape[0])
+            if const_expr(mCuSeqlensQ is not None)
+            else cute.size(mQ.shape[0]) * cute.size(mQ.shape[3]),
+            tile_shape_mn=(self.tile_m, self.tile_n),
+            qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
+            mCuSeqlensQ=mCuSeqlensQ,
+            mSeqUsedQ=mSeqUsedQ,
+            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)
+        fastdiv_mods = utils.compute_fastdiv_mods(mQ, mK, self.qhead_per_kvhead, self.pack_gqa, aux_data.tensors)
 
         self.kernel(
             mQ,
@@ -698,6 +720,10 @@ class FlashAttentionForwardSm80(FlashAttentionForwardBase):
             mV,
             mO,
             mLSE,
+            mCuSeqlensQ,
+            mCuSeqlensK,
+            mSeqUsedQ,
+            mSeqUsedK,
             softmax_scale_log2,
             softmax_scale,
             window_size_left,
@@ -714,7 +740,9 @@ class FlashAttentionForwardSm80(FlashAttentionForwardBase):
             tiled_mma_qk,
             tiled_mma_pv,
             SharedStorage,
-            aux_tensors,
+            tile_sched_params,
+            TileScheduler,
+            aux_data,
             fastdiv_mods,
         ).launch(
             grid=grid_dim,
@@ -731,6 +759,10 @@ class FlashAttentionForwardSm80(FlashAttentionForwardBase):
         mV: cute.Tensor,
         mO: cute.Tensor,
         mLSE: Optional[cute.Tensor],
+        mCuSeqlensQ: Optional[cute.Tensor],
+        mCuSeqlensK: Optional[cute.Tensor],
+        mSeqUsedQ: Optional[cute.Tensor],
+        mSeqUsedK: Optional[cute.Tensor],
         softmax_scale_log2: Float32,
         softmax_scale: Optional[Float32],
         window_size_left: Optional[Int32],
@@ -747,289 +779,320 @@ class FlashAttentionForwardSm80(FlashAttentionForwardBase):
         tiled_mma_qk: cute.TiledMma,
         tiled_mma_pv: cute.TiledMma,
         SharedStorage: cutlass.Constexpr,
-        aux_tensors=None,
+        tile_sched_params,
+        TileScheduler: cutlass.Constexpr[Callable],
+        aux_data: AuxData = AuxData(),
         fastdiv_mods=None,
     ):
         # Thread index, block index
         tidx, _, _ = cute.arch.thread_idx()
-        m_block, num_head, batch_size = cute.arch.block_idx()
-
-        block_info = BlockInfo(
-            self.tile_m,
-            self.tile_n,
-            self.is_causal,
-            self.is_local,
-            False,  # is_split_kv
-            window_size_left,
-            window_size_right,
-            qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
-        )
-        seqlen = SeqlenInfoQK.create(seqlen_q_static=mQ.shape[0], seqlen_k_static=mK.shape[0])
-        n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block)
-        # TODO: return early if n_block_max == 0
-        # if self.is_causal:
-        #     if n_block_max <= 0:
-        #         return
-        n_block = n_block_max - 1
 
-        # ///////////////////////////////////////////////////////////////////////////////
-        # Get the appropriate tiles for this thread block.
-        # ///////////////////////////////////////////////////////////////////////////////
-        blkQ_shape = (self.tile_m, self.tile_hdim)
-        blkK_shape = (self.tile_n, self.tile_hdim)
-        blkV_shape = (self.tile_n, self.tile_hdimv)
-        gQ = cute.local_tile(mQ[None, None, num_head, batch_size], blkQ_shape, (m_block, 0))
-        num_head_kv = num_head // self.qhead_per_kvhead
-        gK = cute.local_tile(mK[None, None, num_head_kv, batch_size], blkK_shape, (None, 0))
-        gV = cute.local_tile(mV[None, None, num_head_kv, batch_size], blkV_shape, (None, 0))
+        tile_scheduler = TileScheduler.create(tile_sched_params)
+        work_tile = tile_scheduler.initial_work_tile_info()
 
-        # ///////////////////////////////////////////////////////////////////////////////
-        # Get shared memory buffer
-        # ///////////////////////////////////////////////////////////////////////////////
-        smem = cutlass.utils.SmemAllocator()
-        storage = smem.allocate(SharedStorage)
-        sQ = storage.sQ.get_tensor(sQ_layout)
-        sK = storage.sK.get_tensor(sK_layout)
-        if const_expr(not self.Q_in_regs):
-            sV = storage.sV.get_tensor(sV_layout)
-        else:
-            sV = cute.make_tensor(cute.recast_ptr(sQ.iterator, dtype=self.dtype), sV_layout)
-        # Transpose view of V to tensor with layout (head_dim_v, tile_n) for tiled mma
-        sVt = layout_utils.transpose_view(sV)
+        if work_tile.is_valid_tile:
+            m_block, num_head, batch_size, _ = work_tile.tile_idx
+
+            block_info = BlockInfo(
+                self.tile_m,
+                self.tile_n,
+                self.is_causal,
+                self.is_local,
+                False,  # is_split_kv
+                window_size_left,
+                window_size_right,
+                qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
+            )
+            seqlen = SeqlenInfoQK.create(
+                batch_idx=batch_size,
+                seqlen_q_static=mQ.shape[0] if const_expr(not self.pack_gqa) else mQ.shape[0][1],
+                seqlen_k_static=mK.shape[0],
+                mCuSeqlensQ=mCuSeqlensQ,
+                mCuSeqlensK=mCuSeqlensK,
+                mSeqUsedQ=mSeqUsedQ,
+                mSeqUsedK=mSeqUsedK,
+            )
+            n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block)
+            # For varlen, wasted grid tiles (where batch_idx >= num_batch) will have
+            # seqlen_q=seqlen_k=0 and n_block_max=0.  Clamp to 0 so we don't use a
+            # negative block index for K/V loads; the load/store predicates already
+            # guard all memory accesses when seqlen is 0.
+            n_block = cutlass.max(n_block_max - 1, 0)
+
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Get the appropriate tiles for this thread block.
+            # ///////////////////////////////////////////////////////////////////////////////
+            blkQ_shape = (self.tile_m, self.tile_hdim)
+            blkK_shape = (self.tile_n, self.tile_hdim)
+            blkV_shape = (self.tile_n, self.tile_hdimv)
+            num_head_kv = num_head if const_expr(self.pack_gqa) else num_head // self.qhead_per_kvhead
+            mQ_cur = seqlen.offset_batch_Q(mQ, batch_size, dim=3)[None, None, num_head]
+            if const_expr(not seqlen.has_cu_seqlens_k):
+                mK_cur = mK[None, None, num_head_kv, batch_size]
+                mV_cur = mV[None, None, num_head_kv, batch_size]
+            else:
+                mK_cur = cute.domain_offset((seqlen.offset_k, 0), mK[None, None, num_head_kv])
+                mV_cur = cute.domain_offset((seqlen.offset_k, 0), mV[None, None, num_head_kv])
+            if const_expr(not self.pack_gqa):
+                gQ = cute.local_tile(mQ_cur, blkQ_shape, (m_block, 0))
+            gK = cute.local_tile(mK_cur, blkK_shape, (None, 0))
+            gV = cute.local_tile(mV_cur, blkV_shape, (None, 0))
 
-        gmem_thr_copy_K = gmem_tiled_copy_K.get_slice(tidx)
-        gmem_thr_copy_V = gmem_tiled_copy_V.get_slice(tidx)
-        # (CPY_Atom, CPY_N, CPY_K, n_block)
-        tKsK, tKgK = gmem_thr_copy_K.partition_D(sK), gmem_thr_copy_K.partition_S(gK)
-        # (CPY_Atom, CPY_N, CPY_K, n_block)
-        tVsV, tVgV = gmem_thr_copy_V.partition_D(sV), gmem_thr_copy_V.partition_S(gV)
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Get shared memory buffer
+            # ///////////////////////////////////////////////////////////////////////////////
+            smem = cutlass.utils.SmemAllocator()
+            storage = smem.allocate(SharedStorage)
+            sQ = storage.sQ.get_tensor(sQ_layout)
+            sK = storage.sK.get_tensor(sK_layout)
+            if const_expr(not self.Q_in_regs):
+                sV = storage.sV.get_tensor(sV_layout)
+            else:
+                sV = cute.make_tensor(cute.recast_ptr(sQ.iterator, dtype=self.dtype), sV_layout)
+            # Transpose view of V to tensor with layout (head_dim_v, tile_n) for tiled mma
+            sVt = layout_utils.transpose_view(sV)
 
-        # ///////////////////////////////////////////////////////////////////////////////
-        # Tile MMA compute thread partitions and allocate accumulators
-        # ///////////////////////////////////////////////////////////////////////////////
-        thr_mma_qk = tiled_mma_qk.get_slice(tidx)
-        thr_mma_pv = tiled_mma_pv.get_slice(tidx)
-        tSrQ = thr_mma_qk.make_fragment_A(thr_mma_qk.partition_A(sQ))
-        tSrK = thr_mma_qk.make_fragment_B(thr_mma_qk.partition_B(sK[None, None, 0]))
-        tOrVt = thr_mma_pv.make_fragment_B(thr_mma_pv.partition_B(sVt[None, None, 0]))
-        acc_shape_O = thr_mma_pv.partition_shape_C((self.tile_m, self.tile_hdimv))
-        acc_O = cute.make_fragment(acc_shape_O, Float32)
-        acc_O.fill(0.0)
+            gmem_thr_copy_K = gmem_tiled_copy_K.get_slice(tidx)
+            gmem_thr_copy_V = gmem_tiled_copy_V.get_slice(tidx)
+            # (CPY_Atom, CPY_N, CPY_K, n_block)
+            tKsK, tKgK = gmem_thr_copy_K.partition_D(sK), gmem_thr_copy_K.partition_S(gK)
+            # (CPY_Atom, CPY_N, CPY_K, n_block)
+            tVsV, tVgV = gmem_thr_copy_V.partition_D(sV), gmem_thr_copy_V.partition_S(gV)
 
-        # ///////////////////////////////////////////////////////////////////////////////
-        # Smem copy atom tiling
-        # ///////////////////////////////////////////////////////////////////////////////
-        smem_copy_atom_QK = cute.make_copy_atom(
-            warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4),
-            self.dtype,
-        )
-        smem_copy_atom_V = cute.make_copy_atom(
-            warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4),
-            self.dtype,
-        )
-        smem_thr_copy_Q = utils.make_tiled_copy_A(smem_copy_atom_QK, tiled_mma_qk).get_slice(tidx)
-        smem_thr_copy_K = utils.make_tiled_copy_B(smem_copy_atom_QK, tiled_mma_qk).get_slice(tidx)
-        smem_thr_copy_V = utils.make_tiled_copy_B(smem_copy_atom_V, tiled_mma_pv).get_slice(tidx)
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Tile MMA compute thread partitions and allocate accumulators
+            # ///////////////////////////////////////////////////////////////////////////////
+            thr_mma_qk = tiled_mma_qk.get_slice(tidx)
+            thr_mma_pv = tiled_mma_pv.get_slice(tidx)
+            tSrQ = thr_mma_qk.make_fragment_A(thr_mma_qk.partition_A(sQ))
+            tSrK = thr_mma_qk.make_fragment_B(thr_mma_qk.partition_B(sK[None, None, 0]))
+            tOrVt = thr_mma_pv.make_fragment_B(thr_mma_pv.partition_B(sVt[None, None, 0]))
+            acc_shape_O = thr_mma_pv.partition_shape_C((self.tile_m, self.tile_hdimv))
+            acc_O = cute.make_rmem_tensor(acc_shape_O, Float32)
+            acc_O.fill(0.0)
 
-        tSsQ = smem_thr_copy_Q.partition_S(sQ)
-        tSsK = smem_thr_copy_K.partition_S(sK)
-        tOsVt = smem_thr_copy_V.partition_S(sVt)
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Smem copy atom tiling
+            # ///////////////////////////////////////////////////////////////////////////////
+            smem_copy_atom_QK = cute.make_copy_atom(
+                warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4),
+                self.dtype,
+            )
+            smem_copy_atom_V = cute.make_copy_atom(
+                warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4),
+                self.dtype,
+            )
+            smem_thr_copy_Q = utils.make_tiled_copy_A(smem_copy_atom_QK, tiled_mma_qk).get_slice(tidx)
+            smem_thr_copy_K = utils.make_tiled_copy_B(smem_copy_atom_QK, tiled_mma_qk).get_slice(tidx)
+            smem_thr_copy_V = utils.make_tiled_copy_B(smem_copy_atom_V, tiled_mma_pv).get_slice(tidx)
 
-        # ///////////////////////////////////////////////////////////////////////////////
-        # Predicate: Mark indices that need to copy when problem_shape isn't a multiple
-        # of tile_shape
-        # ///////////////////////////////////////////////////////////////////////////////
-        # Construct identity layout for KV
-        cK = cute.make_identity_tensor((self.tile_n, self.tile_hdim))
-        tKcK = gmem_thr_copy_K.partition_S(cK)
-        t0KcK = gmem_thr_copy_K.get_slice(0).partition_S(cK)
-        if const_expr(self.tile_hdim == self.tile_hdimv):
-            tVcV = tKcK
-            t0VcV = t0KcK
-        else:
-            cV = cute.make_identity_tensor((self.tile_n, self.tile_hdimv))
-            tVcV = gmem_thr_copy_V.partition_S(cV)
-            t0VcV = gmem_thr_copy_V.get_slice(0).partition_S(cV)
-        # Allocate predicate tensors for m and n, here we only allocate the tile of k, and
-        # use "if" on the mn dimension.
-        # This is to reduce register pressure and gets 2-3% performance gain.
-        tKpK = utils.predicate_k(tKcK, limit=mK.shape[1])
-        if const_expr(self.same_hdim_kv):
-            tVpV = tKpK
-        else:
-            tVpV = utils.predicate_k(tVcV, limit=mV.shape[1])
+            tSsQ = smem_thr_copy_Q.partition_S(sQ)
+            tSsK = smem_thr_copy_K.partition_S(sK)
+            tOsVt = smem_thr_copy_V.partition_S(sVt)
 
-        # shape: (atom_v_m * rest_m)
-        softmax = Softmax.create(
-            softmax_scale_log2,
-            num_rows=acc_O.shape[0][0] * acc_O.shape[1],
-            softmax_scale=softmax_scale,
-        )
-        softmax.reset()
-
-        # group parameters for compute_one_n_block
-        mma_params = SimpleNamespace(
-            thr_mma_qk=thr_mma_qk,
-            thr_mma_pv=thr_mma_pv,
-            tSrQ=tSrQ,
-            tSrK=tSrK,
-            tOrVt=tOrVt,
-            acc_O=acc_O,
-        )
-        smem_copy_params = SimpleNamespace(
-            smem_thr_copy_Q=smem_thr_copy_Q,
-            smem_thr_copy_K=smem_thr_copy_K,
-            smem_thr_copy_V=smem_thr_copy_V,
-            tSsQ=tSsQ,
-            tSsK=tSsK,
-            tOsVt=tOsVt,
-        )
-        load_K = partial(
-            self.load_K, gmem_tiled_copy_K, tKgK, tKsK, tKcK, t0KcK, tKpK, seqlen=seqlen.seqlen_k
-        )
-        load_V = partial(
-            self.load_V, gmem_tiled_copy_V, tVgV, tVsV, tVcV, t0VcV, tVpV, seqlen=seqlen.seqlen_k
-        )
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Predicate: Mark indices that need to copy when problem_shape isn't a multiple
+            # of tile_shape
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Construct identity layout for KV
+            cK = cute.make_identity_tensor((self.tile_n, self.tile_hdim))
+            tKcK = gmem_thr_copy_K.partition_S(cK)
+            t0KcK = gmem_thr_copy_K.get_slice(0).partition_S(cK)
+            if const_expr(self.tile_hdim == self.tile_hdimv):
+                tVcV = tKcK
+                t0VcV = t0KcK
+            else:
+                cV = cute.make_identity_tensor((self.tile_n, self.tile_hdimv))
+                tVcV = gmem_thr_copy_V.partition_S(cV)
+                t0VcV = gmem_thr_copy_V.get_slice(0).partition_S(cV)
+            # Allocate predicate tensors for m and n, here we only allocate the tile of k, and
+            # use "if" on the mn dimension.
+            # This is to reduce register pressure and gets 2-3% performance gain.
+            tKpK = utils.predicate_k(tKcK, limit=mK.shape[1])
+            if const_expr(self.same_hdim_kv):
+                tVpV = tKpK
+            else:
+                tVpV = utils.predicate_k(tVcV, limit=mV.shape[1])
 
-        compute_one_n_block = partial(
-            self.compute_one_n_block,
-            mma_params=mma_params,
-            smem_copy_params=smem_copy_params,
-            softmax=softmax,
-            load_K=load_K,
-            load_V=load_V,
-            score_mod=self.score_mod,
-            batch_idx=batch_size,
-            head_idx=num_head,
-            m_block=m_block,
-            aux_tensors=aux_tensors,
-            fastdiv_mods=fastdiv_mods,
-        )
+            # shape: (atom_v_m * rest_m)
+            softmax = Softmax.create(
+                softmax_scale_log2,
+                num_rows=acc_O.shape[0][0] * acc_O.shape[1],
+                softmax_scale=softmax_scale,
+            )
+            softmax.reset()
+
+            # group parameters for compute_one_n_block
+            mma_params = SimpleNamespace(
+                thr_mma_qk=thr_mma_qk,
+                thr_mma_pv=thr_mma_pv,
+                tSrQ=tSrQ,
+                tSrK=tSrK,
+                tOrVt=tOrVt,
+                acc_O=acc_O,
+            )
+            smem_copy_params = SimpleNamespace(
+                smem_thr_copy_Q=smem_thr_copy_Q,
+                smem_thr_copy_K=smem_thr_copy_K,
+                smem_thr_copy_V=smem_thr_copy_V,
+                tSsQ=tSsQ,
+                tSsK=tSsK,
+                tOsVt=tOsVt,
+            )
+            load_K = partial(
+                self.load_K, gmem_tiled_copy_K, tKgK, tKsK, tKcK, t0KcK, tKpK, seqlen=seqlen.seqlen_k
+            )
+            load_V = partial(
+                self.load_V, gmem_tiled_copy_V, tVgV, tVsV, tVcV, t0VcV, tVpV, seqlen=seqlen.seqlen_k
+            )
 
-        # ///////////////////////////////////////////////////////////////////////////////
-        # Prologue
-        # ///////////////////////////////////////////////////////////////////////////////
-        # Start async loads of the last mn-tile, where we take care of the mn residue
-        gmem_thr_copy_Q = gmem_tiled_copy_Q.get_slice(tidx)
-        self.load_Q(gmem_thr_copy_Q, gQ, sQ, m_block, seqlen=seqlen.seqlen_q, headdim=mQ.shape[1])
-        cute.arch.cp_async_commit_group()
+            compute_one_n_block = partial(
+                self.compute_one_n_block,
+                mma_params=mma_params,
+                smem_copy_params=smem_copy_params,
+                softmax=softmax,
+                load_K=load_K,
+                load_V=load_V,
+                score_mod=self.score_mod,
+                batch_idx=batch_size,
+                head_idx=num_head,
+                m_block=m_block,
+                aux_data=aux_data,
+                fastdiv_mods=fastdiv_mods,
+            )
 
-        def preprocess_Q():
-            cute.arch.cp_async_wait_group(self.num_stages * 2 - 1)
-            if const_expr(self.Q_in_regs):
-                cute.arch.barrier()
-                tSrQ_copy_view = smem_thr_copy_Q.retile(tSrQ)
-                cute.copy(smem_thr_copy_Q, tSsQ, tSrQ_copy_view)
-
-        # If Q_in_regs, we load Q, then load 1 stage of K, then (optionally) rotate Q and
-        # read from smem_q to registers, then load V.
-        # If !Q_in_regs, we load Q, load all stages of K & V, then (optionally) rotate Q.
-        if const_expr(self.Q_in_regs):
-            load_K(n_block, smem_pipe_write=0, need_predicates=True)
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Prologue
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Start async loads of the last mn-tile, where we take care of the mn residue
+            gmem_thr_copy_Q = gmem_tiled_copy_Q.get_slice(tidx)
+            if const_expr(not self.pack_gqa):
+                self.load_Q(gmem_thr_copy_Q, gQ, sQ, m_block, seqlen=seqlen.seqlen_q, headdim=mQ.shape[1])
+            else:
+                pack_gqa = PackGQA(self.tile_m, self.tile_hdim, self.check_hdim_oob, self.qhead_per_kvhead)
+                pack_gqa.load_Q(mQ_cur, sQ, gmem_tiled_copy_Q, tidx, m_block, seqlen.seqlen_q)
             cute.arch.cp_async_commit_group()
-            preprocess_Q()
-            cute.arch.barrier()  # Make sure all threads have read smem_q before loading V
 
-        for stage in cutlass.range_constexpr(self.num_stages):
-            if const_expr(not self.Q_in_regs or stage > 0):
-                if stage == 0 or n_block - stage >= 0:
-                    load_K(n_block - stage, smem_pipe_write=stage, need_predicates=stage == 0)
-                cute.arch.cp_async_commit_group()
-            if const_expr(stage < self.num_stages - 1):
-                if stage == 0 or n_block - stage >= 0:
-                    load_V(n_block - stage, smem_pipe_write=stage, need_predicates=stage == 0)
+            def preprocess_Q():
+                cute.arch.cp_async_wait_group(self.num_stages * 2 - 1)
+                if const_expr(self.Q_in_regs):
+                    cute.arch.barrier()
+                    tSrQ_copy_view = smem_thr_copy_Q.retile(tSrQ)
+                    cute.copy(smem_thr_copy_Q, tSsQ, tSrQ_copy_view)
+
+            # If Q_in_regs, we load Q, then load 1 stage of K, then (optionally) rotate Q and
+            # read from smem_q to registers, then load V.
+            # If !Q_in_regs, we load Q, load all stages of K & V, then (optionally) rotate Q.
+            if const_expr(self.Q_in_regs):
+                load_K(n_block, smem_pipe_write=0, need_predicates=True)
                 cute.arch.cp_async_commit_group()
-        if const_expr(not self.Q_in_regs):
-            preprocess_Q()
+                preprocess_Q()
+                cute.arch.barrier()  # Make sure all threads have read smem_q before loading V
+
+            for stage in cutlass.range_constexpr(self.num_stages):
+                if const_expr(not self.Q_in_regs or stage > 0):
+                    if stage == 0 or n_block - stage >= 0:
+                        load_K(n_block - stage, smem_pipe_write=stage, need_predicates=stage == 0)
+                    cute.arch.cp_async_commit_group()
+                if const_expr(stage < self.num_stages - 1):
+                    if stage == 0 or n_block - stage >= 0:
+                        load_V(n_block - stage, smem_pipe_write=stage, need_predicates=stage == 0)
+                    cute.arch.cp_async_commit_group()
+            if const_expr(not self.Q_in_regs):
+                preprocess_Q()
 
-        # ///////////////////////////////////////////////////////////////////////////////
-        # Mainloop
-        # ///////////////////////////////////////////////////////////////////////////////
-        # Start processing of the first n-block.
-        # For performance reason, we separate out two kinds of iterations:
-        # those that need masking on S, and those that don't.
-        # We need masking on S for the very last block when K and V has length not multiple of tile_n.
-        # We also need masking on S if it's causal, for the last several blocks.
-        mask = AttentionMask(
-            self.tile_m,
-            self.tile_n,
-            seqlen.seqlen_q,
-            seqlen.seqlen_k,
-            window_size_left,
-            window_size_right,
-            self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
-        )
-        mask_fn = partial(
-            mask.apply_mask,
-            m_block=m_block,
-            thr_mma=thr_mma_qk,
-            mask_causal=self.is_causal,
-            mask_local=self.is_local,
-            fastdiv_mods=fastdiv_mods if const_expr(self.mask_mod is not None) else None,
-        )
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Mainloop
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Start processing of the first n-block.
+            # For performance reason, we separate out two kinds of iterations:
+            # those that need masking on S, and those that don't.
+            # We need masking on S for the very last block when K and V has length not multiple of tile_n.
+            # We also need masking on S if it's causal, for the last several blocks.
+            mask = AttentionMask(
+                self.tile_m,
+                self.tile_n,
+                seqlen,
+                window_size_left,
+                window_size_right,
+                self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
+            )
+            mask_fn = partial(
+                mask.apply_mask,
+                batch_idx=batch_size,
+                head_idx=num_head,
+                m_block=m_block,
+                thr_mma=thr_mma_qk,
+                mask_causal=self.is_causal,
+                mask_local=self.is_local,
+                aux_data=aux_data,
+                fastdiv_mods=fastdiv_mods if const_expr(self.mask_mod is not None) else None,
+            )
 
-        # First iteration with seqlen masking
-        smem_pipe_read = Int32(0)
-        smem_pipe_write = Int32(self.num_stages - 1)
-        compute_one_n_block(
-            n_block,
-            smem_pipe_read,
-            smem_pipe_write,
-            is_first_n_block=True,
-            check_inf=True,
-            mask_fn=partial(mask_fn, mask_seqlen=True),
-        )
-        smem_pipe_read = self.advance_pipeline(smem_pipe_read)
-        smem_pipe_write = self.advance_pipeline(smem_pipe_write)
-        # Next couple of iterations with causal masking
-        if const_expr(self.is_causal or self.is_local):
-            n_block_min_causal_local_mask = block_info.get_n_block_min_causal_local_mask(
-                seqlen, m_block, n_block_min
+            # First iteration with seqlen masking
+            smem_pipe_read = Int32(0)
+            smem_pipe_write = Int32(self.num_stages - 1)
+            compute_one_n_block(
+                n_block,
+                smem_pipe_read,
+                smem_pipe_write,
+                is_first_n_block=True,
+                seqlen=seqlen,
+                mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=True),
             )
-            for n_tile in cutlass.range(n_block_max - 1 - n_block_min_causal_local_mask, unroll=1):
-                n_block = n_block_max - 2 - n_tile
+            smem_pipe_read = self.advance_pipeline(smem_pipe_read)
+            smem_pipe_write = self.advance_pipeline(smem_pipe_write)
+            # Next couple of iterations with causal masking
+            if const_expr(self.is_causal or self.is_local):
+                n_block_min_causal_local_mask = block_info.get_n_block_min_causal_local_mask(
+                    seqlen, m_block, n_block_min
+                )
+                for n_tile in cutlass.range(n_block_max - 1 - n_block_min_causal_local_mask, unroll=1):
+                    n_block = n_block_max - 2 - n_tile
+                    compute_one_n_block(
+                        n_block,
+                        smem_pipe_read,
+                        smem_pipe_write,
+                        seqlen=seqlen,
+                        mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=True),
+                    )
+                    smem_pipe_read = self.advance_pipeline(smem_pipe_read)
+                    smem_pipe_write = self.advance_pipeline(smem_pipe_write)
+            # The remaining iterations have no masking
+            for n_tile in cutlass.range(n_block, unroll=1):
                 compute_one_n_block(
-                    n_block,
-                    smem_pipe_read,
-                    smem_pipe_write,
-                    check_inf=True,
-                    mask_fn=partial(mask_fn, mask_seqlen=False),
+                    n_block - n_tile - 1, smem_pipe_read, smem_pipe_write,
+                    seqlen=seqlen, is_first_n_block=False,
+                    mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=False)
                 )
                 smem_pipe_read = self.advance_pipeline(smem_pipe_read)
                 smem_pipe_write = self.advance_pipeline(smem_pipe_write)
-        # The remaining iterations have no masking
-        for n_tile in cutlass.range(n_block, unroll=1):
-            compute_one_n_block(
-                n_block - n_tile - 1, smem_pipe_read, smem_pipe_write, check_inf=True
-            )
-            smem_pipe_read = self.advance_pipeline(smem_pipe_read)
-            smem_pipe_write = self.advance_pipeline(smem_pipe_write)
-        # TODO: local
+            # TODO: local
 
-        # normalize acc_O by row_sum and calculate the lse
-        row_scale = softmax.finalize()
-        softmax.rescale_O(acc_O, row_scale)
+            # normalize acc_O by row_sum and calculate the lse
+            row_scale = softmax.finalize()
+            softmax.rescale_O(acc_O, row_scale)
 
-        # ///////////////////////////////////////////////////////////////////////////////
-        # Epilogue
-        # ///////////////////////////////////////////////////////////////////////////////
-        # reuse sQ's data iterator
-        sO = cute.make_tensor(sQ.iterator, sO_layout)
-        self.epilogue(
-            acc_O,
-            softmax.row_sum,
-            mO,
-            mLSE,
-            sO,
-            seqlen,
-            gmem_tiled_copy_O,
-            None,
-            tiled_mma_pv,
-            tidx,
-            m_block,
-            num_head,
-            batch_size,
-        )
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Epilogue
+            # ///////////////////////////////////////////////////////////////////////////////
+            # reuse sQ's data iterator
+            sO = cute.make_tensor(sQ.iterator, sO_layout)
+            self.epilogue(
+                acc_O,
+                softmax.row_sum,
+                mO,
+                mLSE,
+                sO,
+                seqlen,
+                gmem_tiled_copy_O,
+                None,
+                tiled_mma_pv,
+                tidx,
+                m_block,
+                num_head,
+                batch_size,
+            )
 
     @cute.jit
     def compute_one_n_block(
@@ -1047,7 +1110,7 @@ class FlashAttentionForwardSm80(FlashAttentionForwardBase):
         head_idx: cutlass.Int32,
         m_block: cutlass.Int32,
         seqlen: SeqlenInfoQK,
-        aux_tensors=None,
+        aux_data: AuxData = AuxData(),
         fastdiv_mods=None,
         mask_fn: Optional[Callable] = None,
         is_first_n_block: cutlass.Constexpr = False,
@@ -1064,7 +1127,7 @@ class FlashAttentionForwardSm80(FlashAttentionForwardBase):
             cute.arch.barrier()
 
         acc_shape_S = mma_params.thr_mma_qk.partition_shape_C((self.tile_m, self.tile_n))
-        acc_S = cute.make_fragment(acc_shape_S, Float32)
+        acc_S = cute.make_rmem_tensor(acc_shape_S, Float32)
         acc_S.fill(0.0)
         # wait for smem tile QK before mma calculation for S
         sync()
@@ -1102,9 +1165,9 @@ class FlashAttentionForwardSm80(FlashAttentionForwardBase):
                 m_block,
                 acc_S,
                 n_block,
-                seqlen,
                 softmax_scale=softmax.softmax_scale,
-                aux_tensors=aux_tensors,
+                seqlen=seqlen,
+                aux_data=aux_data,
                 fastdiv_mods=fastdiv_mods,
             )
 
@@ -1142,1285 +1205,45 @@ class FlashAttentionForwardSm80(FlashAttentionForwardBase):
         )
         # if const_expr(self.num_stages > 1):
         #     load_K_next()
-
-
-class FlashAttentionForwardSm90(FlashAttentionForwardBase):
-    arch = 90
-
-    def __init__(
+    @cute.jit
+    def apply_score_mod(
         self,
-        *args,
-        intra_wg_overlap: bool = True,
-        mma_pv_is_rs: bool = True,
-        **kwargs,
+        thr_mma_qk,
+        batch_idx,
+        head_idx,
+        m_block,
+        acc_S,
+        n_block,
+        softmax_scale,
+        seqlen,
+        aux_data: AuxData = AuxData(),
+        fastdiv_mods=None,
     ):
-        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
-
-    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
+        # 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)
 
-    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),  # Might need (1, 2, 1) for hdim 512
-            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,
+        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,
         )
-        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 for Q, 1 for O, self.num_stages*2 for K, self.num_stages*2 for V,
-        mbar_ptr_QO_struct = cute.struct.MemRange[cutlass.Int64, 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: mbar_ptr_QO_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: mbar_ptr_QO_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,
-        stream: cuda.CUstream,
-        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_tensors: Optional[list] = 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)
-            )
-        )
-
-        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_mma_warp_groups = self.num_mma_threads // self.num_threads_per_warp_group
-        self.num_threads = self.num_threads_per_warp_group * (self.num_mma_warp_groups + 1)
-        self.num_producer_threads = 32
-        self.num_Q_load_threads = self.num_mma_threads  # If not TMA_Q, MMA threads load Q
-        self.num_epilogue_threads = self.num_mma_threads
-        self.num_mma_regs = (
-            256
-            if self.num_mma_warp_groups == 1
-            else (240 if self.num_mma_warp_groups == 2 else 160)
-        )
-        self.num_producer_regs = (
-            56 if self.num_mma_warp_groups == 1 else (24 if self.num_mma_warp_groups == 2 else 32)
-        )
-        # self.num_mma_regs = 232
-        # self.num_producer_regs = 40
-        self.use_block_sparsity = cutlass.const_expr(blocksparse_tensors is not None)
-
-        self.use_scheduler_barrier = (
-            (self.num_mma_warp_groups >= 2 and self.tile_hdim <= 128)
-            if const_expr(self.intra_wg_overlap)
-            else (self.num_mma_warp_groups == 2)
-        )
-        self.use_tma_Q = self.arch >= 90 and not (
-            self.pack_gqa and self.tile_m % self.qhead_per_kvhead != 0
-        )
-        self.use_tma_O = (
-            self.arch >= 90 and mCuSeqlensQ is None and mSeqUsedQ is None and not self.pack_gqa
-        )
-        # TODO: rescale_O_before_gemm
-        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()
-
-        if const_expr(self.pack_gqa):
-            shape_Q_packed = (
-                (self.qhead_per_kvhead, mQ.shape[0]),
-                mQ.shape[1],
-                mK.shape[2],
-                *mQ.shape[3:],
-            )
-            stride_Q_packed = (
-                (mQ.stride[2], mQ.stride[0]),
-                mQ.stride[1],
-                mQ.stride[2] * self.qhead_per_kvhead,
-                *mQ.stride[3:],
-            )
-            mQ = cute.make_tensor(
-                mQ.iterator, cute.make_layout(shape_Q_packed, stride=stride_Q_packed)
-            )
-            shape_O_packed = (
-                (self.qhead_per_kvhead, mO.shape[0]),
-                mK.shape[1],
-                mK.shape[2],
-                *mO.shape[3:],
-            )
-            stride_O_packed = (
-                (mO.stride[2], mO.stride[0]),
-                mO.stride[1],
-                mO.stride[2] * self.qhead_per_kvhead,
-                *mO.stride[3:],
-            )
-            mO = cute.make_tensor(
-                mO.iterator, cute.make_layout(shape_O_packed, stride=stride_O_packed)
-            )
-            if const_expr(mLSE is not None):
-                shape_LSE_packed = (
-                    (self.qhead_per_kvhead, mLSE.shape[0]),
-                    mK.shape[2],
-                    *mLSE.shape[2:],
-                )
-                stride_LSE_packed = (
-                    (mLSE.stride[1], mLSE.stride[0]),
-                    mLSE.stride[1] * self.qhead_per_kvhead,
-                    *mLSE.stride[2:],
-                )
-                mLSE = cute.make_tensor(
-                    mLSE.iterator, cute.make_layout(shape_LSE_packed, stride=stride_LSE_packed)
-                )
-
-        # 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),
-            ]
-        }
-        tma_atom_Q, tma_tensor_Q = None, None
-        if const_expr(self.use_tma_Q):
-            tma_atom_Q, tma_tensor_Q = cpasync.make_tiled_tma_atom(
-                gmem_tiled_copy_Q,
-                mQ,
-                self.sQ_layout,
-                (self.tile_m, self.tile_hdim),  # No mcast
-            )
-        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):
-            tma_atom_O, tma_tensor_O = cpasync.make_tiled_tma_atom(
-                gmem_tiled_copy_O,
-                mO,
-                self.sO_layout,
-                (self.tile_m, self.tile_hdimv),  # No mcast
-            )
-        if const_expr(mCuSeqlensQ is not None or mSeqUsedQ is not None):
-            TileScheduler = SingleTileVarlenScheduler
-        else:
-            TileScheduler = (
-                SingleTileScheduler
-                if const_expr(not self.is_causal or self.is_local)
-                else SingleTileLPTScheduler
-            )
-        tile_sched_args = TileSchedulerArguments(
-            cute.ceil_div(cute.size(mQ.shape[0]), self.tile_m),
-            cute.size(mQ.shape[2]),
-            cute.size(mQ.shape[3])
-            if const_expr(mCuSeqlensQ is None)
-            else cute.size(mCuSeqlensQ.shape[0] - 1),
-            1,  # num_splits
-            cute.size(mK.shape[0]),
-            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,
-        )
-        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)
-        if const_expr(self.score_mod is None):
-            softmax_scale_log2 = softmax_scale * LOG2_E
-            softmax_scale = None
-        else:
-            # NB: If a user passes in a score mod, we want to apply the score-mod in the sm_scaled qk
-            # But in the original base 10. We hijack softmax_scale_log2 to just be the change of base
-            # and correctly apply the softmax_scale prior to score_mod in the softmax step
-            softmax_scale_log2 = LOG2_E
-            softmax_scale = softmax_scale
-        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)
-
-        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
-            )
-            seqlen_k = (
-                cute.size(mK.shape[0])
-                if const_expr(mPageTable is None)
-                else mK.shape[0] * mPageTable.shape[1]
-            )
-            seqlen_q_divmod = FastDivmodDivisor(seqlen_q)
-            seqlen_k_divmod = FastDivmodDivisor(seqlen_k)
-            fastdiv_mods = (seqlen_q_divmod, seqlen_k_divmod)
-
-        self.kernel(
-            tma_tensor_Q if const_expr(self.use_tma_Q) else mQ,
-            tma_tensor_K,
-            tma_tensor_V,
-            tma_tensor_O if const_expr(self.use_tma_O) else mO,
-            mLSE,
-            mCuSeqlensQ,
-            mCuSeqlensK,
-            mSeqUsedQ,
-            mSeqUsedK,
-            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_tensors,
-            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],
-        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_tensors=Optional[list[cute.Tensor]],
-        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 init
-        mbar_ptr_Q = storage.mbar_ptr.data_ptr()
-        if warp_idx == 1:
-            # if tidx < 2:
-            #     # barrierO num threads should be self.num_mma_threads
-            #     cute.arch.mbarrier_init(mbar_ptr_Q + tidx, 1 if tidx == 0 else self.num_mma_threads)
-            if const_expr(not self.use_tma_Q):
-                cute.arch.mbarrier_init(mbar_ptr_Q, self.num_Q_load_threads)
-            # cute.arch.mbarrier_init(mbar_ptr_Q + 1, self.num_mma_threads)
-        # We rely on pipeline_k and pipeline_v to initialize the mbarrier fence and sync
-        pipeline_kv_producer_group = cutlass.pipeline.CooperativeGroup(
-            cutlass.pipeline.Agent.Thread
-        )
-        pipeline_kv_consumer_group = cutlass.pipeline.CooperativeGroup(
-            cutlass.pipeline.Agent.Thread, self.num_mma_threads // cute.arch.WARP_SIZE
-        )
-        pipeline_k = pipeline.PipelineTmaAsync.create(
-            barrier_storage=storage.mbar_ptr_K.data_ptr(),
-            num_stages=self.num_stages,
-            producer_group=pipeline_kv_producer_group,
-            consumer_group=pipeline_kv_consumer_group,
-            tx_count=self.tma_copy_bytes["K"],
-            defer_sync=True,
-        )
-        pipeline_v = pipeline.PipelineTmaAsync.create(
-            barrier_storage=storage.mbar_ptr_V.data_ptr(),
-            num_stages=self.num_stages,
-            producer_group=pipeline_kv_producer_group,
-            consumer_group=pipeline_kv_consumer_group,
-            tx_count=self.tma_copy_bytes["V"],
-            defer_sync=False
-        )
-
-        # ///////////////////////////////////////////////////////////////////////////////
-        # 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],
-            mCuSeqlensQ=mCuSeqlensQ,
-            mCuSeqlensK=mCuSeqlensK,
-            mSeqUsedQ=mSeqUsedQ,
-            mSeqUsedK=mSeqUsedK,
-        )
-        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)
-
-        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,
-                mbar_ptr_Q,
-                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,
-                mQ,
-                mO,
-                mLSE,
-                sQ,
-                sK,
-                sVt,
-                sP,
-                sO,
-                learnable_sink,
-                pipeline_k,
-                pipeline_v,
-                mbar_ptr_Q,
-                gmem_tiled_copy_Q,
-                gmem_tiled_copy_O,
-                tma_atom_O,
-                tidx,
-                softmax_scale_log2,
-                softmax_scale,
-                block_info,
-                SeqlenInfoCls,
-                AttentionMaskCls,
-                TileSchedulerCls,
-                blocksparse_tensors,
-                aux_tensors,
-                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: cute.CopyAtom,
-        tma_atom_K: cute.CopyAtom,
-        tma_atom_V: cute.CopyAtom,
-        pipeline_k: cutlass.pipeline.PipelineAsync,
-        pipeline_v: cutlass.pipeline.PipelineAsync,
-        mbar_ptr_Q: cutlass.Pointer,
-        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
-        if warp_idx_in_wg == 0:
-            q_producer_phase = Int32(1)
-            kv_producer_state = pipeline.make_pipeline_state(
-                cutlass.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
-                )
-                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))
-                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
-                    )
-                # TODO: mcast
-                # TODO check warp_idx if we have 128 producer threads
-                load_K, _, _ = copy_utils.tma_get_copy_fn(
-                    tma_atom_K, 0, cute.make_layout(1), gK, sK
-                )
-                load_K = copy_utils.tma_producer_copy_fn(load_K, pipeline_k)
-                load_V, _, _ = copy_utils.tma_get_copy_fn(
-                    tma_atom_V, 0, cute.make_layout(1), gV, sV
-                )
-                load_V = copy_utils.tma_producer_copy_fn(load_V, pipeline_v)
-
-                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)
-                    # First iteration: load both Q & K with the same mbarrier
-                    n_block = n_block_max - 1
-                    pipeline_k.producer_acquire(
-                        kv_producer_state,
-                        extra_tx_count=self.tma_copy_bytes["Q"]
-                        if const_expr(self.use_tma_Q)
-                        else 0,
-                    )
-                    if const_expr(self.use_tma_Q):
-                        load_Q(tma_bar_ptr=pipeline_k.producer_get_barrier(kv_producer_state))
-                    load_K(src_idx=n_block, producer_state=kv_producer_state)
-
-                    if const_expr(not self.intra_wg_overlap):
-                        pipeline_v.producer_acquire(kv_producer_state)
-                        load_V(src_idx=n_block, producer_state=kv_producer_state)
-                        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
-                            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:
-                        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
-                            kv_producer_state_prev = kv_producer_state.clone()
-                            kv_producer_state.advance()
-                            pipeline_k.producer_acquire(kv_producer_state)
-                            load_K(src_idx=n_block, producer_state=kv_producer_state)
-                            pipeline_v.producer_acquire(kv_producer_state_prev)
-                            load_V(src_idx=n_block_prev, producer_state=kv_producer_state_prev)
-                        n_block = n_block_min
-                        pipeline_v.producer_acquire(kv_producer_state)
-                        load_V(src_idx=n_block, producer_state=kv_producer_state)
-                        kv_producer_state.advance()
-                else:
-                    kv_producer_state = produce_block_sparse_loads(
-                        blocksparse_tensors,
-                        batch_idx,
-                        head_idx,
-                        m_block,
-                        kv_producer_state,
-                        load_Q,
-                        load_K,
-                        load_V,
-                        pipeline_k,
-                        pipeline_v,
-                        self.use_tma_Q,
-                        self.tma_copy_bytes["Q"],
-                        self.intra_wg_overlap,
-                        self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
-                        self.q_subtile_factor if self.q_subtile_factor is not None else 1,
-                    )
-
-                tile_scheduler.prefetch_next_work()
-                tile_scheduler.advance_to_next_work()
-                work_tile = tile_scheduler.get_current_work()
-                # End of persistent scheduler loop
-
-    @cute.jit
-    def mma(
-        self,
-        tiled_mma_qk: cute.TiledMma,
-        tiled_mma_pv: cute.TiledMma,
-        # softmax: Softmax,
-        # acc_O: cute.Tensor,
-        mQ: cute.Tensor,
-        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: cutlass.pipeline.PipelineAsync,
-        pipeline_v: cutlass.pipeline.PipelineAsync,
-        mbar_ptr_Q: cutlass.Pointer,
-        gmem_tiled_copy_Q: cute.TiledCopy,
-        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_tensors: Optional[list],
-        fastdiv_mods=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_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, 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()
-
-        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,
-        )
-
-        q_consumer_phase = Int32(0)
-        kv_consumer_state = pipeline.make_pipeline_state(
-            cutlass.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,
-        )
-
-        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,
-            softmax=softmax,
-        )
-        process_last_half_block = partial(
-            self.last_half_block_overlap,
-            pipeline_v=pipeline_v,
-            mma_pv_fn=mma_pv_fn,
-        )
-        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_tensors=aux_tensors,
-                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_tensors=aux_tensors,
-                    fastdiv_mods=fastdiv_mods,
-                )
-            mma_one_n_block = partial(
-                mma_one_n_block_all,
-                seqlen=seqlen,
-                softmax=softmax,
-                score_mod_fn=score_mod_fn,
-            )
-            # Load Q if not TMA_Q
-            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
-                )
-                mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx]
-                # gmem_thr_copy_Q = gmem_tiled_copy_Q.get_slice(tidx)
-                # gQ = cute.local_tile(mQ_cur, (self.tile_m, self.tile_hdim), (m_block, 0))
-                # self.load_Q(gmem_thr_copy_Q, gQ, sQ, m_block, seqlen=seqlen.seqlen_q,
-                #             headdim=mQ.shape[1])
-                pack_gqa.load_Q(mQ_cur, sQ, gmem_tiled_copy_Q, tidx, m_block, seqlen.seqlen_q)
-                cute.arch.cp_async_mbarrier_arrive_noinc(mbar_ptr_Q)
-
-            n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block)
-            if const_expr(not self.use_tma_Q):
-                cute.arch.mbarrier_wait(mbar_ptr_Q, phase=q_consumer_phase)
-            q_consumer_phase ^= 1
-            # 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,
-                    )
-                    # Need to initialize tOrO in the case of RescaleOBeforeGemm where we will scale tOrO even in the 1st iter
-                    # acc_O.fill(0.0)
-                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
-                # 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 if self.q_subtile_factor is not None else 1,
-                )
-
-                # Handle empty case (when no blocks to process)
-                if not processed_any:
-                    softmax.reset()
-                    acc_O.fill(0.0)
-
-            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_fragment_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,
-        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)
-
-        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_fragment_like(tOrP_acc, self.dtype)
-        )
-        tOrP_cur.store(tOrP_acc.load().to(self.dtype))
-
-        # if pv gemm not rs
-        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()
-
-        return kv_consumer_state
-
-    @cute.jit
-    def last_half_block_overlap(
-        self,
-        kv_consumer_state,
-        pipeline_v,
-        mma_pv_fn: Callable,
-        zero_init: bool,
-    ):
-        """Processes the final PV GEMM when using intra-warpgroup-overlap"""
-
-        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: cutlass.pipeline.PipelineState | pipeline.PipelineStateSimple,
-        n_block: Int32,
-        mma_qk_fn: Callable,
-        mma_pv_fn: Callable,
-        pipeline_k: cutlass.pipeline.PipelineAsync,
-        pipeline_v: cutlass.pipeline.PipelineAsync,
-        acc_O: cute.Tensor,
-        tOrP: cute.Tensor,
-        smem_copy_params: SimpleNamespace,
-        softmax: Softmax,
-        seqlen: SeqlenInfoQK,
-        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_fragment_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: cutlass.pipeline.PipelineState | pipeline.PipelineStateSimple,
-        n_block: Int32,
-        mma_qk_fn: Callable,
-        mma_pv_fn: Callable,
-        pipeline_k: cutlass.pipeline.PipelineAsync,
-        pipeline_v: cutlass.pipeline.PipelineAsync,
-        acc_O: cute.Tensor,
-        tOrP: cute.Tensor,
-        smem_copy_params: SimpleNamespace,
-        softmax: Softmax,
-        seqlen: SeqlenInfoQK,
-        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)
-        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_fragment_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)
-        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
-        return smem_pipe_read
-
-    @cute.jit
-    def mma_init(self):
-        warp_group_idx = utils.canonical_warp_group_idx(sync=False)
-        if const_expr(self.use_scheduler_barrier):
-            if warp_group_idx == 1:
-                cute.arch.barrier_arrive(
-                    barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1),
-                    number_of_threads=2 * self.num_threads_per_warp_group,
-                )
-
-    @cute.jit
-    def apply_score_mod(
-        self,
-        thr_mma_qk,
-        batch_idx,
-        head_idx,
-        m_block,
-        acc_S,
-        n_block,
-        softmax_scale,
-        seqlen,
-        aux_tensors: Optional[list] = None,
-        fastdiv_mods=None,
-    ):
-        # Prepare index tensor
-        cS = cute.make_identity_tensor((self.tile_m, self.tile_n))
-        cS = cute.domain_offset((m_block * self.tile_m, n_block * self.tile_n), cS)
-        tScS = thr_mma_qk.partition_C(cS)
-
-        apply_score_mod_inner(
-            acc_S,
-            tScS,
-            self.score_mod,
-            batch_idx,
-            head_idx,
-            softmax_scale,
-            self.vec_size,
-            self.qk_acc_dtype,
-            aux_tensors,
-            fastdiv_mods,
-            seqlen_info=seqlen,
-            constant_q_idx=None,
-            qhead_per_kvhead=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
-        )
-
-    def warp_scheduler_barrier_sync(self):
-        if const_expr(self.use_scheduler_barrier):
-            cute.arch.barrier(
-                barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1)
-                - 1
-                + utils.canonical_warp_group_idx(sync=False),
-                number_of_threads=2 * self.num_threads_per_warp_group,
-            )
-
-    def warp_scheduler_barrier_arrive(self):
-        if const_expr(self.use_scheduler_barrier):
-            assert self.num_mma_warp_groups in [2, 3]
-            cur_wg = utils.canonical_warp_group_idx(sync=False) - 1
-            if const_expr(self.num_mma_warp_groups == 2):
-                next_wg = 1 - cur_wg
-            else:
-                t = cur_wg + 1
-                next_wg = t % self.num_mma_warp_groups
-            cute.arch.barrier_arrive(
-                barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1) + next_wg,
-                number_of_threads=2 * self.num_threads_per_warp_group,
-            )
+# SM90 forward pass moved to flash_fwd_sm90.py; re-export for backward compatibility
+def __getattr__(name):
+    if name == "FlashAttentionForwardSm90":
+        from .flash_fwd_sm90 import FlashAttentionForwardSm90
+        return FlashAttentionForwardSm90
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/build/torch-cuda/flash_fwd_combine.py b/build/torch-cuda/flash_fwd_combine.py
index 478218b11f9cf827d1c86723b791892182cfeb2b..b48eb86569a2ede835af11023d7ec6683882f089 100644
--- a/build/torch-cuda/flash_fwd_combine.py
+++ b/build/torch-cuda/flash_fwd_combine.py
@@ -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 typing import Callable, 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 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
 
 
@@ -24,7 +31,8 @@ class FlashAttentionForwardCombine:
         dtype: Type[cutlass.Numeric],
         dtype_partial: Type[cutlass.Numeric],
         head_dim: int,
-        m_block_size: int = 8,
+        num_head: int,
+        tile_m: 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 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
@@ -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.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
@@ -58,7 +68,7 @@ class FlashAttentionForwardCombine:
         dtype,
         dtype_partial,
         head_dim,
-        m_block_size,
+        tile_m,
         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:
+        if tile_m % 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:
+        if (tile_m * 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
+            if self.tile_m % 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)
-                )
+                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
@@ -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)
+            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.m_block_size, self.k_block_size, self.stages), order=(1, 0, 2)
+            (self.tile_m, 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,
+        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
@@ -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]
+            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
             ]
@@ -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,
+        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,
@@ -303,6 +329,7 @@ class FlashAttentionForwardCombine:
             cu_seqlens,
             seqused,
             num_splits_dynamic_ptr,
+            virtual_batch_idx,
             semaphore_to_reset,
             SharedStorage,
             self.smem_layout_lse,
@@ -314,6 +341,8 @@ class FlashAttentionForwardCombine:
             seqlen_divmod,
             head_divmod,
             varlen,
+            tile_sched_params,
+            TileScheduler,
         ).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],
+        virtual_batch_idx: 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],
+        tile_sched_params: ParamsBase,
+        TileScheduler: cutlass.Constexpr[Callable],
     ):
         # Thread and block indices
         tidx, _, _ = cute.arch.thread_idx()
-        m_block, k_block, batch_idx = cute.arch.block_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
@@ -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,))
+        sMaxValidSplit = storage.sMaxValidSplit.get_tensor((self.tile_m,))
         sO = storage.sO.get_tensor(smem_layout_o)
 
-        # Handle semaphore reset
+        # 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 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
+                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
 
-        # 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,
+        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
 
-            # 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)
+            # 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 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)
+                # ===============================
+                # 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):
-                    scale[m] = sLSE[s, tOcO[0, m, 0][0]]  # Get scale from smem
+                    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 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
+                # 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 the current stage to be ready
+                # Wait for LSE and initial O partial stages to complete
                 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
+                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)
+                            )
 
-                # Accumulate scaled partial results
+                # ===============================
+                # 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 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)
+                    if tOhidx[m] >= 0:
+                        mO_cur_copy = cute.tiled_divide(
+                            mO_cur[tOmidx[m], None, tOhidx[m]], (elems_per_store,)
                         )
-
-            # ===============================
-            # 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])
+                        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,
+        tOpO: Optional[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]:
+                    if const_expr(tOpO is None) or tOpO[k]:
                         cute.copy(
                             gmem_tiled_copy_O_partial,
                             mO_partial_cur_copy[None, k_idx, split],
diff --git a/build/torch-cuda/flash_fwd_mla_sm100.py b/build/torch-cuda/flash_fwd_mla_sm100.py
new file mode 100644
index 0000000000000000000000000000000000000000..3787fb29ba29be396ee7d18b72633d8c17a7f18b
--- /dev/null
+++ b/build/torch-cuda/flash_fwd_mla_sm100.py
@@ -0,0 +1,3160 @@
+# 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, Uint32, Boolean, const_expr
+from cutlass.cute import FastDivmodDivisor
+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
+
+from .pack_gqa import pack_gqa_layout, make_packgqa_tiled_tma_atom
+from .paged_kv import PagedKVManager
+from . import utils as fa_utils
+from .seqlen_info import SeqlenInfoQK
+from .block_info import BlockInfo
+from .mask import AttentionMask
+from . import blackwell_helpers as fa_sm100_utils
+from .softmax import SoftmaxSm100
+from .tile_scheduler import (
+    SchedulerState,
+    SchedulingMode,
+    TileSchedulerArguments,
+    TileSchedulerProtocol,
+    SingleTileScheduler,
+    SingleTileLPTScheduler,
+    SingleTileVarlenScheduler,
+    ParamsBase,
+)
+from .fa_logging import fa_log, fa_printf
+from .utils import smid, get_batch_from_cu_tensor
+
+from .topk_gather_kv import CpasyncGatherKVManager
+
+
+from .named_barrier import NamedBarrierFwdSm100_MLA2CTA
+
+
+class FlashAttentionMLAForwardSm100:
+    def __init__(
+        self,
+        is_causal: bool = False,
+        use_cpasync_load_KV: bool = False,
+        topk_length: int = 2048,
+        is_topk_gather: bool = True,
+        pack_gqa: bool = False,
+        qhead_per_kvhead: int = 1,
+        nheads_kv: int = 1,
+        hdim: int = 64,
+        hdimv: int = 512,
+        has_seqused_q: bool = False,
+        has_cu_seqlens_q: bool = False,
+        disable_bitmask: bool = False,
+        use_clc_scheduler: bool = True,
+        has_qk: bool = True,
+    ):
+        self.is_causal = is_causal
+        self.is_local = False
+        self.pack_gqa = pack_gqa
+        self.qhead_per_kvhead = qhead_per_kvhead
+        assert qhead_per_kvhead <= 128
+        self.nheads_kv = nheads_kv
+        self.use_tma_O = True
+        self.use_cpasync_load_KV = use_cpasync_load_KV
+        self.use_tma_KV = not use_cpasync_load_KV
+        self.topk_length = topk_length
+        self.is_topk_gather = is_topk_gather
+        if is_topk_gather:
+            assert pack_gqa
+            assert qhead_per_kvhead == 128, "require MQA 128 for DSA path"
+            assert use_cpasync_load_KV
+        # user-provided option if topk indices guaranteed in bounds
+        self.disable_bitmask = disable_bitmask
+        self.has_qk = has_qk
+
+        # ==== tile scheduler ====
+        self.is_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
+        )
+
+        self.is_varlen_q = has_seqused_q or has_cu_seqlens_q
+        self.use_packed_varlen_sched = has_cu_seqlens_q and qhead_per_kvhead == 128 and pack_gqa
+        self.use_varlen_scheduler = self.is_varlen_q and not self.use_packed_varlen_sched
+
+        if const_expr(self.use_varlen_scheduler):
+            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 = 32 if use_cpasync_load_KV else 64
+        self.num_relay_threads = 32 if use_cpasync_load_KV else 0
+        self.num_cpasync_load_threads = 128 if use_cpasync_load_KV else 0
+        self.num_threads = (
+            self.num_softmax_threads
+            + self.num_epilogue_threads
+            + self.num_load_threads
+            + self.num_mma_threads
+            + self.num_empty_threads
+            + self.num_relay_threads
+            + self.num_cpasync_load_threads
+        )
+        self.num_warps = self.num_threads // 32
+        assert self.num_warps == 12 or self.num_warps == 16
+        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.empty_warp_ids = tuple(
+            w
+            for w, active in [
+                (self.relay_warp_id, not use_cpasync_load_KV),
+                (self.clc_scheduler_warp_id, not self.use_clc_scheduler),
+            ]
+            if active
+        )
+        self.cpasync_load_warp_indices = (12, 13, 14, 15)
+
+        # ==== register usage ====
+        if self.num_warps == 16:
+            self.num_regs_load = 112
+            self.num_regs_mma = 112
+            self.num_regs_softmax = 192
+            self.num_regs_epilogue = 128
+            self.num_regs_cpasync = 80 if self.use_cpasync_load_KV else 0
+            self.num_regs_other = 48
+        else:
+            self.num_regs_load = 168 - 40
+            self.num_regs_mma = 168 - 40
+            self.num_regs_softmax = 168 + 80
+            self.num_regs_epilogue = 168 - 40
+            self.num_regs_cpasync = 0
+            self.num_regs_other = 48
+
+        self.num_regs_per_thread = 168 if self.num_warps == 12 else 128
+        self.num_regs_total = 504 if self.num_warps == 12 else 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
+        self.hdimv = hdimv
+        self.cta_tile_m = 64
+        self.cluster_tile_m = self.cta_group_size * self.cta_tile_m
+        self.tile_n = 128
+        assert (
+            pack_gqa is False
+            or self.cluster_tile_m % qhead_per_kvhead == 0
+            or qhead_per_kvhead % self.cluster_tile_m == 0
+        )
+        self.num_hdimv_splits = 2  # split hdimv in half for our Qv @ V^T and P @ V mmas.
+        assert hdimv % 32 == 0
+        assert self.topk_length % self.tile_n == 0 or not self.is_topk_gather
+        self.epi_tile = (self.cta_tile_m, self.hdimv // self.num_hdimv_splits)
+        self.tile_P = (self.cta_tile_m, self.tile_n)
+
+        # ==== MMA info ====
+        self.mma_tiler_QK = (
+            self.cluster_tile_m,
+            self.tile_n,
+            self.hdim,
+        )
+        self.mma_tiler_QvV = (
+            self.cluster_tile_m,
+            self.tile_n,
+            self.hdimv // self.num_hdimv_splits,
+        )
+        self.mma_tiler_PVt = (
+            self.cluster_tile_m,
+            self.hdimv // self.num_hdimv_splits,
+            self.tile_n,
+        )
+        self.major_mode_Q = tcgen05.OperandMajorMode.K
+        self.major_mode_Qvi = tcgen05.OperandMajorMode.K
+        self.major_mode_K = tcgen05.OperandMajorMode.K
+        self.major_mode_Vi = tcgen05.OperandMajorMode.K
+        self.major_mode_Vti = tcgen05.OperandMajorMode.MN
+        self.major_mode_P = tcgen05.OperandMajorMode.K
+        self.operand_source_Q = tcgen05.OperandSource.SMEM
+        self.operand_source_Qvi = tcgen05.OperandSource.SMEM
+        self.operand_source_P = tcgen05.OperandSource.SMEM
+
+        # ==== pipeline info ====
+        self.num_stages_Q = 1
+        self.num_stages_K = 1
+        self.num_stages_Qv = 2
+        self.num_stages_V = 4
+        self.num_stages_S = 2
+        # self.num_stages_P = 1 if has_qk else 2
+        self.num_stages_P = 1
+        self.num_stages_Oi = 1
+        self.num_stages_sm_stats = 2
+        self.num_stages_bitmask = 2
+        assert self.num_stages_S == 2, "mainloops expect 2 stages for S"
+
+        # ==== dtype info ====
+        self.dtype_acc = Float32
+
+        # ==== TMEM info ====
+        SM100_TMEM_CAPACITY_COLUMNS = 512
+        self.tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS
+        self.tmem_cols_S = self.tile_n // self.cta_group_size
+        self.tmem_cols_Oi = (self.hdimv // self.num_hdimv_splits) // self.cta_group_size
+        self.tmem_offset_S = [
+            self.tmem_cols_S * stage for stage in range(self.num_stages_S)
+        ]  # allocate 64 TMEM columns for each stage of S
+        self.tmem_offset_O0 = self.tmem_cols_S * self.num_stages_S
+        self.tmem_offset_O1 = self.tmem_offset_O0 + self.tmem_cols_Oi
+        self.tmem_offsets_O = [self.tmem_offset_O0, self.tmem_offset_O1]
+        self.total_tmem = self.tmem_offset_O1 + self.tmem_cols_Oi
+        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, disabled=False):
+            if disabled:
+                return cute.struct.MemRange[dtype, 0]
+            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]
+
+        (sQ_struct, sK_struct, sQv_struct, sV_struct, sP_struct) = (
+            smem_struct_align(dtype, layout, disabled)
+            for dtype, layout, disabled in [
+                (self.dtype_Q, self.sQ_layout_staged, not self.has_qk),
+                (self.dtype_K, self.sK_layout_staged, not self.has_qk),
+                (self.dtype_Qv, self.sQv_layout_staged, False),
+                (self.dtype_V, self.sV_layout_staged, False),
+                (self.dtype_P, self.sP_layout_staged, False),
+            ]
+        )
+        sStats_struct = cute.struct.MemRange[Float32, cute.cosize(self.sStats_layout)]
+        sScale_struct = cute.struct.MemRange[Float32, cute.cosize(self.sScale_layout)]
+        sBitmask_struct = cute.struct.MemRange[Uint32, cute.cosize(self.sBitmask_layout)]
+
+        (
+            mbar_ptr_Q_struct,
+            mbar_ptr_K_struct,
+            mbar_ptr_Qv_struct,
+            mbar_ptr_V_struct,
+            mbar_ptr_S_struct,
+            mbar_ptr_P_struct,
+            mbar_ptr_O0_struct,
+            mbar_ptr_O1_struct,
+            mbar_sm_stats_struct,
+            mbar_bitmask_struct,
+        ) = (
+            mbar_struct(n)
+            for n in [
+                self.num_stages_Q,
+                self.num_stages_K,
+                self.num_stages_Qv,
+                self.num_stages_V,
+                self.num_stages_S,
+                self.num_stages_P,
+                self.num_stages_Oi,
+                self.num_stages_Oi,
+                self.num_stages_sm_stats,
+                self.num_stages_bitmask,
+            ]
+        )
+        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_Q: mbar_ptr_Q_struct
+            mbar_ptr_K: mbar_ptr_K_struct
+            mbar_ptr_Qv: mbar_ptr_Qv_struct
+            mbar_ptr_V: mbar_ptr_V_struct
+            mbar_ptr_S: mbar_ptr_S_struct
+            mbar_ptr_P: mbar_ptr_P_struct
+            mbar_ptr_O0: mbar_ptr_O0_struct
+            mbar_ptr_O1: mbar_ptr_O1_struct
+            mbar_ptr_K_cpasync: mbar_ptr_K_struct
+            mbar_ptr_V_cpasync: mbar_ptr_V_struct
+            mbar_ptr_sm_stats: mbar_sm_stats_struct
+            mbar_ptr_bitmask: mbar_bitmask_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]
+            sO_empty_mbar_ptr: cutlass.Int64
+
+            sRowMax: sStats_struct
+            sRowSum: sStats_struct
+            sScale: sScale_struct
+            sBitmask: sBitmask_struct
+            sQv: sQv_struct
+            sQ: sQ_struct
+            sK: sK_struct
+            sV: sV_struct
+            sP: sP_struct
+
+        # print("smem bytes = ", SharedStorage.size_in_bytes())
+
+        return SharedStorage
+
+    # fmt: off
+    @cute.jit
+    def __call__(
+        self,
+        mQ: Optional[cute.Tensor],    # (b, s_q, h, d)     or (total_q, h, d)    if there is cu_seqlens_q
+        mQv: cute.Tensor,             # (b, s_q, h, dv)    or (total_q, h, d)    if there is cu_seqlens_q
+        mK: Optional[cute.Tensor],    # (b, 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, 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],  # (b, s_q, h)        or (total_q, h)       if there is cu_seqlens_q
+        softmax_scale: Float32,
+        mP: Optional[cute.Tensor] = None,           # (b, s_q, h, topk)            or (total_q, h, topk)           if there is cu_seqlens_q
+        mRowMax: Optional[cute.Tensor] = None,      # (b, s_q, topk // tile_n, h)  or (total_q, topk // tile_n, 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)
+        mIndexTopk: Optional[cute.Tensor] = None,   # (b, s_q, topk)  or (total_q, topk) if there is cu_seqlens_q
+        mPageTable: Optional[cute.Tensor] = None,
+        window_size_left: Int32 | int | None = None,
+        window_size_right: Int32 | int | None = None,
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
+        stream: cuda.CUstream = None,
+    ):
+        # fmt: on
+        self.store_P = mP is not None
+        self.store_row_max = mRowMax is not None
+
+        if const_expr(self.has_qk):
+            assert mQ is not None and mK is not None, "has_qk requires mQ and mK"
+        else:
+            assert mQ is None and mK is None, "not has_qk disallows mQ and mK"
+
+        # ==== dtype info ====
+        self.dtype_Q = mQ.element_type if self.has_qk else cutlass.BFloat16
+        self.dtype_K = mK.element_type if self.has_qk else cutlass.BFloat16
+        self.dtype_Qv = mQv.element_type
+        self.dtype_V = mV.element_type
+        self.dtype_P = mV.element_type
+        self.dtype_O = mO.element_type
+
+        if const_expr(self.store_P):
+            assert mP.element_type == self.dtype_P
+
+        # ==== Prepare Tensors ====
+        new_stride = lambda mX: (
+            *(cute.assume(s, divby=128 // mX.element_type.width) for s in mX.stride[:-1]),
+            mX.stride[-1],
+        )
+        mQ, mQv, mK, mV, mO, mP = [
+            cute.make_tensor(mX.iterator, cute.make_layout(mX.shape, stride=new_stride(mX)))
+            if mX is not None
+            else None
+            for mX in (mQ, mQv, mK, mV, mO, mP)
+        ]
+
+        # (b, s, h, d)  -> (s, d, h, b)  or
+        # (total, h, d) -> (total, d, h) or
+        # (num_pages, page_size, h_k, d) -> (page_size, d, h_k, num_pages)
+        QO_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1]
+        KV_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensK is None) else [0, 2, 1]
+        mQ, mQv, mO, mP = [
+            cute.make_tensor(mX.iterator, cute.select(mX.layout, mode=QO_layout_transpose))
+            if mX is not None
+            else None
+            for mX in (mQ, mQv, mO, mP)
+        ]
+        mK, mV = [
+            cute.make_tensor(mX.iterator, cute.select(mX.layout, mode=KV_layout_transpose))
+            if mX is not None
+            else None
+            for mX in (mK, mV)
+        ]
+        # (s_k, dv, h_k, b)  -> (dv, s_k, h_k, b) or
+        # (total_k, dv, h_k) -> (dv, total_k, h_k)
+        V_layout_transpose = [1, 0, 2, 3] if const_expr(mCuSeqlensK is None) else [1, 0, 2]
+        mVt = cute.make_tensor(mV.iterator, cute.select(mV.layout, mode=V_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)
+            )
+            if mIndexTopk is not None
+            else None
+        )
+        # (b, s_q, h) -> (s_q, h, b) or (total_q, h) -> (total_q, h)
+        LSE_layout_transpose = [1, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 1]
+        mLSE = (
+            cute.make_tensor(mLSE.iterator, cute.select(mLSE.layout, mode=LSE_layout_transpose))
+            if mLSE is not None
+            else None
+        )
+        # (b, s, topk//128, h) => (s, topk//128, h, b) or
+        # (total, topk//128, h) == (total, topk//128, 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 = cute.make_tensor(
+                mRowMax.iterator, cute.select(mRowMax.layout, mode=rowmax_layout_transpose)
+            )
+
+        topk_length_dynamic = mIndexTopk.shape[0] if mIndexTopk is not None else None
+
+        self.o_layout = cutlass.utils.LayoutEnum.from_tensor(mO)
+        self.p_layout = cutlass.utils.LayoutEnum.ROW_MAJOR
+        if const_expr(self.store_P):
+            assert cutlass.utils.LayoutEnum.from_tensor(mP) == self.p_layout
+
+        mO_og = mO
+        mP_og = mP
+        if const_expr(self.pack_gqa):
+            mQ, mQv, mO, mP, mRowMax = [
+                pack_gqa_layout(mX, self.qhead_per_kvhead, self.nheads_kv, head_idx=2)
+                if mX is not None
+                else None
+                for mX in (mQ, mQv, mO, mP, mRowMax)
+            ]
+            if const_expr(mLSE is not None):
+                mLSE = pack_gqa_layout(mLSE, self.qhead_per_kvhead, self.nheads_kv, head_idx=1)
+
+        # ==== Prepare MMAs ====
+        # (local_var, dtype_a, major_a, major_b, mma_tiler, operand_source_a)
+        # fmt: off
+        _mma_specs = [
+            ("tiled_mma_QK",    self.dtype_Q,  self.major_mode_Q,   self.major_mode_K,   self.mma_tiler_QK,    self.operand_source_Q),
+            ("tiled_mma_QvV", self.dtype_Qv, self.major_mode_Qvi, self.major_mode_Vi,  self.mma_tiler_QvV, self.operand_source_Qvi),
+            ("tiled_mma_PVt",  self.dtype_P,  self.major_mode_P,   self.major_mode_Vti, self.mma_tiler_PVt,  self.operand_source_P),
+        ]
+        tiled_mma_QK, tiled_mma_QvV, tiled_mma_PVt = (
+            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
+        )
+        # fmt: on
+
+        # ==== Prepare SMEM layouts and TMAs ====
+        # (attr, make_fn, tiled_mma, mma_tiler, dtype, num_stages)
+        # fmt: off
+        _smem_layout_specs = [
+            ("sQ_layout",  sm100_utils.make_smem_layout_a, tiled_mma_QK,  self.mma_tiler_QK,  self.dtype_Q,  self.num_stages_Q),
+            ("sK_layout",  sm100_utils.make_smem_layout_b, tiled_mma_QK,  self.mma_tiler_QK,  self.dtype_K,  self.num_stages_K),
+            ("sP_layout",  sm100_utils.make_smem_layout_a, tiled_mma_PVt, self.mma_tiler_PVt, self.dtype_P,  self.num_stages_P),
+            ("sQv_layout", sm100_utils.make_smem_layout_a, tiled_mma_QvV, self.mma_tiler_QvV, self.dtype_Qv, self.num_stages_Qv),
+            ("sV_layout",  sm100_utils.make_smem_layout_b, tiled_mma_QvV, self.mma_tiler_QvV, self.dtype_V,  self.num_stages_V),
+            ("sVt_layout", sm100_utils.make_smem_layout_b, tiled_mma_PVt, self.mma_tiler_PVt, self.dtype_V,  self.num_stages_V),
+        ]
+        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]))
+        # fmt: on
+
+        self.sStats_layout = cute.make_layout((self.cta_tile_m, self.cta_group_size))
+        self.sScale_layout = cute.make_layout((self.cta_tile_m, self.num_stages_sm_stats))
+        self.sBitmask_layout = cute.make_layout((self.tile_n // 32, self.num_stages_bitmask))
+
+        # fmt: off
+        for attr, dtype, layout in [
+            ("tma_copy_bytes_Q",   self.dtype_Q,  self.sQ_layout),
+            ("tma_copy_bytes_K",   self.dtype_K,  self.sK_layout),
+            ("tma_copy_bytes_Qvi", self.dtype_Qv, self.sQv_layout),
+            ("tma_copy_bytes_Vi",  self.dtype_V,  self.sV_layout),
+        ]:
+            setattr(self, attr, cute.size_in_bytes(dtype, layout) * self.cta_group_size)
+        # fmt: on
+
+        tma_load_op = cpasync.CopyBulkTensorTileG2SOp(self.cta_group)
+        cta_layout_vmnk = cute.tiled_divide(
+            cute.make_layout(self.cluster_shape_mnk), (tiled_mma_QK.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, kv_only)
+        # fmt: off
+        _tma_specs = [
+            ("tma_atom_Q",  "tma_tensor_Q",  A, mQ,  self.sQ_layout,  self.mma_tiler_QK,  tiled_mma_QK,  False),
+            ("tma_atom_Qv", "tma_tensor_Qv", A, mQv, self.sQv_layout, self.mma_tiler_QvV, tiled_mma_QvV, False),
+            ("tma_atom_K",  "tma_tensor_K",  B, mK,  self.sK_layout,  self.mma_tiler_QK,  tiled_mma_QK,  True),
+            ("tma_atom_V",  "tma_tensor_V",  B, mV,  self.sV_layout,  self.mma_tiler_QvV, tiled_mma_QvV, True),
+            ("tma_atom_Vt", "tma_tensor_Vt", B, mVt, self.sVt_layout, self.mma_tiler_PVt, tiled_mma_PVt, True),
+        ]
+        _tmas = {}
+        for atom_name, tensor_name, make_fn, m, smem_layout, mma_tiler, tiled_mma, kv_only in _tma_specs:
+            _tmas[atom_name], _tmas[tensor_name] = (
+                make_tma(make_fn, m, smem_layout, mma_tiler, tiled_mma)
+                if const_expr((not kv_only or self.use_tma_KV) and m is not None)
+                else (None, None)
+            )
+
+        (tma_atom_Q,  tma_tensor_Q,
+         tma_atom_Qv, tma_tensor_Qv,
+         tma_atom_K,  tma_tensor_K,
+         tma_atom_V,  tma_tensor_V,
+         tma_atom_Vt, tma_tensor_Vt) = _tmas.values()
+        # fmt: on
+
+        tma_store_op = cpasync.CopyBulkTensorTileS2GOp()
+        self.ragged_tma_O = (
+            self.use_tma_O
+            and self.is_varlen_q
+            and self.pack_gqa
+            and self.cta_tile_m % self.qhead_per_kvhead == 0
+        )
+        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.ragged_tma_O)
+            else cpasync.make_tiled_tma_atom
+        )
+
+        # ==== Set up P smem -> gmem tma store ====
+
+        # S<3,4,3> o 0 o ((8,8),(64,2),(1,1)):((64,512),(1,4096),(0,0))
+        sP_layout_out = sm100_utils.make_smem_layout_epi(
+            self.dtype_P, self.p_layout, self.tile_P, self.num_stages_P
+        )
+
+        if const_expr(self.store_P):
+            # TODO: add asserts
+            mP_tma = mP_og if const_expr(self.ragged_tma_O) else mP
+            if const_expr(self.ragged_tma_O):
+                mP_tma = copy_utils.create_ragged_tensor_for_tma(
+                    mP_tma, ragged_dim=0, ptr_shift=True
+                )
+            tma_atom_P, tma_tensor_P = make_tiled_tma_atom_fn(
+                tma_store_op, mP_tma, cute.select(sP_layout_out, mode=[0, 1]), self.tile_P
+            )
+        else:
+            tma_atom_P = None
+            tma_tensor_P = None
+
+        # ==== Set up Oi smem -> gmem tma store ====
+
+        self.overlap_sO_sV = True
+        if const_expr(self.overlap_sO_sV):
+            num_stages_sO = self.num_stages_V
+        else:
+            num_stages_sO = self.num_hdimv_splits
+        sO_layout = sm100_utils.make_smem_layout_epi(
+            self.dtype_O, self.o_layout, self.epi_tile, num_stages_sO
+        )
+
+        if const_expr(self.use_tma_O):
+            mO_tma = mO_og if const_expr(self.ragged_tma_O) else mO
+            if const_expr(self.ragged_tma_O):
+                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(
+                tma_store_op, mO_tma, cute.select(sO_layout, mode=[0, 1]), self.epi_tile
+            )
+        else:
+            tma_atom_O = None
+            tma_tensor_O = None
+
+        # ==== Set up Oi rmem -> gmem copy ====
+        universal_copy_bits = 128
+        atom_universal_copy = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            self.dtype_O,
+            num_bits_per_copy=universal_copy_bits,
+        )
+        thread_layout_O_r2g = cute.make_layout((64, 2), stride=(1, 64))
+        value_layout_O_r2g = cute.make_layout(
+            (1, self.hdimv // self.num_hdimv_splits // self.cta_group_size)
+        )
+        tiled_copy_O_r2g = cute.make_tiled_copy_tv(
+            atom=atom_universal_copy,
+            thr_layout=thread_layout_O_r2g,
+            val_layout=value_layout_O_r2g,
+        )
+
+        # ==== Allocate shared memory ====
+        SharedStorage = self._get_shared_storage_cls()
+
+        # ==== Tile scheduler ====
+
+        TileScheduler = self.TileScheduler
+        
+        batch_size_for_sched = (
+            cute.size(mQv.shape[3]) if const_expr(mCuSeqlensQ is None)
+            else cute.size(mCuSeqlensQ.shape[0] - 1) if self.use_varlen_scheduler
+            else 1
+        )
+
+        tile_sched_args = TileSchedulerArguments(
+            num_block=cute.ceil_div(cute.size(mQv.shape[0]), self.cluster_tile_m),
+            num_head=cute.size(mQv.shape[2]),
+            num_batch=batch_size_for_sched,
+            num_splits=1,  # todo: split_kv
+            seqlen_k=cute.size(mV.shape[0])
+            if const_expr(mPageTable is None)
+            else cute.size(mV.shape[0]) * cute.size(mPageTable.shape[1]),
+            headdim=self.hdim,
+            headdim_v=self.hdimv,
+            total_q=cute.size(mQv.shape[0])
+            if const_expr(mCuSeqlensQ is not None)
+            else cute.size(mQv.shape[0]) * cute.size(mQv.shape[3]),
+            tile_shape_mn=(self.cta_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_K.width // 8,
+            is_persistent=self.is_persistent,
+            # lpt=self.is_causal or self.is_local,
+            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(NamedBarrierFwdSm100_MLA2CTA.Cpasync),
+            num_threads=self.num_cpasync_load_threads,
+        )
+        self.softmax_barrier = cutlass.pipeline.NamedBarrier(
+            barrier_id=int(NamedBarrierFwdSm100_MLA2CTA.Softmax),
+            num_threads=self.num_softmax_threads,
+        )
+        self.epi_barrier = cutlass.pipeline.NamedBarrier(
+            barrier_id=int(NamedBarrierFwdSm100_MLA2CTA.Epilogue),
+            num_threads=self.num_epilogue_threads,
+        )
+        # softmax -> correction
+        self.sm_stats_barrier_full = cutlass.pipeline.NamedBarrier(
+            barrier_id=int(NamedBarrierFwdSm100_MLA2CTA.SoftmaxStatsFull),
+            num_threads=self.num_softmax_threads + self.num_epilogue_threads,
+        )
+        self.sm_stats_barrier_empty = cutlass.pipeline.NamedBarrier(
+            barrier_id=int(NamedBarrierFwdSm100_MLA2CTA.SoftmaxStatsEmpty),
+            num_threads=self.num_softmax_threads + self.num_epilogue_threads,
+        )
+
+        LOG2_E = math.log2(math.e)
+        softmax_scale_log2 = softmax_scale * LOG2_E
+
+        # ==== Launch kernel ====
+        self.kernel(
+            tma_tensor_Q,
+            tma_tensor_Qv,
+            tma_tensor_K if self.use_tma_KV else mK,
+            tma_tensor_V if self.use_tma_KV else mV,
+            tma_tensor_Vt if self.use_tma_KV else mVt,
+            tma_tensor_O if self.use_tma_O else mO,
+            tma_tensor_P,
+            mLSE,
+            mRowMax,
+            mCuSeqlensQ,
+            mCuSeqlensK,
+            mSeqUsedQ,
+            mSeqUsedK,
+            mIndexTopk,
+            mPageTable,
+            tma_atom_Q,
+            tma_atom_Qv,
+            tma_atom_K,
+            tma_atom_V,
+            tma_atom_Vt,
+            tma_atom_O,
+            tma_atom_P,
+            tiled_copy_O_r2g,
+            self.sQ_layout_staged,
+            self.sK_layout_staged,
+            self.sQv_layout_staged,
+            self.sV_layout_staged,
+            self.sVt_layout_staged,
+            self.sP_layout_staged,
+            self.sStats_layout,
+            self.sScale_layout,
+            self.sBitmask_layout,
+            sO_layout,
+            sP_layout_out,
+            tiled_mma_QK,
+            tiled_mma_QvV,
+            tiled_mma_PVt,
+            softmax_scale,
+            softmax_scale_log2,
+            topk_length_dynamic,
+            tile_sched_params,
+            SharedStorage,
+        ).launch(
+            grid=grid_dim,
+            block=(
+                self.num_threads,
+                1,
+                1,
+            ),
+            cluster=self.cluster_shape_mnk,
+            smem=SharedStorage.size_in_bytes(),
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        mQ: Optional[cute.Tensor],
+        mQv: cute.Tensor,
+        mK: Optional[cute.Tensor],
+        mV: cute.Tensor,
+        mVt: cute.Tensor,
+        mO: cute.Tensor,
+        mP: Optional[cute.Tensor],
+        mLSE: Optional[cute.Tensor],
+        mRowMax: Optional[cute.Tensor],
+        mCuSeqlensQ: Optional[cute.Tensor],
+        mCuSeqlensK: Optional[cute.Tensor],
+        mSeqUsedQ: Optional[cute.Tensor],
+        mSeqUsedK: Optional[cute.Tensor],
+        mIndexTopk: Optional[cute.Tensor],
+        mPageTable: Optional[cute.Tensor],
+        tma_atom_Q: cute.CopyAtom,
+        tma_atom_Qv: cute.CopyAtom,
+        tma_atom_K: Optional[cute.CopyAtom],
+        tma_atom_V: Optional[cute.CopyAtom],
+        tma_atom_Vt: Optional[cute.CopyAtom],
+        tma_atom_O: Optional[cute.CopyAtom],
+        tma_atom_P: Optional[cute.CopyAtom],
+        tiled_copy_O_r2g: cute.TiledCopy,
+        sQ_layout_staged: cute.ComposedLayout,
+        sK_layout_staged: cute.ComposedLayout,
+        sQv_layout_staged: cute.ComposedLayout,
+        sV_layout_staged: cute.ComposedLayout,
+        sVt_layout_staged: cute.ComposedLayout,
+        sP_layout_staged: cute.ComposedLayout,
+        sStats_layout: cute.Layout,
+        sScale_layout: cute.Layout,
+        sBitmask_layout: cute.Layout,
+        sO_layout: cute.ComposedLayout,
+        sP_layout_out: cute.ComposedLayout,
+        tiled_mma_QK: cute.TiledMma,
+        tiled_mma_QvV: cute.TiledMma,
+        tiled_mma_PVt: 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_QvV.thr_id.shape,)
+        )
+        cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+        mma_tile_coord_v = cta_rank_in_cluster % cute.size(tiled_mma_QvV.thr_id.shape)
+        is_leader_cta = mma_tile_coord_v == 0
+
+        # ==== Allocate SMEM ====
+        smem = cutlass.utils.SmemAllocator()
+        storage = smem.allocate(SharedStorage)
+
+        # ==== TMEM stuff ====
+        tmem_alloc_barrier = pipeline.NamedBarrier(
+            barrier_id=int(NamedBarrierFwdSm100_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:
+            if const_expr(self.has_qk):
+                cpasync.prefetch_descriptor(tma_atom_Q)
+            cpasync.prefetch_descriptor(tma_atom_Qv)
+            if const_expr(self.use_tma_KV):
+                if const_expr(self.has_qk):
+                    cpasync.prefetch_descriptor(tma_atom_K)
+                cpasync.prefetch_descriptor(tma_atom_V)
+                cpasync.prefetch_descriptor(tma_atom_Vt)
+            if const_expr(self.use_tma_O):
+                cpasync.prefetch_descriptor(tma_atom_O)
+
+        # ==== Construct pipelines ====
+        tma_warp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1)
+        mma_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
+        )
+
+        TmaUmma = pipeline.PipelineTmaUmma
+        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 else {}),
+                **({"tx_count": tx_count} if tx_count is not None else {}),
+            )
+
+        # Unconditional pipelines
+        # fmt: off
+        pipeline_Q = None
+        if const_expr(self.has_qk):
+            pipeline_Q    = make_pipeline(TmaUmma,   storage.mbar_ptr_Q,        self.num_stages_Q,        tma_warp,           mma_warp,           self.tma_copy_bytes_Q)
+        pipeline_Qv       = make_pipeline(TmaUmma,   storage.mbar_ptr_Qv,       self.num_stages_Qv,       tma_warp,           mma_warp,           self.tma_copy_bytes_Qvi)
+        pipeline_S        = make_pipeline(UmmaAsync, storage.mbar_ptr_S,        self.num_stages_S,        mma_warp,           sm_threads_cluster)
+        pipeline_P        = make_pipeline(AsyncUmma, storage.mbar_ptr_P,        self.num_stages_P,        sm_threads_cluster, mma_warp)
+        pipeline_O0       = make_pipeline(UmmaAsync, storage.mbar_ptr_O0,       self.num_stages_Oi,       mma_warp,           epi_threads_cluster)
+        pipeline_O1       = make_pipeline(UmmaAsync, storage.mbar_ptr_O1,       self.num_stages_Oi,       mma_warp,           epi_threads_cluster)
+        pipeline_sm_stats = make_pipeline(Async,     storage.mbar_ptr_sm_stats, self.num_stages_sm_stats, sm_threads,         epi_threads)
+
+        # K/V pipelines: type and producer depend on use_tma_KV
+        if const_expr(self.use_tma_KV):
+            pipeline_K = None
+            if const_expr(self.has_qk):
+                pipeline_K     = make_pipeline(TmaUmma, storage.mbar_ptr_K, self.num_stages_K, tma_warp, mma_warp, self.tma_copy_bytes_K)
+            pipeline_V         = make_pipeline(TmaUmma, storage.mbar_ptr_V, self.num_stages_V, tma_warp, mma_warp, self.tma_copy_bytes_Vi)
+            pipeline_K_cpasync = pipeline_V_cpasync = pipeline_bitmask = None
+        else:
+            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)
+            pipeline_K = pipeline_K_cpasync = None
+            if const_expr(self.has_qk):
+                pipeline_K         = make_pipeline(AsyncUmma, storage.mbar_ptr_K,         self.num_stages_K, relay_warps_cluster,  mma_warp)
+                pipeline_K_cpasync = make_pipeline(Async,     storage.mbar_ptr_K_cpasync, self.num_stages_K, cpasync_load_threads, relay_threads)
+            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)
+            pipeline_bitmask   = (
+                make_pipeline(Async, storage.mbar_ptr_bitmask, self.num_stages_bitmask, cpasync_load_threads, sm_threads)
+                if const_expr(self.is_topk_gather and not self.disable_bitmask) else None
+            )
+        # fmt: on
+
+        sO_empty_mbar_ptr = None
+        if const_expr(self.use_tma_O and self.overlap_sO_sV):
+            sO_empty_mbar_ptr = storage.sO_empty_mbar_ptr
+            if warp_idx == 0:
+                cute.arch.mbarrier_init(sO_empty_mbar_ptr, 1)
+
+        pipeline.pipeline_init_arrive(cluster_shape_mn=cta_layout_vmnk, is_relaxed=True)
+
+        # ==== Get SMEM tensors ====
+        # fmt: off
+        sQ, sK, sQv, sV, sVt, sP, sP_out = (
+            store.get_tensor(layout.outer, swizzle=layout.inner)
+            if const_expr(store._size > 0) else None
+            for store, layout in [
+                (storage.sQ,  sQ_layout_staged),
+                (storage.sK,  sK_layout_staged),
+                (storage.sQv, sQv_layout_staged),
+                (storage.sV,  sV_layout_staged),
+                (storage.sV,  sVt_layout_staged),  # sVt reuses sV storage
+                (storage.sP,  sP_layout_staged),
+                (storage.sP,  sP_layout_out),
+            ]
+        )
+        # fmt: on
+        sRowMax = storage.sRowMax.get_tensor(sStats_layout)
+        sRowSum = storage.sRowSum.get_tensor(sStats_layout)
+        sScale = storage.sScale.get_tensor(sScale_layout)
+        sBitmask = None
+        if const_expr(self.is_topk_gather):
+            sBitmask = storage.sBitmask.get_tensor(sBitmask_layout)
+
+        if const_expr(self.overlap_sO_sV):
+            sO_iterator = sV.iterator
+            assert cute.cosize(sO_layout) <= cute.cosize(sV_layout_staged)
+        else:
+            sO_iterator = sQv.iterator
+            assert cute.cosize(sO_layout) <= cute.cosize(sQv_layout_staged)
+        sO = cute.make_tensor(
+            cute.recast_ptr(sO_iterator, sO_layout.inner, self.dtype_O), sO_layout.outer
+        )
+
+        # ==== Get thread MMAs and accumulator fragments ====
+        thr_mma_QK = tiled_mma_QK.get_slice(mma_tile_coord_v)
+        thr_mma_QvV = tiled_mma_QvV.get_slice(mma_tile_coord_v)
+        thr_mma_PVt = tiled_mma_PVt.get_slice(mma_tile_coord_v)
+
+        acc_shape_S = thr_mma_QvV.partition_shape_C(self.mma_tiler_QvV[:2])
+        tStS_fake = thr_mma_QvV.make_fragment_C(cute.append(acc_shape_S, self.num_stages_S))
+
+        acc_shape_Oi = thr_mma_PVt.partition_shape_C(self.mma_tiler_PVt[:2])
+        tOtO0_fake = thr_mma_PVt.make_fragment_C(acc_shape_Oi)
+        tOtO1_fake = thr_mma_PVt.make_fragment_C(acc_shape_Oi)
+
+        block_info = BlockInfo(
+            self.cta_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=mQv.shape[0] if const_expr(not self.pack_gqa) else mQv.shape[0][1],
+            seqlen_k_static=mV.shape[0]
+            if const_expr(mPageTable is None)
+            else mV.shape[0] * mPageTable.shape[1],
+            tile_m=self.cta_tile_m,
+            tile_n=self.tile_n,
+            mCuSeqlensQ=mCuSeqlensQ,
+            mCuSeqlensK=mCuSeqlensK,
+            mSeqUsedQ=mSeqUsedQ,
+            mSeqUsedK=mSeqUsedK,
+        )
+        AttentionMaskCls = partial(
+            AttentionMask,
+            self.cta_tile_m * self.cta_group_size,
+            self.tile_n,
+            qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
+        )
+
+        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:
+                if const_expr(self.num_regs_other < self.num_regs_per_thread):
+                    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:
+                    if const_expr(self.num_regs_other < self.num_regs_per_thread):
+                        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]:
+                    if const_expr(self.num_regs_other < self.num_regs_per_thread):
+                        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_K,
+                    pipeline_V,
+                    pipeline_K_cpasync,
+                    pipeline_V_cpasync,
+                    sO_empty_mbar_ptr,
+                    topk_length_dynamic,
+                    block_info,
+                    SeqlenInfoCls,
+                    tile_scheduler=tile_scheduler,
+                    mCuSeqlensQ=mCuSeqlensQ,
+                )
+
+            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,
+                    mK,
+                    mV,
+                    mVt,
+                    sK,
+                    sV,
+                    sVt,
+                    sBitmask,
+                    pipeline_K,
+                    pipeline_V,
+                    pipeline_K_cpasync,
+                    pipeline_V_cpasync,
+                    pipeline_bitmask,
+                    sO_empty_mbar_ptr,
+                    topk_length_dynamic,
+                    block_info,
+                    SeqlenInfoCls,
+                    tile_scheduler=tile_scheduler,
+                    mPageTable=mPageTable,
+                    mCuSeqlensQ=mCuSeqlensQ,
+                )
+
+        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(
+                mQ,
+                mK,
+                mQv,
+                mV,
+                mVt,
+                sQ,
+                sK,
+                sQv,
+                sV,
+                sVt,
+                tma_atom_Q,
+                tma_atom_K,
+                tma_atom_Qv,
+                tma_atom_V,
+                tma_atom_Vt,
+                pipeline_Q,
+                pipeline_K,
+                pipeline_Qv,
+                pipeline_V,
+                sO_empty_mbar_ptr,
+                thr_mma_QK,
+                thr_mma_QvV,
+                thr_mma_PVt,
+                topk_length_dynamic,
+                block_info,
+                SeqlenInfoCls,
+                tile_scheduler=tile_scheduler,
+                mPageTable=mPageTable,
+                mCuSeqlensQ=mCuSeqlensQ,
+            )
+
+        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)
+            tStS = cute.make_tensor(tmem_ptr, tStS_fake.layout)
+            tOtO0 = cute.make_tensor(tmem_ptr + self.tmem_offset_O0, tOtO0_fake.layout)
+            tOtO1 = cute.make_tensor(tmem_ptr + self.tmem_offset_O1, tOtO1_fake.layout)
+            self.mma(
+                sQ,
+                sK,
+                sQv,
+                sV,
+                sVt,
+                sP,
+                tStS,
+                tOtO0,
+                tOtO1,
+                tiled_mma_QK,
+                tiled_mma_QvV,
+                tiled_mma_PVt,
+                pipeline_Q,
+                pipeline_K,
+                pipeline_Qv,
+                pipeline_V,
+                pipeline_S,
+                pipeline_P,
+                pipeline_O0,
+                pipeline_O1,
+                sO_empty_mbar_ptr,
+                is_leader_cta,
+                topk_length_dynamic,
+                block_info,
+                SeqlenInfoCls,
+                tile_scheduler=tile_scheduler,
+                mCuSeqlensQ=mCuSeqlensQ,
+            )
+            tmem.relinquish_alloc_permit()
+            tmem_alloc_barrier.arrive_and_wait()
+            tmem.free(tmem_ptr)
+
+        if warp_idx in self.softmax_warp_indices:
+            if const_expr(self.num_regs_softmax > self.num_regs_per_thread):
+                cute.arch.setmaxregister_increase(self.num_regs_softmax)
+            tmem.wait_for_alloc()
+            tmem_ptr = tmem.retrieve_ptr(self.dtype_acc)
+            tStS = cute.make_tensor(tmem_ptr, tStS_fake.layout)
+            self.softmax_loop(
+                softmax_scale,
+                softmax_scale_log2,
+                mLSE,
+                mRowMax,
+                sRowMax,
+                sRowSum,
+                sScale,
+                sBitmask,
+                sP,
+                tStS,
+                thr_mma_QvV,
+                pipeline_S,
+                pipeline_P,
+                pipeline_sm_stats,
+                pipeline_bitmask,
+                sO_empty_mbar_ptr,
+                AttentionMaskCls,
+                topk_length_dynamic,
+                block_info,
+                SeqlenInfoCls,
+                tile_scheduler=tile_scheduler,
+                tma_atom_P=tma_atom_P,
+                mP=mP,
+                sP_out=sP_out,
+                mCuSeqlensQ=mCuSeqlensQ,
+            )
+            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)
+            tOtO0 = cute.make_tensor(tmem_ptr + self.tmem_offset_O0, tOtO0_fake.layout)
+            tOtO1 = cute.make_tensor(tmem_ptr + self.tmem_offset_O1, tOtO1_fake.layout)
+            self.correction_loop(
+                softmax_scale_log2,
+                mO,
+                mLSE,
+                tma_atom_O,
+                sRowMax,
+                sRowSum,
+                sScale,
+                sO,
+                tOtO0,
+                tOtO1,
+                pipeline_O0,
+                pipeline_O1,
+                pipeline_sm_stats,
+                sO_empty_mbar_ptr,
+                tiled_copy_O_r2g,
+                topk_length_dynamic,
+                block_info,
+                SeqlenInfoCls,
+                tile_scheduler=tile_scheduler,
+                mCuSeqlensQ=mCuSeqlensQ,
+            )
+            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()
+            # cluster_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_K: Optional[pipeline.PipelineAsyncUmma],
+        pipeline_V: pipeline.PipelineAsyncUmma,
+        pipeline_K_cpasync: Optional[pipeline.PipelineAsync],
+        pipeline_V_cpasync: pipeline.PipelineAsync,
+        sO_empty_mbar_ptr: Optional[cute.Pointer],
+        topk_length_dynamic: Optional[Int32],
+        block_info: BlockInfo,
+        SeqlenInfoCls: Callable,
+        tile_scheduler: TileSchedulerProtocol,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+    ):
+        # ==== Make pipeline states ====
+        # pipeline_{K,V0,V1} producer
+        # pipeline_{K,V0,V1}_cpasync consumer
+        Producer, Consumer = pipeline.PipelineUserType.Producer, pipeline.PipelineUserType.Consumer
+        relay_K_fn = None
+        if const_expr(self.has_qk):
+            producer_state_K = pipeline.make_pipeline_state(Producer, stages=self.num_stages_K)
+            consumer_state_K = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_K)
+            relay_K_fn = partial(self.relay_inner, pipeline_K_cpasync, pipeline_K)
+
+        producer_state_V = pipeline.make_pipeline_state(Producer, stages=self.num_stages_V)
+        consumer_state_V = pipeline.make_pipeline_state(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:
+            cluster_m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+            if const_expr(self.use_packed_varlen_sched):
+                batch_idx = get_batch_from_cu_tensor(cluster_m_block, mCuSeqlensQ)
+                if const_expr(not self.is_topk_gather):
+                    cluster_m_block -= mCuSeqlensQ[batch_idx]
+
+            seqlen = SeqlenInfoCls(batch_idx)
+            if const_expr(self.is_topk_gather):
+                n_block_min = 0
+                n_block_max = self.topk_length // self.tile_n
+                # n_block_max = topk_length_dynamic // self.tile_n
+            else:
+                n_block_min, n_block_max = block_info.get_n_block_min_max(
+                    seqlen,
+                    cluster_m_block,
+                )
+            num_n_blocks = n_block_max - n_block_min
+
+            # ==== Prologue ====
+            # relay K, V0, V1
+            if const_expr(self.has_qk):
+                consumer_state_K, producer_state_K = relay_K_fn(consumer_state_K, producer_state_K)
+            for _ in cutlass.range_constexpr(self.num_hdimv_splits):
+                consumer_state_V, producer_state_V = relay_V_fn(consumer_state_V, producer_state_V)
+
+            # ==== Mainloop ====
+            for _ in cutlass.range(num_n_blocks - 1, unroll=2):
+                # relay K, V0, V1, Vt0, Vt1
+                if const_expr(self.has_qk):
+                    consumer_state_K, producer_state_K = relay_K_fn(
+                        consumer_state_K, producer_state_K
+                    )
+                for _ in cutlass.range_constexpr(2 * self.num_hdimv_splits):
+                    consumer_state_V, producer_state_V = relay_V_fn(
+                        consumer_state_V, producer_state_V
+                    )
+
+            # ==== Epilogue ===
+            # relay Vt0, Vt1
+            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()
+
+        if const_expr(self.has_qk):
+            pipeline_K.producer_tail(producer_state_K)
+        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,
+        mK: Optional[cute.Tensor],
+        mV: cute.Tensor,
+        mVt: cute.Tensor,
+        sK: Optional[cute.Tensor],
+        sV: cute.Tensor,
+        sVt: cute.Tensor,
+        sBitmask: Optional[cute.Tensor],
+        pipeline_K: Optional[pipeline.PipelineAsyncUmma],
+        pipeline_V: pipeline.PipelineAsyncUmma,
+        pipeline_K_cpasync: Optional[pipeline.PipelineAsync],
+        pipeline_V_cpasync: pipeline.PipelineAsync,
+        pipeline_bitmask: pipeline.PipelineAsync,
+        sO_empty_mbar_ptr: Optional[cute.Pointer],
+        topk_length_dynamic: Optional[Int32],
+        block_info: BlockInfo,
+        SeqlenInfoCls: Callable,
+        tile_scheduler: TileSchedulerProtocol,
+        mPageTable: Optional[cute.Tensor] = None,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+    ):
+        # ==== cpasync load warpgroup ====
+        # Description: loads tiles of K, V, V0, V1 from gmem to smem using cpasync
+        # produces: K, V, V0, V1, bitmask
+        # consumes: -
+
+        # cpasync load is used for both topk gather and paged KV with page_size != tile_n
+        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 = pipeline.PipelineUserType.Producer
+        if const_expr(self.has_qk):
+            producer_state_K = pipeline.make_pipeline_state(Producer, stages=self.num_stages_K)
+        producer_state_V = pipeline.make_pipeline_state(Producer, stages=self.num_stages_V)
+        if const_expr(self.is_topk_gather and not self.disable_bitmask):
+            producer_state_bitmask = pipeline.make_pipeline_state(
+                Producer, stages=self.num_stages_bitmask
+            )
+        if const_expr(self.use_tma_O):
+            producer_phase_O = Int32(1)
+
+        work_tile = tile_scheduler.initial_work_tile_info()
+        while work_tile.is_valid_tile:
+            cluster_m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+            if const_expr(self.use_packed_varlen_sched):
+                batch_idx = get_batch_from_cu_tensor(cluster_m_block, mCuSeqlensQ)
+                if const_expr(not self.is_topk_gather):
+                    cluster_m_block -= mCuSeqlensQ[batch_idx]
+            head_idx_kv = (
+                head_idx // self.qhead_per_kvhead if const_expr(not self.pack_gqa) else head_idx
+            )
+
+            seqlen = SeqlenInfoCls(batch_idx)
+            if const_expr(self.is_topk_gather):
+                n_block_min = 0
+                n_block_max = self.topk_length // self.tile_n
+                # n_block_max = topk_length_dynamic // self.tile_n
+            else:
+                n_block_min, n_block_max = block_info.get_n_block_min_max(
+                    seqlen,
+                    cluster_m_block,
+                )
+            num_n_blocks = n_block_max - n_block_min
+
+            if const_expr(self.is_topk_gather):
+                # ==== Topk gather path ====
+                # cluster_m_block == m_idx under MQA 128 assumption
+                m_idx = cluster_m_block
+                if const_expr(not seqlen.has_cu_seqlens_q):
+                    mIndexTopk_cur = mIndexTopk[None, m_idx, batch_idx]
+                else:
+                    offset_q = seqlen.offset_q if const_expr(not self.use_packed_varlen_sched) else 0
+                    mIndexTopk_cur = mIndexTopk[None, m_idx + offset_q]
+
+                if const_expr(self.is_causal):
+                    m_local_idx = (
+                        m_idx - seqlen.offset_q if const_expr(self.use_packed_varlen_sched) else m_idx
+                    )
+                    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.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,
+                    sBitmask,
+                    pipeline_bitmask,
+                )
+
+                # (seqlen_k, hdim) or (seqlen_k, hdimv)
+                if const_expr(self.has_qk):
+                    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]
+                # (hdimv, seqlen_k)
+                if const_expr(not seqlen.has_cu_seqlens_k):
+                    mVt_cur = mVt[None, None, head_idx_kv, batch_idx]
+                else:
+                    mVt_cur = cute.domain_offset((0, seqlen.offset_k), mVt[None, None, head_idx_kv])
+
+                hdimv_split_per_cta = self.hdimv // self.num_hdimv_splits // self.cta_group_size
+                mVt_cur = cute.tiled_divide(mVt_cur, (hdimv_split_per_cta,))
+                mVt_cur = cute.logical_divide(mVt_cur, (1, self.cta_group_size, 1))
+                mVt_cur = mVt_cur[(0, None), (cta_rank_in_cluster, None), (0, None)]
+                mVt_cur = cute.group_modes(mVt_cur, 0, 2)  # ((hdimv//4, 2), seqlen_k)
+
+                load_K = None
+                if const_expr(self.has_qk):
+                    load_K = partial(
+                        self.cpasync_gather_load_KV,
+                        cpasync_gather_kv_manager,
+                        pipeline_K,
+                        pipeline_K_cpasync,
+                        sK,
+                        False,
+                        "K",
+                        mK_cur,
+                    )
+                load_V = partial(
+                    self.cpasync_gather_load_KV,
+                    cpasync_gather_kv_manager,
+                    pipeline_V,
+                    pipeline_V_cpasync,
+                    sV,
+                    False,
+                    "V",
+                    mV_cur,
+                )
+                load_Vt = partial(
+                    self.cpasync_gather_load_KV,
+                    cpasync_gather_kv_manager,
+                    pipeline_V,
+                    pipeline_V_cpasync,
+                    sVt,
+                    True,
+                    "V",
+                    mVt_cur,
+                )
+
+                # process n_blocks in decreasing order
+                n_block = n_block_max - 1
+
+                # ==== Prologue ====
+                # K, V0, V1
+                cpasync_gather_kv_manager.load_index_topk(n_block, transpose=False)
+                if const_expr(self.has_qk):
+                    producer_state_K = load_K(producer_state_K)
+                for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                    producer_state_V = load_V(producer_state_V, d_offset=split * self.hdimv // 2)
+                if const_expr(not self.disable_bitmask):
+                    producer_state_bitmask = cpasync_gather_kv_manager.compute_bitmask(
+                        producer_state_bitmask
+                    )
+
+                if const_expr(self.use_tma_O and self.overlap_sO_sV):
+                    cute.arch.mbarrier_wait(sO_empty_mbar_ptr, phase=producer_phase_O)
+                    producer_phase_O ^= 1
+
+                # ==== Mainloop ====
+                for _ in cutlass.range(num_n_blocks - 1, unroll=2):
+                    # K, V0, V1
+                    cpasync_gather_kv_manager.load_index_topk(n_block - 1, transpose=False)
+                    if const_expr(self.has_qk):
+                        producer_state_K = load_K(producer_state_K)
+                    for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                        producer_state_V = load_V(
+                            producer_state_V, d_offset=split * self.hdimv // 2
+                        )
+                    if const_expr(not self.disable_bitmask):
+                        producer_state_bitmask = cpasync_gather_kv_manager.compute_bitmask(
+                            producer_state_bitmask
+                        )
+                    # Vt0, Vt1
+                    cpasync_gather_kv_manager.load_index_topk(n_block, transpose=True)
+                    for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                        producer_state_V = load_Vt(
+                            producer_state_V, d_offset=split * hdimv_split_per_cta
+                        )
+                    # advance n_block
+                    n_block -= 1
+
+                # ==== Epilogue ====
+
+                # Vt0, Vt1 for n_block = 0
+                cpasync_gather_kv_manager.load_index_topk(0, transpose=True)
+                for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                    producer_state_V = load_Vt(
+                        producer_state_V, d_offset=split * hdimv_split_per_cta
+                    )
+            else:
+                # ==== Paged KV cp.async path (page_size != tile_n) ====
+                page_size_divmod = FastDivmodDivisor(cute.size(mV.shape[0]))
+                hdimv_split = self.hdimv // self.num_hdimv_splits
+                hdimv_split_per_cta = hdimv_split // self.cta_group_size
+
+                # CTA-split Vt: (dv, page_size, h_k, num_pages) -> ((dv/4, 2), page_size, h_k, num_pages)
+                mVt_cta = cute.tiled_divide(mVt, (hdimv_split_per_cta,))
+                mVt_cta = cute.logical_divide(mVt_cta, (1, self.cta_group_size, 1, 1, 1))
+                mVt_cta = mVt_cta[
+                    (0, None), (cta_rank_in_cluster, None), (0, None), (0, None), (0, None)
+                ]
+                mVt_cta = cute.group_modes(mVt_cta, 0, 2)
+
+                # PagedKVManager for K (hdim=64): uses "K" mode only
+                if const_expr(self.has_qk):
+                    paged_kv_K = PagedKVManager.create(
+                        mPageTable,
+                        mK,
+                        mK,
+                        page_size_divmod,
+                        batch_idx,
+                        head_idx_kv,
+                        tidx,
+                        seqlen.seqlen_k,
+                        0,
+                        self.tile_n,
+                        self.hdim,
+                        self.hdim,
+                        self.num_cpasync_load_threads,
+                        mK.element_type,
+                        arch=100,
+                    )
+                # PagedKVManager for V/Vt: "K" mode → V (non-transposed), "V" mode → Vt (transposed)
+                paged_kv_V = PagedKVManager.create(
+                    mPageTable,
+                    mV,
+                    mVt_cta,
+                    page_size_divmod,
+                    batch_idx,
+                    head_idx_kv,
+                    tidx,
+                    seqlen.seqlen_k,
+                    0,
+                    self.tile_n,
+                    hdimv_split,
+                    hdimv_split_per_cta,
+                    self.num_cpasync_load_threads,
+                    mV.element_type,
+                    arch=100,
+                )
+
+                if const_expr(self.has_qk):
+                    load_K = partial(
+                        self.cpasync_paged_load_KV,
+                        paged_kv_K,
+                        pipeline_K,
+                        pipeline_K_cpasync,
+                        sK,
+                        False,
+                        "K",
+                        cta_rank_in_cluster,
+                    )
+                load_V = partial(
+                    self.cpasync_paged_load_KV,
+                    paged_kv_V,
+                    pipeline_V,
+                    pipeline_V_cpasync,
+                    sV,
+                    False,
+                    "K",
+                    cta_rank_in_cluster,
+                )
+                load_Vt = partial(
+                    self.cpasync_paged_load_KV,
+                    paged_kv_V,
+                    pipeline_V,
+                    pipeline_V_cpasync,
+                    sVt,
+                    True,
+                    "V",
+                    cta_rank_in_cluster,
+                )
+
+                n_block_first = n_block_max - 1
+                n_block = n_block_first
+                safe_n_block_first = n_block_first if num_n_blocks > 0 else 0
+
+                # ==== Prologue ====
+                if const_expr(self.has_qk):
+                    paged_kv_K.load_page_table(safe_n_block_first)
+                    producer_state_K = load_K(n_block_first, producer_state_K)
+                paged_kv_V.load_page_table(safe_n_block_first)
+                for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                    producer_state_V = load_V(
+                        n_block_first, producer_state_V, d_offset=split * self.hdimv // 2
+                    )
+
+                if const_expr(self.use_tma_O and self.overlap_sO_sV):
+                    cute.arch.mbarrier_wait(sO_empty_mbar_ptr, phase=producer_phase_O)
+                    producer_phase_O ^= 1
+
+                # ==== Mainloop ====
+                for n_block_idx in cutlass.range(num_n_blocks - 1, unroll=2):
+                    n_block = n_block_first - n_block_idx
+                    # K, V0, V1 for next block in descending order
+                    if const_expr(self.has_qk):
+                        paged_kv_K.load_page_table(n_block - 1)
+                        producer_state_K = load_K(n_block - 1, producer_state_K)
+                    paged_kv_V.load_page_table(n_block - 1)
+                    for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                        producer_state_V = load_V(
+                            n_block - 1, producer_state_V, d_offset=split * self.hdimv // 2
+                        )
+                    # Vt0, Vt1 for current block
+                    paged_kv_V.load_page_table(n_block)
+                    for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                        producer_state_V = load_Vt(
+                            n_block, producer_state_V, d_offset=split * hdimv_split_per_cta
+                        )
+
+                # ==== Epilogue ====
+                paged_kv_V.load_page_table(n_block_min)
+                for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                    producer_state_V = load_Vt(
+                        n_block_min, producer_state_V, d_offset=split * hdimv_split_per_cta
+                    )
+
+            # Advance to next tile
+            work_tile = tile_scheduler.advance_to_next_work()
+
+        # note: don't use producer tail with pipeline_X_cpasync since we never use its producer_acquire.
+        if const_expr(self.is_topk_gather and not self.disable_bitmask):
+            # pipeline_bitmask invokes producer acquire in gather kv manager,
+            # so we should call its producer tail.
+            pipeline_bitmask.producer_tail(producer_state_bitmask)
+
+    @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 cpasync_paged_load_KV(
+        self,
+        paged_kv_manager: PagedKVManager,
+        pipeline_mma: pipeline.PipelineAsyncUmma,
+        pipeline_cpasync: pipeline.PipelineAsync,
+        sX: cute.Tensor,
+        transpose: bool,
+        K_or_V: str,
+        cta_rank_in_cluster: Int32,
+        n_block: Int32,
+        producer_state: pipeline.PipelineState,
+        d_offset: int = 0,
+    ):
+        """Load one tile of K or V from paged gmem to smem using cp.async.
+
+        Uses PagedKVManager for page table lookups and pointer computation,
+        with smem reshaping via cute.composition (same approach as CpasyncGatherKVManager).
+
+        For non-transposed tensors (K, V0, V1): K_or_V="K", transpose=False
+        For transposed tensors (Vt0, Vt1):      K_or_V="V", transpose=True
+        """
+        stage = producer_state.index
+        pipeline_mma.producer_acquire(producer_state)
+
+        # NOTE: load_page_table() must be called by the caller BEFORE this method.
+        # Calling it here (through a @cute.jit boundary) causes MLIR SSA verification
+        # errors because the rmem tensor writes inside load_page_table's dynamic
+        # cutlass.range loop cross region boundaries. This matches the SM90/SM100
+        # pattern where load_page_table is called directly in the loop body.
+
+        # Compute row pointers from cached page table entries
+        tPrXPtr = paged_kv_manager.compute_X_ptr(K_or_V, d_offset)
+
+        # Reshape smem to flat 2D using composition (matches CpasyncGatherKVManager.load_X)
+        head_dim = (
+            paged_kv_manager.head_dim_v_padded
+            if const_expr(K_or_V == "V")
+            else paged_kv_manager.head_dim_padded
+        )
+        cta_tile_n = self.tile_n if const_expr(transpose) else self.tile_n // self.cta_group_size
+        order = (1, 0) if const_expr(transpose) else (0, 1)
+
+        sX_stage = sX[None, None, None, stage]
+        sX_nd_layout = cute.make_ordered_layout((cta_tile_n, head_dim), order=order)
+        sX_nd = cute.composition(sX_stage, sX_nd_layout)
+
+        cX = cute.make_identity_tensor((cta_tile_n, head_dim))
+        tXsX = paged_kv_manager.gmem_thr_copy_KV.partition_D(sX_nd)
+        tXcX = paged_kv_manager.gmem_thr_copy_KV.partition_S(cX)
+        tXc0X = paged_kv_manager.gmem_thr_copy_KV.get_slice(0).partition_S(cX)
+
+        base_offset = n_block * self.tile_n
+        if const_expr(not transpose):
+            base_offset += cta_tile_n * cta_rank_in_cluster
+        seqlenk_row_limit = (
+            paged_kv_manager.seqlen_k - base_offset - tXcX[0][0] if n_block >= 0 else 0
+        )
+
+        if const_expr(not transpose):
+            offset = cta_rank_in_cluster * (
+                paged_kv_manager.gmem_threads_per_row // self.cta_group_size
+            )
+        else:
+            offset = 0
+
+        for m in cutlass.range_constexpr(cute.size(tXsX, mode=[1])):
+            row_valid = tXc0X[0, m, 0][0] < seqlenk_row_limit
+            should_load = cute.make_fragment_like(tXsX[(0, None), m, 0], cute.Boolean)
+            should_load.fill(row_valid)
+
+            x_ptr_i64 = fa_utils.shuffle_sync(
+                tPrXPtr[m // paged_kv_manager.gmem_threads_per_row],
+                (m + offset) % paged_kv_manager.gmem_threads_per_row,
+                width=paged_kv_manager.gmem_threads_per_row,
+            )
+            x_gmem_ptr = cute.make_ptr(
+                paged_kv_manager.mK_paged.element_type,
+                x_ptr_i64,
+                cute.AddressSpace.gmem,
+                assumed_align=16,
+            )
+            mX_cur = cute.make_tensor(x_gmem_ptr, cute.make_layout((head_dim,)))
+            mX_cur_copy = cute.tiled_divide(mX_cur, (paged_kv_manager.async_copy_elems,))
+
+            for k in cutlass.range_constexpr(cute.size(tXsX, mode=[2])):
+                ki = tXcX[0, 0, k][1] // paged_kv_manager.async_copy_elems
+                mX_cur_copy_ki = mX_cur_copy[None, ki]
+                tXsX_k = tXsX[None, m, k]
+                mX_cur_copy_ki = cute.make_tensor(mX_cur_copy_ki.iterator, tXsX_k.layout)
+                cute.copy(
+                    paged_kv_manager.gmem_tiled_copy_KV,
+                    mX_cur_copy_ki,
+                    tXsX_k,
+                    pred=should_load,
+                )
+
+        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,
+        mQ: Optional[cute.Tensor],
+        mK: Optional[cute.Tensor],
+        mQv: cute.Tensor,
+        mV: cute.Tensor,
+        mVt: cute.Tensor,
+        sQ: Optional[cute.Tensor],
+        sK: Optional[cute.Tensor],
+        sQv: cute.Tensor,
+        sV: cute.Tensor,
+        sVt: cute.Tensor,
+        tma_atom_Q: Optional[cute.CopyAtom],
+        tma_atom_K: Optional[cute.CopyAtom],
+        tma_atom_Qv: cute.CopyAtom,
+        tma_atom_V: cute.CopyAtom,
+        tma_atom_Vt: cute.CopyAtom,
+        pipeline_Q: Optional[pipeline.PipelineAsync],
+        pipeline_K: Optional[pipeline.PipelineAsync],
+        pipeline_Qv: pipeline.PipelineAsync,
+        pipeline_V: pipeline.PipelineAsync,
+        sO_empty_mbar_ptr: Optional[cute.Pointer],
+        thr_mma_QK: cute.ThrMma,
+        thr_mma_QvV: cute.ThrMma,
+        thr_mma_PVt: cute.ThrMma,
+        topk_length_dynamic: Optional[Int32],
+        block_info: BlockInfo,
+        SeqlenInfoCls: Callable,
+        tile_scheduler: TileSchedulerProtocol,
+        mPageTable: Optional[cute.Tensor] = None,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+    ):
+        # ==== Load warp ====
+        # Description: loads tiles of Q, Qv, K, V, V0, V1 from gmem to smem using TMA
+        # produces: Q, Qv, K, V, V0, V1
+        # consumes: -
+
+        # ==== Make pipeline states ====
+        Producer = pipeline.PipelineUserType.Producer
+        if const_expr(self.has_qk):
+            producer_state_Q = pipeline.make_pipeline_state(Producer, stages=self.num_stages_Q)
+        producer_state_Qv = pipeline.make_pipeline_state(Producer, stages=self.num_stages_Qv)
+        if const_expr(self.use_tma_KV):
+            if const_expr(self.has_qk):
+                producer_state_K = pipeline.make_pipeline_state(Producer, stages=self.num_stages_K)
+            producer_state_V = pipeline.make_pipeline_state(Producer, stages=self.num_stages_V)
+        if const_expr(self.use_tma_O):
+            producer_phase_O = Int32(1)
+
+        work_tile = tile_scheduler.initial_work_tile_info()
+        while work_tile.is_valid_tile:
+            cluster_m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+            if const_expr(self.use_packed_varlen_sched):
+                batch_idx = get_batch_from_cu_tensor(cluster_m_block, mCuSeqlensQ)
+                cluster_m_block -= mCuSeqlensQ[batch_idx]
+            head_idx_kv = (
+                head_idx // self.qhead_per_kvhead if const_expr(not self.pack_gqa) else head_idx
+            )
+
+            seqlen = SeqlenInfoCls(batch_idx)
+            if const_expr(self.is_topk_gather):
+                n_block_min = 0
+                n_block_max = self.topk_length // self.tile_n
+                # n_block_max = topk_length_dynamic // self.tile_n
+            else:
+                n_block_min, n_block_max = block_info.get_n_block_min_max(
+                    seqlen,
+                    cluster_m_block,
+                )
+            num_n_blocks = n_block_max - n_block_min
+            even_n_blocks = num_n_blocks % 2 == 0 and num_n_blocks > 0
+            num_n_block_groups = cute.ceil_div(num_n_blocks, self.num_stages_S)
+
+            # ==== Partition GMEM tensors ====
+            # (seqlen_q, hdim or hdimv//2)
+            if const_expr(self.has_qk):
+                mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx]
+                # (mma_tile_m, hdim or hdimv//2)
+                gQ = cute.local_tile(
+                    mQ_cur,
+                    (self.mma_tiler_QK[0], self.mma_tiler_QK[2]),
+                    (cluster_m_block, 0),
+                )
+                tSgQ = thr_mma_QK.partition_A(gQ)
+                tQsQ, tQgQ = cpasync.tma_partition(
+                    atom=tma_atom_Q,
+                    cta_coord=0,
+                    cta_layout=cute.make_layout(1),
+                    smem_tensor=cute.group_modes(sQ, 0, 3),
+                    gmem_tensor=cute.group_modes(tSgQ, 0, 3),
+                )
+            mQv_cur = seqlen.offset_batch_Q(mQv, batch_idx, dim=3)[None, None, head_idx]
+            gQv = cute.local_tile(
+                mQv_cur,
+                (self.mma_tiler_QvV[0], self.mma_tiler_QvV[2]),
+                (cluster_m_block, None),
+            )
+            tSgQv = thr_mma_QvV.partition_A(gQv)
+            tQvsQv, tQvgQv = cpasync.tma_partition(
+                atom=tma_atom_Qv,
+                cta_coord=0,
+                cta_layout=cute.make_layout(1),
+                smem_tensor=cute.group_modes(sQv, 0, 3),
+                gmem_tensor=cute.group_modes(tSgQv, 0, 3),
+            )
+
+            if const_expr(self.use_tma_KV):
+                if const_expr(mPageTable is None):
+                    mPageTable_cur = None
+                    # Non-paged: select batch, tile over seqlen_k
+                    if const_expr(self.has_qk):
+                        # (seqlen_k, hdim)
+                        mK_cur = seqlen.offset_batch_K(mK, batch_idx, dim=3)[
+                            None, None, head_idx_kv
+                        ]
+                        # (tile_n, hdim, num_n_blocks)
+                        gK = cute.local_tile(
+                            mK_cur,
+                            (self.mma_tiler_QK[1], self.mma_tiler_QK[2]),
+                            (None, 0),
+                        )
+                    # (seqlen_k, hdimv)
+                    mV_cur = seqlen.offset_batch_K(mV, batch_idx, dim=3)[None, None, head_idx_kv]
+                    # (hdimv, seqlen_k)
+                    if const_expr(not seqlen.has_cu_seqlens_k):
+                        mVt_cur = mVt[None, None, head_idx_kv, batch_idx]
+                    else:
+                        mVt_cur = cute.domain_offset(
+                            (0, seqlen.offset_k), mVt[None, None, head_idx_kv]
+                        )
+                    # (tile_n, hdimv//4, num_n_blocks, num_d_blocks=4)
+                    gV = cute.local_tile(
+                        mV_cur,
+                        (self.mma_tiler_QvV[1], self.mma_tiler_QvV[2]),
+                        (None, None),
+                    )
+                    # (tile_n, hdimv//4, num_d_blocks=4, num_n_blocks)
+                    gV = cute.make_tensor(gV.iterator, cute.select(gV.layout, mode=[0, 1, 3, 2]))
+                    # (hdimv//4, tile_n, num_d_blocks=4, num_n_blocks)
+                    gVt = cute.local_tile(
+                        mVt_cur,
+                        (self.mma_tiler_PVt[1], self.mma_tiler_PVt[2]),
+                        (None, None),
+                    )
+                else:
+                    mPageTable_cur = mPageTable[batch_idx, None]
+                    # Paged KV: keep pages dim, index by page_idx at load time
+                    # TMA path assumes page_size == tile_n
+                    if const_expr(self.has_qk):
+                        # (page_size, hdim, num_pages)
+                        mK_cur = mK[None, None, head_idx_kv, None]
+                        # (tile_n, hdim, num_pages)
+                        gK = cute.local_tile(
+                            mK_cur,
+                            (self.mma_tiler_QK[1], self.mma_tiler_QK[2]),
+                            (0, 0, None),
+                        )
+                    # (page_size, hdimv, num_pages)
+                    mV_cur = mV[None, None, head_idx_kv, None]
+                    # (hdimv, page_size, num_pages)
+                    mVt_cur = mVt[None, None, head_idx_kv, None]
+                    # (tile_n, hdimv//4, num_d_blocks=4, num_pages)
+                    gV = cute.local_tile(
+                        mV_cur,
+                        (self.mma_tiler_QvV[1], self.mma_tiler_QvV[2]),
+                        (0, None, None),
+                    )
+                    # (hdimv//4, tile_n, num_d_blocks=4, num_pages)
+                    gVt = cute.local_tile(
+                        mVt_cur,
+                        (self.mma_tiler_PVt[1], self.mma_tiler_PVt[2]),
+                        (None, 0, None),
+                    )
+
+                if const_expr(self.has_qk):
+                    tSgK = thr_mma_QK.partition_B(gK)
+                    tKsK, tKgK = cpasync.tma_partition(
+                        atom=tma_atom_K,
+                        cta_coord=0,
+                        cta_layout=cute.make_layout(1),
+                        smem_tensor=cute.group_modes(sK, 0, 3),
+                        gmem_tensor=cute.group_modes(tSgK, 0, 3),
+                    )
+
+                tSgV = thr_mma_QvV.partition_B(gV)
+                tOgVt = thr_mma_PVt.partition_B(gVt)
+                tVsV, tVgV = cpasync.tma_partition(
+                    atom=tma_atom_V,
+                    cta_coord=0,
+                    cta_layout=cute.make_layout(1),
+                    smem_tensor=cute.group_modes(sV, 0, 3),
+                    gmem_tensor=cute.group_modes(tSgV, 0, 3),
+                )
+                tVtsVt, tVtgVt = cpasync.tma_partition(
+                    atom=tma_atom_Vt,
+                    cta_coord=0,
+                    cta_layout=cute.make_layout(1),
+                    smem_tensor=cute.group_modes(sVt, 0, 3),
+                    gmem_tensor=cute.group_modes(tOgVt, 0, 3),
+                )
+
+            if const_expr(self.has_qk):
+                load_Q = partial(self.load_inner, tma_atom_Q, tQgQ, tQsQ, pipeline_Q)
+            load_Qv = partial(self.load_inner, tma_atom_Qv, tQvgQv, tQvsQv, pipeline_Qv)
+
+            if const_expr(self.use_tma_KV):
+                if const_expr(self.has_qk):
+                    load_K = partial(self.load_inner, tma_atom_K, tKgK, tKsK, pipeline_K)
+                load_V = partial(self.load_inner, tma_atom_V, tVgV, tVsV, pipeline_V)
+                load_Vt = partial(self.load_inner, tma_atom_Vt, tVtgVt, tVtsVt, pipeline_V)
+
+            # ==== Load stationary operands ====
+
+            # copy Q, Qvi gmem -> smem
+            if const_expr(self.has_qk):
+                producer_state_Q = load_Q(producer_state_Q)
+            for dv_split in cutlass.range_constexpr(2):
+                producer_state_Qv = load_Qv(producer_state_Qv, block=dv_split)
+
+            if const_expr(self.use_tma_KV):
+                # ==== Prologue ====
+                n_block_first = n_block_max - 1 if n_block_max > 0 else 0
+                block = self._get_block_idx(n_block_first, mPageTable_cur)
+                # copy K gmem -> smem
+                if const_expr(self.has_qk):
+                    producer_state_K = load_K(producer_state_K, block=block)
+                # copy Vi gmem -> smem
+                for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                    producer_state_V = load_V(producer_state_V, block=block, split=split)
+
+                if const_expr(self.use_tma_O and self.overlap_sO_sV):
+                    cute.arch.mbarrier_wait(sO_empty_mbar_ptr, phase=producer_phase_O)
+                    producer_phase_O ^= 1
+
+                # ==== Main loop ====
+                for n_block_group in cutlass.range(num_n_block_groups - 1, unroll=1):
+                    for stage in cutlass.range_constexpr(self.num_stages_S):
+                        n_block = n_block_max - 1 - n_block_group * self.num_stages_S - stage
+                        block_next = self._get_block_idx(n_block - 1, mPageTable_cur)
+                        block = self._get_block_idx(n_block, mPageTable_cur)
+                        if const_expr(self.has_qk):
+                            # copy K gmem -> smem
+                            producer_state_K = load_K(producer_state_K, block=block_next)
+                        # copy Vi gmem -> smem
+                        for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                            producer_state_V = load_V(
+                                producer_state_V, block=block_next, split=split
+                            )
+                        # copy Vti gmem -> smem
+                        for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                            producer_state_V = load_Vt(producer_state_V, block=block, split=split)
+
+                # ==== Epilogue ====
+                num_final_n_blocks = self.num_stages_S if even_n_blocks else self.num_stages_S - 1
+                for stage in cutlass.range(num_final_n_blocks, unroll_full=True):
+                    n_block = num_final_n_blocks - 1 - stage
+                    block = self._get_block_idx(n_block, mPageTable_cur)
+                    if n_block > 0:
+                        block_next = self._get_block_idx(n_block - 1, mPageTable_cur)
+                        if const_expr(self.has_qk):
+                            # copy K gmem -> smem
+                            producer_state_K = load_K(producer_state_K, block=block_next)
+                        # copy Vi gmem -> smem
+                        for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                            producer_state_V = load_V(
+                                producer_state_V, block=block_next, split=split
+                            )
+                    # copy Vti gmem -> smem
+                    for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                        producer_state_V = load_Vt(producer_state_V, block=block, split=split)
+
+            # Advance to next tile
+            work_tile = tile_scheduler.advance_to_next_work()
+
+        if const_expr(self.has_qk):
+            pipeline_Q.producer_tail(producer_state_Q)
+        pipeline_Qv.producer_tail(producer_state_Qv)
+        if const_expr(self.use_tma_KV):
+            if const_expr(self.has_qk):
+                pipeline_K.producer_tail(producer_state_K)
+            pipeline_V.producer_tail(producer_state_V)
+
+    @cute.jit
+    def _get_block_idx(
+        self,
+        n_block,
+        mPageTable_cur: Optional[cute.Tensor],
+    ):
+        if const_expr(mPageTable_cur is not None):
+            return mPageTable_cur[n_block]
+        else:
+            return n_block
+
+    @cute.jit
+    def load_inner(
+        self,
+        tma_atom: cute.CopyAtom,
+        tXgX: cute.Tensor,
+        tXsX: cute.Tensor,
+        load_pipeline: pipeline.PipelineAsync,
+        producer_state: pipeline.PipelineState,
+        block: Optional[Int32] = None,
+        split: Optional[Int32] = None,
+    ):
+        if const_expr(split is not None):
+            tXgX = tXgX[(None, split, None)]
+        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)
+        tma_bar_ptr = load_pipeline.producer_get_barrier(producer_state)
+        cute.copy(tma_atom, tXgX, tXsX, tma_bar_ptr=tma_bar_ptr)
+        producer_state.advance()
+        return producer_state
+
+    @cute.jit
+    def mma(
+        self,
+        sQ: Optional[cute.Tensor],
+        sK: Optional[cute.Tensor],
+        sQv: cute.Tensor,
+        sV: cute.Tensor,
+        sVt: cute.Tensor,
+        sP: cute.Tensor,
+        tStS: cute.Tensor,
+        tOtO0: cute.Tensor,
+        tOtO1: cute.Tensor,
+        tiled_mma_QK: cute.TiledMma,
+        tiled_mma_QvV: cute.TiledMma,
+        tiled_mma_PVt: cute.TiledMma,
+        pipeline_Q: Optional[pipeline.PipelineAsync],
+        pipeline_K: Optional[pipeline.PipelineAsync],
+        pipeline_Qv: pipeline.PipelineAsync,
+        pipeline_V: pipeline.PipelineAsync,
+        pipeline_S: pipeline.PipelineAsync,
+        pipeline_P: pipeline.PipelineAsync,
+        pipeline_O0: pipeline.PipelineAsync,
+        pipeline_O1: pipeline.PipelineAsync,
+        sO_empty_mbar_ptr: Optional[cute.Pointer],
+        is_leader_cta: Boolean,
+        topk_length_dynamic: Optional[Int32],
+        block_info: BlockInfo,
+        SeqlenInfoCls: Callable,
+        tile_scheduler: TileSchedulerProtocol,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+    ):
+        # ==== mma warp ====
+        # Description: Computes Q @ K^T, Qv @ V^T, and P @ V
+        # Produces: S, O
+        # Consumes: Q, K, Qv, V, P
+
+        pipelines_O = [pipeline_O0, pipeline_O1]
+        tOtOs = [tOtO0, tOtO1]
+
+        use_ptx_gemm_QK = True
+        use_ptx_gemm_QvV = True
+        use_ptx_gemm_PVt = True
+
+        # Operands for S = Q @ K^T
+        if const_expr(self.has_qk):
+            tSrQ = tiled_mma_QK.make_fragment_A(sQ)
+            tSrK = tiled_mma_QK.make_fragment_B(sK)
+
+        # Operands for S += Qv @ V^T
+        tSrQv = tiled_mma_QvV.make_fragment_A(sQv)
+        tSrV = tiled_mma_QvV.make_fragment_B(sV)
+
+        # Operands for Oi = P @ Vi
+        tOrP = tiled_mma_PVt.make_fragment_A(sP)
+        tOrVt = tiled_mma_PVt.make_fragment_B(sVt)
+
+        # GEMM functions
+        if const_expr(self.has_qk):
+            if const_expr(use_ptx_gemm_QK):
+                gemm_QK = [
+                    partial(
+                        fa_sm100_utils.gemm_ptx_partial,
+                        tiled_mma_QK.op,
+                        self.tmem_offset_S[stage],
+                        zero_init=True,
+                        cta_group=self.cta_group_size,
+                    )
+                    for stage in range(self.num_stages_S)
+                ]
+            else:
+                gemm_QK = [
+                    partial(
+                        fa_sm100_utils.gemm,
+                        tiled_mma_QK,
+                        tStS[None, None, None, stage],
+                        zero_init=True,
+                    )
+                    for stage in range(self.num_stages_S)
+                ]
+        if const_expr(use_ptx_gemm_QvV):
+            gemm_QvV = [
+                partial(
+                    fa_sm100_utils.gemm_ptx_partial,
+                    tiled_mma_QvV.op,
+                    self.tmem_offset_S[stage],
+                    cta_group=self.cta_group_size,
+                )
+                for stage in range(self.num_stages_S)
+            ]
+        else:
+            gemm_QvV = [
+                partial(
+                    fa_sm100_utils.gemm,
+                    tiled_mma_QvV,
+                    tStS[None, None, None, stage],
+                )
+                for stage in range(self.num_stages_S)
+            ]
+
+        if const_expr(use_ptx_gemm_PVt):
+            gemm_PVt = [
+                partial(
+                    fa_sm100_utils.gemm_ptx_partial,
+                    tiled_mma_PVt.op,
+                    self.tmem_offsets_O[split],
+                    cta_group=self.cta_group_size,
+                )
+                for split in range(self.num_hdimv_splits)
+            ]
+        else:
+            gemm_PVt = [
+                partial(
+                    fa_sm100_utils.gemm,
+                    tiled_mma_PVt,
+                    tOtOs[split],
+                )
+                for split in range(self.num_hdimv_splits)
+            ]
+
+        Consumer, Producer = pipeline.PipelineUserType.Consumer, pipeline.PipelineUserType.Producer
+        if const_expr(self.has_qk):
+            consumer_state_Q = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_Q)
+            consumer_state_K = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_K)
+        consumer_state_Qv = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_Qv)
+        consumer_state_V = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_V)
+        producer_state_S = pipeline.make_pipeline_state(Producer, stages=self.num_stages_S)
+        consumer_state_P = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_P)
+        producer_state_O0 = pipeline.make_pipeline_state(Producer, stages=self.num_stages_Oi)
+        producer_state_O1 = pipeline.make_pipeline_state(Producer, stages=self.num_stages_Oi)
+
+        mma_fn = self.mma_inner
+        if const_expr(self.has_qk):
+            mma_QK = partial(
+                mma_fn, gemm_QK, pipeline_K, tSrQ, sQ, tSrK, sK, use_ptx=use_ptx_gemm_QK
+            )
+        mma_QvV = partial(
+            mma_fn, gemm_QvV, pipeline_V, tSrQv, sQv, tSrV, sV, use_ptx=use_ptx_gemm_QvV
+        )
+        mma_PVt = partial(
+            mma_fn, gemm_PVt, pipeline_V, tOrP, sP, tOrVt, sVt, use_ptx=use_ptx_gemm_PVt
+        )
+
+        work_tile = tile_scheduler.initial_work_tile_info()
+        O_should_accumulate = False
+        while work_tile.is_valid_tile:
+            cluster_m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+            if const_expr(self.use_packed_varlen_sched):
+                batch_idx = get_batch_from_cu_tensor(cluster_m_block, mCuSeqlensQ)
+                cluster_m_block -= mCuSeqlensQ[batch_idx]
+
+            seqlen = SeqlenInfoCls(batch_idx)
+            if const_expr(self.is_topk_gather):
+                n_block_min = 0
+                # n_block_max = self.topk_length // self.tile_n
+                n_block_max = topk_length_dynamic // self.tile_n
+            else:
+                n_block_min, n_block_max = block_info.get_n_block_min_max(
+                    seqlen,
+                    cluster_m_block,
+                )
+            num_n_blocks = n_block_max - n_block_min
+            even_n_blocks = num_n_blocks % 2 == 0 and num_n_blocks > 0
+            num_n_block_groups = cute.ceil_div(num_n_blocks, self.num_stages_S)
+
+            if is_leader_cta:
+                if const_expr(self.has_qk):
+                    pipeline_Q.consumer_wait(consumer_state_Q)
+
+                consumer_wait_state_Qv = consumer_state_Qv.clone()
+                for _ in cutlass.range_constexpr(self.num_hdimv_splits):
+                    pipeline_Qv.consumer_wait(consumer_wait_state_Qv)
+                    consumer_wait_state_Qv.advance()
+
+                producer_states_O = [producer_state_O0, producer_state_O1]
+
+                # ==== Prologue ====
+                pipeline_S.producer_acquire(producer_state_S)
+                if const_expr(self.has_qk):
+                    # S = Q @ K^T
+                    consumer_state_K = mma_QK(consumer_state_K, acc_stage=0)
+                # S += Qvi @ Vi^T
+                for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                    consumer_state_V = mma_QvV(
+                        consumer_state_V,
+                        acc_stage=0,
+                        a_stage=split,
+                        zero_init=split == 0 and not self.has_qk,
+                    )
+                pipeline_S.producer_commit(producer_state_S)
+                producer_state_S.advance()
+
+                # ==== Mainloop ====
+                for _ in cutlass.range(num_n_block_groups - 1, unroll=1):
+                    for stage in cutlass.range_constexpr(self.num_stages_S):
+                        next_stage = const_expr((stage + 1) % self.num_stages_S)
+                        pipeline_S.producer_acquire(producer_state_S)
+                        if const_expr(self.has_qk):
+                            # S = Q @ K^T
+                            consumer_state_K = mma_QK(consumer_state_K, acc_stage=next_stage)
+                        # S += Qvi @ Vi^T
+                        for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                            consumer_state_V = mma_QvV(
+                                consumer_state_V,
+                                acc_stage=next_stage,
+                                a_stage=split,
+                                zero_init=split == 0 and not self.has_qk,
+                            )
+                        pipeline_S.producer_commit(producer_state_S)
+                        producer_state_S.advance()
+                        # Oi += P @ Vi
+                        pipeline_P.consumer_wait(consumer_state_P)
+                        for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                            producer_state_Oi = producer_states_O[split]
+                            pipelines_O[split].producer_acquire(producer_state_Oi)
+                            consumer_state_V = mma_PVt(
+                                consumer_state_V,
+                                acc_stage=split,
+                                a_stage=consumer_state_P.index,
+                                zero_init=not O_should_accumulate,
+                            )
+                            pipelines_O[split].producer_commit(producer_state_Oi)
+                            producer_state_Oi.advance()
+                            producer_states_O[split] = producer_state_Oi
+                        pipeline_P.consumer_release(consumer_state_P)
+                        consumer_state_P.advance()
+                        O_should_accumulate = True
+
+                # ==== Epilogue ====
+                num_final_n_blocks = self.num_stages_S if even_n_blocks else self.num_stages_S - 1
+                for stage in cutlass.range_constexpr(self.num_stages_S):
+                    n_block = num_final_n_blocks - 1 - stage
+                    if const_expr(stage == 0):
+                        if n_block > 0:
+                            pipeline_S.producer_acquire(producer_state_S)
+                            if const_expr(self.has_qk):
+                                # S = Q @ K^T
+                                consumer_state_K = mma_QK(consumer_state_K, acc_stage=stage + 1)
+                            # S += Qvi @ Vi^T
+                            for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                                consumer_state_V = mma_QvV(
+                                    consumer_state_V,
+                                    acc_stage=stage + 1,
+                                    a_stage=split,
+                                    zero_init=split == 0 and not self.has_qk,
+                                )
+                            pipeline_S.producer_commit(producer_state_S)
+                            producer_state_S.advance()
+                    if n_block >= 0:
+                        # Oi += P @ Vi
+                        pipeline_P.consumer_wait(consumer_state_P)
+                        for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                            producer_state_Oi = producer_states_O[split]
+                            pipelines_O[split].producer_acquire(producer_state_Oi)
+                            consumer_state_V = mma_PVt(
+                                consumer_state_V,
+                                acc_stage=split,
+                                a_stage=consumer_state_P.index,
+                                zero_init=not O_should_accumulate,
+                            )
+                            pipelines_O[split].producer_commit(producer_state_Oi)
+                            producer_state_Oi.advance()
+                            producer_states_O[split] = producer_state_Oi
+                        pipeline_P.consumer_release(consumer_state_P)
+                        consumer_state_P.advance()
+                    O_should_accumulate = True
+
+                producer_state_O0, producer_state_O1 = producer_states_O
+
+                if const_expr(self.has_qk):
+                    pipeline_Q.consumer_release(consumer_state_Q)
+                    consumer_state_Q.advance()
+
+                # if we overlap sOi with sQvi for tma store, need to acquire signal
+                if const_expr(self.use_tma_O and not self.overlap_sO_sV):
+                    pipeline_O0.producer_tail(producer_state_O0.clone())
+                    pipeline_O1.producer_tail(producer_state_O1.clone())
+
+                for _ in cutlass.range_constexpr(self.num_hdimv_splits):
+                    pipeline_Qv.consumer_release(consumer_state_Qv)
+                    consumer_state_Qv.advance()
+
+            # Advance to next tile
+            work_tile = tile_scheduler.advance_to_next_work()
+            O_should_accumulate = False
+
+        pipeline_S.producer_tail(producer_state_S)
+        pipeline_O0.producer_tail(producer_state_O0)
+        pipeline_O1.producer_tail(producer_state_O1)
+
+    @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,
+        use_ptx: bool = True,
+    ):
+        if const_expr(acc_stage is not None):
+            gemm = gemm[acc_stage]
+
+        tCrA_cur = tCrA[None, None, None, a_stage]
+        sA_cur = sA[None, None, None, a_stage]
+        b_stage = consumer_state.index
+        tCrB_cur = tCrB[None, None, None, b_stage]
+        sB_cur = sB[None, None, None, b_stage]
+
+        load_pipeline.consumer_wait(consumer_state)
+
+        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
+        gemm(**kwargs)
+
+        load_pipeline.consumer_release(consumer_state)
+        consumer_state.advance()
+        return consumer_state
+
+    @cute.jit
+    def softmax_loop(
+        self,
+        softmax_scale: Float32,
+        softmax_scale_log2: Float32,
+        mLSE: Optional[cute.Tensor],
+        mRowMax: Optional[cute.Tensor],
+        sRowMax: cute.Tensor,
+        sRowSum: cute.Tensor,
+        sScale: cute.Tensor,
+        sBitmask: Optional[cute.Tensor],
+        sP: cute.Tensor,
+        tStS: cute.Tensor,
+        thr_mma_S: cute.ThrMma,
+        pipeline_S: pipeline.PipelineAsync,
+        pipeline_P: pipeline.PipelineAsync,
+        pipeline_sm_stats: pipeline.PipelineAsync,
+        pipeline_bitmask: Optional[pipeline.PipelineAsync],
+        sO_empty_mbar_ptr: Optional[cute.Pointer],
+        AttentionMaskCls: Callable,
+        topk_length_dynamic: Optional[Int32],
+        block_info: BlockInfo,
+        SeqlenInfoCls: Callable,
+        tile_scheduler: TileSchedulerProtocol,
+        tma_atom_P: Optional[cute.CopyAtom] = None,
+        mP: Optional[cute.Tensor] = None,
+        sP_out: Optional[cute.Tensor] = None,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+    ):
+        # ==== softmax warpgroup ====
+        # Description: computes softmax on S and writes the result to P
+        # Produces: P, softmax stats
+        # Consumes: S, bitmask (for topk sparsity)
+
+        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())
+
+        tSAcc = tStS[(None, None), 0, 0, 0]
+        tSAcc_staged = [tStS[(None, None), 0, 0, stage] for stage in range(self.num_stages_S)]
+
+        cS = cute.make_identity_tensor(self.mma_tiler_QK[:2])  # (128, 128)
+        tScS = thr_mma_S.partition_C(cS)[(None, None), 0, 0]  # (64, 128)
+
+        # S tmem -> rmem copy objects
+        tmem_load_atom = cute.make_copy_atom(
+            tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)),
+            self.dtype_acc,
+        )
+        tmem_load_tiled = tcgen05.make_tmem_copy(tmem_load_atom, tSAcc)
+        tmem_load_thr = tmem_load_tiled.get_slice(tidx)
+        # S tmem -> rmem copy operands
+        tStS_t2r = tmem_load_thr.partition_S(tSAcc)  # (((32, 32), 1), 1, 2)
+        tStS_t2r_staged = [
+            tmem_load_thr.partition_S(tSAcc_staged[stage]) for stage in range(self.num_stages_S)
+        ]
+        tScS_t2r = tmem_load_thr.partition_D(tScS)
+        tSrS_t2r = cute.make_rmem_tensor(tScS_t2r.shape, self.dtype_acc)
+
+        # P rmem -> smem copy objects
+        universal_copy_bits = 128
+        smem_store_atom = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            self.dtype_P,
+            num_bits_per_copy=universal_copy_bits,
+        )
+        smem_store_tiled = cute.make_tiled_copy_D(smem_store_atom, tmem_load_tiled)
+        smem_store_thr = smem_store_tiled.get_slice(tidx)
+        # P rmem -> smem copy operands
+        sP_mnp_layout = cute.make_ordered_layout(
+            self.tile_P + (self.num_stages_P,), order=(0, 1, 2)
+        )
+        sP_mnp = cute.composition(sP, sP_mnp_layout)
+        sP_smem_view = smem_store_thr.partition_D(sP_mnp)
+
+        Consumer, Producer = pipeline.PipelineUserType.Consumer, pipeline.PipelineUserType.Producer
+        consumer_state_S = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_S)
+        producer_state_P = pipeline.make_pipeline_state(Producer, stages=self.num_stages_P)
+        producer_state_sm_stats = pipeline.make_pipeline_state(Producer, stages=self.num_stages_sm_stats)
+        consumer_state_bitmask = None
+        if const_expr(self.is_topk_gather and not self.disable_bitmask):
+            consumer_state_bitmask = pipeline.make_pipeline_state(
+                Consumer, stages=self.num_stages_bitmask
+            )
+
+        work_tile = tile_scheduler.initial_work_tile_info()
+        while work_tile.is_valid_tile:
+            cluster_m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+            if const_expr(self.use_packed_varlen_sched):
+                batch_idx = get_batch_from_cu_tensor(cluster_m_block, mCuSeqlensQ)
+                cluster_m_block -= mCuSeqlensQ[batch_idx]
+            cta_m_block = cluster_m_block * self.cta_group_size + cta_rank_in_cluster
+            seqlen = SeqlenInfoCls(batch_idx)
+            if const_expr(self.is_topk_gather):
+                n_block_min = 0
+                n_block_max = self.topk_length // self.tile_n
+                # n_block_max = topk_length_dynamic // self.tile_n
+            else:
+                n_block_min, n_block_max = block_info.get_n_block_min_max(
+                    seqlen,
+                    cluster_m_block,
+                )
+            num_n_blocks = n_block_max - n_block_min
+            even_n_blocks = num_n_blocks % 2 == 0 and num_n_blocks > 0
+            num_n_block_groups = cute.ceil_div(num_n_blocks, self.num_stages_S)
+
+            gRowMax = None
+            if const_expr(mRowMax is not None):
+                # (seqlen_q, {seqlen_k_rounded, topk} / tile_n)
+                if const_expr(not seqlen.has_cu_seqlens_q):
+                    mRowMax_cur = mRowMax[None, None, head_idx, batch_idx]
+                else:
+                    q_offset = (
+                        seqlen.offset_q if const_expr(not self.pack_gqa) else (0, seqlen.offset_q)
+                    )
+                    mRowMax_cur = cute.domain_offset((q_offset, 0), mRowMax[None, None, head_idx])
+                # (cta_tile_m, {seqlen_k_rounded, topk} / tile_n)
+                gRowMax = cute.local_tile(mRowMax_cur, (self.cta_tile_m,), (cta_m_block, None))
+
+            store_P = None
+            if const_expr(self.store_P):
+                # (seqlen_q, seqlen_k)
+                mP_cur = seqlen.offset_batch_Q(mP, batch_idx, dim=3, ragged=self.ragged_tma_O)[
+                    None, None, head_idx
+                ]
+                # (cta_tile_m, tile_n, num_n_blocks)
+                gP = cute.local_tile(mP_cur, self.tile_P, (cta_m_block, None))
+                store_P, tPsP, tPgP = copy_utils.tma_get_copy_fn(
+                    tma_atom_P,
+                    0,
+                    cute.make_layout(1),
+                    sP_out,
+                    gP,
+                )
+
+            mask = AttentionMaskCls(seqlen)
+            mask_fn = partial(
+                mask.apply_mask_sm100,
+                m_block=cluster_m_block,
+                thr_mma=thr_mma_S,
+                thr_tmem_load=tmem_load_thr,
+                mask_causal=self.is_causal,
+                mask_local=self.is_local,
+                batch_idx=batch_idx,
+                head_idx=head_idx,
+                r2p=False,  # TODO: fix r2p for 2cta
+            )
+            disable_mask = self.disable_bitmask and self.is_topk_gather
+
+            softmax = SoftmaxSm100.create(
+                softmax_scale_log2,
+                rescale_threshold=8.0 if const_expr(self.dtype_Q.width == 16) else 0.0,
+                softmax_scale=softmax_scale,
+            )
+            softmax.reset()
+
+            softmax_step_fn = partial(
+                self.softmax_step,
+                softmax,
+                sRowMax,
+                sScale,
+                sBitmask,
+                tStS_t2r_staged,
+                tSrS_t2r,
+                sP_smem_view,
+                tmem_load_thr,
+                smem_store_thr,
+                pipeline_S,
+                pipeline_P,
+                pipeline_sm_stats,
+                pipeline_bitmask,
+                tidx,
+                warp_idx,
+                store_P=store_P,
+                gRowMax=gRowMax,
+            )
+
+            ### first iteration ###
+            n_block = n_block_max - 1
+            (
+                consumer_state_S,
+                producer_state_P,
+                producer_state_sm_stats,
+                consumer_state_bitmask,
+            ) = softmax_step_fn(
+                consumer_state_S,
+                producer_state_P,
+                producer_state_sm_stats,
+                consumer_state_bitmask,
+                0,
+                n_block,
+                mask_fn=partial(mask_fn, mask_seqlen=True)
+                if not const_expr(disable_mask)
+                else None,
+                is_first=True,
+            )
+            n_block -= 1
+
+            ### Separate iterations with causal masking
+            # note: For square mma tile, can mask at most 1 n_block_group
+            if const_expr((self.is_causal or self.is_local) and not self.is_topk_gather):
+                n_block_min_causal_local_mask = block_info.get_n_block_min_causal_local_mask(
+                    seqlen, cluster_m_block, n_block_min
+                )
+                num_masked_n_blocks = n_block_max - 1 - n_block_min_causal_local_mask
+                num_masked_n_block_groups = min(
+                    num_n_block_groups - 1, cute.ceil_div(num_masked_n_blocks, self.num_stages_S)
+                )
+                num_n_block_groups -= num_masked_n_block_groups
+                for _ in cutlass.range(num_masked_n_block_groups, unroll=1):
+                    for stage in cutlass.range_constexpr(self.num_stages_S):
+                        (
+                            consumer_state_S,
+                            producer_state_P,
+                            producer_state_sm_stats,
+                            consumer_state_bitmask,
+                        ) = softmax_step_fn(
+                            consumer_state_S,
+                            producer_state_P,
+                            producer_state_sm_stats,
+                            consumer_state_bitmask,
+                            1 - stage,
+                            n_block,
+                            mask_fn=partial(mask_fn, mask_seqlen=False),
+                        )
+                        n_block -= 1
+
+            ### Mainloop ###
+            for n_block_group in cutlass.range(num_n_block_groups - 1, unroll=1):
+                for stage in cutlass.range_constexpr(self.num_stages_S):
+                    (
+                        consumer_state_S,
+                        producer_state_P,
+                        producer_state_sm_stats,
+                        consumer_state_bitmask,
+                    ) = softmax_step_fn(
+                        consumer_state_S,
+                        producer_state_P,
+                        producer_state_sm_stats,
+                        consumer_state_bitmask,
+                        1 - stage,
+                        n_block,
+                        mask_fn=partial(mask_fn, mask_seqlen=False)
+                        if const_expr(self.is_topk_gather and not self.disable_bitmask)
+                        else None,
+                    )
+                    n_block -= 1
+
+            ### last iteration if even ###
+            # always mask to simplify logic
+            if even_n_blocks:
+                (
+                    consumer_state_S,
+                    producer_state_P,
+                    producer_state_sm_stats,
+                    consumer_state_bitmask,
+                ) = softmax_step_fn(
+                    consumer_state_S,
+                    producer_state_P,
+                    producer_state_sm_stats,
+                    consumer_state_bitmask,
+                    1,
+                    n_block,
+                    mask_fn=partial(mask_fn, mask_seqlen=False)
+                    if not const_expr(disable_mask)
+                    else None,
+                )
+                n_block -= 1
+
+            # write row max and sum to smem
+            sRowSum[tidx % self.cta_tile_m, warp_idx // self.cta_group_size] = softmax.row_sum[0]
+            if const_expr(mLSE is not None):
+                if tidx < self.cta_tile_m:
+                    sRowMax[tidx, 0] = softmax.row_max[0]
+            self.sm_stats_barrier_full.arrive()
+
+            # Advance to next tile
+            work_tile = tile_scheduler.advance_to_next_work()
+            self.sm_stats_barrier_empty.arrive_and_wait()
+
+        pipeline_P.producer_tail(producer_state_P)
+        pipeline_sm_stats.producer_tail(producer_state_sm_stats)
+
+    @cute.jit
+    def softmax_step(
+        self,
+        softmax: SoftmaxSm100,
+        sRowMax: cute.Tensor,
+        sScale: cute.Tensor,
+        sBitmask: Optional[cute.Tensor],
+        tStS_t2r_staged: cute.Tensor,
+        tSrS_t2r: cute.Tensor,
+        sP_smem_view: cute.Tensor,
+        tmem_load_thr: cute.CopyAtom,
+        smem_store_thr: cute.CopyAtom,
+        pipeline_S: pipeline.PipelineAsync,
+        pipeline_P: pipeline.PipelineAsync,
+        pipeline_sm_stats: pipeline.PipelineAsync,
+        pipeline_bitmask: Optional[pipeline.PipelineAsync],
+        tidx: Int32,
+        warp_idx: Int32,
+        consumer_state_S: pipeline.PipelineState,
+        producer_state_P: pipeline.PipelineState,
+        producer_state_sm_stats: pipeline.PipelineState,
+        consumer_state_bitmask: Optional[pipeline.PipelineState],
+        stage: cutlass.Constexpr[Int32],
+        n_block: Int32,
+        mask_fn: Optional[Callable] = None,
+        is_first: Boolean = False,
+        store_P: Optional[Callable] = None,
+        gRowMax: Optional[cute.Tensor] = None,
+    ):
+        leader_warp = warp_idx == 0
+        tSrP = cute.make_rmem_tensor(tSrS_t2r.shape, self.dtype_P)
+        rP_smem_view = smem_store_thr.retile(tSrP)
+
+        pipeline_S.consumer_wait(consumer_state_S)
+        cute.copy(tmem_load_thr, tStS_t2r_staged[stage], tSrS_t2r)
+        cute.arch.fence_view_async_tmem_load()
+        pipeline_S.consumer_release(consumer_state_S)
+
+        rBitmask = None
+        if const_expr(self.is_topk_gather and not self.disable_bitmask):
+            assert pipeline_bitmask is not None
+            assert consumer_state_bitmask is not None
+            pipeline_bitmask.consumer_wait(consumer_state_bitmask)
+            rBitmask = cute.make_rmem_tensor((self.tile_n // 64,), dtype=Uint32)
+            bitmask_col_offset = self.tile_n // 64 if warp_idx >= 2 else 0
+            for i in cutlass.range_constexpr(cute.size(rBitmask)):
+                rBitmask[i] = sBitmask[bitmask_col_offset + i, consumer_state_bitmask.index]
+
+        if const_expr(mask_fn is not None):
+            mask_fn(tSrS_t2r, n_block=n_block, rBitmask=rBitmask)
+
+        # compute threadwise row_max
+        row_max = softmax.compute_row_max_local(tSrS_t2r.load(), is_first)
+        # 2-thread reduce row_max through smem
+        assert self.cta_tile_m * self.cta_group_size == 128
+        sRowMax[tidx % self.cta_tile_m, warp_idx // self.cta_group_size] = row_max
+        self.softmax_barrier.arrive_and_wait()
+        # must release after barrier sync
+        if const_expr(self.is_topk_gather and not self.disable_bitmask):
+            pipeline_bitmask.consumer_release(consumer_state_bitmask)
+        row_max0 = sRowMax[tidx % self.cta_tile_m, 0]
+        row_max1 = sRowMax[tidx % self.cta_tile_m, 1]
+        row_max = max(row_max0, row_max1)
+
+        row_max, acc_scale = softmax.update_row_max_from_local(row_max, is_first)
+
+        if const_expr(gRowMax is not None):
+            if tidx < self.cta_tile_m:
+                gRowMax[tidx, n_block] = row_max
+
+        # note: acc_scales agree for paired threads
+        pipeline_sm_stats.producer_acquire(producer_state_sm_stats)
+        if warp_idx < self.cta_group_size:
+            sScale[tidx % self.cta_tile_m, producer_state_sm_stats.index] = acc_scale
+        pipeline_sm_stats.producer_commit(producer_state_sm_stats)
+
+        # x -> scale_log2*x-rowmax
+        softmax.scale_subtract_rowmax(tSrS_t2r, row_max)
+
+        # x -> exp2(x)
+        softmax.apply_exp2_convert(tSrS_t2r, tSrP)
+
+        if const_expr(self.store_P):
+            if leader_warp:
+                cute.arch.cp_async_bulk_wait_group(self.num_stages_P - 1, read=True)
+            self.softmax_barrier.arrive_and_wait()
+
+        pipeline_P.producer_acquire(producer_state_P)
+        cute.copy(
+            smem_store_thr, rP_smem_view, sP_smem_view[None, None, None, producer_state_P.index]
+        )
+        cute.arch.fence_view_async_shared()
+        pipeline_P.producer_commit(producer_state_P)
+        # unconditionally necessary for sRowMax read to complete before next iter's store
+        self.softmax_barrier.arrive_and_wait()
+
+        if const_expr(self.store_P):
+            if leader_warp:
+                store_P(src_idx=producer_state_P.index, dst_idx=n_block)
+                cute.arch.cp_async_bulk_commit_group()
+
+        consumer_state_S.advance()
+        producer_state_P.advance()
+        producer_state_sm_stats.advance()
+        if const_expr(self.is_topk_gather and not self.disable_bitmask):
+            consumer_state_bitmask.advance()
+
+        softmax.update_row_sum(tSrS_t2r.load(), acc_scale, is_first)
+
+        return consumer_state_S, producer_state_P, producer_state_sm_stats, consumer_state_bitmask
+
+    @cute.jit
+    def correction_loop(
+        self,
+        softmax_scale_log2: Float32,
+        mO: cute.Tensor,
+        mLSE: Optional[cute.Tensor],
+        tma_atom_O: Optional[cute.CopyAtom],
+        sRowMax: cute.Tensor,
+        sRowSum: cute.Tensor,
+        sScale: cute.Tensor,
+        sO: cute.Tensor,
+        tOtO0: cute.Tensor,
+        tOtO1: cute.Tensor,
+        pipeline_O0: pipeline.PipelineAsync,
+        pipeline_O1: pipeline.PipelineAsync,
+        pipeline_sm_stats: pipeline.PipelineAsync,
+        sO_empty_mbar_ptr: Optional[cute.Pointer],
+        tiled_copy_O_r2g: cute.TiledCopy,
+        topk_length_dynamic: Optional[Int32],
+        block_info: BlockInfo,
+        SeqlenInfoCls: Callable,
+        tile_scheduler: TileSchedulerProtocol,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+    ):
+        ### ==== correction/epilogue warpgroup ====
+        # Correction: copy scale smem -> rmem, copy O tmem -> rmem, rescale O, store O rmem -> tmem
+        # Epilogue:   copy O tmem -> rmem, do final scaling of O, store O rmem -> gmem,
+        #             optionally store LSE
+        # Produces: -
+        # Consumes: O, softmax stats
+
+        tidx = cute.arch.thread_idx()[0] % self.num_epilogue_threads
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % (
+            self.num_epilogue_threads // 32
+        )
+        cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+        leader_warp = warp_idx == 0
+
+        tOtO0 = tOtO0[(None, None), 0, 0]  # (64, (128, 2))
+        tOtO1 = tOtO1[(None, None), 0, 0]  # (64, (128, 2))
+        tOtOs = [tOtO0, tOtO1]
+
+        # tuneable parameter
+        corr_tile_size = math.gcd(32, self.tmem_cols_Oi)
+
+        tmem_load_atom_O = cute.make_copy_atom(
+            tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)),
+            self.dtype_acc,
+        )
+        tmem_store_atom_O = cute.make_copy_atom(
+            tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)),
+            self.dtype_acc,
+        )
+        thr_tmem_load_O = tcgen05.make_tmem_copy(tmem_load_atom_O, tOtO0).get_slice(tidx)
+        thr_tmem_store_O = tcgen05.make_tmem_copy(tmem_store_atom_O, tOtO0).get_slice(tidx)
+
+        # ((32,1),1,4)
+        tOtOs_t2r = [
+            thr_tmem_load_O.partition_S(tOtOs[split]) for split in range(self.num_hdimv_splits)
+        ]
+        tOtOs_r2t = [
+            thr_tmem_store_O.partition_D(tOtOs[split]) for split in range(self.num_hdimv_splits)
+        ]
+
+        cOi = cute.make_identity_tensor((self.cta_tile_m, self.hdimv // self.num_hdimv_splits))
+        thr_tiled_copy_O_r2g = tiled_copy_O_r2g.get_slice(tidx)
+        tOicOi = thr_tiled_copy_O_r2g.partition_S(cOi)
+
+        tOicOi_t2r = thr_tmem_load_O.partition_D(tOicOi[(None, None), 0, 0])
+
+        pipelines_O = [pipeline_O0, pipeline_O1]
+
+        Consumer = pipeline.PipelineUserType.Consumer
+        consumer_state_O0 = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_Oi)
+        consumer_state_O1 = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_Oi)
+        consumer_state_sm_stats = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_sm_stats)
+
+        do_correction_rescale = partial(
+            self.correction_rescale,
+            thr_tmem_load_O,
+            thr_tmem_store_O,
+            tOicOi_t2r,
+        )
+
+        work_tile = tile_scheduler.initial_work_tile_info()
+        while work_tile.is_valid_tile:
+            cluster_m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+            if const_expr(self.use_packed_varlen_sched):
+                batch_idx = get_batch_from_cu_tensor(cluster_m_block, mCuSeqlensQ)
+                cluster_m_block -= mCuSeqlensQ[batch_idx]
+            cta_m_block = cluster_m_block * self.cta_group_size + cta_rank_in_cluster
+
+            seqlen = SeqlenInfoCls(batch_idx)
+            if const_expr(self.is_topk_gather):
+                n_block_min = 0
+                n_block_max = self.topk_length // self.tile_n
+                # n_block_max = topk_length_dynamic // self.tile_n
+            else:
+                n_block_min, n_block_max = block_info.get_n_block_min_max(
+                    seqlen,
+                    cluster_m_block,
+                )
+            num_n_blocks = n_block_max - n_block_min
+
+            consumer_states_O = [consumer_state_O0, consumer_state_O1]
+
+            # acquire first signal and release immediately
+            pipeline_sm_stats.consumer_wait(consumer_state_sm_stats)
+            pipeline_sm_stats.consumer_release(consumer_state_sm_stats)
+            consumer_state_sm_stats.advance()
+
+            for _ in cutlass.range(num_n_blocks - 1, unroll=1):
+                pipeline_sm_stats.consumer_wait(consumer_state_sm_stats)
+                scale = sScale[tidx % self.cta_tile_m, consumer_state_sm_stats.index]
+                should_rescale = cute.arch.vote_ballot_sync(scale < 1.0) != 0
+                pipeline_sm_stats.consumer_release(consumer_state_sm_stats)
+                consumer_state_sm_stats.advance()
+
+                for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                    consumer_state_Oi = consumer_states_O[split]
+                    pipelines_O[split].consumer_wait(consumer_state_Oi)
+                    if should_rescale:
+                        do_correction_rescale(
+                            tOtOs_t2r[split],
+                            tOtOs_r2t[split],
+                            scale,
+                        )
+                    pipelines_O[split].consumer_release(consumer_state_Oi)
+                    consumer_state_Oi.advance()
+                    consumer_states_O[split] = consumer_state_Oi
+
+            # (seqlen_q, hdimv)
+            mO_cur = seqlen.offset_batch_Q(mO, batch_idx, dim=3, ragged=self.ragged_tma_O)[
+                None, None, head_idx
+            ]
+            # (cta_tile_m, hdimv//2, 2)
+            gO = cute.local_tile(
+                mO_cur,
+                (self.cta_tile_m, self.hdimv // self.num_hdimv_splits),
+                (cta_m_block, None),
+            )
+            tOgO = thr_tiled_copy_O_r2g.partition_D(gO)
+            # ((32, 1), 1, 4)
+            tOrOs_t2r = [
+                cute.make_rmem_tensor(tOicOi_t2r.shape, self.dtype_acc)
+                for split in range(self.num_hdimv_splits)
+            ]
+            tOrOs_r2g_f32 = [
+                thr_tiled_copy_O_r2g.retile(tOrOs_t2r[split])
+                for split in range(self.num_hdimv_splits)
+            ]
+            tOrOs_r2g = [
+                cute.make_rmem_tensor_like(tOrOs_r2g_f32[split], self.dtype_O)
+                for split in range(self.num_hdimv_splits)
+            ]
+            if const_expr(self.use_tma_O):
+                tOsO = thr_tiled_copy_O_r2g.partition_D(sO)
+                store_O, _, _ = copy_utils.tma_get_copy_fn(
+                    tma_atom_O,
+                    0,
+                    cute.make_layout(1),
+                    sO,
+                    gO,
+                )
+
+            self.sm_stats_barrier_full.arrive_and_wait()
+
+            row_sum0 = sRowSum[tidx % self.cta_tile_m, 0]
+            row_sum1 = sRowSum[tidx % self.cta_tile_m, 1]
+            row_sum = row_sum0 + row_sum1
+            acc_O_mn_row_is_zero_or_nan = row_sum == 0.0 or row_sum != row_sum
+            scale = cute.arch.rcp_approx(row_sum if not acc_O_mn_row_is_zero_or_nan else 1.0)
+
+            row_max = 0.0
+            if const_expr(mLSE is not None):
+                if tidx < self.cta_tile_m:
+                    row_max = sRowMax[tidx, 0]
+
+            self.sm_stats_barrier_empty.arrive()
+
+            seqlen_q = (
+                seqlen.seqlen_q
+                if const_expr(not self.pack_gqa)
+                else seqlen.seqlen_q * self.qhead_per_kvhead
+            )
+
+            # compute and store lse to gmem
+            if const_expr(mLSE is not None):
+                if const_expr(not seqlen.has_cu_seqlens_q):
+                    mLSE_cur = mLSE[None, head_idx, batch_idx]
+                else:
+                    lse_offset = (
+                        seqlen.offset_q if const_expr(not self.pack_gqa) else (0, seqlen.offset_q)
+                    )
+                    mLSE_cur = cute.domain_offset((lse_offset,), mLSE[None, head_idx])
+                gLSE = cute.local_tile(mLSE_cur, (self.cta_tile_m,), (cta_m_block,))
+                if tidx < self.cta_tile_m:
+                    LN2 = math.log(2.0)
+                    lse = (
+                        (row_max * softmax_scale_log2 + cute.math.log2(row_sum, fastmath=True))
+                        * LN2
+                        if not acc_O_mn_row_is_zero_or_nan
+                        else -Float32.inf
+                    )
+                    if tidx < seqlen_q - cta_m_block * self.cta_tile_m:
+                        gLSE[tidx] = lse
+
+            row_idx = cta_m_block * self.cta_tile_m + tOicOi[0][0]
+
+            for split in cutlass.range_constexpr(self.num_hdimv_splits):
+                consumer_state_Oi = consumer_states_O[split]
+                pipelines_O[split].consumer_wait(consumer_state_Oi)
+                # copy Oi tmem -> rmem
+                cute.copy(
+                    thr_tmem_load_O,
+                    tOtOs_t2r[split],
+                    tOrOs_t2r[split],
+                )
+
+                # scale and downcast Oi
+                tOrOs_r2g[split].store((tOrOs_r2g_f32[split].load() * scale).to(self.dtype_O))
+
+                if const_expr(not self.use_tma_O):
+                    # copy Oi rmem -> gmem
+                    if row_idx < seqlen_q:
+                        cute.copy(
+                            thr_tiled_copy_O_r2g,
+                            tOrOs_r2g[split],
+                            tOgO[None, None, None, split],
+                        )
+                else:
+                    # copy Oi rmem -> smem
+                    if const_expr(self.overlap_sO_sV):
+                        # last slot for Vti is always 2, 3
+                        sO_idx = 2 + split
+                    else:
+                        sO_idx = split
+                    cute.copy(
+                        thr_tiled_copy_O_r2g,
+                        tOrOs_r2g[split],
+                        tOsO[None, None, None, sO_idx],
+                    )
+                    cute.arch.fence_view_async_shared()
+                    self.epi_barrier.arrive_and_wait()
+                    # tma store Oi smem -> gmem
+                    if leader_warp:
+                        store_O(src_idx=sO_idx, dst_idx=split)
+                        cute.arch.cp_async_bulk_commit_group()
+                        cute.arch.cp_async_bulk_wait_group(1 - split, read=True)
+                        if const_expr(split == 1 and self.overlap_sO_sV):
+                            with cute.arch.elect_one():
+                                cute.arch.mbarrier_arrive(sO_empty_mbar_ptr)
+
+            consumer_state_O0, consumer_state_O1 = consumer_states_O
+
+            cute.arch.fence_view_async_tmem_load()
+            pipeline_O0.consumer_release(consumer_state_O0)
+            pipeline_O1.consumer_release(consumer_state_O1)
+            consumer_state_O0.advance()
+            consumer_state_O1.advance()
+
+            # Advance to next tile
+            work_tile = tile_scheduler.advance_to_next_work()
+
+    @cute.jit
+    def correction_rescale(
+        self,
+        thr_tmem_load: cute.CopyAtom,
+        thr_tmem_store: cute.CopyAtom,
+        tOcO_t2r: cute.Tensor,
+        tOtO_t2r: cute.Tensor,
+        tOtO_r2t: cute.Tensor,
+        scale: Float32,
+    ):
+        tOrO_t2r_frg = cute.make_rmem_tensor_like(tOcO_t2r[None, None, 0], self.dtype_acc)
+
+        for i in cutlass.range_constexpr(cute.size(tOtO_t2r, mode=[2])):
+            tOtO_t2r_cur = tOtO_t2r[None, None, i]
+            tOtO_r2t_cur = tOtO_r2t[None, None, i]
+
+            cute.copy(thr_tmem_load, tOtO_t2r_cur, tOrO_t2r_frg)
+            for j in cutlass.range(0, cute.size(tOrO_t2r_frg), 2, unroll_full=True):
+                tOrO_t2r_frg[j], tOrO_t2r_frg[j + 1] = cute.arch.mul_packed_f32x2(
+                    (tOrO_t2r_frg[j], tOrO_t2r_frg[j + 1]), (scale, scale)
+                )
+            cute.copy(thr_tmem_store, tOrO_t2r_frg, tOtO_r2t_cur)
+        cute.arch.fence_view_async_tmem_store()
diff --git a/build/torch-cuda/flash_fwd_sm100.py b/build/torch-cuda/flash_fwd_sm100.py
index e9ba465fa4e0a6958f4fcb5350ec7bc876b00bc9..c260f075fe228b41db8ff4862eb206eb8cabdb31 100644
--- a/build/torch-cuda/flash_fwd_sm100.py
+++ b/build/torch-cuda/flash_fwd_sm100.py
@@ -13,9 +13,8 @@
 # https://github.com/NVIDIA/cutlass/tree/main/examples/77_blackwell_fmha
 # https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/blackwell/fmha.py
 
-import enum
 import math
-from typing import Type, Tuple, Callable, Optional, Literal
+from typing import Tuple, Callable, Optional, Literal, NamedTuple
 from functools import partial
 
 import cuda.bindings.driver as cuda
@@ -28,6 +27,7 @@ import cutlass.cute.nvgpu.tcgen05 as tcgen05
 import cutlass.utils.blackwell_helpers as sm100_utils_basic
 from cutlass import pipeline
 from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
+from cutlass.utils import ClcDynamicPersistentTileScheduler
 from cutlass.base_dsl.arch import Arch
 from cutlass.cutlass_dsl import BaseDSL
 
@@ -35,7 +35,9 @@ from .quack import copy_utils, layout_utils
 
 from .paged_kv import PagedKVManager
 from .cute_dsl_utils import assume_tensor_aligned
+from . import utils
 from . import pipeline as pipeline_custom
+import cutlass.pipeline as cutlass_pipeline
 from .mask import AttentionMask
 from .softmax import SoftmaxSm100, apply_score_mod_inner
 from .seqlen_info import SeqlenInfoQK
@@ -47,33 +49,88 @@ from .block_sparse_utils import (
     softmax_block_sparse_sm100,
     handle_block_sparse_empty_tile_correction_sm100,
 )
-from .pack_gqa import PackGQA
+from .pack_gqa import PackGQA, pack_gqa_layout
 from . import mma_sm100_desc as sm100_desc
 from . import blackwell_helpers as sm100_utils
+from .named_barrier import NamedBarrierFwdSm100
 from cutlass.cute import FastDivmodDivisor
 from .quack.cute_dsl_utils import ParamsBase
 from .tile_scheduler import (
+    SchedulerState,
+    SchedulingMode,
     TileSchedulerArguments,
+    TileSchedulerProtocol,
     SingleTileScheduler,
     StaticPersistentTileScheduler,
     SingleTileLPTScheduler,
     SingleTileVarlenScheduler,
+    DynamicPersistentVarlenScheduler,
 )
-
-
-class NamedBarrierFwd(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()
-#     WarpSchedulerWG1 = enum.auto()
-#     WarpSchedulerWG2 = enum.auto()
+from .fa_logging import fa_log, fa_printf
+from .utils import smid
+from .utils import AuxData
+
+# === TUNING KNOBS (agent-editable) ===
+# Keys: (use_2cta_instrs: bool, is_causal: bool, head_dim_padded: int, is_sm103: bool)
+# Values:
+#   ex2_emu_freq: int — how often to use emulated exp2 (0=all hardware exp2, higher=more emulation).
+#                        SM103 has fast native exp2, so set freq=0 there.
+#   ex2_emu_res: int — (hd256 only) number of fragment-pairs per freq period to emulate.
+#   ex2_emu_start_frg: int — fragment index to start emulation from
+#   num_regs_softmax: int — register count for softmax warps (multiple of 8)
+#   num_regs_correction: int — register count for correction warps (multiple of 8)
+#   num_regs_other is derived: 512 - num_regs_softmax * 2 - num_regs_correction
+#                  (hd256 exception: num_regs_other is fixed at 32, not derived)
+
+# Note [Low Precision Scaling]
+# P is in (0, 1] and is cast to the input dtype before P @ V, so scaling it by 2^max_offset
+# spends the dtype's unused upper code points on the probability tail. A positive
+# rescale_threshold lets the row max lag by that many log2 units, so P can reach
+# 2^(max_offset + rescale_threshold); above the dtype max the top probabilities saturate
+# while the FP32 denominator still counts them in full, shrinking the output (#2716).
+
+# log2 of the largest finite value representable in each supported input dtype.
+_LOG2_DTYPE_MAX = {
+    cutlass.Float8E4M3FN: math.log2(448.0),
+    cutlass.Float8E5M2: math.log2(57344.0),
+    cutlass.Float16: math.log2(65504.0),
+    cutlass.BFloat16: math.log2(3.3895313892515355e38),
+}
+
+_TUNING_CONFIG = {
+    (True, False, 128, False): {"ex2_emu_freq": 10, "ex2_emu_start_frg": 1, "num_regs_softmax": 176, "num_regs_correction": 88},
+    (False, True, 128, False): {"ex2_emu_freq": 16, "ex2_emu_start_frg": 1, "num_regs_softmax": 192, "num_regs_correction": 72},
+    (True, False, 192, False): {"ex2_emu_freq": 16, "ex2_emu_start_frg": 0, "num_regs_softmax": 184, "num_regs_correction": 80},
+    (False, True, 192, False): {"ex2_emu_freq": 32, "ex2_emu_start_frg": 1, "num_regs_softmax": 192, "num_regs_correction": 72},
+    (True, False, 128, True): {"ex2_emu_freq": 0, "ex2_emu_start_frg": 0, "num_regs_softmax": 176, "num_regs_correction": 80},
+    (False, True, 128, True): {"ex2_emu_freq": 0, "ex2_emu_start_frg": 0, "num_regs_softmax": 176, "num_regs_correction": 64},
+    (True, False, 192, True): {"ex2_emu_freq": 0, "ex2_emu_start_frg": 0, "num_regs_softmax": 176, "num_regs_correction": 64},
+    (False, True, 192, True): {"ex2_emu_freq": 0, "ex2_emu_start_frg": 0, "num_regs_softmax": 176, "num_regs_correction": 72},
+    (True, False, 256, False): {"ex2_emu_freq": 14, "ex2_emu_res": 6, "ex2_emu_start_frg": 0, "num_regs_softmax": 256, "num_regs_correction": 160},
+    (True, True, 256, False): {"ex2_emu_freq": 14, "ex2_emu_res": 6, "ex2_emu_start_frg": 0, "num_regs_softmax": 256, "num_regs_correction": 160},
+}
+_FP8_TUNING_CONFIG = {
+    (True, False, 128, False): {'ex2_emu_freq': 10, 'ex2_emu_start_frg': 1, 'num_regs_softmax': 160, 'num_regs_correction': 72},
+    # Causal hd128 FP8 previously inherited bf16's freq=16. FP8 fwd is MUFU/ex2-bound, so a more
+    # aggressive emulation freq=8 offloads more exp from MUFU: +3.4%(4k)..+5.5%(16k) on B200,
+    # MHA & GQA, accuracy-neutral (3-run validated, locked clocks). freq=8 would regress non-causal
+    # (0.94x), hence keyed on is_causal=True only.
+    (False, True, 128, False): {'ex2_emu_freq': 8, 'ex2_emu_start_frg': 1},
+}
+_FP8_SMALL_HDIM_REGS = {
+    False: {"num_regs_softmax": 168, "num_regs_correction": 96, "num_regs_other": 80},
+    True: {"num_regs_softmax": 152, "num_regs_correction": 96, "num_regs_other": 112},
+}
+# === END TUNING KNOBS ===
+
+
+class DescaleTensors(NamedTuple):
+    q_descale: Optional[cute.Tensor] = None
+    k_descale: Optional[cute.Tensor] = None
+    v_descale: Optional[cute.Tensor] = None
+
+    def __new_from_mlir_values__(self, values):
+        return DescaleTensors(*((*values, None, None, None)[:3]))
 
 
 class FlashAttentionForwardSm100:
@@ -88,17 +145,21 @@ class FlashAttentionForwardSm100:
         is_local: bool = False,
         is_split_kv: bool = False,
         pack_gqa: bool = False,
-        q_subtile_factor: int | None = None,
+        q_subtile_factor: int = 1,
+        kv_subtile_factor: int = 1,
         m_block_size: int = 128,
         n_block_size: int = 128,
         q_stage: cutlass.Constexpr[int] = 2,
-        is_persistent: bool = True,
+        is_static_persistent: bool = True,
         score_mod: cutlass.Constexpr | None = None,
         mask_mod: cutlass.Constexpr | None = None,
         has_aux_tensors: cutlass.Constexpr = False,
         paged_kv_non_tma: bool = False,
         is_varlen_q: bool = False,
         use_2cta_instrs: bool = False,
+        use_clc_scheduler: bool = False,
+        has_tile_count_semaphore: bool = False,
+        seqlen_k_per_split: Optional[int] = None,
     ):
         self.use_tma_KV = not paged_kv_non_tma
         # self.dtype = dtype
@@ -123,8 +184,13 @@ class FlashAttentionForwardSm100:
         self.split_P_arrive = int(self.split_P_arrive / 32) * 32  # multiple of 32
         assert self.split_P_arrive % 32 == 0
         assert self.split_P_arrive < self.n_block_size
+        assert seqlen_k_per_split is None or seqlen_k_per_split % n_block_size == 0
+        self.num_n_blocks_per_split = (
+            seqlen_k_per_split // n_block_size if seqlen_k_per_split is not None else None
+        )
         self.arch = BaseDSL._get_dsl().get_arch_enum()
-        assert self.arch >= Arch.sm_100 and self.arch <= Arch.sm_110f, "Only SM 10.x and 11.x are supported"
+        assert self.arch.is_family_of(Arch.sm_100f) or self.arch.is_family_of(Arch.sm_110f), \
+            "Only SM 10.x and 11.x are supported"
 
         self.cta_group_size = 2 if self.use_2cta_instrs else 1
         # cta_tiler M includes only 1 CTA, the scheduler will take into account the cluster shape
@@ -136,44 +202,101 @@ class FlashAttentionForwardSm100:
         self.qk_acc_dtype = Float32
         self.pv_acc_dtype = Float32
         self.cluster_shape_mn = (2, 1) if self.use_2cta_instrs else (1, 1)
-        self.is_persistent = is_persistent
+        self.is_static_persistent = is_static_persistent
         self.is_causal = is_causal
         self.is_local = is_local
         self.is_varlen_q = is_varlen_q
-        self.use_correction_warps_for_epi = is_varlen_q
         self.qhead_per_kvhead = qhead_per_kvhead
         self.is_split_kv = is_split_kv
         self.pack_gqa = pack_gqa
+        self.use_tma_O = (
+            not (self.pack_gqa and self.m_block_size % self.qhead_per_kvhead != 0)
+            and not (self.pack_gqa and self.is_split_kv)
+            and not is_varlen_q
+        )
+        self.use_correction_warps_for_epi = not self.use_tma_O
         self.q_subtile_factor = q_subtile_factor
-        if pack_gqa:
-            assert m_block_size % self.qhead_per_kvhead == 0, (
-                "For PackGQA, m_block_size must be divisible by qhead_per_kvhead"
-            )
+        self.kv_subtile_factor = kv_subtile_factor
         assert not (self.is_split_kv and self.head_dim_v_padded >= 192), (
             "SplitKV is not supported for hdim >= 192"
         )
         self.score_mod = score_mod
         self.mask_mod = mask_mod
-        self.vec_size: cutlass.Constexpr = getattr(
+        self.score_vec_size: cutlass.Constexpr = getattr(
             score_mod, "__vec_size__", 1 if cutlass.const_expr(has_aux_tensors) else 2
         )
+        self.mask_vec_size: cutlass.Constexpr = getattr(mask_mod, "__vec_size__", 1)
         # Does S1 need to wait for S0 to finish
         # self.s0_s1_barrier = self.head_dim_padded in [64, 96] and (not self.is_causal and not self.is_local)
-        is_sm103 = self.arch >= Arch.sm_103 and self.arch <= Arch.sm_103f
-        # self.enable_ex2_emu = self.head_dim_padded <= 128 and not is_sm103
-        self.enable_ex2_emu = (self.head_dim_padded <= 128 or (self.head_dim_padded == 192 and self.use_2cta_instrs and not self.is_causal and not self.is_local)) and not is_sm103
+        # NOTE: is_family_of also matches any future sm_10x with x > 3 — intentional.
+        # The flag gates ex2 emulation; sm_103 (B300) has fast hardware ex2 and later
+        # Blackwell variants are assumed to inherit this, so forward-inclusion is correct
+        # despite the literal `is_sm103` name.
+        is_sm103 = self.arch.is_family_of(Arch.sm_103f)
+        self.is_sm103 = is_sm103
+        # SM103 ld.red is profitable except for D32 when scores are unmodified.
+        self.use_ldred_rowmax = (
+            is_sm103
+            and self.score_mod is None
+            and self.mask_mod is None
+            and self.head_dim_padded != 32
+        )
+        # enable_ex2_emu is derived: True if tuning config has freq > 0, else fallback to default logic
+        _default_enable_ex2_emu = (self.head_dim_padded <= 128 or (self.head_dim_padded == 192 and self.use_2cta_instrs and not self.is_causal and not self.is_local)) and not is_sm103
+        self.enable_ex2_emu = _default_enable_ex2_emu
         self.s0_s1_barrier = False
         self.overlap_sO_sQ = (
             (self.head_dim_padded == 192 and self.head_dim_v_padded >= 64) or
             (self.head_dim_v_padded >= 128 and self.is_split_kv)
         )
-        if self.overlap_sO_sQ:
-            self.is_persistent = False
 
         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"
         )
 
+        # ClC does not compose with these other features, so disable even if requested
+        self.use_clc_scheduler = (
+            use_clc_scheduler
+            and self.use_tma_KV
+            and not (has_tile_count_semaphore and is_varlen_q)
+        )
+        self.dynamic_persistent = (
+            has_tile_count_semaphore and is_varlen_q
+        ) or self.use_clc_scheduler
+        self.is_persistent = self.dynamic_persistent or self.is_static_persistent
+        self.sched_stages = 1
+        if self.use_clc_scheduler:
+            assert self.cluster_shape_mn[1] == 1, f"CLC requires cluster N == 1: {self.cluster_shape_mn}"
+            assert self.cluster_shape_mn[0] in (1, 2), f"bad CLC cluster M: {self.cluster_shape_mn}"
+            assert self.cluster_shape_mn[0] == self.cta_group_size, (
+                f"CLC cluster M != cta_group_size: {self.cluster_shape_mn}, {self.cta_group_size}"
+            )
+
+        self.scheduling_mode = (
+            SchedulingMode.CLC if self.use_clc_scheduler
+            else SchedulingMode.DYNAMIC if self.dynamic_persistent
+            else SchedulingMode.STATIC
+        )
+
+        self.use_varlen_scheduler = False
+        if is_varlen_q:
+            if self.dynamic_persistent and not self.use_clc_scheduler:
+                self.use_varlen_scheduler = True
+                self.TileScheduler = DynamicPersistentVarlenScheduler
+            elif self.is_static_persistent and not self.use_clc_scheduler:
+                self.TileScheduler = StaticPersistentTileScheduler
+            else:
+                self.use_varlen_scheduler = True
+                self.TileScheduler = SingleTileVarlenScheduler
+        elif self.is_causal or self.is_local or self.use_clc_scheduler:
+            self.TileScheduler = SingleTileLPTScheduler
+        elif self.is_static_persistent:
+            self.TileScheduler = StaticPersistentTileScheduler
+        else:
+            self.TileScheduler = SingleTileScheduler
+
+        fa_log(1, f"TileScheduler={self.TileScheduler.__name__}, scheduling_mode={self.scheduling_mode.name}, USE_2CTA={self.use_2cta_instrs}")
+
         self.softmax0_warp_ids = (0, 1, 2, 3)
         self.softmax1_warp_ids = (4, 5, 6, 7)
         self.correction_warp_ids = (8, 9, 10, 11)
@@ -195,8 +318,10 @@ class FlashAttentionForwardSm100:
             )
         )
 
+        self.use_tma_Q = not (self.pack_gqa and self.m_block_size % self.qhead_per_kvhead != 0)
+
         if self.q_stage == 1:
-            if not self.use_tma_KV:
+            if not self.use_tma_KV or not self.use_tma_Q:
                 self.empty_warp_ids = self.empty_warp_ids + self.load_warp_ids
                 self.load_warp_ids = self.softmax1_warp_ids
             else:
@@ -209,8 +334,12 @@ class FlashAttentionForwardSm100:
         if self.use_correction_warps_for_epi:
             self.empty_warp_ids = self.empty_warp_ids + self.epilogue_warp_ids
             self.epilogue_warp_ids = self.correction_warp_ids
-        elif self.is_varlen_q: # fallback
-            self.epilogue_warp_ids = (13, 14)
+
+        if self.dynamic_persistent:
+            assert len(self.empty_warp_ids) > 0, (
+                "dynamic persistent scheduling requires an empty warp to serve as the scheduler warp"
+            )
+        self.scheduler_warp_id = self.empty_warp_ids[0] if self.dynamic_persistent else None
 
         self.tmem_s_offset = [0, self.n_block_size]  # e.g., 0, 128
         self.tmem_o_offset = [
@@ -227,31 +356,26 @@ class FlashAttentionForwardSm100:
         # vec buffer for row_max & row_sum
         self.tmem_vec_offset = self.tmem_s_offset
 
+        # Look up tuning config for register counts and ex2_emu params
+        _tune_key = (self.use_2cta_instrs, self.is_causal, self.head_dim_padded, self.is_sm103)
+        self._tune = _TUNING_CONFIG.get(_tune_key, {})
+        if "ex2_emu_freq" in self._tune:
+            self.enable_ex2_emu = self._tune["ex2_emu_freq"] > 0
         if self.head_dim_padded < 96:
             self.num_regs_softmax = 200 if not paged_kv_non_tma else 184
             self.num_regs_correction = 64
             self.num_regs_other = 48 if not paged_kv_non_tma else 80
         else:
-            # self.num_regs_softmax = 192 if self.is_causal or self.is_local else 184
-            if not self.enable_ex2_emu:
-                self.num_regs_softmax = 192 if not paged_kv_non_tma else 184
+            if not paged_kv_non_tma and "num_regs_softmax" in self._tune:
+                self.num_regs_softmax = self._tune["num_regs_softmax"]
+                self.num_regs_correction = self._tune["num_regs_correction"]
+            elif not paged_kv_non_tma:
+                self.num_regs_softmax = 192
+                self.num_regs_correction = 80
             else:
-                # self.num_regs_softmax = 200 if not paged_kv_non_tma else 184
-                self.num_regs_softmax = 192 if not paged_kv_non_tma else 184
-            # self.num_regs_softmax = 176
-            # self.num_regs_correction = 96
-            # self.num_regs_correction = 64 if self.is_causal or self.is_local else 80
-            if not self.enable_ex2_emu:
-                self.num_regs_correction = 80 if not paged_kv_non_tma else 64
-            else:
-                # self.num_regs_correction = 64
-                self.num_regs_correction = 80 if not paged_kv_non_tma else 64
-            # self.num_regs_other = 32
-            # self.num_regs_other = 64
-            # self.num_regs_other = 80
-            self.num_regs_other = 48 if not paged_kv_non_tma else 80
-            # self.num_regs_other = 96 if self.is_causal or self.is_local else 80
-            # self.num_regs_other = 64 if self.is_causal or self.is_local else 80
+                self.num_regs_softmax = 184
+                self.num_regs_correction = 64
+            self.num_regs_other = 512 - self.num_regs_softmax * 2 - self.num_regs_correction
 
         self.buffer_align_bytes = 1024
 
@@ -271,7 +395,10 @@ class FlashAttentionForwardSm100:
         smem_size_k_per_stage = self.n_block_size * self.head_dim_padded * self.k_dtype.width // 8
         smem_size_v_per_stage = self.n_block_size * self.head_dim_v_padded * self.v_dtype.width // 8
         smem_size_kv_per_stage = max(smem_size_k_per_stage, smem_size_v_per_stage) // self.cta_group_size
-        kv_stage = (224 * 1024 - smem_size_q_o) // smem_size_kv_per_stage
+        # Cap small head_dim from over-staging: the 224*1024 budget undercounts
+        # per-stage state, so at hd_padded=16 the unbounded formula picks 52 stages
+        # and overflows the 227 KB SMEM cap. No-op for hd_padded >= 32 (max 26).
+        kv_stage = min((224 * 1024 - smem_size_q_o) // smem_size_kv_per_stage, 32)
         if self.head_dim_padded == 192 and self.head_dim_v_padded == 128 and kv_stage == 2:
             # For hdim 192,128, we can fit 3 stages if we use uneven_kv_smem
              kv_stage = 3
@@ -289,7 +416,7 @@ class FlashAttentionForwardSm100:
             self.head_dim_padded == 192 and self.head_dim_v_padded == 128 and self.kv_stage == 3
         )
         self.uneven_kv_smem_offset = (
-            self.m_block_size * (self.head_dim_padded - self.head_dim_v_padded) // 2
+            self.n_block_size * (self.head_dim_padded - self.head_dim_v_padded) // 2
             if self.uneven_kv_smem
             else 0
         )
@@ -304,7 +431,6 @@ class FlashAttentionForwardSm100:
         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,
-        stream: cuda.CUstream,
         mCuSeqlensQ: Optional[cute.Tensor] = None,
         mCuSeqlensK: Optional[cute.Tensor] = None,
         mSeqUsedQ: Optional[cute.Tensor] = None,
@@ -313,8 +439,19 @@ class FlashAttentionForwardSm100:
         window_size_left: Int32 | int | None = None,
         window_size_right: Int32 | int | None = None,
         learnable_sink: Optional[cute.Tensor] = None,
+        descale_tensors: Optional[DescaleTensors] = None,
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
-        aux_tensors: Optional[list] = None,
+        aux_data: AuxData = AuxData(),
+        num_splits_dynamic_ptr: Optional[cute.Tensor] = None,
+        tile_count_semaphore: Optional[cute.Tensor] = None,
+        virtual_batch_idx_ptr: Optional[cute.Tensor] = None,
+        num_nheads_in_l2_ptr: Optional[cute.Tensor] = None,
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
+        mCuTotalSplitsMBlocks: Optional[cute.Tensor] = None,
+        mBlocksToBatchIdx: Optional[cute.Tensor] = None,
+        max_seqlen_q: Int32 | int | None = None,
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
+        stream: cuda.CUstream = None,
     ):
         """Execute the Fused Multi-Head Attention operation on the provided tensors.
 
@@ -366,23 +503,33 @@ class FlashAttentionForwardSm100:
             raise TypeError(f"Type mismatch: {self.q_dtype} != {self.k_dtype}")
         if const_expr(self.q_dtype != self.v_dtype):
             raise TypeError(f"Type mismatch: {self.q_dtype} != {self.v_dtype}")
+        if const_expr(self.q_dtype.width == 8):
+            paged_kv_non_tma = not self.use_tma_KV
+            if const_expr(self.head_dim_padded < 96):
+                fp8_regs = _FP8_SMALL_HDIM_REGS[paged_kv_non_tma]
+                self.num_regs_softmax = fp8_regs["num_regs_softmax"]
+                self.num_regs_correction = fp8_regs["num_regs_correction"]
+                self.num_regs_other = fp8_regs["num_regs_other"]
+            else:
+                fp8_tune = _FP8_TUNING_CONFIG.get(
+                    (self.use_2cta_instrs, self.is_causal, self.head_dim_padded, self.is_sm103), {}
+                )
+                if const_expr("ex2_emu_freq" in fp8_tune):
+                    self._tune = {**self._tune, **fp8_tune}
+                    self.enable_ex2_emu = self._tune["ex2_emu_freq"] > 0
+                if const_expr(not paged_kv_non_tma and "num_regs_softmax" in fp8_tune):
+                    self.num_regs_softmax = fp8_tune["num_regs_softmax"]
+                    self.num_regs_correction = fp8_tune["num_regs_correction"]
+                    self.num_regs_other = 512 - self.num_regs_softmax * 2 - self.num_regs_correction
         self._setup_attributes()
-        self.use_tma_O = self.arch >= Arch.sm_90 and mCuSeqlensQ is None and mSeqUsedQ is None
-        # This can be tuned
-        # This is currently very ad-hoc, we should tune it systematically
         self.ex2_emu_freq = 0
-        # self.ex2_emu_start_frg = 1 if self.is_causal else 0
-        self.ex2_emu_start_frg = 1
+        self.ex2_emu_start_frg = self._tune.get("ex2_emu_start_frg", 1)
         if const_expr(self.enable_ex2_emu):
-            self.ex2_emu_freq = 16
-            if const_expr(self.head_dim_padded == 128 and self.use_2cta_instrs):
-                self.ex2_emu_freq = 12
+            self.ex2_emu_freq = self._tune.get("ex2_emu_freq", 16)
             if const_expr(
                 self.pack_gqa and self.head_dim_padded > 64 and not self.is_causal and not self.is_local
             ):
-                self.ex2_emu_freq = 32 if mCuSeqlensQ is not None or mSeqUsedQ is not None else 10
-            if const_expr(self.head_dim_padded > 64 and self.is_causal):
-                self.ex2_emu_freq = 10
+                self.ex2_emu_freq = 32 if mCuSeqlensQ is not None or mSeqUsedQ is not None else self._tune.get("ex2_emu_freq", 10)
 
         cta_group = tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE
         q_major_mode = tcgen05.OperandMajorMode.K
@@ -462,50 +609,11 @@ class FlashAttentionForwardSm100:
             )
 
         if const_expr(self.pack_gqa):
-            shape_Q_packed = (
-                (self.qhead_per_kvhead, mQ.shape[0]),
-                mQ.shape[1],
-                mK.shape[2],
-                *mQ.shape[3:],
-            )
-            stride_Q_packed = (
-                (mQ.stride[2], mQ.stride[0]),
-                mQ.stride[1],
-                mQ.stride[2] * self.qhead_per_kvhead,
-                *mQ.stride[3:],
-            )
-            mQ = cute.make_tensor(
-                mQ.iterator, cute.make_layout(shape_Q_packed, stride=stride_Q_packed)
-            )
-            shape_O_packed = (
-                (self.qhead_per_kvhead, mO.shape[0]),
-                mO.shape[1],
-                mK.shape[2],
-                *mO.shape[3:],
-            )
-            stride_O_packed = (
-                (mO.stride[2], mO.stride[0]),
-                mO.stride[1],
-                mO.stride[2] * self.qhead_per_kvhead,
-                *mO.stride[3:],
-            )
-            mO = cute.make_tensor(
-                mO.iterator, cute.make_layout(shape_O_packed, stride=stride_O_packed)
-            )
+            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):
-                shape_LSE_packed = (
-                    (self.qhead_per_kvhead, mLSE.shape[0]),
-                    mK.shape[2],
-                    *mLSE.shape[2:],
-                )
-                stride_LSE_packed = (
-                    (mLSE.stride[1], mLSE.stride[0]),
-                    mLSE.stride[1] * self.qhead_per_kvhead,
-                    *mLSE.stride[2:],
-                )
-                mLSE = cute.make_tensor(
-                    mLSE.iterator, cute.make_layout(shape_LSE_packed, stride=stride_LSE_packed)
-                )
+                mLSE = pack_gqa_layout(mLSE, self.qhead_per_kvhead, nheads_kv, head_idx=1)
 
         self.tma_copy_bytes = {
             name: cute.size_in_bytes(mX.element_type, cute.select(layout, mode=[0, 1, 2]))
@@ -522,14 +630,24 @@ class FlashAttentionForwardSm100:
         tma_load_op = cpasync.CopyBulkTensorTileG2SOp(cta_group)
         tma_store_op = cpasync.CopyBulkTensorTileS2GOp()
 
-        tma_atom_Q, mQ = cute.nvgpu.make_tiled_tma_atom_A(
-            tma_load_op,
-            mQ,
-            cute.select(sQ_layout, mode=[0, 1, 2]),
-            self.mma_tiler_qk,
-            tiled_mma_qk,
-            cta_layout_vmnk.shape,
-        )
+        if const_expr(self.use_tma_Q):
+            tma_atom_Q, mQ = cute.nvgpu.make_tiled_tma_atom_A(
+                tma_load_op,
+                mQ,
+                cute.select(sQ_layout, mode=[0, 1, 2]),
+                self.mma_tiler_qk,
+                tiled_mma_qk,
+                cta_layout_vmnk.shape,
+            )
+            gmem_tiled_copy_Q = None
+        else:
+            tma_atom_Q = None
+            async_copy_elems = 128 // self.q_dtype.width
+            num_load_threads = cute.arch.WARP_SIZE * len(self.load_warp_ids)
+            threads_per_row = math.gcd(self.head_dim_padded // async_copy_elems, num_load_threads)
+            gmem_tiled_copy_Q = copy_utils.tiled_copy_2d(
+                self.q_dtype, threads_per_row, num_load_threads, async_copy_elems, is_async=True
+            )
 
         tma_atom_K = None
         tma_atom_V = None
@@ -578,19 +696,14 @@ class FlashAttentionForwardSm100:
             vO_layout = cute.make_layout((1, async_copy_elems))
             gmem_tiled_copy_O = cute.make_tiled_copy_tv(atom_universal_copy, tO_layout, vO_layout)
 
-        if const_expr(mCuSeqlensQ is not None or mSeqUsedQ is not None):
-            TileScheduler = SingleTileVarlenScheduler
+        if const_expr(max_seqlen_q is None):
+            eff_seqlen_q = cute.size(mQ.shape[0])
         else:
-            if const_expr(self.is_causal or self.is_local):
-                TileScheduler = SingleTileLPTScheduler
-            else:
-                TileScheduler = (
-                    SingleTileScheduler
-                    if const_expr(not self.is_persistent)
-                    else StaticPersistentTileScheduler
-                )
+            eff_seqlen_q = max_seqlen_q if const_expr(not self.pack_gqa) else max_seqlen_q * self.qhead_per_kvhead
+        TileScheduler = self.TileScheduler
+        _num_block_divisor = self.cta_tiler[0] * (self.cta_group_size if not self.is_persistent and self.cta_group_size > 1 else 1)
         tile_sched_args = TileSchedulerArguments(
-            cute.ceil_div(cute.size(mQ.shape[0]), self.cta_tiler[0]),
+            cute.ceil_div(eff_seqlen_q, _num_block_divisor),
             cute.size(mQ.shape[2]),
             cute.size(mQ.shape[3])
             if const_expr(mCuSeqlensQ is None)
@@ -613,8 +726,18 @@ class FlashAttentionForwardSm100:
             lpt=self.is_causal or self.is_local,
             is_split_kv=self.is_split_kv,
             cluster_shape_mn=self.cluster_shape_mn,
+            use_cluster_idx=not self.is_persistent and self.cta_group_size > 1,
+            num_splits_dynamic_ptr=num_splits_dynamic_ptr,
+            virtual_batch_idx_ptr=virtual_batch_idx_ptr,
+            num_nheads_in_l2_ptr=num_nheads_in_l2_ptr,
+            cu_total_m_blocks_ptr=mCuTotalMBlocks,
+            cu_total_splits_m_blocks_ptr=mCuTotalSplitsMBlocks,
+            blocks_to_batch_idx_ptr=mBlocksToBatchIdx,
+            tile_count_semaphore=tile_count_semaphore.iterator if tile_count_semaphore is not None else None,
+        )
+        tile_sched_params = TileScheduler.to_underlying_arguments(
+            tile_sched_args, scheduling_mode=self.scheduling_mode
         )
-        tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
         self.tile_scheduler_cls = TileScheduler
         grid_dim = TileScheduler.get_grid_shape(tile_sched_params)
 
@@ -623,6 +746,12 @@ class FlashAttentionForwardSm100:
             cute.cosize(sQ_layout) if const_expr(not self.overlap_sO_sQ) else
             cutlass.max(cute.cosize(sQ_layout), cute.cosize(sO_layout) * self.o_dtype.width // self.q_dtype.width)
         )
+        # K and V alias the same physical buffer and may have different extents.
+        sKV_size = cutlass.max(cute.cosize(sK_layout), cute.cosize(sV_layout))
+
+        sched_response_size = self.sched_stages * 4 if self.dynamic_persistent else 0
+        sched_mbar_size = self.sched_stages * 2 if self.dynamic_persistent else 0
+        load_epi_mbar_size = 2 if const_expr(self.overlap_sO_sQ) else 0
 
         @cute.struct
         class SharedStorage:
@@ -635,14 +764,23 @@ class FlashAttentionForwardSm100:
             mbar_softmax_stats: cute.struct.MemRange[Int64, self.q_stage * 2]
             # mbar_softmax_stats: cute.struct.MemRange[Int64, self.q_stage * 4 * 2]
             mbar_O_epi: cute.struct.MemRange[Int64, self.q_stage * 2]
+            mbar_load_epi: cute.struct.MemRange[Int64, load_epi_mbar_size]
             mbar_s0_s1_sequence: cute.struct.MemRange[Int64, 2 * 2]
             # Tmem dealloc cluster barrier
-            tmem_dealloc_mbar_ptr: Int64
+            tmem_dealloc_mbar: Int64
             # Tmem holding buffer
             tmem_holding_buf: Int32
             # Smem tensors
             # store row max and row sum
             sScale: cute.struct.MemRange[Float32, self.q_stage * self.m_block_size * 2]
+            # Scheduler buffers placed here to utilize padding before sO's 1024-byte
+            # alignment. This avoids adding bytes at the end when we're at the smem limit.
+            # PipelineClcFetchAsync / PipelineAsync both expect
+            # 2 * sched_stages mbarriers (full + empty). Response is 4 Int32 per stage
+            # (CLC HW response, or work_info written by dynamic persistent producer).
+            sched_mbar_ptr: cute.struct.MemRange[Int64, sched_mbar_size]
+            sched_response: cute.struct.MemRange[Int32, sched_response_size]
+            # Large TMA buffers with 1024-byte alignment
             sO: cute.struct.Align[
                 cute.struct.MemRange[self.o_dtype, sO_size], self.buffer_align_bytes
             ]
@@ -650,42 +788,16 @@ class FlashAttentionForwardSm100:
                 cute.struct.MemRange[self.q_dtype, sQ_size], self.buffer_align_bytes
             ]
             sK: cute.struct.Align[
-                # cute.cosize(sK_layout) is correct even in the case of self.uneven_kv_smem
-                cute.struct.MemRange[self.k_dtype, cute.cosize(sK_layout)],
+                cute.struct.MemRange[self.k_dtype, sKV_size],
                 self.buffer_align_bytes,
             ]
 
         self.shared_storage = SharedStorage
 
-        LOG2_E = math.log2(math.e)
-        if const_expr(self.score_mod is None):
-            softmax_scale_log2 = softmax_scale * LOG2_E
-            softmax_scale = None
-        else:
-            # NB: If a users passes in a score mod, we want to apply the score-mod in the sm_scaled qk
-            # But in the original base 10. We hijack softmax_scale_log2 to just be the change of base
-            # and correctly apply the softmax_scale prior to score_mod in the softmax step
-            softmax_scale_log2 = LOG2_E
-            softmax_scale = softmax_scale
-
-        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)
-
-        fastdiv_mods = None
-        if cutlass.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
-            )
-            seqlen_k = (
-                cute.size(mK.shape[0])
-                if const_expr(mPageTable is None)
-                else mK.shape[0] * mPageTable.shape[1]
-            )
-            seqlen_q_divmod = FastDivmodDivisor(seqlen_q)
-            seqlen_k_divmod = FastDivmodDivisor(seqlen_k)
-            fastdiv_mods = (seqlen_q_divmod, seqlen_k_divmod)
+        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)
 
         head_divmod = None
         if cutlass.const_expr(self.pack_gqa):
@@ -694,6 +806,10 @@ class FlashAttentionForwardSm100:
         self.use_block_sparsity = cutlass.const_expr(blocksparse_tensors is not None)
         if cutlass.const_expr(self.use_block_sparsity and mPageTable is not None):
             raise NotImplementedError("Block sparsity + paged KV not supported on SM100")
+        if cutlass.const_expr(self.use_block_sparsity and self.is_varlen_q):
+            assert const_expr(blocksparse_tensors.cu_total_m_blocks is not None), (
+                "blocksparse_tensors.cu_total_m_blocks must be provided for varlen blocksparsity"
+            )
 
         # Launch the kernel synchronously
         self.kernel(
@@ -716,18 +832,24 @@ class FlashAttentionForwardSm100:
             window_size_left,
             window_size_right,
             learnable_sink,
+            descale_tensors,
             blocksparse_tensors,
             sQ_layout,
             sK_layout,
             tP_layout,
             sV_layout,
             sO_layout,
+            gmem_tiled_copy_Q,
             gmem_tiled_copy_O,
             tiled_mma_qk,
             tiled_mma_pv,
             tile_sched_params,
             num_splits,
-            aux_tensors,
+            num_splits_dynamic_ptr,
+            tile_count_semaphore,
+            virtual_batch_idx_ptr,
+            num_nheads_in_l2_ptr,
+            aux_data,
             fastdiv_mods,
             head_divmod,
         ).launch(
@@ -738,6 +860,18 @@ class FlashAttentionForwardSm100:
             min_blocks_per_mp=1,
         )
 
+    def _generate_attention_mask_cls(self, window_size_left, window_size_right):
+        return partial(
+            AttentionMask,
+            self.m_block_size,
+            self.n_block_size,
+            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
+            ),
+        )
+
     #  GPU device kernel
     @cute.kernel
     def kernel(
@@ -752,7 +886,7 @@ class FlashAttentionForwardSm100:
         mSeqUsedQ: Optional[cute.Tensor],
         mSeqUsedK: Optional[cute.Tensor],
         mPageTable: Optional[cute.Tensor],
-        tma_atom_Q: cute.CopyAtom,
+        tma_atom_Q: Optional[cute.CopyAtom],
         tma_atom_K: Optional[cute.CopyAtom],
         tma_atom_V: Optional[cute.CopyAtom],
         tma_atom_O: Optional[cute.CopyAtom],
@@ -761,18 +895,24 @@ class FlashAttentionForwardSm100:
         window_size_left: Optional[Int32],
         window_size_right: Optional[Int32],
         learnable_sink: Optional[cute.Tensor],
+        descale_tensors: Optional[DescaleTensors],
         blocksparse_tensors: Optional[BlockSparseTensors],
         sQ_layout: cute.ComposedLayout,
         sK_layout: cute.ComposedLayout,
         tP_layout: cute.ComposedLayout,
         sV_layout: cute.ComposedLayout,
         sO_layout: cute.ComposedLayout,
+        gmem_tiled_copy_Q: Optional[cute.TiledCopy],
         gmem_tiled_copy_O: Optional[cute.TiledCopy],
         tiled_mma_qk: cute.TiledMma,
         tiled_mma_pv: cute.TiledMma,
         tile_sched_params: ParamsBase,
         num_splits: Int32,
-        aux_tensors: Optional[list] = None,
+        num_splits_dynamic_ptr: Optional[cute.Tensor] = None,
+        tile_count_semaphore: Optional[cute.Tensor] = None,
+        virtual_batch_idx_ptr: Optional[cute.Tensor] = None,
+        num_nheads_in_l2_ptr: Optional[cute.Tensor] = None,
+        aux_data: AuxData = AuxData(),
         fastdiv_mods=(None, None),
         head_divmod=None,
     ):
@@ -814,7 +954,7 @@ class FlashAttentionForwardSm100:
         storage = smem.allocate(self.shared_storage)
 
         tmem_alloc_barrier = pipeline.NamedBarrier(
-            barrier_id=int(NamedBarrierFwd.TmemPtr),
+            barrier_id=int(NamedBarrierFwdSm100.TmemPtr),
             num_threads=cute.arch.WARP_SIZE * len(
                 (self.mma_warp_id,
                  *self.softmax0_warp_ids,
@@ -824,17 +964,18 @@ class FlashAttentionForwardSm100:
         )
         # Tensor memory dealloc barrier init
         tmem = cutlass.utils.TmemAllocator(
-            storage.tmem_holding_buf,
+            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,
+            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
         )
 
         ThreadCooperativeGroup = partial(pipeline.CooperativeGroup, pipeline.Agent.Thread)
         mma_warp = ThreadCooperativeGroup(len([self.mma_warp_id]))
-        load_warps = ThreadCooperativeGroup(len(self.load_warp_ids))
         tma_warp = ThreadCooperativeGroup(1)
+        load_warps = ThreadCooperativeGroup(len(self.load_warp_ids))
+        load_threads = ThreadCooperativeGroup(len(self.load_warp_ids) * cute.arch.WARP_SIZE)
         softmax_warps = ThreadCooperativeGroup(len(self.softmax0_warp_ids))
         softmax_threads = ThreadCooperativeGroup(cute.arch.WARP_SIZE * len(self.softmax0_warp_ids))
         # softmax_threads = ThreadCooperativeGroup(cute.arch.WARP_SIZE)
@@ -845,6 +986,7 @@ class FlashAttentionForwardSm100:
         softmax_correction_threads = ThreadCooperativeGroup(
             cute.arch.WARP_SIZE * len(self.softmax0_warp_ids + self.correction_warp_ids)
         )
+        epilogue_warps = ThreadCooperativeGroup(len(self.epilogue_warp_ids))
         epilogue_threads = ThreadCooperativeGroup(cute.arch.WARP_SIZE * len(self.epilogue_warp_ids))
         # For UMMA-bridging pipelines: the non-MMA side spans both CTAs in the cluster,
         # so the thread count must include warps from both CTAs.
@@ -857,15 +999,25 @@ class FlashAttentionForwardSm100:
         softmax_correction_threads_cluster = ThreadCooperativeGroup(
             cute.arch.WARP_SIZE * len(self.softmax0_warp_ids + self.correction_warp_ids) * self.cta_group_size
         )
-        pipeline_q = pipeline_custom.PipelineTmaUmma.create(
-            barrier_storage=storage.mbar_load_Q.data_ptr(),
-            num_stages=self.q_stage,
-            producer_group=tma_warp,
-            consumer_group=mma_warp,
-            tx_count=self.tma_copy_bytes["Q"],
-            cta_layout_vmnk=cta_layout_vmnk,
-            defer_sync=True,
-        )
+        if const_expr(self.use_tma_Q):
+            pipeline_q = pipeline_custom.PipelineTmaUmma.create(
+                barrier_storage=storage.mbar_load_Q.data_ptr(),
+                num_stages=self.q_stage,
+                producer_group=tma_warp,
+                consumer_group=mma_warp,
+                tx_count=self.tma_copy_bytes["Q"],
+                cta_layout_vmnk=cta_layout_vmnk,
+                defer_sync=True,
+            )
+        else:
+            pipeline_q = pipeline_custom.PipelineAsyncUmma.create(
+                barrier_storage=storage.mbar_load_Q.data_ptr(),
+                num_stages=self.q_stage,
+                producer_group=load_threads,
+                consumer_group=mma_warp,
+                cta_layout_vmnk=cta_layout_vmnk,
+                defer_sync=True,
+            )
         if const_expr(self.use_tma_KV):
             pipeline_kv = pipeline_custom.PipelineTmaUmma.create(
                 barrier_storage=storage.mbar_load_KV.data_ptr(),
@@ -877,13 +1029,10 @@ class FlashAttentionForwardSm100:
                 defer_sync=True,
             )
         else:
-            cpasync_producer_group = pipeline.CooperativeGroup(
-                pipeline.Agent.Thread, len(self.load_warp_ids) * cute.arch.WARP_SIZE
-            )
             pipeline_kv = pipeline.PipelineAsyncUmma.create(
                 barrier_storage=storage.mbar_load_KV.data_ptr(),
                 num_stages=self.kv_stage,
-                producer_group=cpasync_producer_group,
+                producer_group=load_threads,
                 consumer_group=mma_warp,
                 cta_layout_vmnk=cta_layout_vmnk,
                 defer_sync=True,
@@ -938,7 +1087,7 @@ class FlashAttentionForwardSm100:
         )
         # Should put the NamedBarrier inside the pipeline class so we'll just have pipeline_sm_stats
         sm_stats_barrier = pipeline_custom.NamedBarrier(
-            barrier_id=int(NamedBarrierFwd.SoftmaxStatsW0), num_threads=cute.arch.WARP_SIZE * 2
+            barrier_id=int(NamedBarrierFwdSm100.SoftmaxStatsW0), num_threads=cute.arch.WARP_SIZE * 2
         )
         pipeline_o_epi = None
         if const_expr(not self.use_correction_warps_for_epi):
@@ -950,6 +1099,25 @@ class FlashAttentionForwardSm100:
                 defer_sync=True,
             )
 
+        pipeline_load_epi = None
+        if const_expr(self.overlap_sO_sQ and self.is_persistent):
+            # when overlapping sO and sQ with a persistent kernel, we need this
+            # additional pipeline to ensure sO from the previous work tile is
+            # free for use by sQ in the current one.
+            epi_warps_for_release = (
+                ThreadCooperativeGroup(len(self.correction_warp_ids))
+                if self.use_correction_warps_for_epi
+                else epilogue_warps
+            )
+            pipeline_load_epi = pipeline_custom.PipelineAsync.create(
+                barrier_storage=storage.mbar_load_epi.data_ptr(),
+                num_stages=1,
+                producer_group=epi_warps_for_release,
+                consumer_group=load_warps,
+                defer_sync=True,
+            )
+
+
         # Cluster arrive after barrier init
         pipeline_init_arrive(cluster_shape_mn=cta_layout_vmnk, is_relaxed=True)
 
@@ -999,6 +1167,9 @@ class FlashAttentionForwardSm100:
             window_size_left,
             window_size_right,
             qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
+            num_splits=num_splits,
+            pack_split_idx=num_splits_dynamic_ptr is not None,
+            num_n_blocks_per_split=self.num_n_blocks_per_split,
         )
         SeqlenInfoCls = partial(
             SeqlenInfoQK.create,
@@ -1010,26 +1181,101 @@ class FlashAttentionForwardSm100:
             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
+            ),
         )
-        AttentionMaskCls = partial(
-            AttentionMask,
-            self.m_block_size,
-            self.n_block_size,
-            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,
+        AttentionMaskCls = self._generate_attention_mask_cls(
+            window_size_left, window_size_right
         )
-        TileSchedulerCls = partial(self.tile_scheduler_cls.create, tile_sched_params)
-
         # Cluster wait before tensor memory alloc
         pipeline_init_wait(cluster_shape_mn=cta_layout_vmnk)
 
+        sched_ctx = None
+        if const_expr(self.use_clc_scheduler or self.dynamic_persistent):
+            sched_response_ptr = storage.sched_response.data_ptr()
+            sched_mbar_ptr = storage.sched_mbar_ptr.data_ptr()
+            sched_producer_group = cutlass_pipeline.CooperativeGroup(
+                cutlass_pipeline.Agent.Thread
+            )
+            num_sched_consumer_warps_per_cta = self.threads_per_cta // cute.arch.WARP_SIZE
+            # NB on CTA0 warp15 == scheduler on CTA1 == empty but still both consume
+            num_sched_consumer_warps = num_sched_consumer_warps_per_cta * self.cta_group_size
+            sched_consumer_group = cutlass_pipeline.CooperativeGroup(
+                cutlass_pipeline.Agent.Thread,
+                cute.arch.WARP_SIZE * num_sched_consumer_warps,
+            )
+            if const_expr(self.use_clc_scheduler):
+                _block_idx = cute.arch.block_idx()
+                sched_ctx = SchedulerState.create_clc(
+                    hw_scheduler=ClcDynamicPersistentTileScheduler.create(
+                        self.tile_scheduler_cls.clc_problem_shape(tile_sched_params),
+                        _block_idx,
+                        cute.arch.grid_dim(),
+                        sched_response_ptr,
+                    ),
+                    pipeline=cutlass_pipeline.PipelineClcFetchAsync.create(
+                        barrier_storage=sched_mbar_ptr,
+                        num_stages=self.sched_stages,
+                        producer_group=sched_producer_group,
+                        consumer_group=sched_consumer_group,
+                        tx_count=16,
+                        cta_layout_vmnk=cta_layout_vmnk,
+                    ),
+                    consumer_state=cutlass_pipeline.make_pipeline_state(
+                        cutlass_pipeline.PipelineUserType.Consumer, self.sched_stages
+                    ),
+                    producer_state=cutlass_pipeline.make_pipeline_state(
+                        cutlass_pipeline.PipelineUserType.Producer, self.sched_stages
+                    ),
+                )
+            else:
+                assert tile_count_semaphore is not None
+                sched_ctx = SchedulerState.create_dynamic_persistent(
+                    work_info=storage.sched_response.get_tensor((4,)),
+                    pipeline=cutlass_pipeline.PipelineAsync.create(
+                        barrier_storage=sched_mbar_ptr,
+                        num_stages=self.sched_stages,
+                        producer_group=sched_producer_group,
+                        consumer_group=sched_consumer_group,
+                    ),
+                    consumer_state=cutlass_pipeline.make_pipeline_state(
+                        cutlass_pipeline.PipelineUserType.Consumer, self.sched_stages
+                    ),
+                    producer_state=cutlass_pipeline.make_pipeline_state(
+                        cutlass_pipeline.PipelineUserType.Producer, self.sched_stages
+                    ),
+                )
+        if const_expr(self.use_clc_scheduler or self.dynamic_persistent):
+            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)}"
+
         # ///////////////////////////////////////////////////////////////////////////////
-        #  EMPTY
+        #  EMPTY / SCHEDULER WARP
         # ///////////////////////////////////////////////////////////////////////////////
-        for i in cutlass.range_constexpr(len(self.empty_warp_ids)):
-            if warp_idx == self.empty_warp_ids[i]:
+        if const_expr(self.dynamic_persistent):
+            if warp_idx == self.scheduler_warp_id:
                 cute.arch.setmaxregister_decrease(self.num_regs_other)
+                # CLC: only leader CTA produces.
+                if is_leader_cta:
+                    self.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.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)
 
         # ///////////////////////////////////////////////////////////////////////////////
         #  LOAD
@@ -1049,13 +1295,15 @@ class FlashAttentionForwardSm100:
                 tma_atom_Q,
                 tma_atom_K,
                 tma_atom_V,
+                gmem_tiled_copy_Q,
                 pipeline_q,
                 pipeline_kv,
+                pipeline_load_epi,
                 block_info,
                 num_splits,
                 SeqlenInfoCls,
-                TileSchedulerCls,
                 blocksparse_tensors,
+                tile_scheduler=tile_scheduler,
             )
 
         # ///////////////////////////////////////////////////////////////////////////////
@@ -1085,8 +1333,8 @@ class FlashAttentionForwardSm100:
                 block_info,
                 num_splits,
                 SeqlenInfoCls,
-                TileSchedulerCls,
                 blocksparse_tensors,
+                tile_scheduler=tile_scheduler,
             )
             # Dealloc the tensor memory buffer
             tmem.relinquish_alloc_permit()
@@ -1105,11 +1353,13 @@ class FlashAttentionForwardSm100:
                     gmem_tiled_copy_O,
                     tma_atom_O,
                     pipeline_o_epi,
+                    pipeline_load_epi,
                     block_info,
                     num_splits,
                     SeqlenInfoCls,
-                    TileSchedulerCls,
                     mma_tile_coord_v,
+                    blocksparse_tensors=blocksparse_tensors,
+                    tile_scheduler=tile_scheduler,
                 )
 
         # ///////////////////////////////////////////////////////////////////////////////
@@ -1128,6 +1378,7 @@ class FlashAttentionForwardSm100:
                 self.softmax_loop,
                 softmax_scale_log2=softmax_scale_log2,
                 softmax_scale=softmax_scale,
+                descale_tensors=descale_tensors,
                 thr_mma_qk=thr_mma_qk,
                 sScale=sScale,
                 mLSE=mLSE,
@@ -1141,11 +1392,11 @@ class FlashAttentionForwardSm100:
                 num_splits=num_splits,
                 SeqlenInfoCls=SeqlenInfoCls,
                 AttentionMaskCls=AttentionMaskCls,
-                TileSchedulerCls=TileSchedulerCls,
-                aux_tensors=aux_tensors,
+                aux_data=aux_data,
                 fastdiv_mods=fastdiv_mods,
                 head_divmod=head_divmod,
                 blocksparse_tensors=blocksparse_tensors,
+                tile_scheduler=tile_scheduler,
             )
 
             if const_expr(not self.s0_s1_barrier):
@@ -1182,15 +1433,17 @@ class FlashAttentionForwardSm100:
                 pipeline_sm_stats,
                 sm_stats_barrier,
                 pipeline_o_epi,
+                pipeline_load_epi,
                 learnable_sink,
+                descale_tensors,
                 gmem_tiled_copy_O,
                 tma_atom_O,
                 softmax_scale_log2,
                 block_info,
                 num_splits,
                 SeqlenInfoCls,
-                TileSchedulerCls,
                 blocksparse_tensors,
+                tile_scheduler=tile_scheduler,
             )
             tmem_alloc_barrier.arrive()
 
@@ -1199,8 +1452,8 @@ class FlashAttentionForwardSm100:
     @cute.jit
     def load(
         self,
-        thr_mma_qk: cute.core.ThrMma,
-        thr_mma_pv: cute.core.ThrMma,
+        thr_mma_qk: cute.ThrMma,
+        thr_mma_pv: cute.ThrMma,
         mQ: cute.Tensor,
         mK: cute.Tensor,
         mV: cute.Tensor,
@@ -1208,35 +1461,42 @@ class FlashAttentionForwardSm100:
         sK: cute.Tensor,
         sV: cute.Tensor,
         mPageTable: Optional[cute.Tensor],
-        tma_atom_Q: cute.CopyAtom,
+        tma_atom_Q: Optional[cute.CopyAtom],
         tma_atom_K: Optional[cute.CopyAtom],
         tma_atom_V: Optional[cute.CopyAtom],
+        gmem_tiled_copy_Q: Optional[cute.TiledCopy],
         pipeline_q: pipeline.PipelineAsync,
         pipeline_kv: pipeline.PipelineAsync,
+        pipeline_load_epi: Optional[pipeline.PipelineAsync],
         block_info: BlockInfo,
         num_splits: Int32,
         SeqlenInfoCls: Callable,
-        TileSchedulerCls: Callable,
         blocksparse_tensors: Optional[BlockSparseTensors],
+        tile_scheduler: TileSchedulerProtocol,
     ):
         num_load_threads = len(self.load_warp_ids) * cute.arch.WARP_SIZE
         tidx = cute.arch.thread_idx()[0] % num_load_threads
         warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+        load_epi_consumer_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, 1
+        )
+        issue_kv_for_this_warp = (
+            const_expr(not self.use_tma_KV or len(self.load_warp_ids) == 1) or
+            warp_idx == self.load_warp_ids[0]
+        )
+        issue_q_for_this_warp = (
+            const_expr(not self.use_tma_Q or len(self.load_warp_ids) == 1) or
+            warp_idx == self.load_warp_ids[0]
+        )
         q_producer_phase = Int32(1)
         kv_producer_state = pipeline.make_pipeline_state(
             pipeline.PipelineUserType.Producer, self.kv_stage
         )
-        tile_scheduler = TileSchedulerCls()
         work_tile = tile_scheduler.initial_work_tile_info()
         while work_tile.is_valid_tile:
             m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx
             seqlen = SeqlenInfoCls(batch_idx)
             mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx]
-            tiler_gQ = ((self.mma_tiler_qk[0] * self.q_stage), self.head_dim_padded)
-            gQ = cute.local_tile(mQ_cur, tiler_gQ, (m_block, 0))  # (128 * 2, 128)
-            gQ = layout_utils.select(
-                cute.flat_divide(gQ, (self.mma_tiler_qk[0],)), mode=[0, 2, 1]
-            )  # (128, 128, 2)
 
             head_idx_kv = (
                 head_idx // self.qhead_per_kvhead if const_expr(not self.pack_gqa) else head_idx
@@ -1258,12 +1518,32 @@ class FlashAttentionForwardSm100:
                 gV = cute.local_tile(
                     mV_cur, cute.select(self.mma_tiler_pv, mode=[1, 2]), (0, None, None)
                 )
-            tSgQ = thr_mma_qk.partition_A(gQ)
             tSgK = thr_mma_qk.partition_B(gK)
             tOgV = thr_mma_pv.partition_B(gV)
-            load_Q_fn, _, _ = copy_utils.tma_get_copy_fn(
-                tma_atom_Q, 0, cute.make_layout(1), tSgQ, sQ
-            )
+            if const_expr(self.use_tma_Q):
+                tiler_gQ = ((self.mma_tiler_qk[0] * self.q_stage), self.head_dim_padded)
+                gQ = cute.local_tile(mQ_cur, tiler_gQ, (m_block, 0))  # (128 * 2, 128)
+                gQ = layout_utils.select(
+                    cute.flat_divide(gQ, (self.mma_tiler_qk[0],)), mode=[0, 2, 1]
+                )  # (128, 128, 2)
+                tSgQ = thr_mma_qk.partition_A(gQ)
+                load_Q_fn, _, _ = copy_utils.tma_get_copy_fn(
+                    tma_atom_Q, 0, cute.make_layout(1), tSgQ, sQ
+                )
+                load_Q = partial(self.load_Q, load_Q_fn, pipeline_q=pipeline_q, phase=q_producer_phase)
+            else:
+                assert gmem_tiled_copy_Q is not None
+                load_Q = partial(
+                    self.load_Q_non_tma,
+                    mQ_cur,
+                    sQ,
+                    gmem_tiled_copy_Q,
+                    pipeline_q,
+                    tidx,
+                    seqlen.seqlen_q,
+                    m_block,
+                    phase=q_producer_phase,
+                )
 
             if const_expr(self.use_tma_KV):
                 tKsK, tKgK = cpasync.tma_partition(
@@ -1302,7 +1582,6 @@ class FlashAttentionForwardSm100:
                 tKsK, tKgK = None, None
                 tVsV, tVgV = None, None
 
-            load_Q = partial(self.load_Q, load_Q_fn, pipeline_q=pipeline_q, phase=q_producer_phase)
             load_K = partial(
                 self.load_KV,
                 tma_atom_K,
@@ -1326,9 +1605,14 @@ class FlashAttentionForwardSm100:
 
             if const_expr(not self.use_block_sparsity):
                 n_block_min, n_block_max = block_info.get_n_block_min_max(
-                    seqlen, m_block, split_idx, num_splits
+                    seqlen,
+                    m_block,
+                    split_idx=split_idx,
+                    num_splits=num_splits,
                 )
-                if const_expr(not self.is_split_kv) or n_block_min < n_block_max:
+                if const_expr(self.is_split_kv and block_info.pack_split_idx):
+                    split_idx = split_idx & 0xFFFF
+                if self.process_work_tile(seqlen, n_block_min, n_block_max):
                     n_block_first = n_block_max - 1 if n_block_max > 0 else 0
                     page_idx = (
                         mPageTable[batch_idx, n_block_first]
@@ -1337,24 +1621,19 @@ class FlashAttentionForwardSm100:
                     )
                     if const_expr(not self.use_tma_KV):
                         paged_kv_manager.load_page_table(n_block_first)
-                    load_K(block=n_block_max - 1, producer_state=kv_producer_state, page_idx=page_idx)  # K0
+                    if issue_kv_for_this_warp:
+                        load_K(block=n_block_max - 1, producer_state=kv_producer_state, page_idx=page_idx)  # K0
                     # load_K(block=n_block_max - 1, producer_state=kv_producer_state, page_idx=page_idx, extra_tx_count=self.tma_copy_bytes["Q"])  # K0
-                    if const_expr(len(self.load_warp_ids) == 1) or warp_idx == self.load_warp_ids[0]:
-                        # load_Q(block=0, stage=0)  # Q0
-                        pipeline_q.producer_acquire_w_index_phase(0, q_producer_phase)
-                        # pipeline_q.sync_object_empty.wait(0, q_producer_phase)
-                        tma_bar_ptr = pipeline_q.sync_object_full.get_barrier(0)
-                        # tma_bar_ptr = pipeline_kv.producer_get_barrier(kv_producer_state)
-                        load_Q_fn(src_idx=0, dst_idx=0, tma_bar_ptr=tma_bar_ptr)
-                    kv_producer_state.advance()
-                    if const_expr(self.q_stage == 2) and (const_expr(len(self.load_warp_ids) == 1) or warp_idx == self.load_warp_ids[0]):
-                        # load_Q(block=1, stage=1)  # Q1
-                        pipeline_q.producer_acquire_w_index_phase(1, q_producer_phase)
-                        tma_bar_ptr = pipeline_q.sync_object_full.get_barrier(1)
-                        load_Q_fn(src_idx=1, dst_idx=1, tma_bar_ptr=tma_bar_ptr)
+                    if issue_q_for_this_warp:
+                        load_Q(block=0, stage=0)
+                    if issue_kv_for_this_warp:
+                        kv_producer_state.advance()
+                    if const_expr(self.q_stage == 2) and issue_q_for_this_warp:
+                        load_Q(block=1, stage=1)
                     q_producer_phase ^= 1
-                    load_V(block=n_block_max - 1, producer_state=kv_producer_state, page_idx=page_idx)  # V0
-                    kv_producer_state.advance()
+                    if issue_kv_for_this_warp:
+                        load_V(block=n_block_max - 1, producer_state=kv_producer_state, page_idx=page_idx)  # V0
+                        kv_producer_state.advance()
                     for i in cutlass.range(n_block_max - 1 - n_block_min, unroll=1):
                         n_block = n_block_max - 2 - i
                         page_idx = (
@@ -1365,17 +1644,27 @@ class FlashAttentionForwardSm100:
                         if const_expr(not self.use_tma_KV):
                             paged_kv_manager.load_page_table(n_block)
                     # if cute.arch.thread_idx()[0] % 32 == 0: cute.printf("n_block = {}, page_idx = {}", n_block, page_idx)
-                        load_K(block=n_block, producer_state=kv_producer_state, page_idx=page_idx)  # Ki
-                        kv_producer_state.advance()
-                        load_V(block=n_block, producer_state=kv_producer_state, page_idx=page_idx)  # Vi
-                        kv_producer_state.advance()
+                        if issue_kv_for_this_warp:
+                            load_K(block=n_block, producer_state=kv_producer_state, page_idx=page_idx)  # Ki
+                            kv_producer_state.advance()
+                            load_V(block=n_block, producer_state=kv_producer_state, page_idx=page_idx)  # Vi
+                            kv_producer_state.advance()
 
             else:
+                # Match the dense path (get_n_block_min_max): the scheduler packs the
+                # per-batch dynamic num_splits into the top 16 bits of split_idx.
+                num_splits_dyn = num_splits
+                if const_expr(self.is_split_kv and block_info.pack_split_idx):
+                    num_splits_dyn = split_idx >> 16
+                    split_idx = split_idx & 0xFFFF
                 kv_producer_state, q_producer_phase = produce_block_sparse_loads_sm100(
                     blocksparse_tensors,
                     batch_idx,
                     head_idx,
                     m_block,
+                    seqlen,
+                    split_idx,
+                    num_splits_dyn,
                     kv_producer_state,
                     load_Q,
                     load_K,
@@ -1384,24 +1673,29 @@ class FlashAttentionForwardSm100:
                     self.q_stage,
                     q_producer_phase,
                     self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
-                    self.q_subtile_factor if self.q_subtile_factor is not None else 1,
+                    self.q_subtile_factor,
+                    self.kv_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
 
-        pipeline_kv.producer_tail(kv_producer_state)
-        # This is equivalent to pipeline_q.producer_tail
-        if const_expr(len(self.load_warp_ids) == 1) or warp_idx == self.load_warp_ids[0]:
+            work_tile = tile_scheduler.advance_to_next_work()
+            if const_expr(pipeline_load_epi is not None):
+                pipeline_load_epi.consumer_wait(load_epi_consumer_state)
+                with cute.arch.elect_one():
+                    pipeline_load_epi.consumer_release(load_epi_consumer_state)
+                load_epi_consumer_state.advance()
+
+        if issue_kv_for_this_warp:
+            pipeline_kv.producer_tail(kv_producer_state)
+        # This is equivalent to pipeline_q.producer_tail for the TMA-Q producer warp.
+        if issue_q_for_this_warp:
             pipeline_q.producer_acquire_w_index_phase(self.q_stage - 1, q_producer_phase)
 
     @cute.jit
     def mma(
         self,
-        tiled_mma_qk: cute.core.ThrMma,
-        tiled_mma_pv: cute.core.ThrMma,
+        tiled_mma_qk: cute.ThrMma,
+        tiled_mma_pv: cute.ThrMma,
         sQ: cute.Tensor,
         sK: cute.Tensor,
         sV: cute.Tensor,
@@ -1417,8 +1711,8 @@ class FlashAttentionForwardSm100:
         block_info: BlockInfo,
         num_splits: Int32,
         SeqlenInfoCls: Callable,
-        TileSchedulerCls: Callable,
         blocksparse_tensors: Optional[BlockSparseTensors],
+        tile_scheduler=None,
     ):
         tSrQ = tiled_mma_qk.make_fragment_A(sQ)
         tSrK = tiled_mma_qk.make_fragment_B(sK)
@@ -1430,6 +1724,7 @@ class FlashAttentionForwardSm100:
 
         qk_mma_op, pv_mma_op = tiled_mma_qk.op, tiled_mma_pv.op
         qk_mma_idesc, pv_mma_idesc = sm100_desc.mma_op_to_idesc(qk_mma_op), sm100_desc.mma_op_to_idesc(pv_mma_op)
+        qk_mma_kind = sm100_utils._tcgen05_mma_kind(qk_mma_op)
         q_smem_base = sm100_desc.smem_desc_base_from_tensor(sQ, sm100_desc.Major.K)
         k_smem_base = sm100_desc.smem_desc_base_from_tensor(sK, sm100_desc.Major.K)
         v_smem_base = sm100_desc.smem_desc_base_from_tensor(sV, sm100_desc.Major.MN)
@@ -1456,8 +1751,9 @@ class FlashAttentionForwardSm100:
                 # idesc=qk_mma_idesc,
                 smem_desc_base_b=k_smem_base,
                 tCrB_layout=tSrK[None, None, None, 0].layout,
-                smem_var_name_prefix=f"fa_fwd_q_smem_desc",
-                idesc_var_name=f"fa_fwd_qk_mma_idesc",
+                smem_var_name_prefix="fa_fwd_q_smem_desc",
+                idesc_var_name="fa_fwd_qk_mma_idesc",
+                kind=qk_mma_kind,
                 smem_offset=-sQ_stage_stride if stage == 0 else sQ_stage_stride,
                 zero_init=True,
                 cta_group=self.cta_group_size,
@@ -1507,32 +1803,41 @@ class FlashAttentionForwardSm100:
         )
         P_full_O_rescaled_phase = Int32(0)
 
-        tile_scheduler = TileSchedulerCls()
         work_tile = tile_scheduler.initial_work_tile_info()
         while work_tile.is_valid_tile:
             m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx
             seqlen = SeqlenInfoCls(batch_idx)
-
             block_iter_count = Int32(0)
             process_tile = False
 
             if const_expr(self.use_block_sparsity):
+                # See the load warp: unpack dynamic num_splits packed by the scheduler.
+                num_splits_dyn = num_splits
+                if const_expr(self.is_split_kv and block_info.pack_split_idx):
+                    num_splits_dyn = split_idx >> 16
+                    split_idx = split_idx & 0xFFFF
                 block_iter_count = get_total_block_count(
                     blocksparse_tensors,
                     batch_idx,
                     head_idx,
                     m_block,
+                    split_idx,
+                    num_splits_dyn,
                     self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
-                    self.q_subtile_factor if self.q_subtile_factor is not None else 1,
+                    self.q_subtile_factor,
+                    seqlen_info=seqlen,
+                    kv_subtile_factor=self.kv_subtile_factor,
                 )
                 process_tile = block_iter_count > Int32(0)
             else:
-                n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block, split_idx, num_splits)
+                n_block_min, n_block_max = block_info.get_n_block_min_max(
+                    seqlen,
+                    m_block,
+                    split_idx=split_idx,
+                    num_splits=num_splits,
+                )
                 block_iter_count = n_block_max - n_block_min
-                if const_expr(not self.is_split_kv):
-                    process_tile = True
-                else:
-                    process_tile = n_block_min < n_block_max
+                process_tile = self.process_work_tile(seqlen, n_block_min, n_block_max)
 
             if process_tile and is_leader_cta:
                 for stage in cutlass.range_constexpr(self.q_stage):
@@ -1678,8 +1983,7 @@ class FlashAttentionForwardSm100:
                 # End of GEMM_PV1(i_end) (P1 * Vi_end -> O1)
 
             # Advance to next tile
-            tile_scheduler.advance_to_next_work()
-            work_tile = tile_scheduler.get_current_work()
+            work_tile = tile_scheduler.advance_to_next_work()
         # End of persistent scheduler loop
 
         # We don't need pipeline_s_p_o.producer_tail() since there's no dangling mbarrier at the end
@@ -1688,13 +1992,40 @@ class FlashAttentionForwardSm100:
         # pipeline_o_acc.producer_acquire() inside the loop.
 
     # for both softmax0 and softmax1 warp group
+    @cute.jit
+    def _kv_head_idx(self, head_idx: Int32) -> Int32:
+        """Map query-head tile index -> KV-head index (FA3 descale semantics)."""
+        if cutlass.const_expr(self.pack_gqa):
+            return head_idx
+        return head_idx // self.qhead_per_kvhead
+
+    @cute.jit
+    def _load_effective_descales(
+        self,
+        descale_tensors: Optional[DescaleTensors],
+        batch_idx: Int32,
+        kv_head_idx: Int32,
+    ) -> Tuple[Float32, Float32]:
+        """Load effective QK and V descales, defaulting unspecified tensors to identity."""
+        qk_descale = Float32(1.0)
+        v_descale = Float32(1.0)
+        if cutlass.const_expr(descale_tensors is not None):
+            if cutlass.const_expr(descale_tensors.q_descale is not None):
+                qk_descale = qk_descale * Float32(descale_tensors.q_descale[batch_idx, kv_head_idx])
+            if cutlass.const_expr(descale_tensors.k_descale is not None):
+                qk_descale = qk_descale * Float32(descale_tensors.k_descale[batch_idx, kv_head_idx])
+            if cutlass.const_expr(descale_tensors.v_descale is not None):
+                v_descale = Float32(descale_tensors.v_descale[batch_idx, kv_head_idx])
+        return qk_descale, v_descale
+
     @cute.jit
     def softmax_loop(
         self,
         stage: int | Int32,
         softmax_scale_log2: Float32,
-        softmax_scale: Float32,
-        thr_mma_qk: cute.core.ThrMma,
+        softmax_scale: Float32 | None,
+        descale_tensors: Optional[DescaleTensors],
+        thr_mma_qk: cute.ThrMma,
         tStS: cute.Tensor,  # ((TILE_M, TILE_N), 1, 1, q_stage)
         sScale: cute.Tensor,
         mLSE: Optional[cute.Tensor],
@@ -1708,11 +2039,11 @@ class FlashAttentionForwardSm100:
         num_splits: Int32,
         SeqlenInfoCls: Callable,
         AttentionMaskCls: Callable,
-        TileSchedulerCls: Callable,
-        aux_tensors: Optional[list] = None,
+        aux_data: AuxData = AuxData(),
         fastdiv_mods=(None, None),
         head_divmod=None,
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
+        tile_scheduler=None,
     ):
         """Compute softmax on attention scores from QK matrix multiplication.
 
@@ -1731,6 +2062,7 @@ class FlashAttentionForwardSm100:
             * (len(self.softmax0_warp_ids))
         )
         warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4
+        aux_tensors = aux_data.tensors
 
         cta_qk_tiler = (self.mma_tiler_qk[0] // thr_mma_qk.thr_id.shape, self.mma_tiler_qk[1])
         tSAcc = tStS[(None, None), 0, 0, stage]  # (128, 128)
@@ -1745,9 +2077,12 @@ class FlashAttentionForwardSm100:
         )
         tStP = cute.make_tensor(tSAcc.iterator + self.tmem_s_to_p_offset, tStP_layout)
 
-        tmem_load_atom = cute.make_copy_atom(
-            tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype
+        tmem_load_op = (
+            tcgen05.copy.LdRed32x32bOp(tcgen05.copy.Repetition(32))
+            if const_expr(self.use_ldred_rowmax)
+            else tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32))
         )
+        tmem_load_atom = cute.make_copy_atom(tmem_load_op, self.qk_acc_dtype)
         thr_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tSAcc).get_slice(tidx)
         tStS_t2r = thr_tmem_load.partition_S(tSAcc)  # (((32,32),1),1,4)
 
@@ -1759,7 +2094,10 @@ class FlashAttentionForwardSm100:
         )
         tStScale_r2t = thr_tmem_store_scale.partition_D(tStScale)
         tmem_store_atom = cute.make_copy_atom(
-            tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(16)), Float32
+            tcgen05.copy.St32x32bOp(
+                tcgen05.copy.Repetition(8 if const_expr(self.q_dtype.width == 8) else 16)
+            ),
+            Float32,
         )
         thr_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP).get_slice(tidx)
         tStP_r2t = thr_tmem_store.partition_D(tStP)  # (((16,32),1),1,4)
@@ -1772,12 +2110,19 @@ class FlashAttentionForwardSm100:
 
         warp_idx_in_wg = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4
 
-        tile_scheduler = TileSchedulerCls()
         work_tile = tile_scheduler.initial_work_tile_info()
         while work_tile.is_valid_tile:
             m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx
+            kv_head_idx = self._kv_head_idx(head_idx)
             seqlen = SeqlenInfoCls(batch_idx)
-            n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block, split_idx, num_splits)
+            n_block_min, n_block_max = block_info.get_n_block_min_max(
+                seqlen, m_block, split_idx=split_idx, num_splits=num_splits,
+            )
+            # Keep the dynamic num_splits for the block-sparse helpers below.
+            num_splits_dyn = num_splits
+            if const_expr(self.is_split_kv and block_info.pack_split_idx):
+                num_splits_dyn = split_idx >> 16
+                split_idx = split_idx & 0xFFFF
 
             mask = AttentionMaskCls(seqlen)
             shared_mask_kwargs = dict(
@@ -1788,7 +2133,8 @@ class FlashAttentionForwardSm100:
                 mask_local=self.is_local,
                 batch_idx=batch_idx,
                 head_idx=head_idx,
-                aux_tensors=aux_tensors,
+                aux_data=aux_data,
+                vec_size=self.mask_vec_size,
             )
 
             # Recompute fastdiv_mods if necessary
@@ -1830,10 +2176,28 @@ class FlashAttentionForwardSm100:
             else:
                 mask_fn_none = None
 
+            qk_descale, _ = self._load_effective_descales(descale_tensors, batch_idx, kv_head_idx)
+
+            # See Note [Low Precision Scaling]
+            max_offset = 8 if cutlass.const_expr(self.q_dtype.width == 8) else 0
+            if const_expr(self.score_mod is None):
+                softmax_scale_log2_eff = softmax_scale_log2 * qk_descale
+                softmax_scale_eff = None
+            else:
+                softmax_scale_log2_eff = softmax_scale_log2
+                softmax_scale_eff = softmax_scale * qk_descale
+
+            rescale_threshold = 8.0 if const_expr(self.q_dtype.width == 16) else 0.0
+            # See Note [Low Precision Scaling]
+            assert max_offset + rescale_threshold < _LOG2_DTYPE_MAX[self.q_dtype], (
+                f"max_offset ({max_offset}) + rescale_threshold ({rescale_threshold}) must stay "
+                f"below log2(max {self.q_dtype} value) to avoid saturating P"
+            )
             softmax = SoftmaxSm100.create(
-                softmax_scale_log2,
-                rescale_threshold=8.0 if const_expr(self.q_dtype.width == 16) else 0.0,
-                softmax_scale=softmax_scale,
+                softmax_scale_log2_eff,
+                rescale_threshold=rescale_threshold,
+                softmax_scale=softmax_scale_eff,
+                max_offset=max_offset,
             )
             softmax.reset()
 
@@ -1843,13 +2207,17 @@ class FlashAttentionForwardSm100:
                     batch_idx,
                     head_idx,
                     m_block,
+                    split_idx,
+                    num_splits_dyn,
                     self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
-                    self.q_subtile_factor if self.q_subtile_factor is not None else 1,
+                    self.q_subtile_factor,
+                    seqlen_info=seqlen,
+                    kv_subtile_factor=self.kv_subtile_factor,
                 )
                 has_work = tile_block_count > Int32(0)
             else:
                 tile_block_count = n_block_max - n_block_min
-                has_work = const_expr(not self.is_split_kv) or tile_block_count > Int32(0)
+                has_work = self.process_work_tile(seqlen, n_block_min, n_block_max)
 
             softmax_step = partial(
                 self.softmax_step,
@@ -1872,13 +2240,12 @@ class FlashAttentionForwardSm100:
                 head_idx=head_idx,
                 m_block=(self.q_stage * m_block + stage) * self.cta_group_size,
                 seqlen=seqlen,
-                aux_tensors=aux_tensors,
+                aux_data=aux_data,
                 fastdiv_mods=fastdiv_mods,
                 head_divmod=head_divmod,
             )
 
             if const_expr(self.use_block_sparsity) or has_work:
-                # See block_sparse_utils.py NOTE [SM100 block-sparse empty tiles: mbarrier contract].
                 pipeline_sm_stats.producer_acquire_w_index_phase(stage, sm_stats_producer_phase)
                 sm_stats_producer_phase ^= 1
 
@@ -1901,6 +2268,9 @@ class FlashAttentionForwardSm100:
                     batch_idx,
                     head_idx,
                     m_block,
+                    seqlen,
+                    split_idx,
+                    num_splits_dyn,
                     softmax_step,
                     mask_fn,
                     mask_fn_none,
@@ -1913,7 +2283,8 @@ class FlashAttentionForwardSm100:
                     Int32(stage),
                     check_m_boundary,
                     self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
-                    self.q_subtile_factor if self.q_subtile_factor is not None else 1,
+                    self.q_subtile_factor,
+                    self.kv_subtile_factor,
                 )
                 if not empty_tile:
                     sScale[tidx + stage * self.m_block_size] = softmax.row_sum[0]
@@ -1928,7 +2299,7 @@ class FlashAttentionForwardSm100:
                     sm_stats_barrier.arrive_w_index(index=stage * 4 + warp_idx)
                     # if tidx == 0: cute.printf("softmax row sum stage %d: %f\n", stage, softmax.row_sum[0])
             else:
-                if const_expr(not self.is_split_kv) or tile_block_count > Int32(0):
+                if has_work:
                     mma_si_consumer_phase, sm_stats_producer_phase, s0_s1_sequence_phase = softmax_step(
                         mma_si_consumer_phase,
                         sm_stats_producer_phase,
@@ -2015,8 +2386,7 @@ class FlashAttentionForwardSm100:
             #         gLSE[tidx] = lse
 
             # Advance to next tile
-            tile_scheduler.advance_to_next_work()
-            work_tile = tile_scheduler.get_current_work()
+            work_tile = tile_scheduler.advance_to_next_work()
         # End of persistent scheduler loop
 
         # This is equivalent to pipeline_sm_stats.producer_tail
@@ -2034,7 +2404,7 @@ class FlashAttentionForwardSm100:
         s0_s1_sequence_phase: Int32,
         n_block: Int32,
         softmax: SoftmaxSm100,
-        thr_mma_qk: cute.core.ThrMma,
+        thr_mma_qk: cute.ThrMma,
         pipeline_s_p_o: pipeline.PipelineAsync,
         pipeline_p_lastsplit: pipeline.PipelineAsync,
         pipeline_sm_stats: pipeline.PipelineAsync,
@@ -2052,7 +2422,7 @@ class FlashAttentionForwardSm100:
         head_idx: Int32,
         m_block: Int32,
         seqlen,
-        aux_tensors: Optional[list] = None,
+        aux_data: AuxData = AuxData(),
         fastdiv_mods=(None, None),
         head_divmod=None,
         mask_fn: Optional[Callable] = None,
@@ -2072,6 +2442,8 @@ class FlashAttentionForwardSm100:
         4. Transforming scores using exp2(x*scale - max*scale)
         5. Computing row sums for normalization
         6. Coordinating pipeline synchronization between different processing stages
+
+        A None mask_fn means the tcgen05.ld.red hardware max is valid.
         """
         warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4
         tilePlikeFP32 = self.mma_tiler_qk[1] // Float32.width * self.v_dtype.width
@@ -2084,8 +2456,16 @@ class FlashAttentionForwardSm100:
 
         # Wait for Si
         pipeline_s_p_o.consumer_wait_w_index_phase(stage, mma_si_consumer_phase)
-        tSrS_t2r = cute.make_fragment(thr_tmem_load.partition_D(tScS).shape, self.qk_acc_dtype)
-        cute.copy(thr_tmem_load, tStS_t2r, tSrS_t2r)
+        tSrS_t2r = cute.make_rmem_tensor(thr_tmem_load.partition_D(tScS).shape, self.qk_acc_dtype)
+        hw_row_max = Float32(-Float32.inf)
+        if const_expr(self.use_ldred_rowmax):
+            # ld.red returns each x32 tile's max in an extra register.
+            tSrS_red = cute.make_rmem_tensor(((1, 1), *tSrS_t2r.shape[1:]), self.qk_acc_dtype)
+            cute.copy(thr_tmem_load, tStS_t2r, (tSrS_t2r, tSrS_red))
+            for i in cutlass.range_constexpr(cute.size(tSrS_red.shape)):
+                hw_row_max = cute.arch.fmax(hw_row_max, tSrS_red[i])
+        else:
+            cute.copy(thr_tmem_load, tStS_t2r, tSrS_t2r)
         # tSrS_t2r = copy_utils.load_t2r(thr_tmem_load, tScS_shape, tStS_t2r)
         if cutlass.const_expr(self.score_mod is not None):
             self.apply_score_mod(
@@ -2098,17 +2478,21 @@ class FlashAttentionForwardSm100:
                 n_block,
                 softmax,
                 seqlen,
-                aux_tensors,
+                aux_data,
                 fastdiv_mods,
                 head_divmod,
             )
 
         if const_expr(mask_fn is not None):
             mask_fn(tSrS_t2r, n_block=n_block)
-        row_max, acc_scale = softmax.update_row_max(tSrS_t2r.load(), is_first)
+        # Masked iterations reduce over post-mask values in software.
+        if const_expr(self.use_ldred_rowmax and mask_fn is None):
+            row_max, acc_scale = softmax.update_row_max_precomputed(hw_row_max, is_first)
+        else:
+            row_max, acc_scale = softmax.update_row_max(tSrS_t2r.load(), is_first)
 
         if const_expr(not is_first):
-            # tSrScale_r2t = cute.make_fragment(thr_tmem_store_scale.partition_S(tScScale).shape, Float32)
+            # tSrScale_r2t = cute.make_rmem_tensor(thr_tmem_store_scale.partition_S(tScScale).shape, Float32)
             # tSrScale_r2t[0] = acc_scale
             # cute.copy(thr_tmem_store_scale, tSrScale_r2t, tStScale_r2t)
             # cute.arch.fence_view_async_tmem_store()
@@ -2124,7 +2508,7 @@ class FlashAttentionForwardSm100:
         # Sequence barrier wait
         if const_expr(self.s0_s1_barrier):
             pipeline_s0_s1_sequence.sync_object_full.wait(stage, s0_s1_sequence_phase)
-        tSrP_r2t_f32 = cute.make_fragment(
+        tSrP_r2t_f32 = cute.make_rmem_tensor(
             thr_tmem_store.partition_S(cute.make_identity_tensor(tScP_shape)).shape, Float32
         )
         tSrP_r2t = cute.make_tensor(
@@ -2134,7 +2518,7 @@ class FlashAttentionForwardSm100:
         softmax.apply_exp2_convert(
             tSrS_t2r,
             tSrP_r2t,
-            ex2_emu_freq=self.ex2_emu_freq if const_expr(mask_fn is None) else 0,
+            ex2_emu_freq=self.ex2_emu_freq,
             ex2_emu_start_frg=self.ex2_emu_start_frg,
         )
         # Sequence barrier arrive
@@ -2166,8 +2550,8 @@ class FlashAttentionForwardSm100:
     @cute.jit
     def correction_loop(
         self,
-        thr_mma_qk: cute.core.ThrMma,
-        thr_mma_pv: cute.core.ThrMma,
+        thr_mma_qk: cute.ThrMma,
+        thr_mma_pv: cute.ThrMma,
         tStS: cute.Tensor,
         tOtO: cute.Tensor,
         sScale: cute.Tensor,
@@ -2179,15 +2563,17 @@ class FlashAttentionForwardSm100:
         pipeline_sm_stats: pipeline.PipelineAsync,
         sm_stats_barrier: pipeline.NamedBarrier,
         pipeline_o_epi: pipeline.PipelineAsync,
+        pipeline_load_epi: Optional[pipeline.PipelineAsync],
         learnable_sink: Optional[cute.Tensor],
+        descale_tensors: Optional[DescaleTensors],
         gmem_tiled_copy_O: cute.TiledCopy,
         tma_atom_O: cute.CopyAtom,
         softmax_scale_log2: Float32,
         block_info: BlockInfo,
         num_splits: Int32,
         SeqlenInfoCls: Callable,
-        TileSchedulerCls: Callable,
         blocksparse_tensors: Optional[BlockSparseTensors] = None,
+        tile_scheduler=None,
     ):
         tidx = cute.arch.thread_idx()[0] % (cute.arch.WARP_SIZE * len(self.correction_warp_ids))
         warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4
@@ -2216,24 +2602,47 @@ class FlashAttentionForwardSm100:
         sm_stats_consumer_phase = Int32(0)
         o_corr_consumer_phase = Int32(0)
         corr_epi_producer_phase = Int32(1)
+        load_epi_producer_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, 1
+        )
 
-        tile_scheduler = TileSchedulerCls()
         work_tile = tile_scheduler.initial_work_tile_info()
         while work_tile.is_valid_tile:
             m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx
+            kv_head_idx = self._kv_head_idx(head_idx)
+            qk_descale, v_descale = self._load_effective_descales(descale_tensors, batch_idx, kv_head_idx)
+            if const_expr(self.score_mod is None):
+                softmax_scale_log2_eff = softmax_scale_log2 * qk_descale
+            else:
+                softmax_scale_log2_eff = softmax_scale_log2
+
+            # Must match the softmax warp's max_offset; max_offset_scale = 2^max_offset.
+            max_offset = Float32(8.0) if cutlass.const_expr(self.q_dtype.width == 8) else Float32(0.0)
+            max_offset_scale = (
+                Float32(256.0) if cutlass.const_expr(self.q_dtype.width == 8) else Float32(1.0)
+            )
             seqlen = SeqlenInfoCls(batch_idx)
-            n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block, split_idx, num_splits)
+            n_block_min, n_block_max = block_info.get_n_block_min_max(
+                seqlen, m_block, split_idx=split_idx, num_splits=num_splits,
+            )
+            # Keep the dynamic num_splits for the block-sparse helper below.
+            num_splits_dyn = num_splits
+            if const_expr(self.is_split_kv and block_info.pack_split_idx):
+                num_splits_dyn = split_idx >> 16
+                split_idx = split_idx & 0xFFFF
 
             if const_expr(self.is_split_kv):
                 mO_cur = seqlen.offset_batch_Q(mO, batch_idx, dim=3)[None, None, head_idx, split_idx]
             else:
                 mO_cur = seqlen.offset_batch_Q(mO, batch_idx, dim=3)[None, None, head_idx]
-            tiler_gO = ((self.mma_tiler_pv[0] * self.q_stage), self.head_dim_v_padded)
-            gO = cute.local_tile(mO_cur, tiler_gO, (m_block, 0))  # (128 * 2, 128)
-            gO = layout_utils.select(
-                cute.flat_divide(gO, (self.mma_tiler_pv[0],)), mode=[0, 2, 1]
-            )  # (128, 128, 2)
-            gO = cute.flat_divide(gO, (self.mma_tiler_pv[0] // self.cta_group_size,))[None, mma_tile_coord_v, None, None]
+            gO = None
+            if const_expr(self.use_tma_O or not self.pack_gqa):
+                tiler_gO = ((self.mma_tiler_pv[0] * self.q_stage), self.head_dim_v_padded)
+                gO = cute.local_tile(mO_cur, tiler_gO, (m_block, 0))  # (128 * 2, 128)
+                gO = layout_utils.select(
+                    cute.flat_divide(gO, (self.mma_tiler_pv[0],)), mode=[0, 2, 1]
+                )  # (128, 128, 2)
+                gO = cute.flat_divide(gO, (self.mma_tiler_pv[0] // self.cta_group_size,))[None, mma_tile_coord_v, None, None]
 
             # Default LSE to -inf for invalid split_idx tiles
             stats = [(0.0, -Float32.inf if const_expr(mLSE is not None or learnable_sink is not None) else None, True)] * self.q_stage
@@ -2244,13 +2653,17 @@ class FlashAttentionForwardSm100:
                     batch_idx,
                     head_idx,
                     m_block,
+                    split_idx,
+                    num_splits_dyn,
                     self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
-                    self.q_subtile_factor if self.q_subtile_factor is not None else 1,
+                    self.q_subtile_factor,
+                    seqlen_info=seqlen,
+                    kv_subtile_factor=self.kv_subtile_factor,
                 )
                 has_work = total_block_count > Int32(0)
             else:
                 total_block_count = n_block_max - n_block_min
-                has_work = const_expr(not self.is_split_kv) or total_block_count > Int32(0)
+                has_work = self.process_work_tile(seqlen, n_block_min, n_block_max)
 
             if has_work:
                 # Ignore first signal from softmax as no correction is required
@@ -2262,7 +2675,7 @@ class FlashAttentionForwardSm100:
                     sm_stats_barrier.arrive_and_wait_w_index(index=1 * 4 + warp_idx)
                 sm_stats_consumer_phase ^= 1
 
-                tSrScale_t2r = cute.make_fragment(tSrScale_t2r_shape, Float32)
+                tSrScale_t2r = cute.make_rmem_tensor(tSrScale_t2r_shape, Float32)
                 for i in cutlass.range(total_block_count - 1, unroll=1):
                     for stage in cutlass.range_constexpr(self.q_stage):
                         # wait for S0 / S1
@@ -2321,19 +2734,21 @@ class FlashAttentionForwardSm100:
                         if const_expr(not self.is_split_kv) or split_idx == 0:
                             if row_max == -Float32.inf:
                                 # It's possible to have an empty row with splitKV.
-                                row_max = sink_val * (LOG2_E / softmax_scale_log2)
-                                row_sum = Float32(1.0)
+                                row_max = sink_val * (LOG2_E / softmax_scale_log2_eff)
+                                row_sum = max_offset_scale
                             else:
                                 row_sum += cute.math.exp2(
-                                    sink_val * LOG2_E - row_max * softmax_scale_log2, fastmath=True
+                                    sink_val * LOG2_E - row_max * softmax_scale_log2_eff + max_offset, fastmath=True
                                 )
                     acc_O_mn_row_is_zero_or_nan = row_sum == 0.0 or row_sum != row_sum
                     stats[stage] = (row_sum, row_max, acc_O_mn_row_is_zero_or_nan)
                     scale = cute.arch.rcp_approx(row_sum if not acc_O_mn_row_is_zero_or_nan else 1.0)
+                    scale = scale * v_descale
                     # Wait for the last O to be ready from the MMA warp
                     pipeline_o_acc.consumer_wait_w_index_phase(stage, o_corr_consumer_phase)
                     if const_expr(not self.use_correction_warps_for_epi):
                         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
                     self.correction_epilogue(
                         thr_mma_pv,
                         tOtO[None, None, None, stage],
@@ -2344,7 +2759,7 @@ class FlashAttentionForwardSm100:
                         scale,
                         sO[None, None, stage],
                         mO_cur,
-                        gO[None, None, stage],
+                        gO_stage,
                         gmem_tiled_copy_O,
                     )
                     # Signal for the next work tile that O buffers in tmem are already read, so
@@ -2392,7 +2807,9 @@ class FlashAttentionForwardSm100:
                         sm_stats_consumer_phase,
                         o_corr_consumer_phase,
                         corr_epi_producer_phase,
-                        softmax_scale_log2,
+                        softmax_scale_log2_eff,
+                        max_offset,
+                        max_offset_scale,
                         mO_cur,
                         gO,
                         gmem_tiled_copy_O_for_empty_tile,
@@ -2414,13 +2831,12 @@ class FlashAttentionForwardSm100:
                         mLSE_cur = cute.domain_offset((offset,), mLSE[None, head_idx])
                 for stage in cutlass.range_constexpr(self.q_stage):
                     m_tile_idx = (m_block * self.q_stage + stage) * self.cta_group_size + mma_tile_coord_v
-                    gLSE = cute.local_tile(mLSE_cur, (self.m_block_size,), (m_tile_idx,))
                     row_sum, row_max, acc_O_mn_row_is_zero_or_nan = stats[stage]
                     # if tidx == 0 and stage <= 1:
                     #     cute.printf("row_sum = {}, row_max = {}, acc_O_mn_row_is_zero_or_nan = {}\n", row_sum, row_max, acc_O_mn_row_is_zero_or_nan)
                     LN2 = math.log(2.0)
                     lse = (
-                        (row_max * softmax_scale_log2 + cute.math.log2(row_sum, fastmath=True)) * LN2
+                        (row_max * softmax_scale_log2_eff + (cute.math.log2(row_sum, fastmath=True) - max_offset)) * LN2
                         if not acc_O_mn_row_is_zero_or_nan
                         else -Float32.inf
                     )
@@ -2429,13 +2845,30 @@ class FlashAttentionForwardSm100:
                         if const_expr(not self.pack_gqa)
                         else seqlen.seqlen_q * self.qhead_per_kvhead
                     )
-                    if tidx < seqlen_q - m_tile_idx * self.m_block_size:
-                        # This actually just works with PackGQA too
-                        gLSE[tidx] = lse
+                    if const_expr(not self.pack_gqa or self.m_block_size % self.qhead_per_kvhead == 0):
+                        gLSE = cute.local_tile(mLSE_cur, (self.m_block_size,), (m_tile_idx,))
+                        if tidx < seqlen_q - m_tile_idx * self.m_block_size:
+                            # This actually just works with PackGQA too
+                            gLSE[tidx] = lse
+                    else:
+                        idx = m_tile_idx * self.m_block_size + tidx
+                        if idx < seqlen_q:
+                            m_idx = idx // self.qhead_per_kvhead
+                            h_idx = idx - m_idx * self.qhead_per_kvhead
+                            lse_ptr_i64 = utils.elem_pointer(mLSE_cur, ((h_idx, m_idx),)).toint()
+                            lse_gmem_ptr = cute.make_ptr(
+                                mLSE_cur.element_type, lse_ptr_i64, cute.AddressSpace.gmem, assumed_align=4
+                            )
+                            cute.make_tensor(lse_gmem_ptr, (1,))[0] = lse
+
+            if const_expr(pipeline_load_epi is not None and self.use_correction_warps_for_epi):
+                pipeline_load_epi.producer_acquire(load_epi_producer_state)
+                with cute.arch.elect_one():
+                    pipeline_load_epi.producer_commit(load_epi_producer_state)
+                load_epi_producer_state.advance()
 
             # Advance to next tile
-            tile_scheduler.advance_to_next_work()
-            work_tile = tile_scheduler.get_current_work()
+            work_tile = tile_scheduler.advance_to_next_work()
         # End of persistent scheduler loop
 
         # This is equivalent to pipeline_o_epi.consumer_tail() for the correction warps
@@ -2445,7 +2878,7 @@ class FlashAttentionForwardSm100:
     @cute.jit
     def correction_rescale(
         self,
-        thr_mma: cute.core.ThrMma,
+        thr_mma: cute.ThrMma,
         tOtO: cute.Tensor,
         tidx: Int32,
         scale: Float32,
@@ -2480,9 +2913,9 @@ class FlashAttentionForwardSm100:
         tOtO_r2t = thr_tmem_store.partition_D(tOtO_i)
 
         frg_count = self.head_dim_v_padded // corr_tile_size
-        tOrO_frg = cute.make_fragment((tOrO_t2r_shape, frg_count), self.pv_acc_dtype)
+        tOrO_frg = cute.make_rmem_tensor((tOrO_t2r_shape, frg_count), self.pv_acc_dtype)
         for i in cutlass.range_constexpr(frg_count):
-            tOrO_frg = cute.make_fragment(tOrO_t2r_shape, self.pv_acc_dtype)
+            tOrO_frg = cute.make_rmem_tensor(tOrO_t2r_shape, self.pv_acc_dtype)
             tOtO_t2r_i = cute.make_tensor(tOtO_t2r.iterator + i * corr_tile_size, tOtO_t2r.layout)
             cute.copy(thr_tmem_load, tOtO_t2r_i, tOrO_frg)
             for j in cutlass.range(0, cute.size(tOrO_frg), 2, unroll_full=True):
@@ -2496,7 +2929,7 @@ class FlashAttentionForwardSm100:
     @cute.jit
     def correction_epilogue(
         self,
-        thr_mma: cute.core.ThrMma,
+        thr_mma: cute.ThrMma,
         tOtO: cute.Tensor,
         tidx: Int32,
         stage: Int32,
@@ -2522,7 +2955,7 @@ class FlashAttentionForwardSm100:
         5. Preparation for efficient TMA store operations
 
         :param thr_mma: Thread MMA operation for the computation
-        :type thr_mma: cute.core.ThrMma
+        :type thr_mma: cute.ThrMma
         :param tOtO: Tensor containing accumulated attention output
         :type tOtO: cute.Tensor
         :param scale: Final scaling factor to apply to the output
@@ -2562,7 +2995,7 @@ class FlashAttentionForwardSm100:
         for i in cutlass.range(self.head_dim_v_padded // corr_tile_size, unroll_full=True):
             tOtO_t2r_i = tOtO_t2r[None, 0, 0, i]
             tOsO_r2s_i = tOsO_s2r[None, 0, 0, i]
-            tOrO_frg = cute.make_fragment(tOcO_t2r[None, 0, 0, i].shape, self.pv_acc_dtype)
+            tOrO_frg = cute.make_rmem_tensor(tOcO_t2r[None, 0, 0, i].shape, self.pv_acc_dtype)
             cute.copy(tiled_tmem_load, tOtO_t2r_i, tOrO_frg)
             for j in cutlass.range(0, cute.size(tOrO_frg), 2, unroll_full=True):
                 tOrO_frg[j], tOrO_frg[j + 1] = cute.arch.mul_packed_f32x2(
@@ -2574,7 +3007,7 @@ class FlashAttentionForwardSm100:
         if const_expr(self.use_correction_warps_for_epi):
             assert(not self.use_tma_O)
             assert(gmem_tiled_copy_O is not None)
-            cute.arch.barrier(barrier_id=int(NamedBarrierFwd.Epilogue),
+            cute.arch.barrier(barrier_id=int(NamedBarrierFwdSm100.Epilogue),
                               number_of_threads=len(self.epilogue_warp_ids) * cute.arch.WARP_SIZE)
             mma_tile_coord_v = thr_mma.thr_idx
             m_tile_idx = (m_block * self.q_stage + stage) * self.cta_group_size + mma_tile_coord_v
@@ -2586,7 +3019,7 @@ class FlashAttentionForwardSm100:
     def _store_O_to_gmem(
         self,
         sO_stage: cute.Tensor,
-        gO: cute.Tensor,
+        gO: Optional[cute.Tensor],
         mO_cur: cute.Tensor,
         gmem_tiled_copy_O: cute.TiledCopy,
         tidx: Int32,
@@ -2597,7 +3030,6 @@ class FlashAttentionForwardSm100:
         gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx)
         tOsO = gmem_thr_copy_O.partition_S(sO_stage)
         cO = cute.make_identity_tensor((self.m_block_size, self.head_dim_v_padded))
-        tOgO = gmem_thr_copy_O.partition_D(gO)
         tOcO = gmem_thr_copy_O.partition_S(cO)
         t0OcO = gmem_tiled_copy_O.get_slice(0).partition_S(cO)
         tOpO = copy_utils.predicate_k(tOcO, limit=mO_cur.shape[1])
@@ -2613,6 +3045,8 @@ class FlashAttentionForwardSm100:
         cute.autovec_copy(tOsO, tOrO)
         # copy acc O from rmem to gmem
         if const_expr(not self.pack_gqa):
+            assert gO is not None
+            tOgO = gmem_thr_copy_O.partition_D(gO)
             for rest_m in cutlass.range_constexpr(cute.size(tOrO.shape[1])):
                 if (
                     t0OcO[0, rest_m, 0][0] < seqlen_q - m_tile_idx * self.m_block_size - tOcO[0][0]
@@ -2638,31 +3072,43 @@ class FlashAttentionForwardSm100:
         gmem_tiled_copy_O: cute.TiledCopy,
         tma_atom_O: Optional[cute.CopyAtom],
         pipeline_o_epi: pipeline.PipelineAsync,
+        pipeline_load_epi: Optional[pipeline.PipelineAsync],
         block_info: BlockInfo,
         num_splits: int,
         SeqlenInfoCls: Callable,
-        TileSchedulerCls: Callable,
         mma_tile_coord_v: Int32 = 0,
+        blocksparse_tensors: Optional[BlockSparseTensors] = None,
+        tile_scheduler=None,
     ):
         epi_consumer_phase = Int32(0)
-        tile_scheduler = TileSchedulerCls()
+        load_epi_producer_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Producer, 1
+        )
         work_tile = tile_scheduler.initial_work_tile_info()
         while work_tile.is_valid_tile:
             m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx
             seqlen = SeqlenInfoCls(batch_idx)
-            n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block, split_idx, num_splits)
+            n_block_min, n_block_max = block_info.get_n_block_min_max(
+                seqlen, m_block, split_idx=split_idx, num_splits=num_splits,
+            )
+            if const_expr(self.is_split_kv and block_info.pack_split_idx):
+                split_idx = split_idx & 0xFFFF
 
-            if const_expr(not self.is_split_kv) or n_block_min < n_block_max:
+            if const_expr(self.use_block_sparsity) or self.process_work_tile(
+                seqlen, n_block_min, n_block_max
+            ):
                 if const_expr(self.is_split_kv):
                     mO_cur = seqlen.offset_batch_Q(mO, batch_idx, dim=3)[None, None, head_idx, split_idx]
                 else:
                     mO_cur = seqlen.offset_batch_Q(mO, batch_idx, dim=3)[None, None, head_idx]
-                tiler_gO = ((self.mma_tiler_pv[0] * self.q_stage), self.head_dim_v_padded)
-                gO = cute.local_tile(mO_cur, tiler_gO, (m_block, 0))  # (128 * 2, 128)
-                gO = layout_utils.select(
-                    cute.flat_divide(gO, (self.mma_tiler_pv[0],)), mode=[0, 2, 1]
-                )  # (128, 128, 2)
-                gO = cute.flat_divide(gO, (self.mma_tiler_pv[0] // self.cta_group_size,))[None, mma_tile_coord_v, None, None]
+                gO = None
+                if const_expr(self.use_tma_O or not self.pack_gqa):
+                    tiler_gO = ((self.mma_tiler_pv[0] * self.q_stage), self.head_dim_v_padded)
+                    gO = cute.local_tile(mO_cur, tiler_gO, (m_block, 0))  # (128 * 2, 128)
+                    gO = layout_utils.select(
+                        cute.flat_divide(gO, (self.mma_tiler_pv[0],)), mode=[0, 2, 1]
+                    )  # (128, 128, 2)
+                    gO = cute.flat_divide(gO, (self.mma_tiler_pv[0] // self.cta_group_size,))[None, mma_tile_coord_v, None, None]
 
                 if const_expr(self.use_tma_O):
                     store_O, _, _ = copy_utils.tma_get_copy_fn(
@@ -2689,17 +3135,56 @@ class FlashAttentionForwardSm100:
                         pipeline_o_epi.consumer_wait_w_index_phase(stage, epi_consumer_phase)
                         # 2. copy O0 / O1 to gmem
                         m_tile_idx = (m_block * self.q_stage + stage) * self.cta_group_size + mma_tile_coord_v
+                        gO_stage = gO[None, None, stage] if const_expr(gO is not None) else None
                         self._store_O_to_gmem(
-                            sO[None, None, stage], gO[None, None, stage], mO_cur, gmem_tiled_copy_O,
+                            sO[None, None, stage], gO_stage, mO_cur, gmem_tiled_copy_O,
                             tidx, seqlen.seqlen_q, m_tile_idx,
                         )
                         pipeline_o_epi.consumer_release_w_index(stage)
 
                 epi_consumer_phase ^= 1
 
+            if const_expr(pipeline_load_epi is not None):
+                pipeline_load_epi.producer_acquire(load_epi_producer_state)
+                with cute.arch.elect_one():
+                    pipeline_load_epi.producer_commit(load_epi_producer_state)
+                load_epi_producer_state.advance()
+
             # Advance to next tile
-            tile_scheduler.advance_to_next_work()
-            work_tile = tile_scheduler.get_current_work()
+            work_tile = tile_scheduler.advance_to_next_work()
+
+    @cute.jit
+    def 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()
+            if cute.arch.thread_idx()[0] == self.scheduler_warp_id * cute.arch.WARP_SIZE:
+                prefix_str = "[CLC] query " if const_expr(self.use_clc_scheduler) else "[DYNAMIC] info "
+                fa_printf(
+                    3,
+                    prefix_str + "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()
 
     def load_Q(
         self,
@@ -2712,6 +3197,39 @@ class FlashAttentionForwardSm100:
         pipeline_q.producer_acquire_w_index_phase(stage, phase)
         load_Q_fn(src_idx=block, dst_idx=stage, tma_bar_ptr=pipeline_q.sync_object_full.get_barrier(stage))
 
+    def load_Q_non_tma(
+        self,
+        mQ: cute.Tensor,
+        sQ: cute.Tensor,
+        gmem_tiled_copy_Q: cute.TiledCopy,
+        pipeline_q: pipeline.PipelineAsync,
+        tidx: Int32,
+        seqlen_q: Int32,
+        m_block: Int32,
+        block: Int32,
+        stage: int,
+        phase: Int32,
+    ):
+        assert self.cta_group_size == 1, "cta_group_size must be 1 for non-tma Q load"
+        pipeline_q.producer_acquire_w_index_phase(stage, phase)
+        pack_gqa = PackGQA(
+            self.m_block_size,
+            self.head_dim_padded,
+            self.check_hdim_oob,
+            self.qhead_per_kvhead,
+        )
+        sQ_stage = sQ[None, None, None, stage]
+        sQ_pi = cute.make_tensor(
+            sQ_stage.iterator,
+            cute.make_layout(
+                (sQ_stage.shape[0][0], (sQ_stage.shape[0][1], sQ_stage.shape[2])),
+                stride=(sQ_stage.stride[0][0], (sQ_stage.stride[0][1], sQ_stage.stride[2])),
+            ),
+        )
+        pack_gqa.load_Q(mQ, sQ_pi, gmem_tiled_copy_Q, tidx, m_block * self.q_stage + block, seqlen_q)
+        cute.arch.cp_async_commit_group()
+        pipeline_q.sync_object_full.arrive_cp_async_mbarrier(stage)
+
     @cute.jit
     def load_KV(
         self,
@@ -2754,7 +3272,10 @@ class FlashAttentionForwardSm100:
         else:
             assert paged_kv_manager is not None
             assert extra_tx_count is None
-            paged_kv_manager.load_KV(block, sX[None, None, None, stage], K_or_V)
+            sX_cur = sX[None, None, None, stage]
+            if const_expr(self.uneven_kv_smem):
+                sX_cur = self.offset_kv_smem(sX_cur, stage, phase ^ 1)
+            paged_kv_manager.load_KV(block, sX_cur, K_or_V)
             cute.arch.cp_async_commit_group()
             pipeline_kv.sync_object_full.arrive_cp_async_mbarrier(stage)
 
@@ -2765,6 +3286,9 @@ class FlashAttentionForwardSm100:
             # (smem_large + smem_small) // 2. So for stage == 1, move right by offset if
             # phase == 0, or left by offset if phase == 1.
             offset = 0 if stage != 1 else self.uneven_kv_smem_offset * (1 - 2 * phase)
+            # Hint that the offset is 128-bit aligned so that
+            # ptr + offset preserves the alignment needed by cp.async.
+            offset = cute.assume(offset, divby=128 // self.k_dtype.width)
             return cute.make_tensor(sX.iterator + offset, sX.layout)
         else:
             return sX
@@ -2774,12 +3298,12 @@ class FlashAttentionForwardSm100:
     #     warp_group_idx = utils.canonical_warp_group_idx(sync=False)
     #     if warp_group_idx == 0:
     #         cute.arch.barrier_arrive(
-    #             barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1), number_of_threads=2 * 128,
+    #             barrier_id=int(NamedBarrierFwdSm100.WarpSchedulerWG1), number_of_threads=2 * 128,
     #         )
 
     # def warp_scheduler_barrier_sync(self):
     #     cute.arch.barrier(
-    #         barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1) + utils.canonical_warp_group_idx(sync=False),
+    #         barrier_id=int(NamedBarrierFwdSm100.WarpSchedulerWG1) + utils.canonical_warp_group_idx(sync=False),
     #         number_of_threads=2 * 128
     #     )
 
@@ -2787,7 +3311,7 @@ class FlashAttentionForwardSm100:
     #     cur_wg = utils.canonical_warp_group_idx(sync=False)
     #     next_wg = 1 - cur_wg
     #     cute.arch.barrier_arrive(
-    #         barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1) + next_wg, number_of_threads=2 * 128,
+    #         barrier_id=int(NamedBarrierFwdSm100.WarpSchedulerWG1) + next_wg, number_of_threads=2 * 128,
     #     )
 
     @cute.jit
@@ -2802,7 +3326,7 @@ class FlashAttentionForwardSm100:
         n_block,
         softmax,
         seqlen: SeqlenInfoQK,
-        aux_tensors=None,
+        aux_data: AuxData = AuxData(),
         fastdiv_mods=(None, None),
         head_divmod=None,
     ):
@@ -2825,7 +3349,7 @@ class FlashAttentionForwardSm100:
             q_idx_logical, head_offset = divmod(q_physical, head_divmod)
             head_idx = head_idx * self.qhead_per_kvhead + head_offset
 
-        if cutlass.const_expr(aux_tensors is not None):
+        if cutlass.const_expr(aux_data.tensors is not None):
             seqlen_q_divmod, _ = fastdiv_mods
             _, q_idx_logical = divmod(q_idx_logical, seqlen_q_divmod)
 
@@ -2836,11 +3360,26 @@ class FlashAttentionForwardSm100:
             batch_idx,
             head_idx,
             softmax.softmax_scale,
-            self.vec_size,
+            self.score_vec_size,
             self.qk_acc_dtype,
-            aux_tensors,
+            aux_data,
             fastdiv_mods,
             seqlen_info=seqlen,
             constant_q_idx=q_idx_logical,
             qhead_per_kvhead=self.qhead_per_kvhead if cutlass.const_expr(self.pack_gqa) else 1,
         )
+
+    @cute.jit
+    def process_work_tile(
+        self,
+        seqlen_info: SeqlenInfoQK,
+        n_block_min: Int32,
+        n_block_max: Int32,
+    ):
+        is_varlen_q = seqlen_info.has_cu_seqlens_q or seqlen_info.has_seqused_q
+        process_work_tile_k = const_expr(not self.is_split_kv) or n_block_min < n_block_max
+        if const_expr(is_varlen_q and not self.use_varlen_scheduler):
+            process_work_tile_q = seqlen_info.seqlen_q > 0
+        else:
+            process_work_tile_q = True
+        return process_work_tile_k and process_work_tile_q
diff --git a/build/torch-cuda/flash_fwd_sm120.py b/build/torch-cuda/flash_fwd_sm120.py
new file mode 100644
index 0000000000000000000000000000000000000000..31719122b212e0edddb0b30bcbbe1e33a6db2530
--- /dev/null
+++ b/build/torch-cuda/flash_fwd_sm120.py
@@ -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
diff --git a/build/torch-cuda/flash_fwd_sm90.py b/build/torch-cuda/flash_fwd_sm90.py
new file mode 100644
index 0000000000000000000000000000000000000000..3c9c8e26946df72fd2d15767b7e5839839577f26
--- /dev/null
+++ b/build/torch-cuda/flash_fwd_sm90.py
@@ -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,
+            )
diff --git a/build/torch-cuda/interface.py b/build/torch-cuda/interface.py
index cc1c9cc48cd62b171ebb1100f3b181297972f569..bb0147bee0b9e8fcf46ff3767cb8f3f7d8e68bcb 100644
--- a/build/torch-cuda/interface.py
+++ b/build/torch-cuda/interface.py
@@ -1,36 +1,20 @@
 # Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
 # [2025-07-04] Version in Cute-DSL, for Hopper and Blackwell. You'll need install nvidia-cutlass-dsl==4.2.0.
 
-# Supported features:
-# - BF16 & FP16 dtype
-# - noncausal & causal attention
-# - MHA, GQA, MQA
-# - hdim 64, 96, 128.
-# - (hdim_qk, hdim_v) = (192, 128) for Blackwell (i.e. DeepSeek shape)
-# - varlen
-# - sliding window
-# - bwd pass for Ampere (will also run on Hopper/Blackwell, but will be slow)
-
-# Features not supported yet:
-# - split (i.e. FlashDecoding)
-# - tuned block sizes
-# - paged KV
-# - append KV to existing KV cache
-# - FP8
-# - bwd pass optimized for Hopper/Blackwell
-
 import os
 import math
+from dataclasses import dataclass
 from functools import lru_cache
 from typing import Optional, Tuple, Callable
 
 import torch
 
 
-import cuda.bindings.driver as cuda
 
 import cutlass
 import cutlass.cute as cute
+from cutlass import Int32, Float32
+from .quack.compile_utils import make_fake_tensor as fake_tensor
 from .cache_utils import get_jit_cache
 from .testing import is_fake_mode
 
@@ -43,33 +27,233 @@ if os.environ.get("CUTE_DSL_PTXAS_PATH", None) is not None:
 
 
 from . import utils
+from . import fa_logging
 from .cute_dsl_utils import (
-    to_cute_tensor, to_cute_aux_tensor, get_aux_tensor_metadata, get_broadcast_dims,
+    get_aux_tensor_metadata,
+    get_broadcast_dims,
+    get_num_sms_for_selection,
+    maybe_contiguous,
+    to_cute_aux_tensor,
+    to_cute_tensor,
+    validate_output_layout,
 )
-from .flash_fwd import FlashAttentionForwardSm90
-from .flash_fwd_sm100 import FlashAttentionForwardSm100
+from .flash_fwd import FlashAttentionForwardSm80
+from .flash_fwd_sm90 import FlashAttentionForwardSm90
+from .flash_fwd_sm100 import FlashAttentionForwardSm100, DescaleTensors
+from .flash_fwd_sm120 import FlashAttentionForwardSm120
 from .flash_bwd_preprocess import FlashAttentionBackwardPreprocess
 from .flash_bwd import FlashAttentionBackwardSm80
 from .flash_bwd_sm90 import FlashAttentionBackwardSm90
 from .flash_bwd_sm100 import FlashAttentionBackwardSm100
-from .flash_bwd_postprocess import FlashAttentionBackwardPostprocess
+from .flash_bwd_sm120 import FlashAttentionBackwardSm120
+from .flash_bwd_postprocess import (
+    FlashAttentionBackwardPostprocess,
+    LearnableSinkBwdTensors,
+)
 from .flash_fwd_combine import FlashAttentionForwardCombine
-
+from .flash_fwd_mla_sm100 import FlashAttentionMLAForwardSm100
+from .prepare_scheduler import FlashPrepareScheduler, SchedulerMetadataTensorsTorch
+from .cu_blocks_kernel import CuSeqlensToBlocksKernel, CuBlocksToBatchKernel
+from .flash_bwd_mla_sm100 import FlashAttentionSparseMLABackwardSm100
+from .flash_bwd_mla_dq_dqv_sm100 import dQdQvGemmKernel
+from .flash_bwd_mla_dk_sm100 import dKGemmKernel
+
+# SM100 head_dim=256 2CTA kernel imports
+from .sm100_hd256_2cta_fmha_forward import BlackwellFusedMultiHeadAttentionForward
+from .sm100_hd256_2cta_fmha_backward import BlackwellFusedMultiHeadAttentionBackward
+
+from .utils import AuxData
 from .block_sparsity import (
     BlockSparseTensorsTorch,
+    block_sparse_bwd_supports_2cta,
+    get_kv_subtile_factor,
+    get_sparse_q_block_size,
     to_cute_block_sparse_tensors,
     normalize_block_sparse_config,
     normalize_block_sparse_config_bwd,
 )
 
+BIN_BATCH_SEARCH_THRESH = 256  # above this batch size SingleTileVarlenScheduler gets a batch-lookup aid
+# Where the cu hint applies, use an O(1) flat-block -> batch lookup instead of the binary search.
+USE_BLOCKS_TO_BATCH: bool = True
+
+
+def _parse_arch_str(arch_str):
+    """Parse arch string (e.g. 'sm_80', 'sm_90a', '80', '100') to int (e.g. 80, 90, 100)."""
+    import re
+    match = re.match(r"^(?:sm_?|SM_?)?(\d+)(\d)([af]?)$", arch_str)
+    if not match:
+        raise ValueError(f"Invalid arch format: {arch_str}")
+    major, minor, _ = match.groups()
+    return int(major) * 10 + int(minor)
+
+
 @lru_cache(maxsize=None)
 def _get_device_arch():
-    """Cached device arch check."""
+    """Cached device arch check.
+
+    Override with FLASH_ATTENTION_ARCH (e.g. 'sm_80' or '80') to select which
+    kernel path to use (SM80/SM90/SM100/SM120) independently of the compilation
+    target (CUTE_DSL_ARCH).
+
+    For CPU-only compilation (no GPU), set:
+      FLASH_ATTENTION_ARCH=sm_80  (kernel selection)
+      CUTE_DSL_ARCH=sm_80         (compilation target)
+      FLASH_ATTENTION_NUM_SMS=132 (target-SKU selector metadata)
+    """
+    arch_override = os.environ.get("FLASH_ATTENTION_ARCH", None)
+    if arch_override is not None:
+        return _parse_arch_str(arch_override)
     major, minor = torch.cuda.get_device_capability()
-    return major * 10 + minor
+    return major * 10 + int(minor)
+
+
+@lru_cache(maxsize=None)
+def _validate_head_dims(head_dim: int, head_dim_v: int, compute_capability: int, alignment: int) -> None:
+    """Validate head dimension constraints based on compute capability."""
+    is_deepseek_shape = head_dim == 192 and head_dim_v == 128
+    is_deepseek_mla_absorbed_shape = (head_dim == 64 or head_dim == head_dim_v) and head_dim_v == 512
+    is_dedicate_kernel_shape = head_dim == 256 and head_dim_v == 256
+    is_standard_range = 8 <= head_dim <= 128 and 8 <= head_dim_v <= 128
+
+    is_sm90_range = 8 <= head_dim <= 256 and 8 <= head_dim_v <= 256
+    if compute_capability == 9:
+        assert is_sm90_range and head_dim % alignment == 0 and head_dim_v % alignment == 0, (
+            f"(head_dim, head_dim_v)=({head_dim}, {head_dim_v}) is not supported on SM90. "
+            f"head_dim and head_dim_v must be between 8 and 256 and divisible by {alignment}."
+        )
+    elif compute_capability in [10, 11]:
+        assert (is_standard_range or is_deepseek_shape or is_deepseek_mla_absorbed_shape or is_dedicate_kernel_shape) and head_dim % alignment == 0 and head_dim_v % alignment == 0, (
+            f"(head_dim, head_dim_v)=({head_dim}, {head_dim_v}) is not supported on SM100/SM110. "
+            f"head_dim and head_dim_v must be between 8 and 128 and divisible by {alignment}, or (192, 128) for DeepSeek, or (256, 256) for hd256."
+        )
+
+
+@dataclass(frozen=True)
+class FwdConfig:
+    m_block_size: int
+    n_block_size: int
+    mma_pv_is_rs: bool
+    intra_wg_overlap: bool
+    q_stage: int = 1
+    num_splits: int = 1
+
+
+def _tile_size_fwd_sm90(head_dim, head_dim_v, is_causal, is_local, sparse_block_size_q=None):
+    """Return FwdConfig for SM90 forward.
+
+    Tile sizes and flags based on tile_size_fwd_sm90 in hopper/tile_size.h, adjusted
+    for the Python kernel's different register/smem tradeoffs (benchmarked on H100 SXM).
+
+    When sparse_block_size_q is set, tile_m must divide it. For head_dim <= 96 the
+    optimal tile_m=192 is used when compatible, otherwise we fall back to 128.
+    """
+    if head_dim <= 64:
+        # C++: 192×192 non-causal, 192×128 causal/local.
+        # Python: 192×128 RS+OL is consistently best across seqlens.
+        if sparse_block_size_q is not None and sparse_block_size_q % 192 != 0:
+            return FwdConfig(128, 128, True, True)
+        return FwdConfig(192, 128, True, True)
+    elif head_dim <= 96:
+        # C++: 192×144 noRS+OL for all cases.
+        # Python: RS is catastrophic with 192× tiles (~300 vs ~600 TFLOPS).
+        # noRS+OL is always required. Causal: 192×128 slightly better short seqlen.
+        if sparse_block_size_q is not None and sparse_block_size_q % 192 != 0:
+            return FwdConfig(128, 128, False, True)
+        if is_causal or is_local:
+            return FwdConfig(192, 128, False, True)
+        else:
+            return FwdConfig(192, 144, False, True)
+    elif head_dim <= 128:
+        return FwdConfig(128, 128, True, True)
+    elif head_dim <= 192:
+        tile_n = 96 if is_local else (128 if head_dim_v <= 128 else 112)
+        return FwdConfig(128, tile_n, True, True)
+    else:  # hdim 256
+        tile_n = 64 if is_local else 80
+        return FwdConfig(128, tile_n, True, True)
+
+@dataclass(frozen=True)
+class BwdConfig:
+    m_block_size: int
+    n_block_size: int
+    num_stages_Q: int
+    num_stages_dO: int
+    num_stages_PdS: int
+    SdP_swapAB: bool
+    dKV_swapAB: bool
+    dQ_swapAB: bool
+    AtomLayoutMSdP: int
+    AtomLayoutNdKV: int
+    AtomLayoutMdQ: int
+    num_wg: int = 2  # MMA warp groups (total threads = (num_wg + 1) * 128)
+    dQ_single_wg: bool = False
+
+
+def _tile_size_bwd_sm90(head_dim, head_dim_v, causal, local, sparse_block_size_q=None):
+    """Return BwdConfig for SM90.
+
+    Configs based on C++ FA3 hopper/flash_bwd_launch_template.h,
+    benchmarked on H100 SXM.
+    """
+    if head_dim <= 64:
+        # C++ FA3: 128, 128, 64, ..., 2, 2, true, false, false, 2, 1, 2, 2
+        return BwdConfig(
+            m_block_size=128, n_block_size=128,
+            num_stages_Q=2, num_stages_dO=2, num_stages_PdS=2,
+            SdP_swapAB=True, dKV_swapAB=False, dQ_swapAB=False,
+            AtomLayoutMSdP=1, AtomLayoutNdKV=2, AtomLayoutMdQ=2,
+        )
+    elif head_dim <= 96:
+        # C++ FA3: 64, 128, 96, dQ_swapAB=False
+        return BwdConfig(
+            m_block_size=64, n_block_size=128,
+            num_stages_Q=2, num_stages_dO=2, num_stages_PdS=2,
+            SdP_swapAB=True, dKV_swapAB=False, dQ_swapAB=False,
+            AtomLayoutMSdP=1, AtomLayoutNdKV=2, AtomLayoutMdQ=1,
+            dQ_single_wg=True,
+        )
+    elif head_dim <= 128:
+        # C++ FA3: causal/local: 64, 128; non-causal: 80, 128 with dQ_swapAB
+        is_causal_or_local = causal or local
+        m_block_size = 64 if is_causal_or_local else 80
+        if sparse_block_size_q is not None and sparse_block_size_q % m_block_size != 0:
+            m_block_size = 64
+        return BwdConfig(
+            m_block_size=m_block_size,
+            n_block_size=128,
+            num_stages_Q=2, num_stages_dO=2, num_stages_PdS=2,
+            SdP_swapAB=True, dKV_swapAB=False,
+            dQ_swapAB=m_block_size % 64 != 0,
+            AtomLayoutMSdP=1, AtomLayoutNdKV=2, AtomLayoutMdQ=1,
+        )
+    elif head_dim <= 192:
+        hdimv128 = head_dim_v <= 128
+        if hdimv128:
+            return BwdConfig(
+                m_block_size=64, n_block_size=96,
+                num_stages_Q=2, num_stages_dO=2, num_stages_PdS=1,
+                SdP_swapAB=False, dKV_swapAB=True, dQ_swapAB=False,
+                AtomLayoutMSdP=1, AtomLayoutNdKV=2, AtomLayoutMdQ=1,
+                num_wg=2,
+            )
+        else:
+            return BwdConfig(
+                m_block_size=64, n_block_size=96,
+                num_stages_Q=2, num_stages_dO=1, num_stages_PdS=1,
+                SdP_swapAB=False, dKV_swapAB=True, dQ_swapAB=False,
+                AtomLayoutMSdP=1, AtomLayoutNdKV=2, AtomLayoutMdQ=1,
+                num_wg=2,
+            )
+    else:
+        # hdim 256
+        return BwdConfig(
+            m_block_size=64, n_block_size=64,
+            num_stages_Q=1, num_stages_dO=1, num_stages_PdS=1,
+            SdP_swapAB=False, dKV_swapAB=False, dQ_swapAB=False,
+            AtomLayoutMSdP=1, AtomLayoutNdKV=1, AtomLayoutMdQ=1,
+        )
 
-def maybe_contiguous(x):
-    return x.contiguous() if x is not None and x.stride(-1) != 1 else x
 
 
 def _validate_tensor(t, name, expected_shape, expected_dtype, expected_device):
@@ -78,34 +262,278 @@ def _validate_tensor(t, name, expected_shape, expected_dtype, expected_device):
     assert t.device == expected_device, f"{name} device {t.device} != expected {expected_device}"
     assert t.is_cuda, f"{name} must be on CUDA"
 
-
 torch2cute_dtype_map = {
     torch.float16: cutlass.Float16,
     torch.bfloat16: cutlass.BFloat16,
     torch.float32: cutlass.Float32,
+    torch.float8_e4m3fn: cutlass.Float8E4M3FN,
+    torch.float8_e5m2: cutlass.Float8E5M2,
 }
 
+_LEARNABLE_SINK_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
+
 
 def num_splits_heuristic(total_mblocks, num_SMs, num_n_blocks, max_splits):
     # If num_n_blocks is too small, use 1 split. For example, we never split for hdim = 128 and seqlen_k = 512.
     if num_n_blocks <= 4:
         return 1
+    # Avoid ZeroDivisionError when batch_size or seqlen_q is 0. The empty-Q
+    # early-exit in _flash_attn_fwd handles correctness for those shapes; this
+    # guard just keeps the heuristic safe if called in other contexts.
+    if total_mblocks == 0:
+        return 1
 
     # NOTE: We should revisit this heuristic after persistence is supported for split KV.
     # Sometimes, it's ideal to over-schedule splits for better efficiency.
     return min(num_SMs // total_mblocks, max_splits, num_n_blocks)
 
 
+def _get_fwd_config(
+    *,
+    arch: int,
+    head_dim: int,
+    head_dim_v: int,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    num_head_kv: int,
+    qhead_per_kvhead: int,
+    pack_gqa: bool,
+    batch_size: int,
+    causal: bool,
+    local: bool,
+    window_size_left: Optional[int],
+    window_size_right: Optional[int],
+    num_splits: int,
+    device,
+    seqlen_q: Optional[int] = None,
+    tile_mn: Optional[Tuple[int, int]] = None,
+    block_sparse_tensors: Optional[BlockSparseTensorsTorch] = None,
+    mma_pv_is_rs: Optional[bool] = None,
+    intra_wg_overlap: Optional[bool] = None,
+) -> FwdConfig:
+    if seqlen_q is None:
+        seqlen_q = max_seqlen_q
+
+    # Base tile sizes and flags: explicit override, else per-arch heuristic.
+    cfg = FwdConfig(128, 128, True, True)
+    if tile_mn is None:
+        if arch // 10 == 12:
+            # SM120 tile sizes tuned for 99 KB SMEM capacity:
+            # D<=64:  128x128 → 48 KB (good occupancy)
+            # D>64:   128x64  → 64 KB (128x128 would use 96 KB, hurting occupancy)
+            if head_dim > 64:
+                cfg = FwdConfig(128, 64, True, True)
+        elif arch // 10 == 8:
+            cfg = FwdConfig(128, 64, True, True)  # SM80, should tune
+        elif arch // 10 == 9:
+            sparse_q = get_sparse_q_block_size(block_sparse_tensors, seqlen_q)
+            cfg = _tile_size_fwd_sm90(
+                head_dim, head_dim_v, causal, local, sparse_block_size_q=sparse_q
+            )
+    else:
+        cfg = FwdConfig(tile_mn[0], tile_mn[1], cfg.mma_pv_is_rs, cfg.intra_wg_overlap)
+
+    tile_m, tile_n = cfg.m_block_size, cfg.n_block_size
+    if mma_pv_is_rs is None:
+        mma_pv_is_rs = cfg.mma_pv_is_rs
+    if intra_wg_overlap is None:
+        intra_wg_overlap = cfg.intra_wg_overlap
+
+    seqlen_q_packgqa = max_seqlen_q * (qhead_per_kvhead if pack_gqa else 1)
+    if arch // 10 in [10, 11]:
+        q_stage = 2 if seqlen_q_packgqa > tile_m else 1
+    else:
+        q_stage = 1
+
+    m_block_size_effective = q_stage * tile_m
+    # Only None is unbounded; preserve 0 (e.g. the right bound of a causal window).
+    window_right_loaded = max_seqlen_k if window_size_right is None else window_size_right
+    window_left_loaded = max_seqlen_k if window_size_left is None else window_size_left
+    seqlen_k_loaded = (
+        max_seqlen_k
+        if not local
+        else max(
+            0,
+            min(
+                max_seqlen_k,
+                window_right_loaded
+                + window_left_loaded
+                + 1
+                + tile_m,
+            ),
+        )
+    )
+    num_m_blocks = (seqlen_q_packgqa + m_block_size_effective - 1) // m_block_size_effective
+    total_mblocks = batch_size * num_head_kv * num_m_blocks
+    num_n_blocks = (seqlen_k_loaded + tile_n - 1) // tile_n
+    num_SMs = None
+    if arch // 10 == 12:
+        assert num_splits == 1, "SM120 forward only supports num_splits=1"
+    elif num_splits < 1:
+        num_SMs = get_num_sms_for_selection(device.index, arch)
+        num_splits = num_splits_heuristic(total_mblocks, num_SMs, num_n_blocks, 128)
+
+    # SplitKV uses float32 partial output, which doubles the O buffer size
+    # in shared memory, causing OOM for diff-headdim (192, 128)
+    if arch // 10 in [10, 11] and head_dim != head_dim_v and num_splits > 1:
+        if num_n_blocks >= 64 and head_dim_v != 512:
+            tile_n = 64
+            num_n_blocks = (seqlen_k_loaded + tile_n - 1) // tile_n
+            if num_SMs is None:
+                num_SMs = get_num_sms_for_selection(device.index, arch)
+            num_splits = num_splits_heuristic(total_mblocks, num_SMs, num_n_blocks, 128)
+        else:
+            num_splits = 1
+
+    return FwdConfig(tile_m, tile_n, mma_pv_is_rs, intra_wg_overlap, q_stage, num_splits)
+
+
+def _resolve_causal_local_window(causal, window_size_left, window_size_right, mask_mod=None):
+    """Resolve causal/local/window settings into canonical form.
+
+    Returns (causal, local, window_size_left, window_size_right).
+    """
+    if mask_mod is not None:
+        return False, False, window_size_left, window_size_right
+    if (window_size_left is not None and window_size_right is not None
+            and window_size_left < 0 and window_size_right < 0):
+        window_size_left = None
+        window_size_right = None
+    if causal:
+        window_size_right = 0
+    if window_size_left is not None or window_size_right is not None:
+        if window_size_left is None and window_size_right == 0:
+            causal, local = True, False
+            window_size_right = None
+        else:
+            causal, local = False, True
+    else:
+        local = False
+    return causal, local, window_size_left, window_size_right
+
+
+def _compute_tile_cumsum(
+    *,
+    num_m_blocks: Optional[torch.Tensor] = None,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    seqused: Optional[torch.Tensor] = None,
+    num_splits_dynamic: Optional[torch.Tensor] = None,
+    virtual_batch_idx: Optional[torch.Tensor] = None,
+    tile_size: int = 1,
+    q_stage: int = 1,
+    cluster_shape_m: int = 1,
+    qhead_per_kvhead: int = 1,
+    pack_gqa: bool = False,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """(cu_total_m_blocks, cu_total_splits_m_blocks), int32, (num_batch + 1,).
+
+    cu_total_splits_m_blocks is None when num_splits_dynamic is None.
+    """
+    assert num_m_blocks is not None or cu_seqlens is not None or seqused is not None, (
+        "_compute_tile_cumsum requires num_m_blocks, cu_seqlens, or seqused"
+    )
+    if num_m_blocks is not None:
+        # num_m_blocks is already in tile_size units; feed it through the seqused slot.
+        seqused = num_m_blocks
+        tile = q_stage * cluster_shape_m
+        seqlen_q_multiplier = 1
+    else:
+        tile = tile_size * q_stage * cluster_shape_m
+        seqlen_q_multiplier = qhead_per_kvhead if pack_gqa and qhead_per_kvhead > 1 else 1
+    batch_size = seqused.shape[0] if seqused is not None else cu_seqlens.shape[0] - 1
+    device = seqused.device if seqused is not None else cu_seqlens.device
+    cu_total_m_blocks = torch.empty(batch_size + 1, dtype=torch.int32, device=device)
+    cu_total_splits_m_blocks = (
+        torch.empty(batch_size + 1, dtype=torch.int32, device=device)
+        if num_splits_dynamic is not None
+        else None
+    )
+    compile_key = (
+        tile,
+        seqlen_q_multiplier,
+        cu_seqlens is not None,
+        seqused is not None,
+        num_splits_dynamic is not None,
+        virtual_batch_idx is not None,
+    )
+    if compile_key not in _compute_tile_cumsum.compile_cache:
+        cute_tensors = [
+            to_cute_tensor(t, assumed_align=4, leading_dim=0) if t is not None else None
+            for t in (
+                cu_total_m_blocks,
+                cu_total_splits_m_blocks,
+                cu_seqlens,
+                seqused,
+                num_splits_dynamic,
+                virtual_batch_idx,
+            )
+        ]
+        _compute_tile_cumsum.compile_cache[compile_key] = cute.compile(
+            CuSeqlensToBlocksKernel(tile=tile, seqlen_q_multiplier=seqlen_q_multiplier),
+            *cute_tensors,
+            cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+            options="--enable-tvm-ffi",
+        )
+    if not is_fake_mode():
+        _compute_tile_cumsum.compile_cache[compile_key](
+            cu_total_m_blocks,
+            cu_total_splits_m_blocks,
+            cu_seqlens,
+            seqused,
+            num_splits_dynamic,
+            virtual_batch_idx,
+        )
+    return cu_total_m_blocks, cu_total_splits_m_blocks
+
+
+_compute_tile_cumsum.compile_cache = get_jit_cache("tile_cumsum")
+
+
+def _blocks_to_batch_size(total_q, num_batch, tile_m, qhead_per_kvhead, pack_gqa):
+    """Upper bound on number of m_blocks in a given varlen invocation"""
+    seqlen_mult = qhead_per_kvhead if pack_gqa and qhead_per_kvhead > 1 else 1
+    return (total_q * seqlen_mult + num_batch * (tile_m - 1)) // tile_m + 1
+
+
+def _compute_blocks_to_batch(cu_total_blocks, num_blocks, device):
+    """Inverted index of _compute_tile_cumsum: flat scheduler block -> batch, int32, (num_blocks,).
+
+    Blocks past the last batch's range map to batch_size (invalid).
+    """
+    blocks_to_batch = torch.empty(num_blocks, dtype=torch.int32, device=device)
+    compile_key = ()
+    if compile_key not in _compute_blocks_to_batch.compile_cache:
+        cu_total_blocks_tensor, blocks_to_batch_tensor = [
+            to_cute_tensor(t, assumed_align=4, leading_dim=0)
+            for t in (cu_total_blocks, blocks_to_batch)
+        ]
+        _compute_blocks_to_batch.compile_cache[compile_key] = cute.compile(
+            CuBlocksToBatchKernel(),
+            cu_total_blocks_tensor,
+            blocks_to_batch_tensor,
+            cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+            options="--enable-tvm-ffi",
+        )
+    if not is_fake_mode():
+        _compute_blocks_to_batch.compile_cache[compile_key](cu_total_blocks, blocks_to_batch)
+    return blocks_to_batch
+
+
+_compute_blocks_to_batch.compile_cache = get_jit_cache("blocks_to_batch")
+
+
 def _flash_attn_fwd(
-    q: torch.Tensor,
-    k: torch.Tensor,
+    q: Optional[torch.Tensor],
+    k: Optional[torch.Tensor],
     v: torch.Tensor,
+    qv: Optional[torch.Tensor] = 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,
     max_seqlen_q: Optional[int] = None,
     max_seqlen_k: Optional[int] = None,
+    min_seqlen_k: Optional[int] = None,
     page_table: Optional[torch.Tensor] = None,
     softmax_scale: Optional[float] = None,
     causal: bool = False,
@@ -113,11 +541,9 @@ def _flash_attn_fwd(
     window_size_left: Optional[int] = None,
     window_size_right: Optional[int] = None,
     learnable_sink: Optional[torch.Tensor] = None,
-    # m_block_size: int = 128,
-    # n_block_size: int = 64,
-    # num_threads: int = 128,
-    m_block_size: int = 128,
-    n_block_size: int = 128,
+    tile_mn: Optional[Tuple[int, int]] = None,
+    mma_pv_is_rs: Optional[bool] = None,
+    intra_wg_overlap: Optional[bool] = None,
     num_threads: int = 384,
     num_splits: int = 1,
     pack_gqa: Optional[bool] = None,
@@ -129,7 +555,15 @@ def _flash_attn_fwd(
     out: Optional[torch.Tensor] = None,
     lse: Optional[torch.Tensor] = None,
     aux_tensors: Optional[list[torch.Tensor]] = None,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+    aux_scalars: Optional[tuple] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+    gather_kv_indices: Optional[torch.Tensor] = None,
+    scheduler_metadata: Optional[SchedulerMetadataTensorsTorch] = None,
+    seqlen_k_per_split: Optional[int] = None,
+    disable_scheduler_metadata: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
     """Forward pass for FlashAttention.
 
     Args:
@@ -138,42 +572,57 @@ def _flash_attn_fwd(
         mask_mod: A callable that takes token position information and selectively masks
         block_sparse_tensors: A tuple of tensors used for block sparsity.
         return_lse: Whether to return the log softmax of the attention scores. If set to True will always calculate
-            Note: the returned LSE currently does not support taking gradient.
+            The returned LSE supports taking gradient.
         out: Optional pre-allocated output tensor. If None, will be allocated internally.
         lse: Optional pre-allocated log-sum-exp tensor. If None, will be allocated when needed.
         aux_tensors: Some score_mods will want to read from global aux_tensors. This is how we thread them through to the inner kernel.
+        aux_scalars: Runtime scalar captures used by score_mod or mask_mod.
     """
-    q, k, v = [maybe_contiguous(t) for t in (q, k, v)]
-    num_head, head_dim = q.shape[-2:]
+    aux_scalars = tuple(aux_scalars) if aux_scalars else None
+    requires_grad = any(
+        t is not None and t.requires_grad for t in (q, k, v, qv, learnable_sink)
+    )
+    fake_mode = is_fake_mode()
+    q, k, v, qv = [maybe_contiguous(t) for t in (q, k, v, qv)]
+    assert q is not None or qv is not None
+    assert v is not None
+    q_descale, k_descale, v_descale = [
+        maybe_contiguous(t, align_bytes=4) for t in (q_descale, k_descale, v_descale)
+    ]
+    page_table = maybe_contiguous(page_table, align_bytes=4)
+    learnable_sink = maybe_contiguous(learnable_sink, align_bytes=4)
+    gather_kv_indices = maybe_contiguous(gather_kv_indices, align_bytes=16)
+    q_shape = q.shape if q is not None else qv.shape
+    num_head, head_dim = q_shape[-2:]
     if cu_seqlens_q is None:
-        batch_size, seqlen_q = q.shape[:2]
+        batch_size, seqlen_q = q_shape[:2]
         total_q = batch_size * seqlen_q
     else:
         batch_size = cu_seqlens_q.shape[0] - 1
         seqlen_q = None
-        total_q = q.shape[0]
+        total_q = q_shape[0]
     if page_table is not None:
         assert cu_seqlens_k is None, "page_table is not supported with cu_seqlens_k"
         assert page_table.dtype == torch.int32, "page_table must be int32"
         assert page_table.stride(-1) == 1, "page_table must be contiguous in the last dimension"
         max_num_pages_per_seq = page_table.shape[1]
         assert page_table.shape == (batch_size, max_num_pages_per_seq)
-        num_pages, page_size = k.shape[:2]
+        num_pages, page_size = v.shape[:2]
         seqlen_k = num_pages * page_size
     else:
         num_pages, page_size = None, None
-        seqlen_k = k.shape[-3]
-    num_head_kv = k.shape[-2]
+        seqlen_k = v.shape[-3]
+    num_head_kv = v.shape[-2]
     head_dim_v = v.shape[-1]
     if cu_seqlens_k is None:
         if page_table is None:
-            assert k.shape == (batch_size, seqlen_k, num_head_kv, head_dim)
+            assert k is None or k.shape == (batch_size, seqlen_k, num_head_kv, head_dim)
             assert v.shape == (batch_size, seqlen_k, num_head_kv, head_dim_v)
         else:
-            assert k.shape == (num_pages, page_size, num_head_kv, head_dim)
+            assert k is None or k.shape == (num_pages, page_size, num_head_kv, head_dim)
             assert v.shape == (num_pages, page_size, num_head_kv, head_dim_v)
     else:
-        assert k.shape == (seqlen_k, num_head_kv, head_dim)
+        assert k is None or k.shape == (seqlen_k, num_head_kv, head_dim)
         assert v.shape == (seqlen_k, num_head_kv, head_dim_v)
         assert cu_seqlens_k.shape == (batch_size + 1,), (
             "cu_seqlens_k must have shape (batch_size + 1,)"
@@ -189,8 +638,16 @@ def _flash_attn_fwd(
     assert seqused_k is None or seqused_k.shape == (batch_size,), (
         "seqused_k must have shape (batch_size,)"
     )
-    assert q.dtype in [torch.float16, torch.bfloat16], "inputs must be float16 or bfloat16"
-    assert q.dtype == k.dtype == v.dtype, "inputs must have the same dtype"
+    assert v.dtype in [torch.float16, torch.bfloat16, torch.float8_e4m3fn, torch.float8_e5m2], (
+        "inputs must be float16, bfloat16, fp8 e4m3fn, or fp8 e5m2"
+    )
+    
+    assert all(t is None or t.dtype == v.dtype for t in (q, k, qv)), (
+        "q, k, v, and qv must have the same dtype"
+    )
+
+    q_dtype = q.dtype if q is not None else qv.dtype
+
     for t in [cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k]:
         if t is not None:
             assert t.dtype == torch.int32, (
@@ -201,47 +658,72 @@ def _flash_attn_fwd(
             )
     if learnable_sink is not None:
         assert learnable_sink.shape == (num_head,)
-        assert learnable_sink.dtype == torch.bfloat16, "learnable_sink must be bfloat16"
-
-    assert all(
-        t is None or t.is_cuda
-        for t in (
-            q,
-            k,
-            v,
-            cu_seqlens_q,
-            cu_seqlens_k,
-            seqused_q,
-            seqused_k,
-            page_table,
-            learnable_sink,
+        assert learnable_sink.dtype in _LEARNABLE_SINK_DTYPES, (
+            "learnable_sink must be float16, bfloat16, or float32"
         )
-    ), "inputs must be on CUDA device"
+
+    if not fake_mode:
+        assert all(
+            t is None or t.is_cuda
+            for t in (
+                q,
+                k,
+                v,
+                qv,
+                q_descale,
+                k_descale,
+                v_descale,
+                cu_seqlens_q,
+                cu_seqlens_k,
+                seqused_q,
+                seqused_k,
+                page_table,
+                learnable_sink,
+            )
+        ), "inputs must be on CUDA device"
+    arch = _get_device_arch() if _arch is None else _arch
+    assert arch // 10 in [8, 9, 10, 11, 12], "Unsupported compute capability. Supported: 8.x, 9.x, 10.x, 11.x, 12.x"
     assert num_head % num_head_kv == 0, "num_head must be divisible by num_head_kv"
-    assert head_dim <= 256, "head_dim must be less than or equal to 256"
-    alignment = 16 // q.element_size()
-    assert head_dim % alignment == 0, f"head_dim must be divisible by {alignment}"
-    assert head_dim_v % alignment == 0, f"head_dim_v must be divisible by {alignment}"
+    alignment = 16 // v.element_size()
+    if arch // 10 not in [8, 12]:
+        _validate_head_dims(head_dim, head_dim_v, arch // 10, alignment)
     if softmax_scale is None:
-        softmax_scale = 1.0 / math.sqrt(head_dim)
+        softmax_scale = (
+            1.0 / math.sqrt(head_dim) if qv is None or q is None
+            else 1.0 / math.sqrt(head_dim + head_dim_v)
+        )
     if softcap == 0.0:
         softcap = None
     qhead_per_kvhead = num_head // num_head_kv
     if pack_gqa is None:
         pack_gqa = qhead_per_kvhead > 1
 
-    out_torch_dtype = q.dtype
-    device = q.device
+    is_fp8 = v.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
+    if is_fp8 and requires_grad:
+        raise NotImplementedError("FA4 CuTe FP8 backward is not supported yet (forward-only).")
+    out_torch_dtype = torch.bfloat16 if is_fp8 else q_dtype
+    device = v.device
     q_batch_seqlen_shape = (batch_size, seqlen_q) if cu_seqlens_q is None else (total_q,)
-    lse_shape = (batch_size, num_head, seqlen_q) if cu_seqlens_q is None else (num_head, total_q)
-    requires_grad = q.requires_grad or k.requires_grad or v.requires_grad
+
+    if qv is None:
+        lse_shape = (batch_size, num_head, seqlen_q) if cu_seqlens_q is None else (num_head, total_q)
+    else:
+        # num_head contiguous better for MQA in MLA absorbed
+        lse_shape = (batch_size, seqlen_q, num_head) if cu_seqlens_q is None else (total_q, num_head)
 
     if out is None:
         out = torch.empty(
             *q_batch_seqlen_shape, num_head, head_dim_v, dtype=out_torch_dtype, device=device
         )
     else:
-        _validate_tensor(out, "out", (*q_batch_seqlen_shape, num_head, head_dim_v), out_torch_dtype, device)
+        _validate_tensor(
+            out,
+            "out",
+            (*q_batch_seqlen_shape, num_head, head_dim_v),
+            out_torch_dtype,
+            device,
+        )
+        validate_output_layout(out, "out", align_bytes=16)
 
     if lse is None:
         lse = (
@@ -251,81 +733,126 @@ def _flash_attn_fwd(
         )
     elif lse is not None:
         _validate_tensor(lse, "lse", lse_shape, torch.float32, device)
+        validate_output_layout(lse, "lse", align_bytes=4)
 
-    dtype = torch2cute_dtype_map[q.dtype]
-    arch = _get_device_arch() if _arch is None else _arch
-
-    assert arch // 10 in [9, 10, 11], "Unsupported compute capability. Supported: 9.x, 10.x, 11.x"
-
-    use_block_sparsity = block_sparse_tensors is not None
-
-    if mask_mod is None:
-        if causal:
-            window_size_right = 0
-        if window_size_left is not None and window_size_right is not None and window_size_left + window_size_right < 0:
-            window_size_left = None
-            window_size_right = None
-        local = window_size_left is not None or window_size_right is not None
-        if window_size_left is not None or window_size_right is not None:
-            if window_size_left is None and window_size_right == 0:
-                causal, local = True, False
-                window_size_right = None
+    if seqlen_k == 0 or total_q == 0:
+        out.zero_()
+        if lse is not None:
+            if learnable_sink is None:
+                lse.fill_(float("-inf"))
             else:
-                causal, local = False, True
+                assert qv is None
+                lse.copy_(
+                    learnable_sink[None, :, None]
+                    if cu_seqlens_q is None
+                    else learnable_sink[:, None]
+                )
+        return out, lse, None, None
+
+    if is_fp8:
+        for t, name in ((q_descale, "q_descale"), (k_descale, "k_descale"), (v_descale, "v_descale")):
+            if t is not None:
+                _validate_tensor(
+                    t,
+                    name,
+                    (batch_size, num_head_kv),
+                    torch.float32,
+                    device,
+                )
     else:
-        causal, local = False, False
+        assert q_descale is None and k_descale is None and v_descale is None, (
+            "q_descale/k_descale/v_descale are only supported for FP8 inputs"
+        )
+
+    dtype = torch2cute_dtype_map[q_dtype]
+    if is_fp8:
+        assert arch // 10 == 10, "FP8 is only supported on SM100 (compute capability 10.x) for FA4 CuTe."
+    use_block_sparsity = block_sparse_tensors is not None
 
-    current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+    causal, local, window_size_left, window_size_right = _resolve_causal_local_window(
+        causal, window_size_left, window_size_right, mask_mod
+    )
 
-    if arch // 10 == 9:  # TODO: tune block size according to hdim.
-        if head_dim == head_dim_v == 128 and not causal and not local and not use_block_sparsity:
-            n_block_size = 192
+    requested_use_clc_scheduler = utils._get_use_clc_scheduler_default()
+    requested_disable_2cta = utils._get_disable_2cta_default(is_fwd=True)
 
-    if arch // 10 in [10, 11]:
-        if (
-            pack_gqa
-            and (128 % qhead_per_kvhead != 0)
-        ):
-            pack_gqa = False
-        # TODO: fix GQA + SplitKV + non-varlen
-        if pack_gqa and num_splits != 1 and cu_seqlens_q is None:
-            pack_gqa = False
+    # SM80/SM120: uses SM80 MMA, 128 threads (4 warps)
+    if arch // 10 in [8, 12]:
+        num_threads = 128
 
     if max_seqlen_q is None:
         max_seqlen_q = seqlen_q if cu_seqlens_q is None else total_q
     if max_seqlen_k is None:
         max_seqlen_k = seqlen_k
-    seqlen_q_packgqa = max_seqlen_q * qhead_per_kvhead
-    if arch // 10 == 10:
-        q_stage = 2 if seqlen_q_packgqa > m_block_size else 1
-    else:
-        q_stage = 1
+    if cu_seqlens_k is None and seqused_k is None:
+        min_seqlen_k = seqlen_k
+
+    fwd_cfg = _get_fwd_config(
+        arch=arch,
+        head_dim=head_dim,
+        head_dim_v=head_dim_v,
+        causal=causal,
+        local=local,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        qhead_per_kvhead=qhead_per_kvhead,
+        pack_gqa=pack_gqa,
+        batch_size=batch_size,
+        num_head_kv=num_head_kv,
+        num_splits=num_splits,
+        device=device,
+        seqlen_q=seqlen_q,
+        tile_mn=tile_mn,
+        block_sparse_tensors=block_sparse_tensors,
+        mma_pv_is_rs=mma_pv_is_rs,
+        intra_wg_overlap=intra_wg_overlap,
+    )
+    tile_m, tile_n = fwd_cfg.m_block_size, fwd_cfg.n_block_size
+    q_stage = fwd_cfg.q_stage
+    num_splits = fwd_cfg.num_splits
+    mma_pv_is_rs = fwd_cfg.mma_pv_is_rs
+    intra_wg_overlap = fwd_cfg.intra_wg_overlap
 
-    if num_splits < 1:
-        m_block_size_effective = q_stage * m_block_size
-        seqlen_k_loaded = max_seqlen_k if not local else max(0, min(max_seqlen_k, window_size_right + window_size_left + 1 + m_block_size))
-        num_n_blocks = (seqlen_k_loaded + n_block_size - 1) // n_block_size
-        num_m_blocks = (seqlen_q_packgqa + m_block_size_effective - 1) // m_block_size_effective
-        total_mblocks = batch_size * num_head_kv * num_m_blocks
-        num_splits = num_splits_heuristic(
-            total_mblocks,
-            torch.cuda.get_device_properties(device).multi_processor_count,
-            num_n_blocks,
-            128,
-        )
+    seqlen_q_packgqa = max_seqlen_q * (qhead_per_kvhead if pack_gqa else 1)
+    max_m_blocks_leq_one = seqlen_q_packgqa <= q_stage * tile_m
 
     is_split_kv = num_splits > 1
     if is_split_kv:
         out_partial = torch.empty(num_splits, *q_batch_seqlen_shape, num_head, head_dim_v, dtype=torch.float32, device=device)
         lse_partial = torch.empty(num_splits, *lse_shape, dtype=torch.float32, device=device)
 
-    # hash score and mask mods for compile cache
-    score_mod_hash = utils.hash_callable(score_mod) if score_mod is not None else False
-    mask_mod_hash = utils.hash_callable(mask_mod) if mask_mod is not None else False
+    use_2cta_instrs = (
+        arch // 10 in [10, 11]
+        and not requested_disable_2cta
+        and not causal
+        and not local
+        and not is_split_kv
+        and cu_seqlens_q is None
+        and seqused_q is None
+        and not use_block_sparsity
+        and page_size in [None, 128]
+        and int(math.ceil(head_dim / 16) * 16) in [128, 192]
+        and int(math.ceil(head_dim_v / 16) * 16) == 128
+        and seqlen_q_packgqa > 2 * tile_m
+        and (tile_m % qhead_per_kvhead == 0 or not pack_gqa)
+    )
+
+    # hd=256 2CTA forward uses dedicated kernel (Blackwell family)
+    use_dedicated_hd256_kernel = arch // 10 in [10, 11] and head_dim == 256 and head_dim_v == 256
+    use_2cta_instrs = use_2cta_instrs or use_dedicated_hd256_kernel
 
     if softcap is not None:
         assert score_mod is None, "softcap and score_mod cannot be used together"
         score_mod = utils.create_softcap_scoremod(softcap)
+    elif score_mod is not None:
+        if arch // 10 == 8:
+            raise NotImplementedError("Custom user-provided score_mod is not supported on SM8x architectures.")
+        
+    # hash score and mask mods for compile cache
+    score_mod_hash = utils.hash_callable(score_mod) if score_mod is not None else False
+    mask_mod_hash = utils.hash_callable(mask_mod) if mask_mod is not None else False
 
     is_varlen = (
         cu_seqlens_q is not None
@@ -334,50 +861,249 @@ def _flash_attn_fwd(
         or seqused_k is not None
     )
 
-    if mask_mod is not None:
-        if is_varlen:
-            raise NotImplementedError(
-                "mask_mod with aux_tensors is not yet supported for varlen sequences. This will be fixed in a future PR."
-            )
+    # CLC regressed for varlen MHA and dense noncausal. Imbalanced varlen shapes
+    # keep more K/V blocks in flight and hurt L2; dense noncausal mostly just
+    # pays work-stealing overhead.
+    is_varlen_mha = is_varlen and qhead_per_kvhead == 1
+    is_dense_noncausal = not is_varlen and not causal and not local
+    use_clc_scheduler = requested_use_clc_scheduler and not is_varlen_mha and not is_dense_noncausal
 
     if use_block_sparsity:
-        if is_varlen:
-            raise NotImplementedError(
-                "Block sparsity is not yet supported for varlen sequences. This will be fixed in a future PR."
-            )
         # NB: pack_gqa requires block sparse head dim == 1 (broadcasted)
-        if pack_gqa and block_sparse_tensors.mask_block_cnt.shape[1] != 1:
+        head_dim_idx = 0 if block_sparse_tensors.mask_block_cnt.ndim == 2 else 1
+        if pack_gqa and block_sparse_tensors.mask_block_cnt.shape[head_dim_idx] != 1:
             pack_gqa = False
-        if is_split_kv:
-            raise NotImplementedError(
-                "Block sparsity is not yet supported with SplitKV. TODO: partition sparse block lists per split."
+        if cu_seqlens_q is not None:
+            assert block_sparse_tensors.cu_total_m_blocks is not None, (
+                "Varlen block sparsity requires block_sparse_tensors.cu_total_m_blocks."
             )
+            if (
+                block_sparse_tensors.cu_block_idx_offsets is None
+                and (cu_seqlens_k is not None or seqused_k is not None)
+            ):
+                raise ValueError(
+                    "Varlen block sparsity with cu_seqlens_k or seqused_k requires "
+                    "block_sparse_tensors.cu_block_idx_offsets."
+                )
 
     # See get_broadcast_dims for why this is needed in compile key
     block_sparse_broadcast_pattern = None
     normalized_block_sparse_tensors = None
-    q_subtile_factor = None
+    q_subtile_factor = 1
+    kv_subtile_factor = 1
     if block_sparse_tensors is not None:
-        if seqlen_q is None:
-            raise ValueError("Block sparsity requires fixed-length sequences (seqlen_q must be known).")
-        (
-            normalized_block_sparse_tensors,
-            block_sparse_broadcast_pattern,
-            q_subtile_factor,
-        ) = normalize_block_sparse_config(
+        block_sparse_config = normalize_block_sparse_config(
             block_sparse_tensors,
             batch_size=batch_size,
             num_head=num_head,
             seqlen_q=seqlen_q,
             seqlen_k=seqlen_k,
-            block_size=(m_block_size, n_block_size),
+            block_size=(tile_m, tile_n),
             q_stage=q_stage,
+            allow_kv_subtile=arch // 10 in [10, 11],
         )
-    if aux_tensors is not None:    
+        normalized_block_sparse_tensors = block_sparse_config.tensors
+        block_sparse_broadcast_pattern = block_sparse_config.broadcast_pattern
+        q_subtile_factor = block_sparse_config.q_subtile_factor
+        kv_subtile_factor = block_sparse_config.kv_subtile_factor
+    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
+
+    if qv is not None:
+        assert arch // 10 in [10, 11], "only support Blackwell arch with qv"
+        assert q is None or qv.shape[:-1] == q.shape[:-1]
+        assert qv.shape[-1] == head_dim_v
+        assert head_dim_v == 512
+        assert q is None or head_dim == 64
+        assert not local, "local not yet supported with qv"
+        assert q_descale is None and k_descale is None and v_descale is None, (
+            "q_descale/k_descale/v_descale are not yet supported with qv"
+        )
+        assert tile_n == 128
+
+        assert not is_split_kv, "split kv not supported with qv"
+        assert learnable_sink is None
+        assert softcap is None
+        assert score_mod is None
+        assert mask_mod is None
+
+        if page_table is not None:
+            assert gather_kv_indices is None, "paged KV + topk sparsity not yet supported together"
+        
+        qv = maybe_contiguous(qv)
+
+        gather_kv_length = 2048  # dummy value
+        sparse_kv = gather_kv_indices is not None
+        # always use kv bitmask by default (handles -1 sentinel)
+        disable_sparse_kv_bitmask = False
+        if sparse_kv:
+            assert gather_kv_indices.shape[:-1] == qv.shape[:-2]
+            gather_kv_length = gather_kv_indices.shape[-1]
+            assert gather_kv_length % 128 == 0
+            # if min_seqlen_k is None or causal:
+            #     disable_sparse_kv_bitmask = False
+            # else:
+            #     # seqlen_k_boundary = min_seqlen_k - max_seqlen_q + 1 if causal else min_seqlen_k
+            #     seqlen_k_boundary = min_seqlen_k
+            #     disable_sparse_kv_bitmask = seqlen_k_boundary >= gather_kv_length
+        
+        if requires_grad and sparse_kv:
+            if cu_seqlens_q is None:
+                p = torch.empty(batch_size, seqlen_q, num_head, gather_kv_length, dtype=q_dtype, device=device)
+                row_max = torch.empty(batch_size, seqlen_q, gather_kv_length//128, num_head, dtype=torch.float32, device=device)
+            else:
+                p = torch.empty(total_q, num_head, gather_kv_length, dtype=q_dtype, device=device)
+                row_max = torch.empty(total_q, gather_kv_length//128, num_head, dtype=torch.float32, device=device)
+        else:
+            p = row_max = None
+    else:
+        assert gather_kv_indices is None, "gather_kv_indices is only supported with qv"
+        gather_kv_length = None
+        sparse_kv = None
+        disable_sparse_kv_bitmask = None
+        p = row_max = None
+
+
+    reuse_scheduler_metadata = scheduler_metadata is not None
+    is_varlen_q = cu_seqlens_q is not None or seqused_q is not None
+    cluster_shape_m = 2 if use_2cta_instrs else 1
+    if use_dedicated_hd256_kernel:
+        # The hd=256 2CTA fwd kernel does not support the dynamic-persistent scheduler.
+        scheduler_metadata = None
+        reuse_scheduler_metadata = False
+    if (
+        is_split_kv
+        and is_varlen_q
+        and scheduler_metadata is None
+        and not disable_scheduler_metadata
+        and not use_dedicated_hd256_kernel
+    ):
+        scheduler_metadata = _get_scheduler_metadata(
+            num_batch=batch_size,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            nheads=num_head,
+            nheads_kv=num_head_kv,
+            headdim=head_dim,
+            num_splits=num_splits,
+            tile_m=tile_m,
+            tile_n=tile_n,
+            headdim_v=head_dim_v,
+            pack_gqa=pack_gqa,
+            causal=causal,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            seqused_q=seqused_q,
+            seqused_k=seqused_k,
+            seqlen_k_per_split=seqlen_k_per_split,
+            q_stage=q_stage,
+            cluster_shape_m=cluster_shape_m,
+            total_q=total_q if cu_seqlens_q is not None else None,
+            use_clc_scheduler=use_clc_scheduler,
+        )
+
+    has_scheduler_metadata = scheduler_metadata is not None and not disable_scheduler_metadata
+    if has_scheduler_metadata:
+        num_m_blocks = scheduler_metadata.num_m_blocks_ptr
+        num_splits_dynamic = scheduler_metadata.num_splits_dynamic_ptr
+        virtual_batch_idx = scheduler_metadata.virtual_batch_idx_ptr
+        num_nheads_in_l2 = scheduler_metadata.num_nheads_in_l2_ptr
+        tile_count_semaphore = scheduler_metadata.tile_count_semaphore
+        assert all(
+            t is None or t.is_cuda
+            for t in scheduler_metadata
+        ), "scheduler metadata must be on CUDA device"
+        assert all(
+            t is None or t.shape == (batch_size,)
+            for t in (
+                num_m_blocks,
+                num_splits_dynamic,
+                virtual_batch_idx,
+                num_nheads_in_l2,
+            )
+        ), "these scheduler metadata tensors must have shape (batch_size,)"
+        if tile_count_semaphore is not None:
+            assert tile_count_semaphore.shape == (1,), "semaphore must have size 1"
+    else:
+        num_m_blocks = None
+        num_splits_dynamic = None
+        virtual_batch_idx = None
+        num_nheads_in_l2 = None
+        tile_count_semaphore = None
+
+    # use binary batch search in SingleTileVarlenScheduler to avoid
+    # O(N^2) lookup; observed to be faster only for batch_size > BIN_BATCH_SEARCH_THRESH; this is tunable
+    cu_total_m_blocks = None
+    cu_total_splits_m_blocks = None
+    blocks_to_batch_idx = None
+    use_single_tile_varlen_scheduler = tile_count_semaphore is None
+    use_cu_hint = (
+        is_varlen_q
+        and use_single_tile_varlen_scheduler
+        and batch_size > BIN_BATCH_SEARCH_THRESH
+        and not use_dedicated_hd256_kernel
+    )
+    if (
+        use_cu_hint
+        and has_scheduler_metadata
+        and scheduler_metadata.cu_total_m_blocks is not None
+    ):
+        cu_total_m_blocks = scheduler_metadata.cu_total_m_blocks
+        cu_total_splits_m_blocks = scheduler_metadata.cu_total_splits_m_blocks
+        blocks_to_batch_idx = scheduler_metadata.blocks_to_batch_idx
+    elif use_cu_hint:
+        cu_total_m_blocks, cu_total_splits_m_blocks = _compute_tile_cumsum(
+            num_m_blocks=num_m_blocks,
+            cu_seqlens=cu_seqlens_q,
+            seqused=seqused_q,
+            num_splits_dynamic=num_splits_dynamic,
+            virtual_batch_idx=virtual_batch_idx,
+            tile_size=tile_m,
+            q_stage=q_stage,
+            cluster_shape_m=cluster_shape_m,
+            qhead_per_kvhead=qhead_per_kvhead,
+            pack_gqa=pack_gqa,
+        )
+    if blocks_to_batch_idx is None and USE_BLOCKS_TO_BATCH and cu_total_m_blocks is not None:
+        blocks_to_batch_idx = _compute_blocks_to_batch(
+            cu_total_m_blocks,
+            _blocks_to_batch_size(total_q, batch_size, tile_m, qhead_per_kvhead, pack_gqa),
+            cu_total_m_blocks.device,
+        )
+
+    # Tensor max_seqlen values (e.g. HF varlen) must not leak into the compile key:
+    # tensor identity changes on every call and defeats the JIT cache.
+    is_static_persistent = (
+        not causal
+        and not local
+        and cu_seqlens_q is None
+        and seqused_q is None
+        and not is_split_kv
+    ) or (
+        not torch.is_tensor(max_m_blocks_leq_one)
+        and max_m_blocks_leq_one
+        and not is_split_kv
+    )
+
+    # CuTe keeps stride-zero modes static when marking layouts dynamic.
+    tensor_broadcast_patterns = tuple(
+        get_broadcast_dims(tensor) if tensor is not None else None
+        for tensor in (
+            q,
+            k,
+            v,
+            qv,
+            page_table,
+            q_descale,
+            k_descale,
+            v_descale,
+            gather_kv_indices,
+        )
+    )
+
     compile_key = (
         dtype,
         head_dim,
@@ -388,7 +1114,9 @@ def _flash_attn_fwd(
         mask_mod_hash,
         use_block_sparsity,
         block_sparse_broadcast_pattern,
+        tensor_broadcast_patterns,
         aux_tensor_metadata,
+        aux_scalar_metadata,
         lse is None,
         cu_seqlens_q is None,
         cu_seqlens_k is None,
@@ -397,18 +1125,51 @@ def _flash_attn_fwd(
         page_table is not None,
         window_size_left is not None,
         window_size_right is not None,
-        learnable_sink is not None,
-        m_block_size,
-        n_block_size,
+        (
+            torch2cute_dtype_map[learnable_sink.dtype]
+            if learnable_sink is not None
+            else None
+        ),
+        q_descale is not None,
+        k_descale is not None,
+        v_descale is not None,
+        block_sparse_tensors is None or block_sparse_tensors.cu_total_m_blocks is None,
+        block_sparse_tensors is None or block_sparse_tensors.cu_block_idx_offsets is None,
+        tile_m,
+        tile_n,
         q_stage,
         num_threads,
         is_split_kv,
         pack_gqa,
         arch,
-        page_size not in [None, 128],  # paged KV non-TMA
+        page_size not in [None, tile_n],  # paged KV non-TMA
+        use_2cta_instrs,
         q_subtile_factor,
+        kv_subtile_factor,
+        mma_pv_is_rs,
+        intra_wg_overlap,
+        use_clc_scheduler,
+        num_splits_dynamic is not None,
+        virtual_batch_idx is not None,
+        num_nheads_in_l2 is not None,
+        tile_count_semaphore is not None,
+        cu_total_m_blocks is not None,
+        cu_total_splits_m_blocks is not None,
+        blocks_to_batch_idx is not None,
+        seqlen_k_per_split,
+        is_static_persistent,
+        q is not None,
+        qv is not None,
+        p is not None,
+        row_max is not None,
+        gather_kv_length,
+        sparse_kv,
+        disable_sparse_kv_bitmask,
+        fa_logging.get_fa_log_level(),
     )
+
     if compile_key not in _flash_attn_fwd.compile_cache:
+        current_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
         (
             cu_seqlens_q_tensor,
             cu_seqlens_k_tensor,
@@ -431,10 +1192,24 @@ def _flash_attn_fwd(
         ]
         if is_split_kv:
             lse_tensor = to_cute_tensor(lse_partial, assumed_align=4)
-        elif lse is not None:
-            lse_tensor = to_cute_tensor(lse, assumed_align=4)
         else:
-            lse_tensor = None
+            lse_tensor = to_cute_tensor(lse, assumed_align=4)
+
+        q_descale_tensor, k_descale_tensor, v_descale_tensor = (
+            to_cute_tensor(t, assumed_align=4, leading_dim=1)
+            for t in (q_descale, k_descale, v_descale)
+        )
+        descale_tensors_tensor = (
+            DescaleTensors(
+                q_descale=q_descale_tensor,
+                k_descale=k_descale_tensor,
+                v_descale=v_descale_tensor,
+            )
+            if q_descale_tensor is not None
+            or k_descale_tensor is not None
+            or v_descale_tensor is not None
+            else None
+        )
 
         sparse_tensors = None
         if normalized_block_sparse_tensors is not None:
@@ -445,10 +1220,54 @@ def _flash_attn_fwd(
         if aux_tensors is not None:
             cute_aux_tensors = [to_cute_aux_tensor(buf) for buf in aux_tensors]
 
-        if arch // 10 == 9:
-            assert page_table is None, "paged KV not supported on SM 9.0"
+        (
+            num_splits_dynamic_tensor,
+            tile_count_semaphore_tensor,
+            virtual_batch_idx_tensor,
+            num_nheads_in_l2_tensor,
+            cu_total_m_blocks_tensor,
+            cu_total_splits_m_blocks_tensor,
+            blocks_to_batch_idx_tensor,
+        ) = [
+            to_cute_tensor(t, assumed_align=4, leading_dim=0)
+            for t in (
+                num_splits_dynamic,
+                tile_count_semaphore,
+                virtual_batch_idx,
+                num_nheads_in_l2,
+                cu_total_m_blocks,
+                cu_total_splits_m_blocks,
+                blocks_to_batch_idx,
+            )
+        ]
+
+        qv_tensor = to_cute_tensor(qv)
+        gather_kv_indices_tensor = to_cute_tensor(gather_kv_indices)
+        p_tensor = to_cute_tensor(p)
+        row_max_tensor = to_cute_tensor(row_max)
+
+        if arch // 10 == 8:
+            assert page_table is None, "paged KV not supported on SM 8.0"
+            assert not is_split_kv, "SplitKV not supported on SM 8.0"
+            fa_fwd = FlashAttentionForwardSm80(
+                dtype,
+                head_dim,
+                head_dim_v,
+                qhead_per_kvhead,
+                is_causal=causal,
+                is_local=local,
+                pack_gqa=pack_gqa,
+                tile_m=tile_m,
+                tile_n=tile_n,
+                num_stages=1,
+                num_threads=num_threads,
+                Q_in_regs=False,
+                score_mod=score_mod,
+                mask_mod=mask_mod,
+                has_aux_tensors=aux_tensors is not None,
+            )
+        elif arch // 10 == 9:
             assert not is_split_kv, "SplitKV not supported on SM 9.0"
-            # fa_fwd = FlashAttentionForwardSm80(
             fa_fwd = FlashAttentionForwardSm90(
                 dtype,
                 head_dim,
@@ -457,108 +1276,272 @@ def _flash_attn_fwd(
                 is_causal=causal,
                 is_local=local,
                 pack_gqa=pack_gqa,
-                tile_m=m_block_size,
-                tile_n=n_block_size,
+                tile_m=tile_m,
+                tile_n=tile_n,
                 # num_stages=1,
                 num_stages=2,
                 num_threads=num_threads,
                 Q_in_regs=False,
-                intra_wg_overlap=True,
-                mma_pv_is_rs=True,
+                intra_wg_overlap=intra_wg_overlap,
+                mma_pv_is_rs=mma_pv_is_rs,
                 mask_mod=mask_mod,
                 score_mod=score_mod,
                 has_aux_tensors=aux_tensors is not None,
                 q_subtile_factor=q_subtile_factor,
+                paged_kv_non_tma=page_size not in [None, tile_n],
             )
         elif arch // 10 in [10, 11]:
-            head_dim_padded = int(math.ceil(head_dim / 16) * 16)
-            head_dim_v_padded = int(math.ceil(head_dim / 16) * 16)
-            use_2cta_instrs = (
-                not causal
-                and not local
-                and not is_split_kv
-                and cu_seqlens_q is None
-                and seqused_q is None
-                and not use_block_sparsity
-                and page_size in [None, 128]
-                and head_dim_padded == 128
-                and head_dim_v_padded == 128
-            )
-            fa_fwd = FlashAttentionForwardSm100(
+            if qv is not None:
+                paged_kv_cpasync = page_table is not None and page_size != tile_n
+                has_qk = q is not None
+                fa_fwd = FlashAttentionMLAForwardSm100(
+                    is_causal=causal,
+                    use_cpasync_load_KV=sparse_kv or paged_kv_cpasync,
+                    topk_length=gather_kv_length,
+                    is_topk_gather=sparse_kv,
+                    pack_gqa=pack_gqa,
+                    qhead_per_kvhead=qhead_per_kvhead,
+                    nheads_kv=num_head_kv,
+                    has_seqused_q=seqused_q is not None,
+                    has_cu_seqlens_q=cu_seqlens_q is not None,
+                    disable_bitmask=disable_sparse_kv_bitmask,
+                    has_qk=has_qk,
+                )
+            else:
+                if use_dedicated_hd256_kernel:
+                    # hd=256 2CTA forward: check for currently unsupported features
+                    assert softcap is None, "SM100 forward with head_dim=256 does not support softcap"
+                    assert not use_block_sparsity, \
+                        "SM100 forward with head_dim=256 does not support block sparsity"
+                    assert learnable_sink is None, \
+                        "SM100 forward with head_dim=256 does not support learnable_sink"
+                    assert seqused_q is None and seqused_k is None, \
+                        "SM100 forward with head_dim=256 does not support seqused_q/seqused_k"
+                    if page_table is not None:
+                        assert max_seqlen_k % page_size == 0, (
+                            f"SM100 hd256 2CTA paged KV requires max_seqlen_k divisible by "
+                            f"page_size ({page_size}), got max_seqlen_k={max_seqlen_k}"
+                        )
+                        assert page_table.shape[1] == max_seqlen_k // page_size, (
+                            f"SM100 hd256 2CTA paged KV requires page_table.shape[1] == "
+                            f"max_seqlen_k // page_size ({max_seqlen_k} // {page_size} = "
+                            f"{max_seqlen_k // page_size}), got {page_table.shape[1]}; "
+                            f"pass page_table[:, :{max_seqlen_k // page_size}] to slice to "
+                            f"the actual sequence length"
+                        )
+                    # pack_gqa is an auto-selected optimization; disable it for hd256 kernel
+                    pack_gqa = False
+
+                flash_fwd_obj_cls = (
+                    BlackwellFusedMultiHeadAttentionForward
+                    if use_dedicated_hd256_kernel
+                    else FlashAttentionForwardSm100
+                )
+
+                fa_fwd_kwargs = dict(
+                    qhead_per_kvhead=qhead_per_kvhead,
+                    is_causal=causal,
+                    is_local=local,
+                    is_split_kv=is_split_kv,
+                    pack_gqa=pack_gqa,
+                    m_block_size=tile_m,
+                    n_block_size=tile_n,
+                    q_stage=q_stage,
+                    is_static_persistent=is_static_persistent,
+                    score_mod=score_mod,
+                    mask_mod=mask_mod,
+                    has_aux_tensors=aux_tensors is not None,
+                    paged_kv_non_tma=page_size not in [None, tile_n],
+                    is_varlen_q=cu_seqlens_q is not None or seqused_q is not None,
+                    q_subtile_factor=q_subtile_factor,
+                    kv_subtile_factor=kv_subtile_factor,
+                    use_2cta_instrs=use_2cta_instrs,
+                    use_clc_scheduler=use_clc_scheduler,
+                    seqlen_k_per_split=seqlen_k_per_split,
+                )
+                if not use_dedicated_hd256_kernel:
+                    fa_fwd_kwargs["has_tile_count_semaphore"] = tile_count_semaphore is not None
+                fa_fwd = flash_fwd_obj_cls(head_dim, head_dim_v, **fa_fwd_kwargs)
+        elif arch // 10 == 12:
+            # SM120 (Blackwell GeForce / DGX Spark): uses SM80 MMA with SM120 SMEM capacity
+            assert not use_block_sparsity, "Block sparsity not supported on SM 12.0"
+            assert page_table is None, "Paged KV not supported on SM 12.0 in this PR"
+            assert not is_split_kv, "SplitKV not supported on SM 12.0 in this PR"
+            fa_fwd = FlashAttentionForwardSm120(
+                dtype,
                 head_dim,
                 head_dim_v,
-                qhead_per_kvhead=qhead_per_kvhead,
+                qhead_per_kvhead,
                 is_causal=causal,
                 is_local=local,
-                is_split_kv=is_split_kv,
                 pack_gqa=pack_gqa,
-                m_block_size=m_block_size,
-                n_block_size=n_block_size,
-                q_stage=q_stage,
-                is_persistent=not causal
-                    and not local
-                    and cu_seqlens_q is None
-                    and seqused_q is None
-                    and not is_split_kv,
+                tile_m=tile_m,
+                tile_n=tile_n,
+                num_stages=1,
+                num_threads=num_threads,
+                Q_in_regs=False,
                 score_mod=score_mod,
                 mask_mod=mask_mod,
                 has_aux_tensors=aux_tensors is not None,
-                paged_kv_non_tma=page_size not in [None, 128],
-                is_varlen_q=cu_seqlens_q is not None or seqused_q is not None,
-                q_subtile_factor=q_subtile_factor,
-                use_2cta_instrs=use_2cta_instrs,
             )
         else:
             raise ValueError(
-                f"Unsupported compute capability: {arch}. Supported: 9.x, 10.x, 11.x"
+                f"Unsupported compute capability: {arch}. Supported: 8.x, 9.x, 10.x, 11.x, 12.x"
             )
         # TODO: check @can_implement
-        _flash_attn_fwd.compile_cache[compile_key] = cute.compile(
-            fa_fwd,
-            q_tensor,
-            k_tensor,
-            v_tensor,
-            o_tensor,
-            lse_tensor,
-            softmax_scale,
-            current_stream,
-            cu_seqlens_q_tensor,
-            cu_seqlens_k_tensor,
-            seqused_q_tensor,
-            seqused_k_tensor,
-            page_table_tensor,
-            window_size_left,
-            window_size_right,
-            learnable_sink_tensor,
-            sparse_tensors,
-            cute_aux_tensors,
-            options="--enable-tvm-ffi",
-        )
-
-    # In "fake mode", we will take torch fake tensors as input and the expected behaviors are:
-    # - Use those fake metadata to populate compilation cache
-    # - Return "fake" output tensors, which could be needed in follow-up fake operations
-    # Thus, we skip the actual kernel invocation here.
-    if not is_fake_mode():
-        _flash_attn_fwd.compile_cache[compile_key](
-            q.detach(),
-            k.detach(),
-            v.detach(),
-            out.detach() if not is_split_kv else out_partial,
-            lse_partial if is_split_kv else lse,
-            softmax_scale,
-            current_stream,
-            cu_seqlens_q,
-            cu_seqlens_k,
-            seqused_q,
-            seqused_k,
-            page_table,
-            window_size_left,
-            window_size_right,
-            learnable_sink,
-            normalized_block_sparse_tensors[:4] if normalized_block_sparse_tensors is not None else None,
-            aux_tensors,
+        if qv is not None:
+            _flash_attn_fwd.compile_cache[compile_key] = cute.compile(
+                fa_fwd,
+                q_tensor,
+                qv_tensor,
+                k_tensor,
+                v_tensor,
+                o_tensor,
+                lse_tensor,
+                softmax_scale,
+                p_tensor,
+                row_max_tensor,
+                cu_seqlens_q_tensor,
+                cu_seqlens_k_tensor,
+                seqused_q_tensor,
+                seqused_k_tensor,
+                gather_kv_indices_tensor,
+                page_table_tensor,
+                window_size_left,
+                window_size_right,
+                current_stream,
+                options="--enable-tvm-ffi",
+            )
+        else:
+            compile_args = [
+                fa_fwd,
+                q_tensor,
+                k_tensor,
+                v_tensor,
+                o_tensor,
+                lse_tensor,
+                softmax_scale,
+                cu_seqlens_q_tensor,
+                cu_seqlens_k_tensor,
+                seqused_q_tensor,
+                seqused_k_tensor,
+                page_table_tensor,
+                window_size_left,
+                window_size_right,
+                learnable_sink_tensor,
+            ]
+            if arch // 10 in [10, 11]:
+                compile_args.append(descale_tensors_tensor)
+            compile_args.extend([
+                sparse_tensors,
+                AuxData(cute_aux_tensors, aux_scalars),
+            ])
+            if arch // 10 in [10, 11] and not use_dedicated_hd256_kernel:
+                compile_args.extend([
+                    num_splits_dynamic_tensor,
+                    tile_count_semaphore_tensor,
+                    virtual_batch_idx_tensor,
+                    num_nheads_in_l2_tensor,
+                    cu_total_m_blocks_tensor,
+                    cu_total_splits_m_blocks_tensor,
+                    blocks_to_batch_idx_tensor,
+                    max_seqlen_q,
+                ])
+            elif arch // 10 in [8, 9, 12]:
+                compile_args.extend([
+                    cu_total_m_blocks_tensor,
+                    cu_total_splits_m_blocks_tensor,
+                ])
+            compile_args.append(current_stream)
+            _flash_attn_fwd.compile_cache[compile_key] = cute.compile(*compile_args, options="--enable-tvm-ffi")
+
+    if not fake_mode:
+        q_call, k_call, v_call, qv_call = [
+            t.detach() if t is not None else None
+            for t in (q, k, v, qv)
+        ]
+        if is_fp8:
+            # need uint8 workaround until we pin torch >= 2.11.0 where fp8 export is supported
+            q_call, k_call, v_call, qv_call = [
+                t.view(torch.uint8) if t is not None else None
+                for t in (q_call, k_call, v_call, qv_call)
+            ]
+        descale_tensors = (
+            DescaleTensors(q_descale=q_descale, k_descale=k_descale, v_descale=v_descale)
+            if q_descale is not None or k_descale is not None or v_descale is not None
+            else None
         )
+        if qv is not None:
+            _flash_attn_fwd.compile_cache[compile_key](
+                q_call,
+                qv_call,
+                k_call,
+                v_call,
+                out.detach(),
+                lse,
+                softmax_scale,
+                p,
+                row_max,
+                cu_seqlens_q,
+                cu_seqlens_k,
+                seqused_q,
+                seqused_k,
+                gather_kv_indices,
+                page_table,
+                window_size_left,
+                window_size_right,
+            )
+        else:
+            call_args = [
+                q_call,
+                k_call,
+                v_call,
+                out.detach() if not is_split_kv else out_partial,
+                lse_partial if is_split_kv else lse,
+                softmax_scale,
+                cu_seqlens_q,
+                cu_seqlens_k,
+                seqused_q,
+                seqused_k,
+                page_table,
+                window_size_left,
+                window_size_right,
+                learnable_sink,
+            ]
+            if arch // 10 in [10, 11]:
+                call_args.append(descale_tensors)
+            call_args.extend([
+                (
+                    normalized_block_sparse_tensors.mask_block_cnt,
+                    normalized_block_sparse_tensors.mask_block_idx,
+                    normalized_block_sparse_tensors.full_block_cnt,
+                    normalized_block_sparse_tensors.full_block_idx,
+                    normalized_block_sparse_tensors.cu_total_m_blocks,
+                    normalized_block_sparse_tensors.cu_block_idx_offsets,
+                    normalized_block_sparse_tensors.dq_write_order,
+                    normalized_block_sparse_tensors.dq_write_order_full,
+                )
+                if normalized_block_sparse_tensors is not None
+                else None,
+                AuxData(aux_tensors, aux_scalars),
+            ])
+            if arch // 10 in [10, 11] and not use_dedicated_hd256_kernel:
+                call_args.extend([
+                    num_splits_dynamic,
+                    tile_count_semaphore,
+                    virtual_batch_idx,
+                    num_nheads_in_l2,
+                    cu_total_m_blocks,
+                    cu_total_splits_m_blocks,
+                    blocks_to_batch_idx,
+                    max_seqlen_q,
+                ])
+            elif arch // 10 in [8, 9, 12]:
+                call_args.extend([
+                    cu_total_m_blocks,
+                    cu_total_splits_m_blocks,
+                ])
+            _flash_attn_fwd.compile_cache[compile_key](*call_args)
     if is_split_kv:
         _flash_attn_fwd_combine(
             out_partial,
@@ -567,13 +1550,272 @@ def _flash_attn_fwd(
             lse.transpose(-1, -2) if lse is not None else None,
             cu_seqlens_q,
             seqused_q,
+            num_splits_dynamic_ptr=num_splits_dynamic if has_scheduler_metadata else None,
+            virtual_batch_idx=virtual_batch_idx if has_scheduler_metadata else None,
+            _arch=arch,
         )
-    return out, lse
+    if reuse_scheduler_metadata and tile_count_semaphore is not None:
+        # TODO: pass tile_count_semaphore to the combine kernel and zero it there when
+        # is_split_kv (using CTA 0, since a later CTA may have exited prematurely), so
+        # that this host-side zeroing is only needed when is_split_kv=False.
+        tile_count_semaphore.zero_()
+    return out, lse, p, row_max
 
 
 _flash_attn_fwd.compile_cache = get_jit_cache("fwd")
 
 
+def make_fake_bwd_tensors(dtype, has_gqa, varlen_q, varlen_k, nheads_major=False):
+    sym = cute.sym_int
+    # divisibility in elements: assumed_align_bytes = divisibility * dtype.width // 8
+    # For 16-byte align: fp16/bf16 → divisibility=8, float32 → divisibility=4
+    div = 128 // dtype.width  # 8 for fp16/bf16
+    # Shared sym_ints for dimensions that must match across tensors
+    b, seqlen_q, seqlen_k, h_q, d, d_v = sym(), sym(), sym(), sym(), sym(), sym()
+    topk = sym()
+    h_kv = h_q if not has_gqa else sym()
+    seqlen_q_rounded, seqlen_k_rounded = sym(), sym()
+    seqlen_q_d_rounded, seqlen_k_d_rounded, seqlen_k_dv_rounded = sym(), sym(), sym()
+    total_q, total_k, total_q_rounded, total_k_rounded = sym(), sym(), sym(), sym()
+    total_q_d_rounded, total_k_d_rounded, total_k_dv_rounded = sym(), sym(), sym()
+    b_seqlenq = (b, seqlen_q) if not varlen_q else (total_q,)
+    b_seqlenk = (b, seqlen_k) if not varlen_k else (total_k,)
+    mQ = fake_tensor(dtype, (*b_seqlenq, h_q, d), divisibility=div)
+    mO = fake_tensor(dtype, (*b_seqlenq, h_q, d_v), divisibility=div)
+    mdO = fake_tensor(dtype, (*b_seqlenq, h_q, d_v), divisibility=div)
+    mK = fake_tensor(dtype, (*b_seqlenk, h_kv, d), divisibility=div)
+    mV = fake_tensor(dtype, (*b_seqlenk, h_kv, d_v), divisibility=div)
+    mdQ = fake_tensor(dtype, (*b_seqlenq, h_q, d), divisibility=div)
+    mdK = fake_tensor(dtype, (*b_seqlenk, h_kv, d), divisibility=div)
+    mdV = fake_tensor(dtype, (*b_seqlenk, h_kv, d_v), divisibility=div)
+
+    sq    = seqlen_q         if not varlen_q else total_q
+    sq_r  = seqlen_q_rounded if not varlen_q else total_q_rounded
+    sq_dr = seqlen_q_d_rounded if not varlen_q else total_q_d_rounded
+
+    def shape(*dims):
+        batch = (b,) if not varlen_q else ()
+        return (*batch, h_q, *dims) if not nheads_major else (*batch, *dims, h_q)
+
+    mLSE     = fake_tensor(Float32, shape(sq),       divisibility=1)
+    mLSElog2 = fake_tensor(Float32, shape(sq_r),     divisibility=4)
+    mPdPsum  = fake_tensor(Float32, shape(sq_r),     divisibility=4)
+    dQaccum  = fake_tensor(Float32, shape(sq_dr),    divisibility=4)
+    mScaleP  = fake_tensor(Float32, shape(sq, topk), divisibility=4)
+
+    if not has_gqa:
+        mdKaccum, mdVaccum = None, None
+    else:
+        if not varlen_k:
+            mdKaccum = fake_tensor(Float32, (b, h_kv, seqlen_k_rounded), divisibility=4)
+            mdVaccum = fake_tensor(Float32, (b, h_kv, seqlen_k_dv_rounded), divisibility=4)
+        else:
+            mdKaccum = fake_tensor(Float32, (h_kv, total_k_rounded), divisibility=4)
+            mdVaccum = fake_tensor(Float32, (h_kv, total_k_dv_rounded), divisibility=4)
+    return mQ, mK, mV, mO, mdO, mdQ, mdK, mdV, mLSE, mLSElog2, mPdPsum, dQaccum, mdKaccum, mdVaccum, mScaleP
+
+
+def _compile_bwd_preprocess(
+    dtype,
+    head_dim,
+    head_dim_v,
+    m_block_size,
+    has_cuseqlens_q,
+    has_seqused_q,
+    has_dlse,
+    has_dq_accum,
+    has_scaleP,
+    use_padded_offsets,
+    nheads_major,
+    pack_gqa,
+    qhead_per_kvhead,
+    nheads_kv,
+    has_cu_total_m_blocks,
+    hdim_multiple_of,
+):
+    """Compile bwd preprocess kernel using cute fake tensors (no real GPU tensors needed)."""
+    mQ, mK, mV, mO, mdO, mdQ, mdK, mdV, mLSE, mLSElog2, mPdPsum, mdQaccum, mdKaccum, mdVaccum, mScaleP = make_fake_bwd_tensors(
+        dtype, has_gqa=True, varlen_q=has_cuseqlens_q, varlen_k=False, nheads_major=nheads_major,
+    )
+    batch = mQ.shape[0] if not has_cuseqlens_q else cute.sym_int()
+    batchp1 = cute.sym_int()
+    mCuSeqlensQ = fake_tensor(Int32, (batchp1,), divisibility=1) if has_cuseqlens_q else None
+    mSequsedQ = fake_tensor(Int32, (batch,), divisibility=1) if has_seqused_q else None
+    mdLSE = fake_tensor(Float32, mLSE.shape, divisibility=1) if has_dlse else None
+    mLSElog2 = None if has_scaleP else mLSElog2
+    mdQaccum = mdQaccum if has_dq_accum else None
+    mRowMax = fake_tensor(Float32, mScaleP.shape, divisibility=1) if has_scaleP else None
+    mScaleP = fake_tensor(Float32, mScaleP.shape, divisibility=1) if has_scaleP else None
+    softmax_scale = Float32(1.0)
+    mCuTotalMBlocks = fake_tensor(Int32, (batchp1,), divisibility=1) if has_cu_total_m_blocks else None
+    fa_bwd_pre = FlashAttentionBackwardPreprocess(
+        dtype, head_dim, head_dim_v, m_block_size,
+        use_padded_offsets=use_padded_offsets,
+        nheads_major=nheads_major,
+        pack_gqa=pack_gqa,
+        qhead_per_kvhead=qhead_per_kvhead,
+        nheads_kv=nheads_kv,
+        hdim_multiple_of=hdim_multiple_of,
+    )
+    return cute.compile(
+        fa_bwd_pre, mO, mdO, mPdPsum, mLSE, mLSElog2, mdQaccum, mCuSeqlensQ, mSequsedQ, mdLSE,
+        mRowMax, mScaleP, softmax_scale, mCuTotalMBlocks,
+        cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+        options="--enable-tvm-ffi",
+    )
+
+
+def _bwd_preprocess(
+    out, dout, dpsum, lse, lse_log2, dq_accum,
+    cu_seqlens_q, seqused_q, dlse,
+    dtype, head_dim, head_dim_v, m_block_size,
+    row_max=None,
+    scale_p=None,
+    use_padded_offsets=True,
+    nheads_major=False,
+    pack_gqa=False,
+    qhead_per_kvhead=1,  # only used with pack_gqa
+    nheads_kv=1,         # only used with pack_gqa
+    softmax_scale=1.0,   # only used with scale_p
+    cu_total_m_blocks=None,
+    hdim_multiple_of=32,
+    *,
+    fake_mode,
+):
+    """Backward preprocess: compute (o * dout).sum(dim=-1) - dLSE, lse * log2_e, and zero out dq_accum."""
+    if row_max is not None:
+        assert scale_p is not None
+    is_varlen = cu_seqlens_q is not None or seqused_q is not None
+    if is_varlen:
+        batch_size = (cu_seqlens_q.shape[0] - 1) if cu_seqlens_q is not None else seqused_q.shape[0]
+    else:
+        batch_size = 0
+    if cu_total_m_blocks is None and is_varlen and batch_size > BIN_BATCH_SEARCH_THRESH:
+        cu_total_m_blocks, _ = _compute_tile_cumsum(
+            cu_seqlens=cu_seqlens_q,
+            seqused=seqused_q,
+            tile_size=m_block_size,
+            qhead_per_kvhead=qhead_per_kvhead,
+            pack_gqa=pack_gqa,
+        )
+    compile_key = (
+        dtype, head_dim, head_dim_v, m_block_size,
+        cu_seqlens_q is not None,
+        seqused_q is not None,
+        dlse is not None,
+        dq_accum is not None,
+        row_max is not None,
+        use_padded_offsets,
+        nheads_major,
+        pack_gqa,
+        qhead_per_kvhead,
+        nheads_kv,
+        cu_total_m_blocks is not None,
+        hdim_multiple_of,
+    )
+    if compile_key not in _bwd_preprocess.compile_cache:
+        _bwd_preprocess.compile_cache[compile_key] = _compile_bwd_preprocess(*compile_key)
+    if not fake_mode:
+        _bwd_preprocess.compile_cache[compile_key](
+            out, dout, dpsum, lse, lse_log2, dq_accum, cu_seqlens_q, seqused_q, dlse,
+            row_max, scale_p, softmax_scale, cu_total_m_blocks,
+        )
+
+
+_bwd_preprocess.compile_cache = get_jit_cache("bwd_pre")
+
+
+def _compile_bwd_postprocess(
+    dtype, hdim, block_size, num_threads, atom_layout, swap_ab,
+    has_cuseqlens_q, has_seqused_q,
+    use_2cta_instrs, cluster_size, arch,
+    has_cu_total_m_blocks,
+    learnable_sink_dtype,
+    hdim_multiple_of,
+):
+    """Compile bwd postprocess kernel using cute fake tensors."""
+    mQ, mK, mV, mO, mdO, mdQ, mdK, mdV, mLSE, mLSElog2, mPdPsum, mdQaccum, mdKaccum, mdVaccum, mScaleP = make_fake_bwd_tensors(
+        dtype, has_gqa=True, varlen_q=has_cuseqlens_q, varlen_k=False
+    )
+    batch = mQ.shape[0] if not has_cuseqlens_q else cute.sym_int()
+    batchp1 = cute.sym_int()
+    mCuSeqlensQ = fake_tensor(Int32, (batchp1,), divisibility=1) if has_cuseqlens_q else None
+    mSeqUsedQ = fake_tensor(Int32, (batch,), divisibility=1) if has_seqused_q else None
+    mCuTotalMBlocks = fake_tensor(Int32, (batchp1,), divisibility=1) if has_cu_total_m_blocks else None
+    sink_tensors = (
+        LearnableSinkBwdTensors(
+            mPdPsum,
+            mLSE,
+            fake_tensor(learnable_sink_dtype, (mQ.shape[-2],), divisibility=1),
+            fake_tensor(learnable_sink_dtype, (mQ.shape[-2],), divisibility=1),
+        )
+        if learnable_sink_dtype is not None
+        else None
+    )
+    fa_bwd_post = FlashAttentionBackwardPostprocess(
+        dtype, hdim, arch, block_size, num_threads, atom_layout, swap_ab,
+        use_2cta_instrs=use_2cta_instrs,
+        cluster_size=cluster_size,
+        hdim_multiple_of=hdim_multiple_of,
+    )
+    return cute.compile(
+        fa_bwd_post, mdQaccum, mdQ, Float32(0.0), mCuSeqlensQ, mSeqUsedQ,
+        sink_tensors,
+        mCuTotalMBlocks,
+        cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+        options="--enable-tvm-ffi",
+    )
+
+
+def _bwd_postprocess_convert(
+    accum, output, scale,
+    cu_seqlens, seqused,
+    arch, dtype, hdim, block_size, num_threads,
+    atom_layout, swap_ab,
+    use_2cta_instrs=False, cluster_size=1,
+    cu_total_m_blocks=None,
+    sink_tensors=None,
+    hdim_multiple_of=32,
+    *,
+    fake_mode,
+):
+    """Backward postprocess: convert float32 accumulator to bf16/fp16 output."""
+    is_varlen = cu_seqlens is not None or seqused is not None
+    if is_varlen:
+        batch_size = (cu_seqlens.shape[0] - 1) if cu_seqlens is not None else seqused.shape[0]
+    else:
+        batch_size = 0
+    if cu_total_m_blocks is None and is_varlen and batch_size > BIN_BATCH_SEARCH_THRESH:
+        cu_total_m_blocks, _ = _compute_tile_cumsum(
+            cu_seqlens=cu_seqlens,
+            seqused=seqused,
+            tile_size=block_size,
+        )
+    compile_key = (
+        dtype, hdim, block_size, num_threads, atom_layout, swap_ab,
+        cu_seqlens is not None, seqused is not None,
+        use_2cta_instrs, cluster_size, arch,
+        cu_total_m_blocks is not None,
+        (
+            torch2cute_dtype_map[sink_tensors.sink.dtype]
+            if sink_tensors is not None
+            else None
+        ),
+        hdim_multiple_of,
+    )
+    if compile_key not in _bwd_postprocess_convert.compile_cache:
+        _bwd_postprocess_convert.compile_cache[compile_key] = _compile_bwd_postprocess(*compile_key)
+    if not fake_mode:
+        _bwd_postprocess_convert.compile_cache[compile_key](
+            accum, output, scale, cu_seqlens, seqused,
+            sink_tensors,
+            cu_total_m_blocks,
+        )
+
+
+_bwd_postprocess_convert.compile_cache = get_jit_cache("bwd_post")
+
+
 def _flash_attn_bwd(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -613,48 +1855,99 @@ def _flash_attn_bwd(
     score_mod_bwd: Optional[Callable] = None,
     mask_mod: Optional[Callable] = None,
     aux_tensors: Optional[list[torch.Tensor]] = None,
+    aux_scalars: Optional[tuple] = None,
     block_sparse_tensors: Optional[BlockSparseTensorsTorch] = None,
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    dlse: Optional[torch.Tensor] = None,
+    learnable_sink: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, ...]:
+    aux_scalars = tuple(aux_scalars) if aux_scalars else None
+    fake_mode = is_fake_mode()
     arch = _get_device_arch()
-    assert arch // 10 in [9, 10, 11], "Unsupported compute capability. Supported: 9.x, 10.x, 11.x"
+    assert arch // 10 in [9, 10, 11, 12], "Unsupported compute capability. Supported: 9.x, 10.x, 11.x, 12.x"
+    if block_sparse_tensors is not None:
+        assert (
+            cu_seqlens_q is None
+            and cu_seqlens_k is None
+            and seqused_q is None
+            and seqused_k is None
+        ), "Varlen backward with block sparsity is not yet supported"
+    if learnable_sink is not None:
+        assert arch // 10 in [9, 10, 11], "Learnable sink backward is supported on SM90 and SM100/SM110"
+        assert lse is not None, "learnable_sink backward requires LSE"
+        if q.numel() == 0 or k.numel() == 0:
+            dq = torch.zeros_like(q) if dq is None else dq.zero_()
+            dk = torch.zeros_like(k) if dk is None else dk.zero_()
+            dv = torch.zeros_like(v) if dv is None else dv.zero_()
+            dsink = (
+                dlse.sum(dim=(0, 2) if dlse.ndim == 3 else 1).to(learnable_sink.dtype)
+                if dlse is not None
+                else torch.zeros_like(learnable_sink)
+            )
+            return dq, dk, dv, dsink
+    sparse_q = None
+    kv_subtile_factor = 1
+    if block_sparse_tensors is not None:
+        if block_sparse_tensors.block_size is not None:
+            sparse_q = block_sparse_tensors.block_size[0]
+        elif arch // 10 == 9:
+            sparse_q = 128
 
     num_head, head_dim = q.shape[-2:]
+    head_dim_v = v.shape[-1]
 
-    if causal:
-        window_size_right = 0
-    if window_size_left is not None and window_size_right is not None and window_size_left + window_size_right < 0:
-        window_size_left = None
-        window_size_right = None
-    local = window_size_left is not None or window_size_right is not None
-    if local:
-        if window_size_left is None and window_size_right == 0:
-            causal, local = True, False
-            window_size_right = None
-        else:
-            causal, local = False, True
+    window_size = [window_size_left, window_size_right]
+    causal, local, window_size_left, window_size_right = _resolve_causal_local_window(
+        causal, window_size_left, window_size_right
+    )
 
-    if arch // 10 == 9:
-        m_block_size = 80 if not causal else 64
-        n_block_size = 128
-        num_stages_Q = 2
-        num_stages_dO = 2
-        num_stages_PdS = 2
-        SdP_swapAB = True
+    if arch // 10 == 12:
+        # SM120: uses SM80 MMA with 99 KB SMEM, 128 threads (4 warps).
+        m_block_size = 64
+        n_block_size = 64
+        if head_dim <= 64:
+            num_stages_Q = 2
+            num_stages_dO = 2
+        else:
+            num_stages_Q = 1
+            num_stages_dO = 1
+        SdP_swapAB = False
         dKV_swapAB = False
-        dQ_swapAB = not causal
-        AtomLayoutMSdP = 1
-        AtomLayoutNdKV = 2
-        AtomLayoutMdQ = 1
+        dQ_swapAB = False
+        AtomLayoutMSdP = 4
+        AtomLayoutNdKV = 4
+        AtomLayoutMdQ = 4
+        V_in_regs = False
+        dQ_single_wg = False
+        cluster_size = 1
+        use_2cta_instrs = False
+        num_threads = 128
+        assert not (block_sparse_tensors is not None), "Block sparsity backward not supported on SM 12.0"
+        assert score_mod is None and score_mod_bwd is None, "score_mod backward not supported on SM 12.0"
+        assert mask_mod is None, "mask_mod backward not supported on SM 12.0"
+        assert deterministic is False, "deterministic backward not supported on SM 12.0"
+    elif arch // 10 == 9:
+        cfg = _tile_size_bwd_sm90(
+            head_dim,
+            head_dim_v,
+            causal,
+            local,
+            sparse_block_size_q=sparse_q,
+        )
+        m_block_size = cfg.m_block_size
+        n_block_size = cfg.n_block_size
+        num_stages_Q = cfg.num_stages_Q
+        num_stages_dO = cfg.num_stages_dO
+        num_stages_PdS = cfg.num_stages_PdS
+        SdP_swapAB = cfg.SdP_swapAB
+        dKV_swapAB = cfg.dKV_swapAB
+        dQ_swapAB = cfg.dQ_swapAB
+        AtomLayoutMSdP = cfg.AtomLayoutMSdP
+        AtomLayoutNdKV = cfg.AtomLayoutNdKV
+        AtomLayoutMdQ = cfg.AtomLayoutMdQ
+        num_threads = (cfg.num_wg + 1) * 128
+        dQ_single_wg = cfg.dQ_single_wg
         cluster_size = 1
         use_2cta_instrs = False
-        assert window_size_left is None and window_size_right is None, "local not supported yet on 9.x"
-        is_varlen = (
-            cu_seqlens_q is not None
-            or cu_seqlens_k is not None
-            or seqused_q is not None
-            or seqused_k is not None
-        )
-        assert not is_varlen, "varlen backward is not yet supported on sm90"
     else:
         m_block_size = 128
         n_block_size = 128
@@ -662,18 +1955,55 @@ def _flash_attn_bwd(
         dKV_swapAB = False
         AtomLayoutMdQ = 1
         AtomLayoutNdKV = 1
-        disable_2cta = (
-            local
-            or score_mod is not None
-            or score_mod_bwd is not None
-            or mask_mod is not None
-        )
-        cluster_size = 2 if head_dim >= 128 and not disable_2cta else 1
-        use_2cta_instrs = cluster_size==2
-    
-    q, k, v, out, dout, lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = [
+        requested_disable_2cta = utils._get_disable_2cta_default()
+        kv_subtile_factor = get_kv_subtile_factor(block_sparse_tensors, n_block_size)
+        use_2cta_instrs = (
+            head_dim >= 128
+            and not requested_disable_2cta
+            and block_sparse_bwd_supports_2cta(block_sparse_tensors, n_block_size)
+        )
+        if block_sparse_tensors is not None and head_dim == 192 and not use_2cta_instrs:
+            reason = (
+                "2CTA was disabled by request"
+                if requested_disable_2cta
+                else (
+                    f"sparse_block_size[1] must cover an even number of tile_n={n_block_size} "
+                    f"tiles; got factor {kv_subtile_factor}"
+                )
+            )
+            raise ValueError(
+                f"SM100 block-sparse backward with head_dim=192 requires 2CTA; {reason}."
+            )
+        cluster_size = 2 if use_2cta_instrs else 1
+
+    use_dedicated_hd256_kernel = arch // 10 in [10, 11] and head_dim == 256 and head_dim_v == 256
+    if use_dedicated_hd256_kernel:
+        assert learnable_sink is None, (
+            "SM100 backward with head_dim=256 does not support learnable_sink"
+        )
+    use_2cta_instrs = use_2cta_instrs or use_dedicated_hd256_kernel
+    is_varlen = (
+        cu_seqlens_q is not None
+        or cu_seqlens_k is not None
+        or seqused_q is not None
+        or seqused_k is not None
+    )
+
+    q, k, v, out, dout, lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, learnable_sink = [
         maybe_contiguous(t)
-        for t in (q, k, v, out, dout, lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k)
+        for t in (
+            q,
+            k,
+            v,
+            out,
+            dout,
+            lse,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqused_q,
+            seqused_k,
+            learnable_sink,
+        )
     ]
     if cu_seqlens_q is None:
         batch_size, seqlen_q = q.shape[:2]
@@ -692,25 +2022,30 @@ def _flash_attn_bwd(
         seqlen_k = max_seqlen_k if max_seqlen_k is not None else total_k
 
     num_head_kv = k.shape[-2]
-    head_dim_v = v.shape[-1]
 
     use_block_sparsity = block_sparse_tensors is not None
-
-    # SM90 block-sparse backward: tile_m=64 is the GCD between a m_block_size that fits,
-    # the base block_m of 128 from forward, and block-sparse size for subtiling.
-    if arch // 10 == 9 and use_block_sparsity:
-        m_block_size = 64
-        # dQ_swapAB tuning: use False when m_block_size=64 (same as causal case)
-        dQ_swapAB = False
-
-    # NB: this could be derived from the block_sparse_tensors but for now we hardcode it to 2
-    subtile_factor = 2
+    if sparse_q is not None and (sparse_q <= 0 or sparse_q % m_block_size != 0):
+        raise ValueError(
+            "Block sparsity requires sparse_block_size[0] to be a multiple of "
+            f"tile_m={m_block_size}; got {sparse_q}."
+        )
+    q_subtile_factor = sparse_q // m_block_size if sparse_q is not None else 2
     seqlen_q_rounded = (seqlen_q + m_block_size - 1) // m_block_size * m_block_size
     seqlen_k_rounded = (seqlen_k + n_block_size - 1) // n_block_size * n_block_size
     num_n_blocks = seqlen_k_rounded // n_block_size
     if cluster_size == 2 and num_n_blocks % cluster_size != 0:
         seqlen_k_rounded = seqlen_k_rounded + n_block_size
 
+    # The single-block specialization below only guards against TVM stride poisoning,
+    # which is a host-side branch predicate that selects a kernel variant. When
+    # max_seqlen is passed as a tensor (e.g. HF/TE varlen), seqlen_*_rounded are tensors,
+    # so `seqlen_*_rounded // block == 1` would leak a tensor into the compile key. Its
+    # pickle hash differs every call, forcing a recompile per step. Only specialize when
+    # the seqlen is already a host scalar; tensor callers fall back to the multi-block
+    # default, keeping the key stable with no device sync.
+    single_q_block = (not torch.is_tensor(seqlen_q_rounded)) and (seqlen_q_rounded // m_block_size == 1)
+    single_k_block = (not torch.is_tensor(seqlen_k_rounded)) and (seqlen_k_rounded // n_block_size == 1)
+
     if cu_seqlens_k is None:
         assert k.shape == (batch_size, seqlen_k, num_head_kv, head_dim)
         assert v.shape == (batch_size, seqlen_k, num_head_kv, head_dim_v)
@@ -744,14 +2079,22 @@ def _flash_attn_bwd(
         if t is not None:
             assert t.dtype == torch.int32, "cu_seqlens_q, cu_seqlens_k must be int32"
     assert lse.dtype == torch.float32, "lse must be float32"
-    assert all(
-        t is None or t.is_cuda for t in (q, k, v, out, dout, lse, cu_seqlens_q, cu_seqlens_k)
-    ), "inputs must be on CUDA device"
+    if dlse is not None:
+        dlse = maybe_contiguous(dlse)
+    if learnable_sink is not None:
+        assert learnable_sink.shape == (num_head,)
+        assert learnable_sink.dtype in _LEARNABLE_SINK_DTYPES, (
+            "learnable_sink must be float16, bfloat16, or float32"
+        )
+    if not fake_mode:
+        assert all(
+            t is None or t.is_cuda
+            for t in (q, k, v, out, dout, lse, cu_seqlens_q, cu_seqlens_k, learnable_sink)
+        ), "inputs must be on CUDA device"
     assert num_head % num_head_kv == 0, "num_head must be divisible by num_head_kv"
-    assert head_dim <= 256, "head_dim must be less than or equal to 256"
     alignment = 16 // q.element_size()
-    assert head_dim % alignment == 0, f"head_dim must be divisible by {alignment}"
-    assert head_dim_v % alignment == 0, f"head_dim_v must be divisible by {alignment}"
+    if arch // 10 != 12:
+        _validate_head_dims(head_dim, head_dim_v, arch // 10, alignment)
     if softmax_scale is None:
         softmax_scale = 1.0 / math.sqrt(head_dim)
     qhead_per_kvhead = num_head // num_head_kv
@@ -759,15 +2102,17 @@ def _flash_attn_bwd(
         pack_gqa = qhead_per_kvhead > 1
     # pack_gqa backward not yet supported in bwd
     pack_gqa = False
-    if arch // 10 not in [10, 11]:
-        assert deterministic is False, "bwd deterministic only supported for sm100/sm110 for now"
-
+    
+    if softcap != 0.0:
+        assert score_mod is None and score_mod_bwd is None, (
+            "softcap and score_mod/score_mod_bwd cannot be used together"
+        )
+        score_mod = utils.create_softcap_scoremod(softcap)
+        score_mod_bwd = utils.create_softcap_scoremod_bwd(softcap)
     if score_mod is not None:
         assert score_mod_bwd is not None, "score_mod_bwd is required when score_mod is provided"
-        assert softcap == 0.0, "softcap and score_mod are mutually exclusive (different log2 scaling)"
-        assert cu_seqlens_q is None and cu_seqlens_k is None, (
-            "varlen + score_mod not supported in bwd yet"
-        )
+        if arch // 10 == 8:
+            raise NotImplementedError("Custom user-provided score_mod is not supported on SM8x architectures.")
 
     device = q.device
     out_torch_dtype = q.dtype
@@ -787,15 +2132,21 @@ def _flash_attn_bwd(
     else:
         _validate_tensor(dv, "dv", v.shape, out_torch_dtype, device)
 
-    head_dim_rounded = (head_dim + 32 - 1) // 32 * 32
+    # Keep accumulator allocation, zeroing, and readback aligned with SM90's swapped MMA.
+    hdim_multiple_of = 64 if arch // 10 == 9 and dKV_swapAB else 32
+    head_dim_rounded = (head_dim + hdim_multiple_of - 1) // hdim_multiple_of * hdim_multiple_of
 
     if cu_seqlens_q is None:
-        dq_accum = torch.empty(
-            batch_size,
-            num_head,
-            seqlen_q_rounded * head_dim_rounded,
-            dtype=torch.float32,
-            device=device,
+        dq_accum = (
+            None
+            if use_dedicated_hd256_kernel
+            else torch.empty(
+                batch_size,
+                num_head,
+                seqlen_q_rounded * head_dim_rounded,
+                dtype=torch.float32,
+                device=device,
+            )
         )
         dpsum = torch.empty(
             batch_size, num_head, seqlen_q_rounded, dtype=torch.float32, device=device
@@ -807,13 +2158,21 @@ def _flash_attn_bwd(
         total_q_rounded_padded = (
             (total_q + cu_seqlens_q.shape[0] * m_block_size - 1) // m_block_size * m_block_size
         )
-        dq_accum = torch.empty(
-            num_head, total_q_rounded_padded * head_dim_rounded, dtype=torch.float32, device=device
+        dq_accum = (
+            None
+            if use_dedicated_hd256_kernel
+            else torch.empty(
+                num_head, total_q_rounded_padded * head_dim_rounded, dtype=torch.float32, device=device
+            )
         )
         dpsum = torch.empty(num_head, total_q_rounded_padded, dtype=torch.float32, device=device)
         lse_log2 = torch.empty(num_head, total_q_rounded_padded, dtype=torch.float32, device=device)
 
-    dKV_postprocess = qhead_per_kvhead > 1
+    # GQA (qhead_per_kvhead > 1) needs dK/dV accum+postprocess since multiple Q heads
+    # accumulate into the same dK/dV. SM90 varlen_k with qhead_per_kvhead==1 now uses
+    # ragged TMA tensors for direct store, so no longer needs accum+postprocess.
+    # hd=256 2CTA backward has its own internal postprocess for dK/dV.
+    dKV_postprocess = qhead_per_kvhead > 1 and not use_dedicated_hd256_kernel
     if dKV_postprocess:
         head_dim_v_rounded = (head_dim_v + 32 - 1) // 32 * 32
         if cu_seqlens_k is None:
@@ -850,96 +2209,64 @@ def _flash_attn_bwd(
             )
 
     dtype = torch2cute_dtype_map[q.dtype]
-    current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
 
     if deterministic:
-        dQ_semaphore = torch.zeros(batch_size, num_head, seqlen_q_rounded // m_block_size, cluster_size, dtype=torch.int32, device="cuda")
+        dQ_semaphore = torch.zeros(batch_size, num_head, seqlen_q_rounded // m_block_size, cluster_size, dtype=torch.int32, device=device)
     else:
         dQ_semaphore = None
 
     if deterministic and qhead_per_kvhead > 1:
-        dK_semaphore = torch.zeros(batch_size, num_head_kv, seqlen_k_rounded // n_block_size, 2, dtype=torch.int32, device="cuda")
-        dV_semaphore = torch.zeros(batch_size, num_head_kv, seqlen_k_rounded // n_block_size, 2, dtype=torch.int32, device="cuda")
+        dK_semaphore = torch.zeros(batch_size, num_head_kv, seqlen_k_rounded // n_block_size, 2, dtype=torch.int32, device=device)
+        dV_semaphore = torch.zeros(batch_size, num_head_kv, seqlen_k_rounded // n_block_size, 2, dtype=torch.int32, device=device)
     else:
         dK_semaphore = None
         dV_semaphore = None
 
-    # Preprocess kernel: compute (o * dout).sum(dim=-1), lse * log2_e, and zero out dq_accum.
-    compile_key_pre = (
-        arch,
-        dtype,
-        head_dim,
-        head_dim_v,
-        m_block_size,
-        num_threads,
-        cu_seqlens_q is None,
-        seqused_q is None,
-        get_broadcast_dims(out),
-        get_broadcast_dims(dout),
-    )
-    if compile_key_pre not in _flash_attn_bwd.compile_cache_pre:
-        o_tensor, do_tensor = [to_cute_tensor(t) for t in (out, dout)]
-        dq_accum_tensor, dpsum_tensor, lse_log2_tensor = [
-            to_cute_tensor(t) for t in (dq_accum, dpsum, lse_log2)
-        ]
-        lse_tensor = to_cute_tensor(lse, assumed_align=4)
-        cu_seqlens_q_tensor, seqused_q_tensor = [
-            to_cute_tensor(t, assumed_align=4) if t is not None else None
-            for t in (cu_seqlens_q, seqused_q)
-        ]
-        fa_bwd_pre = FlashAttentionBackwardPreprocess(
-            dtype,
-            head_dim,
-            head_dim_v,
-            arch,
-            m_block_size,
-            num_threads=num_threads,
+    # SingleTileVarlenScheduler batch-lookup aid, above BIN_BATCH_SEARCH_THRESH;
+    # shared across preprocess, main bwd, and the three postprocess calls.
+    cu_total_m_blocks_q = None
+    cu_total_m_blocks_k = None
+    if is_varlen and batch_size > BIN_BATCH_SEARCH_THRESH and not use_dedicated_hd256_kernel:
+        cu_total_m_blocks_q, _ = _compute_tile_cumsum(
+            cu_seqlens=cu_seqlens_q,
+            seqused=seqused_q,
+            tile_size=m_block_size,
         )
-        # TODO: check @can_implement
-        _flash_attn_bwd.compile_cache_pre[compile_key_pre] = cute.compile(
-            fa_bwd_pre,
-            o_tensor,
-            do_tensor,
-            dpsum_tensor,
-            lse_tensor,
-            lse_log2_tensor,
-            dq_accum_tensor,
-            cu_seqlens_q_tensor,
-            seqused_q_tensor,
-            current_stream,
-            options="--enable-tvm-ffi",
-        )
-    if not is_fake_mode():
-        _flash_attn_bwd.compile_cache_pre[compile_key_pre](
-            out,
-            dout,
-            dpsum,
-            lse,
-            lse_log2,
-            dq_accum,
-            cu_seqlens_q,
-            seqused_q,
-            current_stream,
+        cu_total_m_blocks_k, _ = _compute_tile_cumsum(
+            cu_seqlens=cu_seqlens_k,
+            seqused=seqused_k,
+            tile_size=n_block_size,
+            cluster_shape_m=cluster_size,
         )
 
-    # NB num_threads application for 3 kernels
-    # There are pre, main, post processing kernels, currenlty num_threads is only actually
-    # used for the pre proc, and then we hard code to 384 for the main and post proc, and we do
-    # before cache key gen
-    num_threads = 384
+    dsink = torch.empty_like(learnable_sink) if learnable_sink is not None else None
+
+    # Preprocess kernel: compute (o * dout).sum(dim=-1) - dLSE, lse * log2_e, and zero out dq_accum.
+    # For hd=256 dedicated path, dq_accum is None so preprocess only fills dpsum/lse_log2.
+    _bwd_preprocess(
+        out, dout, dpsum, lse, lse_log2, dq_accum,
+        cu_seqlens_q, seqused_q, dlse,
+        dtype, head_dim, head_dim_v, m_block_size,
+        cu_total_m_blocks=cu_total_m_blocks_q,
+        fake_mode=fake_mode,
+        hdim_multiple_of=hdim_multiple_of,
+    )
+    # num_threads: SM90 derives from BwdConfig.num_wg, SM120 is set to 128 above,
+    # SM100/SM110 uses default from function signature (384).
+    if arch // 10 not in [9, 12]:
+        num_threads = 384
 
     # Backward kernel: compute dk, dv, dq_accum.
     score_mod_hash = utils.hash_callable(score_mod) if score_mod else False
     score_mod_bwd_hash = utils.hash_callable(score_mod_bwd) if score_mod_bwd else False
     mask_mod_hash = utils.hash_callable(mask_mod) if mask_mod else False
     num_aux_tensors = len(aux_tensors) if aux_tensors else 0
-    cute_aux_tensors = None
-    if aux_tensors is not None:
-        cute_aux_tensors = [to_cute_tensor(buf, assumed_align=None, fully_dynamic=True) for buf in aux_tensors]
+    aux_tensor_metadata = get_aux_tensor_metadata(aux_tensors) if aux_tensors is not None else None
+    aux_scalar_metadata = tuple(type(s) for s in aux_scalars) if aux_scalars is not None else None
 
     block_sparse_broadcast_pattern = None
     normalized_block_sparse_tensors = None
-    if block_sparse_tensors is not None:
+    if use_block_sparsity:
         (
             normalized_block_sparse_tensors,
             block_sparse_broadcast_pattern,
@@ -950,10 +2277,35 @@ def _flash_attn_bwd(
             seqlen_q=seqlen_q,
             seqlen_k=seqlen_k,
             block_size=(m_block_size, n_block_size),
-            subtile_factor=subtile_factor,
+            q_subtile_factor=q_subtile_factor,
+            kv_subtile_factor=kv_subtile_factor,
         )
+        if deterministic:
+            if normalized_block_sparse_tensors.dq_write_order is None:
+                raise ValueError(
+                    "deterministic block-sparse backward requires dq_write_order in block_sparse_tensors"
+                )
+            if (
+                normalized_block_sparse_tensors.full_block_cnt is not None
+                and normalized_block_sparse_tensors.dq_write_order_full is None
+            ):
+                raise ValueError(
+                    "deterministic block-sparse backward requires dq_write_order_full when full blocks are present"
+                )
+            if normalized_block_sparse_tensors.spt is None:
+                raise ValueError(
+                    "deterministic block-sparse backward requires block_sparse_tensors.spt "
+                    "to match dq_write_order direction"
+                )
+    if (
+        normalized_block_sparse_tensors is not None
+        and normalized_block_sparse_tensors.spt is not None
+    ):
+        spt = normalized_block_sparse_tensors.spt and deterministic
+    else:
+        spt = (causal or local) and deterministic
 
-    if arch // 10 == 9:
+    if arch // 10 in [8, 9, 12]:
         compile_key = (
             arch,
             dtype,
@@ -961,7 +2313,8 @@ def _flash_attn_bwd(
             head_dim_v,
             qhead_per_kvhead,
             causal,
-            softcap != 0.0,
+            window_size_left is not None,
+            window_size_right is not None,
             m_block_size,
             n_block_size,
             num_threads,
@@ -975,6 +2328,8 @@ def _flash_attn_bwd(
             AtomLayoutNdKV,
             AtomLayoutMdQ,
             V_in_regs,
+            dQ_single_wg,
+            deterministic,
             cu_seqlens_q is None,
             cu_seqlens_k is None,
             seqused_q is None,
@@ -983,12 +2338,19 @@ def _flash_attn_bwd(
             score_mod_bwd_hash,
             mask_mod_hash,
             num_aux_tensors,
+            aux_tensor_metadata,
+            aux_scalar_metadata,
             use_block_sparsity,
+            q_subtile_factor,
             block_sparse_broadcast_pattern,
             get_broadcast_dims(q),
             get_broadcast_dims(k),
             get_broadcast_dims(v),
             get_broadcast_dims(dout),
+            # Prevent TVM stride poisoning when only one block is present.
+            single_q_block,
+            single_k_block,
+            cu_total_m_blocks_k is not None,
         )
     else:
         compile_key = (
@@ -1000,18 +2362,22 @@ def _flash_attn_bwd(
             causal,
             window_size_left is not None,
             window_size_right is not None,
-            softcap != 0.0,
             m_block_size,
             n_block_size,
             num_threads,
             pack_gqa,
             cluster_size,
             use_2cta_instrs,
+            q_subtile_factor,
+            kv_subtile_factor,
             deterministic,
+            spt,
             score_mod_hash,
             score_mod_bwd_hash,
             mask_mod_hash,
             num_aux_tensors,
+            aux_tensor_metadata,
+            aux_scalar_metadata,
             use_block_sparsity,
             block_sparse_broadcast_pattern,
             cu_seqlens_q is None,
@@ -1022,99 +2388,151 @@ def _flash_attn_bwd(
             get_broadcast_dims(k),
             get_broadcast_dims(v),
             get_broadcast_dims(dout),
+            # Prevent TVM stride poisoning when only one block is present.
+            single_q_block,
+            single_k_block,
+            cu_total_m_blocks_k is not None,
         )
+
     if compile_key not in _flash_attn_bwd.compile_cache:
+        current_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+        cute_aux_tensors = (
+            [to_cute_aux_tensor(buf) for buf in aux_tensors]
+            if aux_tensors is not None
+            else None
+        )
         q_tensor, k_tensor, v_tensor, do_tensor, dq_tensor, dk_tensor, dv_tensor = [
             to_cute_tensor(t) for t in (q, k, v, dout, dq, dk, dv)
         ]
-        dq_accum_tensor, dpsum_tensor, lse_log2_tensor = [
-            to_cute_tensor(t) for t in (dq_accum, dpsum, lse_log2)
-        ]
+        lse_log2_tensor, dpsum_tensor = [to_cute_tensor(t) for t in (lse_log2, dpsum)]
+        dq_accum_tensor = to_cute_tensor(dq_accum) if dq_accum is not None else None
         if dKV_postprocess:
             dk_accum_tensor, dv_accum_tensor = [
                 to_cute_tensor(t) for t in (dk_accum, dv_accum)
             ]
-        cu_seqlens_q_tensor, cu_seqlens_k_tensor, seqused_q_tensor, seqused_k_tensor = [
+        cu_seqlens_q_tensor, cu_seqlens_k_tensor, seqused_q_tensor, seqused_k_tensor, cu_total_m_blocks_k_tensor = [
             to_cute_tensor(t, assumed_align=4) if t is not None else None
-            for t in (cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k)
+            for t in (cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, cu_total_m_blocks_k)
         ]
         dQ_semaphore_tensor, dK_semaphore_tensor, dV_semaphore_tensor = [
             utils.convert_from_dlpack_leading_static(t.detach(), leading_dim=3, alignment=4, stride_order=t.dim_order())
             if t is not None else None
             for t in (dQ_semaphore, dK_semaphore, dV_semaphore)
         ]
-        fa_bwd_sm80 = FlashAttentionBackwardSm80(
-            dtype,
-            head_dim,
-            head_dim_v,
-            qhead_per_kvhead,
-            m_block_size,
-            n_block_size,
-            num_stages_Q,
-            num_stages_dO,
-            num_threads,
-            pack_gqa,
-            causal,
-            SdP_swapAB,
-            dKV_swapAB,
-            dQ_swapAB,
-            AtomLayoutMSdP,
-            AtomLayoutNdKV,
-            AtomLayoutMdQ,
-            V_in_regs=V_in_regs,
-        )
-        if arch // 10 == 9:
-            fa_bwd_obj = FlashAttentionBackwardSm90(
+        if arch // 10 in [8, 12]:
+            flash_bwd_obj_cls = FlashAttentionBackwardSm120 if arch // 10 == 12 else FlashAttentionBackwardSm80
+            fa_bwd_obj = flash_bwd_obj_cls(
                 dtype,
                 head_dim,
                 head_dim_v,
                 qhead_per_kvhead,
-                causal,
                 m_block_size,
                 n_block_size,
                 num_stages_Q,
                 num_stages_dO,
-                num_stages_PdS,
+                num_threads,
+                pack_gqa,
+                causal,
+                local,
                 SdP_swapAB,
                 dKV_swapAB,
                 dQ_swapAB,
                 AtomLayoutMSdP,
                 AtomLayoutNdKV,
                 AtomLayoutMdQ,
-                num_threads,
                 V_in_regs=V_in_regs,
                 score_mod=score_mod,
                 score_mod_bwd=score_mod_bwd,
-                mask_mod=mask_mod,
-                has_aux_tensors=aux_tensors is not None,
-                subtile_factor=subtile_factor,
             )
-        else:
-            fa_bwd_obj = FlashAttentionBackwardSm100(
+        elif arch // 10 == 9:
+            fa_bwd_obj = FlashAttentionBackwardSm90(
+                dtype,
                 head_dim,
                 head_dim_v,
-                is_causal=causal,
+                qhead_per_kvhead,
+                causal,
                 is_local=local,
-                qhead_per_kvhead=qhead_per_kvhead,
+                deterministic=deterministic,
                 tile_m=m_block_size,
                 tile_n=n_block_size,
-                cluster_size=cluster_size,
-                use_2cta_instrs=use_2cta_instrs,
-                deterministic=deterministic,
+                Q_stage=num_stages_Q,
+                dO_stage=num_stages_dO,
+                PdS_stage=num_stages_PdS,
+                SdP_swapAB=SdP_swapAB,
+                dKV_swapAB=dKV_swapAB,
+                dQ_swapAB=dQ_swapAB,
+                AtomLayoutMSdP=AtomLayoutMSdP,
+                AtomLayoutNdKV=AtomLayoutNdKV,
+                AtomLayoutMdQ=AtomLayoutMdQ,
+                num_threads=num_threads,
+                V_in_regs=V_in_regs,
                 score_mod=score_mod,
                 score_mod_bwd=score_mod_bwd,
                 mask_mod=mask_mod,
                 has_aux_tensors=aux_tensors is not None,
-                subtile_factor=subtile_factor,
+                q_subtile_factor=q_subtile_factor,
+                dQ_single_wg=dQ_single_wg,
             )
+        else:
+            if use_dedicated_hd256_kernel:
+                assert softcap == 0.0, "SM100 backward with head_dim=256 does not support softcap"
+                assert block_sparse_tensors is None, \
+                    "SM100 backward with head_dim=256 does not support block sparsity"
+                assert dlse is None, \
+                    "SM100 backward with head_dim=256 does not support dlse"
+                assert seqused_q is None and seqused_k is None, \
+                    "SM100 backward with head_dim=256 does not support seqused_q/seqused_k"
+
+                dq_tile_mn = (128, 128)
+                dkdv_tile_mn = (128, 64)
+                fa_bwd_obj = BlackwellFusedMultiHeadAttentionBackward(
+                    head_dim,
+                    head_dim_v,
+                    is_causal=causal,
+                    is_local=local,
+                    qhead_per_kvhead=qhead_per_kvhead,
+                    is_persistent=False,
+                    deterministic=deterministic,
+                    cluster_size=cluster_size,
+                    use_2cta_instrs=use_2cta_instrs,
+                    score_mod=score_mod,
+                    score_mod_bwd=score_mod_bwd,
+                    mask_mod=mask_mod,
+                    has_aux_tensors=aux_tensors is not None,
+                    q_subtile_factor=q_subtile_factor,
+                    tile_m_dq=dq_tile_mn[0],
+                    tile_n_dq=dq_tile_mn[1],
+                    tile_m_dkdv=dkdv_tile_mn[0],
+                    tile_n_dkdv=dkdv_tile_mn[1],
+                )
+            else:
+                fa_bwd_obj = FlashAttentionBackwardSm100(
+                    head_dim,
+                    head_dim_v,
+                    is_causal=causal,
+                    is_local=local,
+                    qhead_per_kvhead=qhead_per_kvhead,
+                    tile_m=m_block_size,
+                    tile_n=n_block_size,
+                    cluster_size=cluster_size,
+                    use_2cta_instrs=use_2cta_instrs,
+                    deterministic=deterministic,
+                    spt=spt,
+                    score_mod=score_mod,
+                    score_mod_bwd=score_mod_bwd,
+                    mask_mod=mask_mod,
+                    has_aux_tensors=aux_tensors is not None,
+                    q_subtile_factor=q_subtile_factor,
+                    kv_subtile_factor=kv_subtile_factor,
+                )
 
         # Block sparse tensors for backward use Q-direction indexing (transposed from forward).
         sparse_tensors_compile = None
         if normalized_block_sparse_tensors is not None:
             sparse_tensors_compile = to_cute_block_sparse_tensors(normalized_block_sparse_tensors)
+        dq_accum_tensor = dq_tensor if use_dedicated_hd256_kernel else dq_accum_tensor
 
-        # TODO: check @can_implement
-        _flash_attn_bwd.compile_cache[compile_key] = cute.compile(
+        compile_args = [
             fa_bwd_obj,
             q_tensor,
             k_tensor,
@@ -1126,23 +2544,29 @@ def _flash_attn_bwd(
             dk_tensor if not dKV_postprocess else dk_accum_tensor,
             dv_tensor if not dKV_postprocess else dv_accum_tensor,
             softmax_scale,
-            current_stream,
             cu_seqlens_q_tensor,
             cu_seqlens_k_tensor,
             seqused_q_tensor,
             seqused_k_tensor,
-            None,  # softcap - not yet supported in backward
             window_size_left,
             window_size_right,
             dQ_semaphore_tensor,
             dK_semaphore_tensor,
             dV_semaphore_tensor,
-            cute_aux_tensors,
+            AuxData(cute_aux_tensors, aux_scalars),
             sparse_tensors_compile,
-            options="--enable-tvm-ffi",
+        ]
+        if not use_dedicated_hd256_kernel:
+            compile_args.append(cu_total_m_blocks_k_tensor)
+        compile_args.append(current_stream)
+
+        # TODO: check @can_implement
+        _flash_attn_bwd.compile_cache[compile_key] = cute.compile(
+            *compile_args, options="--enable-tvm-ffi"
         )
-    if not is_fake_mode():
-        _flash_attn_bwd.compile_cache[compile_key](
+    if not fake_mode:
+        dq_accum = dq if use_dedicated_hd256_kernel else dq_accum
+        call_args = [
             q.detach(),
             k.detach(),
             v.detach(),
@@ -1153,172 +2577,504 @@ def _flash_attn_bwd(
             dk if not dKV_postprocess else dk_accum,
             dv if not dKV_postprocess else dv_accum,
             softmax_scale,
-            current_stream,
             cu_seqlens_q,
             cu_seqlens_k,
             seqused_q,
             seqused_k,
-            None,  # softcap - not yet supported in backward
             window_size_left,
             window_size_right,
             dQ_semaphore,
             dK_semaphore,
             dV_semaphore,
-            aux_tensors,
-            normalized_block_sparse_tensors[:4] if normalized_block_sparse_tensors is not None else None,
+            AuxData(aux_tensors, aux_scalars),
+            (
+                normalized_block_sparse_tensors.mask_block_cnt,
+                normalized_block_sparse_tensors.mask_block_idx,
+                normalized_block_sparse_tensors.full_block_cnt,
+                normalized_block_sparse_tensors.full_block_idx,
+                normalized_block_sparse_tensors.cu_total_m_blocks,
+                normalized_block_sparse_tensors.cu_block_idx_offsets,
+                normalized_block_sparse_tensors.dq_write_order,
+                normalized_block_sparse_tensors.dq_write_order_full,
+            )
+            if normalized_block_sparse_tensors is not None
+            else None,
+        ]
+        if not use_dedicated_hd256_kernel:
+            call_args.append(cu_total_m_blocks_k)
+        _flash_attn_bwd.compile_cache[compile_key](*call_args)
+    # Postprocess: convert dq_accum from float32 to dq in bf16/fp16
+    # hd=256 2CTA backward has its own internal postprocess, skip here.
+    if not use_dedicated_hd256_kernel:
+        if arch // 10 == 9:
+            # dQ postprocess: match main kernel's MMA WG count, unless dQ_single_wg
+            num_threads_post_dQ = 128 if dQ_single_wg else cfg.num_wg * 128
+            num_threads_post_dKV = cfg.num_wg * 128
+        else:
+            num_threads_post_dQ = 128
+            num_threads_post_dKV = 128
+
+        _bwd_postprocess_convert(
+            dq_accum, dq, softmax_scale,
+            cu_seqlens_q, seqused_q,
+            arch, dtype, head_dim, m_block_size, num_threads_post_dQ,
+            AtomLayoutMdQ, dQ_swapAB,
+            use_2cta_instrs=use_2cta_instrs, cluster_size=1,
+            cu_total_m_blocks=cu_total_m_blocks_q,
+            sink_tensors=(
+                LearnableSinkBwdTensors(dpsum, lse, learnable_sink, dsink)
+                if learnable_sink is not None
+                else None
+            ),
+            fake_mode=fake_mode,
+            hdim_multiple_of=hdim_multiple_of,
         )
 
-    num_threads = 256 if arch // 10 == 9 else 128
-    # Postprocess kernel: convert dq_accum from float32 to dq in bf16/fp16
-    compile_key_post = (
-        arch,
+        if dKV_postprocess:
+            # Postprocess: convert dk_accum from float32 to dk in bf16/fp16
+            _bwd_postprocess_convert(
+                dk_accum, dk, softmax_scale,
+                cu_seqlens_k, seqused_k,
+                arch, dtype, head_dim, n_block_size, num_threads_post_dKV,
+                AtomLayoutNdKV, dKV_swapAB,
+                cluster_size=cluster_size,
+                cu_total_m_blocks=cu_total_m_blocks_k if cluster_size == 1 else None,
+                fake_mode=fake_mode,
+                hdim_multiple_of=hdim_multiple_of,
+            )
+            # Postprocess: convert dv_accum from float32 to dv in bf16/fp16
+            _bwd_postprocess_convert(
+                dv_accum, dv, 1.0,
+                cu_seqlens_k, seqused_k,
+                arch, dtype, head_dim_v, n_block_size, num_threads_post_dKV,
+                AtomLayoutNdKV, dKV_swapAB,
+                cluster_size=cluster_size,
+                cu_total_m_blocks=cu_total_m_blocks_k if cluster_size == 1 else None,
+                fake_mode=fake_mode,
+                hdim_multiple_of=hdim_multiple_of,
+            )
+
+    return (dq, dk, dv) if learnable_sink is None else (dq, dk, dv, dsink)
+
+
+_flash_attn_bwd.compile_cache = get_jit_cache("bwd")
+
+
+def _flash_attn_bwd_sparse_mla(
+    q: Optional[torch.Tensor],
+    k: Optional[torch.Tensor],
+    v: torch.Tensor,
+    qv: torch.Tensor,
+    out: torch.Tensor,
+    dout: torch.Tensor,
+    lse: torch.Tensor,
+    p: torch.Tensor,
+    row_max: torch.Tensor,
+    gather_kv_indices: torch.Tensor,
+    learnable_sink: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    m_block_size: int = 128,
+    n_block_size: int = 64,
+    num_threads: int = 256,
+    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,
+    max_seqlen_q: Optional[int] = None,
+    max_seqlen_k: Optional[int] = None,
+    min_seqlen_k: Optional[int] = None,
+    deterministic: bool = False,
+    dq: Optional[torch.Tensor] = None,
+    dk: Optional[torch.Tensor] = None,
+    dv: Optional[torch.Tensor] = None,
+    dqv: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    fake_mode = is_fake_mode()
+    arch = _get_device_arch()
+    assert arch // 10 in [10, 11], "Unsupported compute capability. Supported: 10.x, 11.x"
+    assert gather_kv_indices is not None, "require gather kv indices for backward"
+
+    q_shape = q.shape if q is not None else qv.shape
+    nheads, head_dim = q_shape[-2:]
+    nheads_kv, head_dim_v = v.shape[-2:]
+    qhead_per_kvhead = nheads // nheads_kv
+    gather_kv_length = gather_kv_indices.shape[-1]
+    assert nheads_kv == 1 and qhead_per_kvhead == 128, f"sparse MLA bwd: only MQA 128 supported for now"
+    assert gather_kv_length % 128 == 0, f"sparse MLA bwd: {gather_kv_length=} must be divisible by 128"
+    assert deterministic is False, "sparse MLA bwd: deterministic mode not yet supported"
+    assert learnable_sink is None, "sparse MLA bwd: learnable sink not yet supported"
+    assert seqused_q is None and seqused_k is None, "sparse MLA bwd: seqused_q,k not yet supported"
+
+    if softmax_scale is None:
+        softmax_scale = (
+            1.0 / math.sqrt(head_dim) if qv is None or q is None
+            else 1.0 / math.sqrt(head_dim + head_dim_v)
+        )
+
+    q, k, v, qv, out, dout, lse, p, row_max = [
+        maybe_contiguous(t)
+        for t in (q, k, v, qv, out, dout, lse, p, row_max)
+    ]
+    gather_kv_indices, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, learnable_sink = [
+        maybe_contiguous(t)
+        for t in (gather_kv_indices, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, learnable_sink)
+    ]
+    device = v.device
+
+    varlen_q = cu_seqlens_q is not None or seqused_q is not None
+    if cu_seqlens_q is None:
+        batch_size, seqlen_q = q_shape[:2]
+        total_q = batch_size * seqlen_q
+        p_shape = (batch_size, seqlen_q, nheads, gather_kv_length)
+    else:
+        batch_size = cu_seqlens_q.shape[0] - 1
+        total_q = q_shape[0]
+        seqlen_q = max_seqlen_q if max_seqlen_q is not None else total_q
+        p_shape = (total_q, nheads, gather_kv_length)
+
+    varlen_k = cu_seqlens_k is not None or seqused_k is not None
+    if cu_seqlens_k is None:
+        batch_size, seqlen_k = v.shape[:2]
+        total_k = batch_size * seqlen_k
+    else:
+        batch_size = cu_seqlens_k.shape[0] - 1
+        total_k = v.shape[0]
+        seqlen_k = max_seqlen_k if max_seqlen_k is not None else total_k
+    if not varlen_k:
+        min_seqlen_k = seqlen_k 
+
+    assert varlen_q == varlen_k, "sparse MLA bwd: either q and k are both varlen or not"
+
+    # always use kv bitmask by default (handles -1 sentinel)
+    disable_sparse_kv_bitmask = False
+    # if min_seqlen_k is None or causal:
+    #     disable_sparse_kv_bitmask = False
+    # else:
+    #     disable_sparse_kv_bitmask = min_seqlen_k >= gather_kv_length
+
+    prealloc_dq = dq is not None
+    prealloc_dk = dk is not None
+    prealloc_dqv = dqv is not None
+    prealloc_dv = dv is not None
+    if not prealloc_dq and q is not None:
+        dq = torch.empty_like(q)
+    if not prealloc_dk and k is not None:
+        dk = torch.zeros_like(k, dtype=torch.float32)
+    if not prealloc_dv:
+        dv = torch.zeros_like(v, dtype=torch.float32)
+    if not prealloc_dqv:
+        dqv = torch.empty_like(qv)
+    ds = torch.empty_like(p)
+
+    device = v.device
+    dtype = v.dtype
+    if q is not None:
+        _validate_tensor(dq, "dq", q.shape, dtype, device)
+    if k is not None:
+        _validate_tensor(dk, "dk", k.shape, torch.float32, device)
+    _validate_tensor(dv, "dv", v.shape, torch.float32, device)
+    _validate_tensor(dqv, "dqv", qv.shape, dtype, device)
+    _validate_tensor(p, "p", p_shape, dtype, device)
+
+    if cu_seqlens_q is None:
+        dpsum = torch.empty(batch_size, seqlen_q, nheads, dtype=torch.float32, device=device)
+    else:
+        dpsum = torch.empty(total_q, nheads, dtype=torch.float32, device=device)
+    scale_p = torch.empty_like(row_max)
+
+    dtype = torch2cute_dtype_map[dout.dtype]
+
+    # Preprocess kernel: compute (o * dout).sum(dim=-1), scale_p.
+    _bwd_preprocess(
+        out, dout, dpsum, lse, None, None,
+        cu_seqlens_q, seqused_q, None,
+        dtype, head_dim, head_dim_v, m_block_size,
+        row_max=row_max,
+        scale_p=scale_p,
+        use_padded_offsets=False,
+        nheads_major=True,
+        pack_gqa=True,
+        qhead_per_kvhead=qhead_per_kvhead,
+        nheads_kv=nheads_kv,
+        softmax_scale=softmax_scale,
+        fake_mode=fake_mode,
+    )
+
+    compile_key = (
         dtype,
         head_dim,
-        m_block_size,
-        num_threads,
-        AtomLayoutMdQ,
-        dQ_swapAB,
+        head_dim_v,
+        qhead_per_kvhead,
+        causal,
         cu_seqlens_q is None,
+        cu_seqlens_k is None,
         seqused_q is None,
-        use_2cta_instrs,
-        1, # no cluster for tile_m
-        get_broadcast_dims(dq_accum),
-        get_broadcast_dims(dq),
+        seqused_k is None,
+        q is not None,
+        gather_kv_length,
+        learnable_sink is not None,
+        disable_sparse_kv_bitmask,
     )
-    if compile_key_post not in _flash_attn_bwd.compile_cache_post:
-        dq_accum_tensor = to_cute_tensor(dq_accum)
-        dq_tensor = to_cute_tensor(dq)
-        cu_seqlens_q_tensor, seqused_q_tensor = [
-            to_cute_tensor(t, assumed_align=4) if t is not None else None
-            for t in (cu_seqlens_q, seqused_q)
+
+    if compile_key not in _flash_attn_bwd_sparse_mla.compile_cache:
+        current_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+        (
+            cu_seqlens_q_tensor,
+            cu_seqlens_k_tensor,
+            seqused_q_tensor,
+            seqused_k_tensor,
+            learnable_sink_tensor,
+        ) = [
+            to_cute_tensor(t, assumed_align=4, leading_dim=0)
+            for t in (cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, learnable_sink)
+        ]
+        (
+            v_tensor,
+            qv_tensor,
+            do_tensor,
+            p_tensor,
+            scale_p_tensor,
+            dpsum_tensor,
+            ds_tensor,
+            dv_tensor,
+            gather_kv_indices_tensor,
+         ) = [
+            to_cute_tensor(t) for t in (v, qv, dout, p, scale_p, dpsum, ds, dv, gather_kv_indices)
         ]
-        fa_bwd_post = FlashAttentionBackwardPostprocess(
-            dtype, head_dim, arch, m_block_size, num_threads, AtomLayoutMdQ, dQ_swapAB,
-            use_2cta_instrs=use_2cta_instrs,
+
+        fa_bwd_obj = FlashAttentionSparseMLABackwardSm100(
+            is_causal=causal,
+            topk_length=gather_kv_length,
+            qhead_per_kvhead=qhead_per_kvhead,
+            nheads_kv=nheads_kv,
+            has_seqused_q=seqused_q is not None,
+            disable_bitmask=disable_sparse_kv_bitmask,
         )
-        # TODO: check @can_implement
-        _flash_attn_bwd.compile_cache_post[compile_key_post] = cute.compile(
-            fa_bwd_post,
-            dq_accum_tensor,
-            dq_tensor,
+        fa_bwd_kernel = cute.compile(
+            fa_bwd_obj,
+            do_tensor,
+            v_tensor,
+            qv_tensor,
+            p_tensor,
+            dv_tensor,
+            ds_tensor,
+            gather_kv_indices_tensor,
             softmax_scale,
+            scale_p_tensor,
+            dpsum_tensor,
             cu_seqlens_q_tensor,
+            cu_seqlens_k_tensor,
             seqused_q_tensor,
+            seqused_k_tensor,
             current_stream,
             options="--enable-tvm-ffi",
         )
+        _flash_attn_bwd_sparse_mla.compile_cache[compile_key] = fa_bwd_kernel
 
-    if not is_fake_mode():
-        _flash_attn_bwd.compile_cache_post[compile_key_post](
-            dq_accum,
-            dq,
+    if not fake_mode:
+        _flash_attn_bwd_sparse_mla.compile_cache[compile_key](
+            dout,
+            v,
+            qv,
+            p,
+            dv,
+            ds,
+            gather_kv_indices,
             softmax_scale,
+            scale_p,
+            dpsum,
             cu_seqlens_q,
+            cu_seqlens_k,
             seqused_q,
-            current_stream,
+            seqused_k,
         )
 
-    if dKV_postprocess:
-        # Postprocess kernel: convert dk_accum & dv_accum from float32 to bf16/fp16
-        compile_key_post = (
-            arch,
-            dtype,
-            head_dim,
-            n_block_size,
-            num_threads,
-            AtomLayoutNdKV,
-            dKV_swapAB,
-            cu_seqlens_k is None,
-            seqused_k is None,
-            False, # even for 2cta, is split along hdim, so always False
-            cluster_size, # cluster is for tile_n
-            get_broadcast_dims(dk_accum),
-            get_broadcast_dims(dk),
-        )
-        if compile_key_post not in _flash_attn_bwd.compile_cache_post:
-            dk_accum_tensor = to_cute_tensor(dk_accum)
-            dk_tensor = to_cute_tensor(dk)
-            cu_seqlens_k_tensor, seqused_k_tensor = [
-                to_cute_tensor(t, assumed_align=4) if t is not None else None
-                for t in (cu_seqlens_k, seqused_k)
-            ]
-            fa_bwd_post = FlashAttentionBackwardPostprocess(
-                dtype, head_dim, arch, n_block_size, num_threads, AtomLayoutNdKV, dKV_swapAB,
-                cluster_size=cluster_size,
-            )
-            # TODO: check @can_implement
-            _flash_attn_bwd.compile_cache_post[compile_key_post] = cute.compile(
-                fa_bwd_post,
-                dk_accum_tensor,
-                dk_tensor,
-                softmax_scale,
-                cu_seqlens_k_tensor,
-                seqused_k_tensor,
-                current_stream,
-                options="--enable-tvm-ffi",
-            )
-        if not is_fake_mode():
-            _flash_attn_bwd.compile_cache_post[compile_key_post](
-                dk_accum,
-                dk,
-                softmax_scale,
-                cu_seqlens_k,
-                seqused_k,
-                current_stream,
-            )
-        compile_key_post = (
-            arch,
-            dtype,
-            head_dim_v,
-            n_block_size,
-            num_threads,
-            AtomLayoutNdKV,
-            dKV_swapAB,
-            cu_seqlens_k is None,
-            seqused_k is None,
-            False,
-            cluster_size,
-            get_broadcast_dims(dv_accum),
-            get_broadcast_dims(dv),
+    v = v.squeeze(-2)
+    if k is not None:
+        k = k.squeeze(-2)
+    
+    _sparse_mla_dq_dqv(
+        ds, k, v, dq, dqv, gather_kv_indices, cu_seqlens_q, cu_seqlens_k,
+    )
+
+    if k is not None:
+        dk = dk.squeeze(-2)
+        _sparse_mla_dk(ds, gather_kv_indices, q, dk, cu_seqlens_q, cu_seqlens_k)
+        dk = dk.unsqueeze(-2)
+    
+    # return dk, dv in float32: all-reduce across sequence-parallel ranks must happen
+    # before downcasting to avoid rounding error during inter-rank grad accumulation
+    return dq, dk, dv, dqv
+
+_flash_attn_bwd_sparse_mla.compile_cache = get_jit_cache("bwd_dsa")
+
+
+def _compile_sparse_mla_dq_dqv(
+    dtype, nheads, head_dim, head_dim_v, top_k, varlen_q, varlen_k, compute_dq,
+):
+    sym = cute.sym_int 
+    b, b_plus_1, seqlen_q, seqlen_k = sym(), sym(), sym(), sym()
+    total_q, total_k = sym(), sym()
+    b_seqlenq = (b, seqlen_q) if not varlen_q else (total_q,)
+    b_seqlenk = (b, seqlen_k) if not varlen_k else (total_k,)
+    
+    div = 128 // dtype.width  # 8 for fp16/bf16
+    
+    mdS = fake_tensor(dtype, (*b_seqlenq, nheads, top_k), divisibility=div)
+    mK = fake_tensor(dtype, (*b_seqlenk, head_dim), divisibility=div)
+    mV = fake_tensor(dtype, (*b_seqlenk, head_dim_v), divisibility=div)
+    mdQ = fake_tensor(dtype, (*b_seqlenq, nheads, head_dim), divisibility=div)
+    mdQv = fake_tensor(dtype, (*b_seqlenq, nheads, head_dim_v), divisibility=div)
+    mIdxTopK = fake_tensor(Int32, (*b_seqlenq, top_k), divisibility=div)
+    
+    mCuSeqlensQ = fake_tensor(Int32, (b_plus_1,), divisibility=1) if varlen_q else None 
+    mCuSeqlensK = fake_tensor(Int32, (b_plus_1,), divisibility=1) if varlen_k else None 
+    
+    dq_dqv_gemm = dQdQvGemmKernel(
+        acc_dtype=Float32,
+        nheads=nheads,
+        head_dim_k=head_dim,
+        head_dim_v=head_dim_v,
+        top_k=top_k,
+    )
+    
+    return cute.compile(
+        dq_dqv_gemm,
+        mdS,
+        mK if compute_dq else None,
+        mV,
+        mdQ if compute_dq else None,
+        mdQv,
+        mIdxTopK,
+        mCuSeqlensQ,
+        mCuSeqlensK,
+        cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+        options="--enable-tvm-ffi",
+    )
+
+
+def _sparse_mla_dq_dqv(
+    ds, k, v, dq, dqv, gather_kv_indices, cu_seqlens_q, cu_seqlens_k,
+):
+    """Compute dQ = dS @ K and dQv = dS @ V"""
+    *_, nheads, gather_kv_length = ds.shape
+    
+    head_dim_v = v.shape[-1]
+    head_dim = k.shape[-1] if k is not None else 0
+    
+    dtype = ds.dtype
+    dtype_cute = torch2cute_dtype_map[dtype]
+    
+    varlen_q = cu_seqlens_q is not None
+    varlen_k = cu_seqlens_k is not None
+    
+    compile_key = (
+        dtype_cute, nheads, head_dim, head_dim_v, gather_kv_length, varlen_q, varlen_k, k is not None,
+    )
+    if compile_key not in _sparse_mla_dq_dqv.compile_cache:
+        _sparse_mla_dq_dqv.compile_cache[compile_key] = _compile_sparse_mla_dq_dqv(
+            *compile_key
+        )
+    if not is_fake_mode():
+        _sparse_mla_dq_dqv.compile_cache[compile_key](
+            ds, k, v, dq, dqv, gather_kv_indices, cu_seqlens_q, cu_seqlens_k
         )
-        if compile_key_post not in _flash_attn_bwd.compile_cache_post:
-            dv_accum_tensor = to_cute_tensor(dv_accum)
-            dv_tensor = to_cute_tensor(dv)
-            cu_seqlens_k_tensor, seqused_k_tensor = [
-                to_cute_tensor(t, assumed_align=4) if t is not None else None
-                for t in (cu_seqlens_k, seqused_k)
-            ]
-            fa_bwd_post = FlashAttentionBackwardPostprocess(
-                dtype, head_dim_v, arch, n_block_size, num_threads, AtomLayoutNdKV, dKV_swapAB,
-                cluster_size=cluster_size,
-            )
-            # TODO: check @can_implement
-            _flash_attn_bwd.compile_cache_post[compile_key_post] = cute.compile(
-                fa_bwd_post,
-                dv_accum_tensor,
-                dv_tensor,
-                cutlass.Float32(1.0),
-                cu_seqlens_k_tensor,
-                seqused_k_tensor,
-                current_stream,
-                options="--enable-tvm-ffi",
-            )
-        if not is_fake_mode():
-            _flash_attn_bwd.compile_cache_post[compile_key_post](
-                dv_accum,
-                dv,
-                1.0,
-                cu_seqlens_k,
-                seqused_k,
-                current_stream,
-            )
 
-    return dq, dk, dv
+_sparse_mla_dq_dqv.compile_cache = get_jit_cache("dq_dqv_gemm")
 
 
-_flash_attn_bwd.compile_cache_pre = get_jit_cache("bwd_pre")
-_flash_attn_bwd.compile_cache = get_jit_cache("bwd")
-_flash_attn_bwd.compile_cache_post = get_jit_cache("bwd_post")
+def _compile_sparse_mla_dk(
+    dtype,
+    dtype_acc,
+    nheads: int,
+    head_dim: int,
+    topk: int,
+    varlen: bool,
+):
+    kernel = dKGemmKernel(
+        topk,
+        nheads,
+        head_dim,
+        varlen,
+    )
+    # Check if configuration can be implemented
+    kernel.check_can_implement()
+
+    div = 128 // dtype.width
+
+    sym = cute.sym_int
+    batch_fake = sym()
+    batchp1_fake = sym()
+    seqlen_q_fake = sym()
+    seqlen_k_fake = sym()
+    total_q_fake = (batch_fake, seqlen_q_fake) if not varlen else (sym(),)
+    total_k_fake = (batch_fake, seqlen_k_fake) if not varlen else (sym(),)
+
+    mdS = fake_tensor(dtype, (*total_q_fake, nheads, topk), divisibility=div)
+    mI = fake_tensor(Int32, (*total_q_fake, topk), divisibility=div)
+    mQ = fake_tensor(dtype, (*total_q_fake, nheads, head_dim), divisibility=div)
+    mdK = fake_tensor(dtype_acc, (*total_k_fake, head_dim), divisibility=div)
+    mCuSeqlensQ = fake_tensor(Int32, (batchp1_fake,), divisibility=1) if varlen else None
+    mCuSeqlensK = fake_tensor(Int32, (batchp1_fake,), divisibility=1) if varlen else None
+    
+    return cute.compile(
+        kernel,
+        mdS,
+        mI,
+        mQ,
+        mdK,
+        mCuSeqlensQ,
+        mCuSeqlensK,
+        cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+        options="--enable-tvm-ffi",
+    )
+
+
+def _sparse_mla_dk(
+    dS: torch.Tensor,
+    index_topk: torch.Tensor,
+    q: torch.Tensor,
+    dk: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor],
+    cu_seqlens_k: Optional[torch.Tensor],
+):
+    """Compute dKaccum = scatter(dS'^T @ Q, I).
+
+    Args:
+      dS:          (*total_q, heads, topk), bf16
+      index_topk:  (*total_q, topk), int32
+      Q:           (*total_q, heads, dim), bf16
+      dK:          (*total_q, dim), fp32
+      cuSeqlensQ:  (batch + 1,), int32, omit for non-varlen
+      cuSeqlensK:  (batch + 1,), int32, omit for non-varlen
+
+    Accumulates in place on top of dK.
+
+    For varlen, total_q and total_k are 1-dimensional, and the seqlen indices per batch are
+    determined using the cuSeqlensQ and cuSeqlensK tensors.
+    For non-varlen, total_q and total_k are (batch, seqlen_q) and (batch, seqlen_k).
+    """
+    dtype = dS.dtype
+    dtype_cute = torch2cute_dtype_map[dtype]
+    dtype_acc = dk.dtype
+    dtype_acc_cute = torch2cute_dtype_map[dtype_acc]
+
+    varlen = cu_seqlens_q is not None
+    nheads, topk = dS.shape[-2], dS.shape[-1]
+    head_dim = q.shape[-1] if q is not None else 0
+
+    compile_key = (
+        dtype_cute, dtype_acc_cute, nheads, head_dim, topk, varlen,
+    )
+
+    if compile_key not in _sparse_mla_dk.compile_cache:
+        _sparse_mla_dk.compile_cache[compile_key] = _compile_sparse_mla_dk(*compile_key)
+
+    if not is_fake_mode():
+        _sparse_mla_dk.compile_cache[compile_key](dS, index_topk, q, dk, cu_seqlens_q, cu_seqlens_k)
+    
+_sparse_mla_dk.compile_cache = get_jit_cache("dk_gemm")
 
 
 class FlashAttnFunc(torch.autograd.Function):
@@ -1328,6 +3084,8 @@ class FlashAttnFunc(torch.autograd.Function):
         q: torch.Tensor,
         k: torch.Tensor,
         v: torch.Tensor,
+        qv: Optional[torch.Tensor] = None,
+        gather_kv_indices: Optional[torch.Tensor] = None,
         softmax_scale: Optional[float] = None,
         causal: bool = False,
         window_size: Tuple[Optional[int], Optional[int]] = (None, None),
@@ -1336,28 +3094,28 @@ class FlashAttnFunc(torch.autograd.Function):
         num_splits: int = 1,
         pack_gqa: Optional[bool] = None,
         deterministic: bool = False,
+        score_mod: Optional[Callable] = None,
+        score_mod_bwd: Optional[Callable] = None,
         mask_mod: Optional[Callable] = None,
-        full_block_cnt: Optional[torch.Tensor] = None,
-        full_block_idx: Optional[torch.Tensor] = None,
-        mask_block_cnt: Optional[torch.Tensor] = None,
-        mask_block_idx: Optional[torch.Tensor] = None,
-        block_size: Optional[Tuple[int, int]] = None,
+        aux_tensors: Optional[list] = None,
+        aux_scalars: Optional[tuple] = None,
+        block_sparse_tensors: Optional[BlockSparseTensorsTorch] = None,
+        block_sparse_tensors_bwd: Optional[BlockSparseTensorsTorch] = None,
         return_lse: bool = False,
     ):
-        # Only create block sparse tensors if at least one block sparse parameter is provided
-        block_sparse_tensors = None
-        if any(t is not None for t in [full_block_cnt, full_block_idx, mask_block_cnt, mask_block_idx]):
-            block_sparse_tensors = BlockSparseTensorsTorch(
-                full_block_cnt=full_block_cnt,
-                full_block_idx=full_block_idx,
-                mask_block_cnt=mask_block_cnt,
-                mask_block_idx=mask_block_idx,
-                block_size=block_size,
-            )
-        out, lse = _flash_attn_fwd(
+        aux_scalars = tuple(aux_scalars) if aux_scalars else None
+        shared_kv = k is v
+        if shared_kv and v.shape[-1] == 512:
+            # specialize MLA attention formula
+            # O = softmax(Q @ K.T + Qv @ V.T) @ V
+            # by setting q, k to None
+            qv = q if qv is None else qv
+            q = k = None
+        out, lse, p, row_max = _flash_attn_fwd(
             q,
             k,
             v,
+            qv=qv,
             softmax_scale=softmax_scale,
             causal=causal,
             window_size_left=window_size[0],
@@ -1366,54 +3124,104 @@ class FlashAttnFunc(torch.autograd.Function):
             softcap=softcap,
             num_splits=num_splits,
             pack_gqa=pack_gqa,
+            score_mod=score_mod,
             mask_mod=mask_mod,
+            aux_tensors=aux_tensors,
+            aux_scalars=aux_scalars,
             block_sparse_tensors=block_sparse_tensors,
             return_lse=return_lse,
+            gather_kv_indices=gather_kv_indices,
         )
-        ctx.save_for_backward(q, k, v, out, lse)
+        ctx.save_for_backward(q, k, v, qv, out, lse, p, row_max, gather_kv_indices, learnable_sink, *(aux_tensors or ()))
+        ctx.shared_kv = shared_kv
         ctx.softmax_scale = softmax_scale
         ctx.causal = causal
         ctx.window_size = window_size
         ctx.softcap = softcap
         ctx.deterministic = deterministic
-        # LSE gradient is not supported yet
-        if lse is not None:
-            ctx.mark_non_differentiable(lse)
+        ctx.return_lse = return_lse
+        ctx.score_mod = score_mod 
+        ctx.score_mod_bwd = score_mod_bwd 
+        ctx.mask_mod = mask_mod
+        ctx.aux_scalars = aux_scalars
+        ctx.block_sparse_tensors_bwd = block_sparse_tensors_bwd
+        ctx.set_materialize_grads(False)
         return out, lse
 
     @staticmethod
-    def backward(ctx, dout, *args):
-        q, k, v, out, lse = ctx.saved_tensors
-        dq, dk, dv = _flash_attn_bwd(
-            q,
-            k,
-            v,
-            out,
-            dout,
-            lse,
-            ctx.softmax_scale,
-            ctx.causal,
-            ctx.softcap,
-            window_size_left=ctx.window_size[0],
-            window_size_right=ctx.window_size[1],
-            deterministic=ctx.deterministic,
-        )
-        return dq, dk, dv, *((None,) * 20)  # Extra Nones is fine
+    def backward(ctx, dout, dlse):
+        q, k, v, qv, out, lse, p, row_max, gather_kv_indices, learnable_sink, *aux = ctx.saved_tensors
+        aux_tensors = aux if aux else None
+        if not ctx.return_lse:
+            dlse = None
+        if dout is None:
+            dout = torch.zeros_like(out)
+        if qv is not None:
+            dq, dk, dv, dqv = _flash_attn_bwd_sparse_mla(
+                q,
+                k,
+                v,
+                qv,
+                out,
+                dout,
+                lse,
+                p,
+                row_max,
+                gather_kv_indices,
+                softmax_scale=ctx.softmax_scale,
+                causal=ctx.causal,
+            )
+            if ctx.shared_kv:
+                return dqv, dv, None, None, *((None,) * 30)
+            else:
+                return dq, dk, dv, dqv, *((None,) * 30)
+        else:
+            bwd_result = _flash_attn_bwd(
+                q,
+                k,
+                v,
+                out,
+                dout,
+                lse,
+                ctx.softmax_scale,
+                ctx.causal,
+                ctx.softcap,
+                window_size_left=ctx.window_size[0],
+                window_size_right=ctx.window_size[1],
+                deterministic=ctx.deterministic,
+                score_mod=ctx.score_mod,
+                score_mod_bwd=ctx.score_mod_bwd,
+                mask_mod=ctx.mask_mod,
+                aux_tensors=aux_tensors,
+                aux_scalars=ctx.aux_scalars,
+                block_sparse_tensors=ctx.block_sparse_tensors_bwd,
+                dlse=dlse,
+                learnable_sink=learnable_sink,
+            )
+            if learnable_sink is None:
+                dq, dk, dv = bwd_result
+                dsink = None
+            else:
+                dq, dk, dv, dsink = bwd_result
+            return dq, dk, dv, None, None, None, None, None, dsink, *((None,) * 12)
 
 
 class FlashAttnVarlenFunc(torch.autograd.Function):
     @staticmethod
     def forward(
         ctx,
-        q: torch.Tensor,
-        k: torch.Tensor,
+        q: Optional[torch.Tensor],
+        k: Optional[torch.Tensor],
         v: torch.Tensor,
-        cu_seqlens_q: Optional[torch.Tensor],
-        cu_seqlens_k: Optional[torch.Tensor],
+        qv: Optional[torch.Tensor] = 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,
         max_seqlen_q: Optional[int] = None,
         max_seqlen_k: Optional[int] = None,
+        min_seqlen_k: Optional[int] = None,
+        gather_kv_indices: Optional[torch.Tensor] = None,
         page_table: Optional[torch.Tensor] = None,
         softmax_scale: Optional[float] = None,
         causal: bool = False,
@@ -1424,19 +3232,36 @@ class FlashAttnVarlenFunc(torch.autograd.Function):
         pack_gqa: Optional[bool] = None,
         deterministic: bool = False,
         score_mod: Optional[Callable] = None,
+        score_mod_bwd: Optional[Callable] = None,
+        mask_mod: Optional[Callable] = None,
+        block_sparse_tensors: Optional[list] = None,
         aux_tensors: Optional[list] = None,
+        aux_scalars: Optional[tuple] = None,
         return_lse: bool = False,
+        scheduler_metadata: Optional["SchedulerMetadataTensorsTorch"] = None,
+        seqlen_k_per_split: Optional[int] = None,
+        disable_scheduler_metadata: bool = False,
     ):
-        out, lse = _flash_attn_fwd(
+        aux_scalars = tuple(aux_scalars) if aux_scalars else None
+        shared_kv = k is v
+        if shared_kv and v.shape[-1] == 512:
+            # specialize MLA attention formula
+            # O = softmax(Q @ K.T + Qv @ V.T) @ V
+            # by setting q, k to None
+            qv = q if qv is None else qv
+            q = k = None
+        out, lse, p, row_max = _flash_attn_fwd(
             q,
             k,
             v,
-            cu_seqlens_q,
-            cu_seqlens_k,
-            seqused_q,
-            seqused_k,
+            qv=qv,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            seqused_q=seqused_q,
+            seqused_k=seqused_k,
             max_seqlen_q=max_seqlen_q,
             max_seqlen_k=max_seqlen_k,
+            min_seqlen_k=min_seqlen_k,
             page_table=page_table,
             softmax_scale=softmax_scale,
             causal=causal,
@@ -1447,10 +3272,34 @@ class FlashAttnVarlenFunc(torch.autograd.Function):
             num_splits=num_splits,
             pack_gqa=pack_gqa,
             score_mod=score_mod,
+            mask_mod=mask_mod,
+            block_sparse_tensors=block_sparse_tensors,
             aux_tensors=aux_tensors,
+            aux_scalars=aux_scalars,
             return_lse=return_lse,
+            gather_kv_indices=gather_kv_indices,
+            scheduler_metadata=scheduler_metadata,
+            seqlen_k_per_split=seqlen_k_per_split,
+            disable_scheduler_metadata=disable_scheduler_metadata,
+        )
+        ctx.save_for_backward(
+            q,
+            k,
+            v,
+            qv,
+            out,
+            lse,
+            p,
+            row_max,
+            gather_kv_indices,
+            learnable_sink,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqused_q,
+            seqused_k,
+            *(aux_tensors or ()),
         )
-        ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k)
+        ctx.shared_kv = shared_kv
         ctx.softmax_scale = softmax_scale
         ctx.causal = causal
         ctx.window_size = window_size
@@ -1458,43 +3307,91 @@ class FlashAttnVarlenFunc(torch.autograd.Function):
         ctx.deterministic = deterministic
         ctx.max_seqlen_q = max_seqlen_q
         ctx.max_seqlen_k = max_seqlen_k
-        # LSE gradient is not supported yet
-        if lse is not None:
-            ctx.mark_non_differentiable(lse)
+        ctx.min_seqlen_k = min_seqlen_k
+        ctx.return_lse = return_lse
+        ctx.score_mod = score_mod
+        ctx.score_mod_bwd = score_mod_bwd
+        ctx.mask_mod = mask_mod
+        ctx.aux_scalars = aux_scalars
+        ctx.set_materialize_grads(False)
         return out, lse
 
     @staticmethod
-    def backward(ctx, dout, *args):
-        q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors
-        assert ctx.softcap == 0.0
-        dq, dk, dv = _flash_attn_bwd(
-            q,
-            k,
-            v,
-            out,
-            dout,
-            lse,
-            ctx.softmax_scale,
-            ctx.causal,
-            ctx.softcap,
-            window_size_left=ctx.window_size[0],
-            window_size_right=ctx.window_size[1],
-            cu_seqlens_q=cu_seqlens_q,
-            cu_seqlens_k=cu_seqlens_k,
-            seqused_q=seqused_q,
-            seqused_k=seqused_k,
-            max_seqlen_q=ctx.max_seqlen_q,
-            max_seqlen_k=ctx.max_seqlen_k,
-            deterministic=ctx.deterministic,
-        )
-
-        return dq, dk, dv, *((None,) * 20)
+    def backward(ctx, dout, dlse):
+        q, k, v, qv, out, lse, p, row_max, gather_kv_indices, learnable_sink, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, *aux = ctx.saved_tensors
+        aux_tensors = aux if aux else None
+        if not ctx.return_lse:
+            dlse = None
+        if dout is None:
+            dout = torch.zeros_like(out)
+        if qv is not None:
+            dq, dk, dv, dqv = _flash_attn_bwd_sparse_mla(
+                q,
+                k,
+                v,
+                qv,
+                out,
+                dout,
+                lse,
+                p,
+                row_max,
+                gather_kv_indices,
+                softmax_scale=ctx.softmax_scale,
+                causal=ctx.causal,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                seqused_q=seqused_q,
+                seqused_k=seqused_k,
+                max_seqlen_q=ctx.max_seqlen_q,
+                max_seqlen_k=ctx.max_seqlen_k,
+                min_seqlen_k=ctx.min_seqlen_k,
+            )
+            if ctx.shared_kv:
+                return dqv, dv, None, None, *((None,) * 31)
+            else:
+                return dq, dk, dv, dqv, *((None,) * 31)
+        else:
+            bwd_result = _flash_attn_bwd(
+                q,
+                k,
+                v,
+                out,
+                dout,
+                lse,
+                ctx.softmax_scale,
+                ctx.causal,
+                ctx.softcap,
+                window_size_left=ctx.window_size[0],
+                window_size_right=ctx.window_size[1],
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                seqused_q=seqused_q,
+                seqused_k=seqused_k,
+                max_seqlen_q=ctx.max_seqlen_q,
+                max_seqlen_k=ctx.max_seqlen_k,
+                deterministic=ctx.deterministic,
+                score_mod=ctx.score_mod,
+                score_mod_bwd=ctx.score_mod_bwd,
+                aux_tensors=aux_tensors,
+                aux_scalars=ctx.aux_scalars,
+                mask_mod=ctx.mask_mod,
+                dlse=dlse,
+                learnable_sink=learnable_sink,
+            )
+            if learnable_sink is None:
+                dq, dk, dv = bwd_result
+                dsink = None
+            else:
+                dq, dk, dv, dsink = bwd_result
+            return dq, dk, dv, None, *((None,) * 12), dsink, *((None,) * 14)
 
 
 def flash_attn_func(
     q: torch.Tensor,
     k: torch.Tensor,
     v: torch.Tensor,
+    qv: Optional[torch.Tensor] = None,
+    gather_kv_indices: Optional[torch.Tensor] = None,
     softmax_scale: Optional[float] = None,
     causal: bool = False,
     window_size: Tuple[Optional[int], Optional[int]] = (None, None),
@@ -1503,18 +3400,21 @@ def flash_attn_func(
     num_splits: int = 1,
     pack_gqa: Optional[bool] = None,
     deterministic: bool = False,
+    score_mod: Optional[Callable] = None,
+    score_mod_bwd: Optional[Callable] = None,
     mask_mod: Optional[Callable] = None,
-    full_block_cnt: Optional[torch.Tensor] = None,
-    full_block_idx: Optional[torch.Tensor] = None,
-    mask_block_cnt: Optional[torch.Tensor] = None,
-    mask_block_idx: Optional[torch.Tensor] = None,
-    block_size: Optional[Tuple[int, int]] = None,
+    aux_tensors: Optional[list] = None,
+    aux_scalars: Optional[tuple] = None,
+    block_sparse_tensors: Optional[BlockSparseTensorsTorch] = None,
+    block_sparse_tensors_bwd: Optional[BlockSparseTensorsTorch] = None,
     return_lse: bool = False,
 ):
     return FlashAttnFunc.apply(
         q,
         k,
         v,
+        qv,
+        gather_kv_indices,
         softmax_scale,
         causal,
         window_size,
@@ -1523,12 +3423,13 @@ def flash_attn_func(
         num_splits,
         pack_gqa,
         deterministic,
+        score_mod,
+        score_mod_bwd,
         mask_mod,
-        full_block_cnt,
-        full_block_idx,
-        mask_block_cnt,
-        mask_block_idx,
-        block_size,
+        aux_tensors,
+        aux_scalars,
+        block_sparse_tensors,
+        block_sparse_tensors_bwd,
         return_lse,
     )
 
@@ -1537,12 +3438,15 @@ def flash_attn_varlen_func(
     q: torch.Tensor,
     k: torch.Tensor,
     v: torch.Tensor,
+    qv: Optional[torch.Tensor] = None,
     cu_seqlens_q: Optional[torch.Tensor] = None,
     cu_seqlens_k: Optional[torch.Tensor] = None,
     max_seqlen_q: Optional[int] = None,
     max_seqlen_k: Optional[int] = None,
+    min_seqlen_k: Optional[int] = None,
     seqused_q: Optional[torch.Tensor] = None,
     seqused_k: Optional[torch.Tensor] = None,
+    gather_kv_indices: Optional[torch.Tensor] = None,
     page_table: Optional[torch.Tensor] = None,
     softmax_scale: Optional[float] = None,
     causal: bool = False,
@@ -1553,19 +3457,69 @@ def flash_attn_varlen_func(
     pack_gqa: Optional[bool] = None,
     deterministic: bool = False,
     score_mod: Optional[Callable] = None,
+    score_mod_bwd: Optional[Callable] = None,
+    mask_mod: Optional[Callable] = None,
+    block_sparse_tensors: Optional[BlockSparseTensorsTorch] = None,
     aux_tensors: Optional[list] = None,
+    aux_scalars: Optional[tuple] = None,
     return_lse: bool = False,
+    scheduler_metadata: Optional[SchedulerMetadataTensorsTorch] = None,
+    seqlen_k_per_split: Optional[int] = None,
+    disable_scheduler_metadata: bool = False,
 ):
+    """
+    Tensor arguments:
+        q:  (total_q, nheads,   hdim)   or (batch, seqlen_q, nheads,   hdim)
+        k:  (total_k, nheads_k, hdim)   or (batch, seqlen_k, nheads_k, hdim)
+        v:  (total_k, nheads_k, hdim_v) or (batch, seqlen_k, nheads_k, hdim_v)
+        qv: (total_q, nheads,   hdim_v) or (batch, seqlen_q, nheads,   hdim_v)
+        cu_seqlens_q: (batch + 1)       or seqused_q: (batch)
+        cu_seqlens_k: (batch + 1)       or seqused_k: (batch)
+        gather_kv_indices: (total_q, gather_kv_length) or
+                           (batch, seqlen_q, gather_kv_length)
+        page_table: (batch, max_num_pages_per_seq)
+    
+    Return:
+       out: (total_q, nheads, hdim) or (batch, seqlen_q, nheads, hdim)
+       lse: (nheads, total_q)       or (batch, nheads, seqlen_q) if not has_qv (standard)
+            (total_q, nheads)       or (batch, seqlen_q, nheads) if has_qv
+
+    Explanation of some optional arguments & decisions:
+
+    qv: we write the MLA weight absorbed formula as
+        O = softmax(scale * (Q @ K.T + Qv @ V.T)) @ V
+        where Q = q_pe, Qv = q_nope, K = pe_cache, V = kv_cache.
+
+    lse return shape: with Qv, MQA with nheads at least divisible by 4 is typical,
+        so we arrange for nheads as the contiguous mode for better vectorization.
+
+    gather_kv_indices: used for topk sparsity with MLA absorption kernel.
+
+    min_seqlen_k: for varlen, specifies the minimum kv sequence length for any batch.
+        Used with gather_kv_indices to determine if we need oob masking.
+
+    scheduler_metadata: optional tensors used by certain tile schedulers, for optimization
+        and functionality. computed in get_scheduler_metadata.
+
+    seqlen_k_per_split: when using dynamic (per-batch) num_splits, can set a fixed seqlen_k to be
+        covered per split for bitwise reproducibility.
+
+    disable_scheduler_metadata: if True, ignores scheduler_metadata if it is passed and skips
+        computing metadata fresh.
+    """
     return FlashAttnVarlenFunc.apply(
         q,
         k,
         v,
+        qv,
         cu_seqlens_q,
         cu_seqlens_k,
         seqused_q,
         seqused_k,
         max_seqlen_q,
         max_seqlen_k,
+        min_seqlen_k,
+        gather_kv_indices,
         page_table,
         softmax_scale,
         causal,
@@ -1576,8 +3530,74 @@ def flash_attn_varlen_func(
         pack_gqa,
         deterministic,
         score_mod,
+        score_mod_bwd,
+        mask_mod,
+        block_sparse_tensors,
         aux_tensors,
+        aux_scalars,
         return_lse,
+        scheduler_metadata,
+        seqlen_k_per_split,
+        disable_scheduler_metadata,
+    )
+
+
+def _compile_fwd_combine(
+    _arch, dtype, dtype_partial, head_dim, num_head, tile_m, k_block_size, log_max_splits,
+    has_cu_seqlens, has_seqused, has_lse, has_virtual_batch_idx,
+    has_num_splits_dynamic, has_semaphore_to_reset,
+):
+    """Compile fwd combine kernel using cute fake tensors (no real GPU tensors needed)."""
+    sym = cute.sym_int
+    div = 128 // dtype_partial.width  # 16-byte alignment in elements
+
+    fa_combine = FlashAttentionForwardCombine(
+        dtype=dtype,
+        dtype_partial=dtype_partial,
+        head_dim=head_dim,
+        num_head=num_head,
+        tile_m=tile_m,
+        k_block_size=k_block_size,
+        log_max_splits=log_max_splits,
+    )
+    if not fa_combine.can_implement(
+        dtype, dtype_partial, head_dim, tile_m, k_block_size, log_max_splits,
+        num_threads=256,
+    ):
+        raise RuntimeError(
+            "FlashAttention combine kernel cannot be implemented with given parameters"
+        )
+
+    if has_cu_seqlens:
+        # Varlen: (num_splits, total_q, nheads, headdim)
+        num_splits, total_q, nheads = sym(), sym(), sym()
+        mO_partial = fake_tensor(dtype_partial, (num_splits, total_q, nheads, head_dim), divisibility=div)
+        mLSE_partial = fake_tensor(Float32, (num_splits, total_q, nheads), divisibility=1, leading_dim=1)
+        mO = fake_tensor(dtype, (total_q, nheads, head_dim), divisibility=div)
+        mLSE = fake_tensor(Float32, (total_q, nheads), divisibility=1, leading_dim=0) if has_lse else None
+    else:
+        # Batched: (num_splits, batch, seqlen, nheads, headdim)
+        num_splits, batch, seqlen, nheads = sym(), sym(), sym(), sym()
+        mO_partial = fake_tensor(dtype_partial, (num_splits, batch, seqlen, nheads, head_dim), divisibility=div)
+        mLSE_partial = fake_tensor(Float32, (num_splits, batch, seqlen, nheads), divisibility=1, leading_dim=2)
+        mO = fake_tensor(dtype, (batch, seqlen, nheads, head_dim), divisibility=div)
+        mLSE = fake_tensor(Float32, (batch, seqlen, nheads), divisibility=1, leading_dim=1) if has_lse else None
+        batch = mO_partial.shape[1]
+
+    batch_for_1d = batch if not has_cu_seqlens else sym()
+    batchp1 = sym()
+    mCuSeqlens = fake_tensor(Int32, (batchp1,), divisibility=1) if has_cu_seqlens else None
+    mSeqused = fake_tensor(Int32, (batch_for_1d,), divisibility=1) if has_seqused else None
+    mNumSplitsDynamic = fake_tensor(Int32, (batch_for_1d,), divisibility=1) if has_num_splits_dynamic else None
+    mVirtualBatchIdx = fake_tensor(Int32, (batch_for_1d,), divisibility=1) if has_virtual_batch_idx else None
+    mSemaphore = fake_tensor(Int32, (1,), divisibility=1) if has_semaphore_to_reset else None
+
+    return cute.compile(
+        fa_combine,
+        mO_partial, mLSE_partial, mO, mLSE,
+        mCuSeqlens, mSeqused, mNumSplitsDynamic, mVirtualBatchIdx, mSemaphore,
+        cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+        options="--enable-tvm-ffi",
     )
 
 
@@ -1589,7 +3609,10 @@ def _flash_attn_fwd_combine(
     cu_seqlens: Optional[torch.Tensor] = None,
     seqused: Optional[torch.Tensor] = None,
     num_splits_dynamic_ptr: Optional[torch.Tensor] = None,
+    virtual_batch_idx: Optional[torch.Tensor] = None,
     semaphore_to_reset: Optional[torch.Tensor] = None,
+    *,
+    _arch: Optional[int] = None,
 ) -> None:
     """Forward combine kernel for split attention computation.
 
@@ -1612,39 +3635,25 @@ def _flash_attn_fwd_combine(
     Returns:
         None
     """
-    # Input validation
-    assert out_partial.dim() in [4, 5], "out_partial must have 4 or 5 dimensions"
-    assert lse_partial.dim() in [3, 4], "lse_partial must have 3 or 4 dimensions"
+    fake_mode = is_fake_mode()
     assert out_partial.dtype in [torch.float16, torch.bfloat16, torch.float32], (
         "out_partial must be fp16, bf16, or fp32"
     )
-    assert lse_partial.dtype == torch.float32, "lse_partial must be fp32"
-    assert out_partial.is_cuda and lse_partial.is_cuda, "tensors must be on CUDA device"
-    assert out_partial.stride(-1) == 1, "out_partial must be contiguous in the last dimension"
-    assert lse_partial.stride(-2) == 1, "lse_partial must be contiguous in the seqlen dimension"
-    assert lse_partial.shape == out_partial.shape[:-1]
-
-    # Determine if this is variable length based on dimensions
-    is_varlen = out_partial.dim() == 4
-
-    # Validate output tensor shapes and types
-    assert out.shape == out_partial.shape[1:], "out shape mismatch"
-    if lse is not None:
-        assert lse.shape == lse_partial.shape[1:], "lse shape mismatch"
-        assert lse.dtype == torch.float32, "lse must be fp32"
-
-    # Validate optional tensors
-    for t, name in [
+    if not fake_mode:
+        assert out_partial.is_cuda and lse_partial.is_cuda, "tensors must be on CUDA device"
+    for tensor, name in (
         (cu_seqlens, "cu_seqlens"),
         (seqused, "seqused"),
         (num_splits_dynamic_ptr, "num_splits_dynamic_ptr"),
-    ]:
-        if t is not None:
-            assert t.dtype == torch.int32, f"{name} must be int32"
-            assert t.is_cuda, f"{name} must be on CUDA device"
-            assert t.is_contiguous(), f"{name} must be contiguous"
-
+        (virtual_batch_idx, "virtual_batch_idx"),
+        (semaphore_to_reset, "semaphore_to_reset"),
+    ):
+        if tensor is not None:
+            if not fake_mode:
+                assert tensor.is_cuda, f"{name} must be on CUDA device"
+            assert tensor.is_contiguous(), f"{name} must be contiguous"
     head_dim = out_partial.shape[-1]
+    num_head = out_partial.shape[-2]
     num_splits = out_partial.shape[0]
     assert num_splits <= 256
     # If hdim is 96 or 192, it's faster to round them to 128 or 256 respectively
@@ -1652,101 +3661,41 @@ def _flash_attn_fwd_combine(
     k_block_size = 64 if head_dim <= 64 else 128
     # We want kBlockM to be as small as possible to maximize parallelism.
     # E.g., if hdim is 64, we want kBlockM to be 16 so that we can use 256 threads, each reading 4 elements (floats).
-    m_block_size = 8 if k_block_size % 128 == 0 else (16 if k_block_size % 64 == 0 else 32)
+    tile_m = 8 if k_block_size % 128 == 0 else (16 if k_block_size % 64 == 0 else 32)
     log_max_splits = max(math.ceil(math.log2(num_splits)), 4)
-    if m_block_size == 8:
+    if tile_m == 8:
         # If kBlockM == 8 then the minimum number of splits is 32.
         # TODO: we can deal w this by using 128 threads instead
         log_max_splits = max(log_max_splits, 5)
 
-    current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
-
     # Create combine kernel configuration
     dtype = torch2cute_dtype_map[out.dtype]
     dtype_partial = torch2cute_dtype_map[out_partial.dtype]
-
     compile_key = (
+        _get_device_arch() if _arch is None else _arch,
         dtype,
         dtype_partial,
         head_dim,
-        m_block_size,
+        num_head,
+        tile_m,
         k_block_size,
         log_max_splits,
         cu_seqlens is not None,
         seqused is not None,
         lse is not None,
+        virtual_batch_idx is not None,
+        num_splits_dynamic_ptr is not None,
+        semaphore_to_reset is not None,
     )
-
     if compile_key not in _flash_attn_fwd_combine.compile_cache:
-        out_partial_tensor = to_cute_tensor(
-            out_partial, leading_dim=4 if not is_varlen else 3
-        )
-        lse_partial_tensor = to_cute_tensor(
-            lse_partial, assumed_align=4, leading_dim=lse_partial.ndim - 2
-        )
-        out_tensor = to_cute_tensor(out, leading_dim=3 if not is_varlen else 2)
-        lse_tensor = (
-            to_cute_tensor(lse, assumed_align=4, leading_dim=lse.ndim - 2)
-            if lse is not None
-            else None
-        )
-
-        optional_tensors = [
-            to_cute_tensor(t, assumed_align=4, leading_dim=0)
-            if t is not None
-            else None
-            for t in (cu_seqlens, seqused, num_splits_dynamic_ptr, semaphore_to_reset)
-        ]
-        cu_seqlens_tensor, seqused_tensor, num_splits_dynamic_tensor, semaphore_tensor = (
-            optional_tensors
-        )
-        fa_combine = FlashAttentionForwardCombine(
-            dtype=dtype,
-            dtype_partial=dtype_partial,
-            head_dim=head_dim,
-            m_block_size=m_block_size,
-            k_block_size=k_block_size,
-            log_max_splits=log_max_splits,
-        )
-
-        # Check if implementation is supported
-        if not fa_combine.can_implement(
-            dtype,
-            dtype_partial,
-            head_dim,
-            m_block_size,
-            k_block_size,
-            log_max_splits,
-            num_threads=256,
-        ):
-            raise RuntimeError(
-                "FlashAttention combine kernel cannot be implemented with given parameters"
-            )
-
-        _flash_attn_fwd_combine.compile_cache[compile_key] = cute.compile(
-            fa_combine,
-            out_partial_tensor,
-            lse_partial_tensor,
-            out_tensor,
-            lse_tensor,
-            cu_seqlens_tensor,
-            seqused_tensor,
-            num_splits_dynamic_tensor,
-            semaphore_tensor,
-            current_stream,
-            options="--enable-tvm-ffi",
+        _flash_attn_fwd_combine.compile_cache[compile_key] = _compile_fwd_combine(
+            *compile_key
         )
-    if not is_fake_mode():
+    if not fake_mode:
         _flash_attn_fwd_combine.compile_cache[compile_key](
-            out_partial,
-            lse_partial,
-            out,
-            lse,
-            cu_seqlens,
-            seqused,
-            num_splits_dynamic_ptr,
+            out_partial, lse_partial, out, lse,
+            cu_seqlens, seqused, num_splits_dynamic_ptr, virtual_batch_idx,
             semaphore_to_reset,
-            current_stream,
         )
 
 
@@ -1760,6 +3709,7 @@ def flash_attn_combine(
     out_dtype: Optional[torch.dtype] = None,
     cu_seqlens: Optional[torch.Tensor] = None,
     seqused: Optional[torch.Tensor] = None,
+    virtual_batch_idx: Optional[torch.Tensor] = None,
     return_lse: bool = True,
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
     """Flash Attention combine function for split attention computation.
@@ -1779,6 +3729,9 @@ def flash_attn_combine(
         out_dtype: Optional output dtype. If None, will use fp16/bf16 based on input.
         cu_seqlens: Cumulative sequence lengths for variable length sequences
         seqused: Used sequence lengths for each batch
+        virtual_batch_idx: Optional mapping from virtual batch index to real batch index
+            (int32 tensor of shape (batch_size,)). Used by persistent tile schedulers
+            that reorder batch processing for load balancing.
         return_lse: Whether to return the combined LSE tensor. Default is True.
 
     Returns:
@@ -1795,32 +3748,19 @@ def flash_attn_combine(
     """
     # Input validation
     assert out_partial.dim() in [4, 5], "out_partial must have 4 or 5 dimensions"
-    assert lse_partial.dim() in [3, 4], "lse_partial must have 3 or 4 dimensions"
-    assert out_partial.dtype == torch.float32, "out_partial must be fp32 (from accumulation)"
-    assert lse_partial.dtype == torch.float32, "lse_partial must be fp32"
-
     # Determine if this is variable length based on dimensions
     is_varlen = out_partial.dim() == 4
-
     if is_varlen:
         # Variable length: (num_splits, total_q, num_heads, head_size)
         num_splits, total_q, num_heads, head_size = out_partial.shape
-        assert lse_partial.shape == (num_splits, total_q, num_heads), (
-            "lse_partial shape mismatch for varlen"
-        )
         batch_size = 1  # Treat as single batch for varlen
         seqlen = total_q
     else:
         # Regular batched: (num_splits, batch_size, seqlen, num_heads, head_size)
         num_splits, batch_size, seqlen, num_heads, head_size = out_partial.shape
-        assert lse_partial.shape == (num_splits, batch_size, seqlen, num_heads), (
-            "lse_partial shape mismatch"
-        )
-
     # Determine output dtype
     if out_dtype is None:
         out_dtype = out_partial.dtype
-
     # Create output if not provided
     device = out_partial.device
     if out is None:
@@ -1830,20 +3770,15 @@ def flash_attn_combine(
             out = torch.empty(
                 batch_size, seqlen, num_heads, head_size, dtype=out_dtype, device=device
             )
-
     # Create lse output only if requested
     if return_lse:
         if is_varlen:
-            lse = torch.empty(num_heads, total_q, dtype=torch.float32, device=device).transpose(
-                0, 1
-            )
+            lse = torch.empty(num_heads, total_q, dtype=torch.float32, device=device)
         else:
-            lse = torch.empty(
-                batch_size, num_heads, seqlen, dtype=torch.float32, device=device
-            ).transpose(1, 2)
+            lse = torch.empty(batch_size, num_heads, seqlen, dtype=torch.float32, device=device)
+        lse = lse.transpose(-1, -2)
     else:
         lse = None
-
     _flash_attn_fwd_combine(
         out_partial,
         lse_partial,
@@ -1851,5 +3786,376 @@ def flash_attn_combine(
         lse,
         cu_seqlens,
         seqused,
+        virtual_batch_idx=virtual_batch_idx,
     )
     return out, lse
+
+
+def _get_scheduler_metadata(
+    num_batch: int,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    nheads: int,
+    nheads_kv: int,
+    headdim: int,
+    num_splits: int,
+    tile_m: int,
+    tile_n: int,
+    headdim_v: Optional[int] = None,
+    pack_gqa: Optional[bool] = False,
+    q_stage: int = 1,
+    cluster_shape_m: int = 1,
+    causal: bool = False,
+    enable_pdl: bool = False,
+    sort: bool = False,
+    seqlen_k_new: int = 0,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    cu_seqlens_k_new: Optional[torch.Tensor] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    seqlen_k_per_split: Optional[int] = None,
+    zfill_padded_output: bool = True,
+    total_q: Optional[int] = None,
+    use_clc_scheduler: bool = False,
+) -> SchedulerMetadataTensorsTorch:
+    device = None
+    for t in [cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k]:
+        if t is not None:
+            device = t.device
+            break
+    if device is None:
+        raise ValueError(
+            "At least one of cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k must be provided on device"
+        )
+    if headdim_v is None:
+        headdim_v = headdim
+
+    # Override enable_pdl (not supported yet)
+    enable_pdl = False
+
+    assert not sort, "LPT batch sort not yet implemented"
+
+    if seqlen_k_per_split is not None:
+        assert seqlen_k_per_split % tile_n == 0, "seqlen per split must be divisible by tile_n"
+        n_blocks_per_split = seqlen_k_per_split // tile_n
+        n_blocks_total = (max_seqlen_k + seqlen_k_new + tile_n - 1) // tile_n
+        splits_needed = (n_blocks_total + n_blocks_per_split - 1) // n_blocks_per_split
+        assert num_splits >= splits_needed, (
+            f"seqlen_k_per_split={seqlen_k_per_split} needs num_splits>={splits_needed}, "
+            f"got {num_splits}"
+        )
+    else:
+        n_blocks_per_split = None
+
+    is_split_kv = num_splits > 1
+    needs_prepare_kernel = is_split_kv or causal or sort
+
+    if needs_prepare_kernel:
+        num_m_blocks = torch.empty(num_batch, dtype=torch.int32, device=device)
+        num_splits_dynamic = torch.empty(num_batch, dtype=torch.int32, device=device)
+        virtual_batch_idx = (
+            torch.empty(num_batch, dtype=torch.int32, device=device) if sort else None
+        )
+        num_nheads_in_l2 = (
+            torch.empty(num_batch, dtype=torch.int32, device=device) if causal else None
+        )
+        tile_count_semaphore = (
+            torch.empty(1, dtype=torch.int32, device=device) if not use_clc_scheduler else None
+        )
+
+        num_warps = min((num_batch + 30) // 31, 32)
+        num_warps = 1 << (num_warps - 1).bit_length()
+
+        cache_key = (
+            num_warps,
+            tile_m,
+            tile_n,
+            nheads,
+            nheads_kv,
+            headdim,
+            headdim_v,
+            causal,
+            pack_gqa,
+            enable_pdl,
+            sort,
+            cu_seqlens_q is not None,
+            cu_seqlens_k is not None,
+            cu_seqlens_k_new is not None,
+            seqused_q is not None,
+            seqused_k is not None,
+            leftpad_k is not None,
+            num_m_blocks is not None,
+            num_splits_dynamic is not None,
+            virtual_batch_idx is not None,
+            num_nheads_in_l2 is not None,
+            tile_count_semaphore is not None,
+            n_blocks_per_split is not None,
+            zfill_padded_output,
+        )
+
+        if cache_key not in _get_scheduler_metadata.compile_cache:
+            (
+                num_m_blocks_cute,
+                num_splits_dynamic_cute,
+                virtual_batch_idx_cute,
+                num_nheads_in_l2_cute,
+                tile_count_semaphore_cute,
+                cu_seqlens_q_cute,
+                cu_seqlens_k_cute,
+                cu_seqlens_k_new_cute,
+                seqused_q_cute,
+                seqused_k_cute,
+                leftpad_k_cute,
+            ) = [
+                to_cute_tensor(t, assumed_align=4) if t is not None else None
+                for t in (
+                    num_m_blocks,
+                    num_splits_dynamic,
+                    virtual_batch_idx,
+                    num_nheads_in_l2,
+                    tile_count_semaphore,
+                    cu_seqlens_q,
+                    cu_seqlens_k,
+                    cu_seqlens_k_new,
+                    seqused_q,
+                    seqused_k,
+                    leftpad_k,
+                )
+            ]
+            scheduler = FlashPrepareScheduler(
+                num_warps,
+                tile_m,
+                tile_n,
+                nheads,
+                nheads_kv,
+                headdim,
+                headdim_v,
+                causal,
+                packgqa=pack_gqa,
+                sort=sort,
+                zfill_padded_output=zfill_padded_output,
+            )
+            _get_scheduler_metadata.compile_cache[cache_key] = cute.compile(
+                scheduler,
+                max_seqlen_q,
+                max_seqlen_k,
+                seqlen_k_new,
+                cu_seqlens_q_cute,
+                cu_seqlens_k_cute,
+                cu_seqlens_k_new_cute,
+                seqused_q_cute,
+                seqused_k_cute,
+                leftpad_k_cute,
+                num_batch,
+                num_splits,
+                tile_count_semaphore_cute,
+                num_m_blocks_cute,
+                num_splits_dynamic_cute,
+                virtual_batch_idx_cute,
+                num_nheads_in_l2_cute,
+                n_blocks_per_split,
+                cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+                options="--enable-tvm-ffi",
+            )
+
+        if not is_fake_mode():
+            _get_scheduler_metadata.compile_cache[cache_key](
+                max_seqlen_q,
+                max_seqlen_k,
+                seqlen_k_new,
+                cu_seqlens_q,
+                cu_seqlens_k,
+                cu_seqlens_k_new,
+                seqused_q,
+                seqused_k,
+                leftpad_k,
+                num_batch,
+                num_splits,
+                tile_count_semaphore,
+                num_m_blocks,
+                num_splits_dynamic,
+                virtual_batch_idx,
+                num_nheads_in_l2,
+                n_blocks_per_split,
+            )
+    else:
+        num_m_blocks = None
+        num_splits_dynamic = None
+        virtual_batch_idx = None
+        num_nheads_in_l2 = None
+        tile_count_semaphore = None
+
+    qhead_per_kvhead = nheads // nheads_kv
+    # binary-search hint; only consumed by the single-tile scheduler above this batch
+    has_varlen_info = (
+        cu_seqlens_q is not None or seqused_q is not None
+    )
+    needs_compute_tile_cumsum = (
+        has_varlen_info
+        and num_batch > BIN_BATCH_SEARCH_THRESH
+        and tile_count_semaphore is None
+    )
+    if needs_compute_tile_cumsum:
+        cu_total_m_blocks, cu_total_splits_m_blocks = _compute_tile_cumsum(
+            num_m_blocks=num_m_blocks,
+            cu_seqlens=cu_seqlens_q,
+            seqused=seqused_q,
+            num_splits_dynamic=num_splits_dynamic,
+            virtual_batch_idx=virtual_batch_idx,
+            tile_size=tile_m,
+            q_stage=q_stage,
+            cluster_shape_m=cluster_shape_m,
+            qhead_per_kvhead=qhead_per_kvhead,
+            pack_gqa=bool(pack_gqa),
+        )
+    else:
+        cu_total_m_blocks, cu_total_splits_m_blocks = None, None
+
+    blocks_to_batch_idx = None
+    if USE_BLOCKS_TO_BATCH and cu_total_m_blocks is not None:
+        blocks_to_batch_idx = _compute_blocks_to_batch(
+            cu_total_m_blocks,
+            _blocks_to_batch_size(
+                total_q if total_q is not None else num_batch * max_seqlen_q,
+                num_batch, tile_m, qhead_per_kvhead, pack_gqa,
+            ),
+            cu_total_m_blocks.device,
+        )
+
+    return SchedulerMetadataTensorsTorch(
+        num_m_blocks_ptr=num_m_blocks,
+        num_splits_dynamic_ptr=num_splits_dynamic,
+        virtual_batch_idx_ptr=virtual_batch_idx,
+        num_nheads_in_l2_ptr=num_nheads_in_l2,
+        tile_count_semaphore=tile_count_semaphore,
+        cu_total_m_blocks=cu_total_m_blocks,
+        cu_total_splits_m_blocks=cu_total_splits_m_blocks,
+        blocks_to_batch_idx=blocks_to_batch_idx,
+    )
+
+
+_get_scheduler_metadata.compile_cache = get_jit_cache("scheduler_metadata")
+
+
+def get_scheduler_metadata(
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    nheads: int,
+    nheads_kv: int,
+    headdim: int,
+    num_splits: int,
+    headdim_v: Optional[int] = None,
+    pack_gqa: Optional[int] = None,
+    causal: bool = False,
+    window_size_left: Optional[int] = None,
+    window_size_right: Optional[int] = None,
+    seqlen_k_new: int = 0,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    cu_seqlens_k_new: Optional[torch.Tensor] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    seqlen_k_per_split: Optional[int] = None,
+    _arch: Optional[int] = None,
+) -> SchedulerMetadataTensorsTorch:
+    """Prepares metadata tensors used by varlen tile schedulers (SingleTileVarlenScheduler
+    and DynamicPersistentVarlenScheduler)
+
+    Explanation of selected args:
+        num_splits: maximum number of splits per batch entry that the prepare kernel can emit
+        seqlen_k_per_split: for bitwise reproducibility between forward and backward, can fix
+            an exact seqlen_k per split; num_splits is calculated accordingly.
+
+    Returns
+        SchedulerMetadataTensorsTorch, a named tuple including:
+        - num_splits_dynamic_ptr: per-batch num_splits
+        - num_nheads_in_l2_ptr: used for head swizzle to avoid l2 cache thrashing
+        - tile_count_semaphore: the global semaphore used by DynamicPersistentVarlenScheduler atomic incrementation
+        - cu_total_m_blocks: cumsum tensor counting total m_blocks, used for binary batch search with large batch_size
+        - cu_total_splits_m_blocks: complementary cumsum tensor used for binary batch search and to
+            extract dynamic num splits in the absense of num_splits_dynamic_ptr
+    """
+    arch = _get_device_arch() if _arch is None else _arch
+    if headdim_v is None:
+        headdim_v = headdim
+
+    batch_sizes = {}
+    if cu_seqlens_q is not None:
+        batch_sizes["cu_seqlens_q"] = cu_seqlens_q.shape[0] - 1
+    if cu_seqlens_k is not None:
+        batch_sizes["cu_seqlens_k"] = cu_seqlens_k.shape[0] - 1
+    if seqused_q is not None:
+        batch_sizes["seqused_q"] = seqused_q.shape[0]
+    if seqused_k is not None:
+        batch_sizes["seqused_k"] = seqused_k.shape[0]
+    assert batch_sizes, (
+        "get_scheduler_metadata requires at least one of "
+        "cu_seqlens_q/cu_seqlens_k/seqused_q/seqused_k"
+    )
+    num_batch = next(iter(batch_sizes.values()))
+    assert all(b == num_batch for b in batch_sizes.values()), (
+        f"inconsistent batch size across inputs: {batch_sizes}"
+    )
+    device = next(
+        t.device for t in (cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) if t is not None
+    )
+
+    causal, local, window_size_left, window_size_right = _resolve_causal_local_window(
+        causal, window_size_left, window_size_right
+    )
+
+    qhead_per_kvhead = nheads // nheads_kv
+    if pack_gqa is None:
+        pack_gqa = qhead_per_kvhead > 1
+
+    fwd_cfg = _get_fwd_config(
+        arch=arch,
+        head_dim=headdim,
+        head_dim_v=headdim_v,
+        causal=causal,
+        local=local,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        qhead_per_kvhead=qhead_per_kvhead,
+        pack_gqa=pack_gqa,
+        batch_size=num_batch,
+        num_head_kv=nheads_kv,
+        num_splits=num_splits,
+        device=device,
+    )
+    tile_m, tile_n = fwd_cfg.m_block_size, fwd_cfg.n_block_size
+    q_stage = fwd_cfg.q_stage
+    num_splits = fwd_cfg.num_splits
+
+    return _get_scheduler_metadata(
+        num_batch,
+        max_seqlen_q,
+        max_seqlen_k,
+        nheads,
+        nheads_kv,
+        headdim,
+        num_splits,
+        tile_m,
+        tile_n,
+        headdim_v=headdim_v,
+        pack_gqa=pack_gqa,
+        q_stage=q_stage,
+        causal=causal,
+        enable_pdl=False,  # pdl not yet enabled
+        sort=False,  # LPT batch sort not yet enabled
+        seqlen_k_new=seqlen_k_new,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        cu_seqlens_k_new=cu_seqlens_k_new,
+        seqused_q=seqused_q,
+        seqused_k=seqused_k,
+        leftpad_k=leftpad_k,
+        seqlen_k_per_split=seqlen_k_per_split,
+        zfill_padded_output=True,
+        use_clc_scheduler=utils._get_use_clc_scheduler_default(),
+    )
diff --git a/build/torch-cuda/mask.py b/build/torch-cuda/mask.py
index 8da835700ab3260686bb927e0f4269fe2cf43834..008f08acaa125ab6acefece5c4832c98063073e1 100644
--- a/build/torch-cuda/mask.py
+++ b/build/torch-cuda/mask.py
@@ -1,108 +1,157 @@
 # Copyright (c) 2025, Tri Dao.
 
-from typing import Optional, Callable
+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, const_expr
+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
+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(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
+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))
-    # 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
+    # 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(24, ncol - s * 24)):
-            in_bound = cutlass.Boolean(mask & (1 << i))
-            c = s * 24 + i
+        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
-                # 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)
+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 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)
+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.
 
-    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))
+    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.
 
-    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)
+    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.
 
-        # 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)
+    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)
 
-        # 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
+@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
 
 
@@ -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,
+        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"
@@ -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)
+                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):
@@ -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)
+                    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
@@ -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)
+                has_fastdiv and mask_seqlen and const_expr(aux_data.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(
+                    mask_value = call_mask_mod(
+                        mask_mod,
                         batch_idx_ssa,
                         head_idx_ssa,
                         q_idx_ssa,
                         kv_idx_ssa,
                         self.seqlen_info,
-                        aux_tensors,
+                        aux_data,
                     )
                     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)
+                            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
@@ -284,6 +339,7 @@ class AttentionMask:
                         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
@@ -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
+                        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]
@@ -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
+                            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
                         )
-                        # 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_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,
@@ -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,
+        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)
@@ -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(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):
@@ -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)
+                    mask_r2p_lambda(
+                        acc_S,
+                        lambda s: r2p_bitmask_below(seqlenk_col_limit, s),
+                        rank1=True,
+                    )
 
         elif const_expr(not mask_causal and not mask_local and mask_mod is not None):
-            # Block sparse case w/ mask_mod
-            has_fastdiv = const_expr(
-                fastdiv_mods is not None
-                and fastdiv_mods[0] is not None
-                and fastdiv_mods[1] is not None
+            # 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"
             )
-            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,
+            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,
                 )
-                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
+            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
+                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:
@@ -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)
+                    mask_r2p_lambda(
+                        acc_S,
+                        lambda s: r2p_bitmask_below(col_limit_right, s),
+                        rank1=True,
+                    )
             else:
                 local_row_offset_right = (
-                    causal_row_offset + self.window_size_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 - 1 - self.window_size_left
+                    causal_row_offset - 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)
+                    # 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(
@@ -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,
+        aux_data: AuxData = AuxData(),
         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:
+        Coordinate convention:
         - 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)
+                    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)
@@ -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(
+                    mask_value = call_mask_mod(
+                        mask_mod,
                         batch_idx_ssa,
                         head_idx_ssa,
                         q_idx_ssa,
                         kv_idx_ssa,
                         self.seqlen_info,
-                        aux_tensors,
+                        aux_data,
                     )
                     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)
+                    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
@@ -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]
+                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
diff --git a/build/torch-cuda/metadata.json b/build/torch-cuda/metadata.json
index dcb291053bd17ba509fe37d12a4e017761f980b7..c9f695f9c1b2ec7cf3ddefa46c243412a5dfaaad 100644
--- a/build/torch-cuda/metadata.json
+++ b/build/torch-cuda/metadata.json
@@ -1,9 +1,93 @@
 {
-  "version": 0,
+  "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
+    }
+  }
 }
\ No newline at end of file
diff --git a/build/torch-cuda/metadata.json.sigstore b/build/torch-cuda/metadata.json.sigstore
new file mode 100644
index 0000000000000000000000000000000000000000..64e1d9ce321f86b74454004e9a3bb9394b93ed72
--- /dev/null
+++ b/build/torch-cuda/metadata.json.sigstore
@@ -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="}}
\ No newline at end of file
diff --git a/build/torch-cuda/mma_sm100_desc.py b/build/torch-cuda/mma_sm100_desc.py
index ab8dd098b9246e2215ee8b7959b101296065b79e..9a2adfca8d7952746e61b874117244305c5047e1 100644
--- a/build/torch-cuda/mma_sm100_desc.py
+++ b/build/torch-cuda/mma_sm100_desc.py
@@ -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:
+    if cutlass_type is cutlass.Float8E4M3FN:
         return MXF8F6F4Format.E4M3
-    if cutlass_type is cutlass.FloatE5M2:
+    if cutlass_type is cutlass.Float8E5M2:
         return MXF8F6F4Format.E5M2
     raise TypeError(f"Unsupported CUTLASS scalar type for A/B: {cutlass_type!r}")
 
diff --git a/build/torch-cuda/named_barrier.py b/build/torch-cuda/named_barrier.py
index eadac4b926cc4a40b8b3a08971dbb84d0f78ea29..8eb3ac61193413599efcf6b6c39c1a10702ceefd 100644
--- a/build/torch-cuda/named_barrier.py
+++ b/build/torch-cuda/named_barrier.py
@@ -12,6 +12,19 @@ class NamedBarrierFwd(enum.IntEnum):
     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()
@@ -20,8 +33,10 @@ class NamedBarrierBwd(enum.IntEnum):
     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):
@@ -30,3 +45,19 @@ 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()
diff --git a/build/torch-cuda/pack_gqa.py b/build/torch-cuda/pack_gqa.py
index 8b0fd735e8bcc26db6b471e35e2787a15aaa4b5a..28a950e5de520eef47d7027cb74a52ddb16b0064 100644
--- a/build/torch-cuda/pack_gqa.py
+++ b/build/torch-cuda/pack_gqa.py
@@ -1,25 +1,123 @@
 # 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
+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:
-    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
+    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(
@@ -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)
+        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
diff --git a/build/torch-cuda/paged_kv.py b/build/torch-cuda/paged_kv.py
index e76e9a0a3a8cc9cc80a60f99dbdd64d333b5999c..227db132afccef75f12cb934a50848ee2be20c05 100644
--- a/build/torch-cuda/paged_kv.py
+++ b/build/torch-cuda/paged_kv.py
@@ -28,6 +28,9 @@ class PagedKVManager(ParamsBase):
     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
@@ -55,7 +58,11 @@ class PagedKVManager(ParamsBase):
         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
@@ -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
+        # 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)
@@ -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])
+            # 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,
@@ -111,6 +121,8 @@ class PagedKVManager(ParamsBase):
             num_threads,
             head_dim_padded,
             head_dim_v_padded,
+            arch,
+            v_gmem_transposed,
             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):
+    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(K_or_V == "K"):
-                tPrXPtr[i] = utils.elem_pointer(self.mK_paged, (page_offset, 0, page)).toint()
+            if const_expr(transposed):
+                tPrXPtr[i] = utils.elem_pointer(mX, (d_offset, page_offset, page)).toint()
             else:
-                tPrXPtr[i] = utils.elem_pointer(self.mV_paged, (0, page_offset, page)).toint()
+                tPrXPtr[i] = utils.elem_pointer(mX, (page_offset, d_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).
+    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(
@@ -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]))
+        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))
@@ -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,
-                )
+            self._copy_row_async(tXsX, tXcX, mX_paged_cur_copy, m, should_load)
diff --git a/build/torch-cuda/pipeline.py b/build/torch-cuda/pipeline.py
index e45284ff4272ed6a16e9f7fe7eac959c0a20e76d..f8fdc1e802885416ed94a0a84f7d9260bea84f67 100644
--- a/build/torch-cuda/pipeline.py
+++ b/build/torch-cuda/pipeline.py
@@ -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 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.
@@ -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."
 
 
+# ── 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):
-    @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
+    create = _override_create(NamedBarrierOg, None)  # patched below
 
     @dsl_user_op
     def arrive_w_index(self, index: Int32, *, loc=None, ip=None) -> None:
@@ -134,72 +186,121 @@ class NamedBarrier(NamedBarrierOg):
         )
 
 
+NamedBarrier.create = _override_create(NamedBarrierOg, NamedBarrier)
+
+
+# ── PipelineAsync ────────────────────────────────────────────────────────────
+
+
 @dataclass(frozen=True)
-class PipelineAsync(PipelineAsyncOg):
+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, **kwargs):
+    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)
-        # Can't assign to __class__ directly since the dataclass is frozen
-        # obj.__class__ = PipelineAsync
         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_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,
+    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 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,
+    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,
         )
 
-    @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)
+    # _w_index variants inherited from _PipelineIndexPhaseMixin, which delegate
+    # to producer_commit / consumer_release above.
+
+
+# ── PipelineCpAsync ──────────────────────────────────────────────────────────
 
 
 @dataclass(frozen=True)
-class PipelineTmaAsync(PipelineTmaAsyncOg):
-    """
-    Override producer_acquire to take in extra_tx_count parameter.
-    """
+class PipelineCpAsync(_PipelineIndexPhaseMixin, PipelineCpAsyncOg):
+    _elect_one_release: bool = False
+    _syncwarp_before_release: bool = True
 
     @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)
+    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,
@@ -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.
-    """
+PipelineTmaAsync.create = _override_create(PipelineTmaAsyncOg, PipelineTmaAsync)
+
+
+# ── PipelineTmaUmma ─────────────────────────────────────────────────────────
 
-    @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
+
+@dataclass(frozen=True)
+class PipelineTmaUmma(_PipelineIndexPhaseMixin, PipelineTmaUmmaOg):
+    """Override producer_acquire to take in extra_tx_count parameter."""
 
     @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,
-        )
+PipelineTmaUmma.create = _override_create(PipelineTmaUmmaOg, PipelineTmaUmma)
 
-    @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)
+
+# ── PipelineUmmaAsync ───────────────────────────────────────────────────────
 
 
 @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
+class PipelineUmmaAsync(_PipelineIndexPhaseMixin, PipelineUmmaAsyncOg):
+    pass
 
-    @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)
+PipelineUmmaAsync.create = _override_create(PipelineUmmaAsyncOg, PipelineUmmaAsync)
 
-    @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)
+# ── PipelineAsyncUmma ───────────────────────────────────────────────────────
 
 
 @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
+class PipelineAsyncUmma(_PipelineIndexPhaseMixin, PipelineAsyncUmmaOg):
+    pass
 
-    @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)
+PipelineAsyncUmma.create = _override_create(PipelineAsyncUmmaOg, PipelineAsyncUmma)
diff --git a/build/torch-cuda/prepare_scheduler.py b/build/torch-cuda/prepare_scheduler.py
new file mode 100644
index 0000000000000000000000000000000000000000..eef256edbabf543222d03647def4a4c00262fbcb
--- /dev/null
+++ b/build/torch-cuda/prepare_scheduler.py
@@ -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)
diff --git a/build/torch-cuda/quack/__init__.py b/build/torch-cuda/quack/__init__.py
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..1e049cca6d1a4dd8621bb71710ba56341c762656 100644
--- a/build/torch-cuda/quack/__init__.py
+++ b/build/torch-cuda/quack/__init__.py
@@ -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
diff --git a/build/torch-cuda/quack/activation.py b/build/torch-cuda/quack/activation.py
deleted file mode 100644
index f316dd4159b65ff06690dc4ac95edd19d0348421..0000000000000000000000000000000000000000
--- a/build/torch-cuda/quack/activation.py
+++ /dev/null
@@ -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,
-}
diff --git a/build/torch-cuda/quack/copy_utils.py b/build/torch-cuda/quack/copy_utils.py
index 2c3363899cfbe962c3a44f05ececd9f3e3c3c1e3..fcac27a85686e178460e2718f7da8b60020c28df 100644
--- a/build/torch-cuda/quack/copy_utils.py
+++ b/build/torch-cuda/quack/copy_utils.py
@@ -1,22 +1,118 @@
-# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
+# Copyright (c) 2025-2026, QuACK team.
 
-from typing import Optional, Type, Tuple, Callable, Sequence
+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.cute.nvgpu import cpasync, warp, warpgroup
+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.dialects import llvm
 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
@@ -41,6 +137,30 @@ def cvt_copy(
     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)
@@ -48,6 +168,13 @@ def load_s2r(src: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor:
     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,
@@ -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
+    tiled_copy: cute.TiledCopy,
+    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)
+    """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
 
 
@@ -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:
+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.core.ThrCopy, tensor: cute.Tensor
-) -> cute.Tensor:
+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,
@@ -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):
+    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,
@@ -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):
+    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,
@@ -394,26 +525,35 @@ def get_smem_load_atom(
 
 
 def get_smem_store_C(
-    tiled_mma: cute.TiledMma,
+    tiled_mma: cute.TiledMma | cute.TiledCopy,
     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)
+    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, **new_kwargs):
-        dst_tensor = tRS_sC if const_expr(dst_idx is None) else tRS_sC[None, None, None, dst_idx]
+    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
 
@@ -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)
+    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(arch, dtype, transpose)
+    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
 
@@ -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
-    )
+    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
 
@@ -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
+    copy_atom = get_smem_store_atom(dtype, transpose)
     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
+    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 == warpgroup.OperandMajorMode.MN
-    copy_atom = get_smem_store_atom(arch, dtype, transpose)
+    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):
@@ -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)
+    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):
@@ -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()
+    smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip)
     # 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",
+    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",
-        has_side_effects=True,
-        is_align_stack=False,
-        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
     )
 
 
@@ -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]
+    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).ir_value()
-    dst_addr = dst_smem_ptr.toint(loc=loc, ip=ip).ir_value()
+    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 = mbar_addr.ir_value()
+    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).ir_value()
+        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 = (
-        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];"
+        "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}];"
     )
 
-    llvm.inline_asm(
-        None,
-        [
+    cute.arch.inline_ptx(
+        ptx,
+        read_only_args=[
             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 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(
@@ -787,6 +1080,18 @@ def tma_producer_copy_fn(copy: Callable, pipeline: cutlass.pipeline.PipelineAsyn
     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,
@@ -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)
+    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_shape_mk[0] % thr_copy_A.tiler_mn[0].shape == 0
+    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_shape_mk[0])
+        limit_m = min(limit_m, tile_M)
     elems_per_load = cute.size(tAsA.shape[0][0])
-    cA = cute.make_identity_tensor(tile_shape_mk)
+    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]
@@ -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]))
+    mA_k = cute.logical_divide(mA, (None, tile_K))
 
-    def copy_fn(src_idx, dst_idx, pred: bool = False):
+    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_shape_mk[1]
+            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)]
@@ -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):
+        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[None, m, None, dst_idx].iterator
+            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)
diff --git a/build/torch-cuda/quack/cute_dsl_utils.py b/build/torch-cuda/quack/cute_dsl_utils.py
index 044a45a13428550d581268c48516a0db55aa9a17..2090fb978d102295dd2be38e41a56c264e7f6c1d 100644
--- a/build/torch-cuda/quack/cute_dsl_utils.py
+++ b/build/torch-cuda/quack/cute_dsl_utils.py
@@ -4,6 +4,9 @@ from typing import Tuple, get_origin
 from functools import lru_cache
 from dataclasses import dataclass, fields
 
+import os
+import re
+
 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)
+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(device: torch.device = None) -> Tuple[int, int]:
+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)}
@@ -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
diff --git a/build/torch-cuda/quack/dsl/__init__.py b/build/torch-cuda/quack/dsl/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..9e8689b6e44049319a5d375dd08ad2eb11a7ca28
--- /dev/null
+++ b/build/torch-cuda/quack/dsl/__init__.py
@@ -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
diff --git a/build/torch-cuda/quack/dsl/cute_tensor_indexing.py b/build/torch-cuda/quack/dsl/cute_tensor_indexing.py
new file mode 100644
index 0000000000000000000000000000000000000000..d6342665d680f71a5f3242005471fa1856df6040
--- /dev/null
+++ b/build/torch-cuda/quack/dsl/cute_tensor_indexing.py
@@ -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"]
diff --git a/build/torch-cuda/quack/layout_utils.py b/build/torch-cuda/quack/layout_utils.py
index 9a47480423b34a84862f04390b97c890208bf24c..77def95de62ff3adfac1810330cb840def01cf83 100644
--- a/build/torch-cuda/quack/layout_utils.py
+++ b/build/torch-cuda/quack/layout_utils.py
@@ -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 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:])
@@ -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))
+        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], l.shape[2][0]),
+                (l.shape[0][0], l.shape[0][1], l.shape[2][0]),
                 l.shape[1],
                 l.shape[2][1],
             ),
             stride=(
-                (l.stride[0], l.stride[2][0]),
+                (l.stride[0][0], l.stride[0][1], 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
+    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
@@ -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
+    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
@@ -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]
+
+
+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
diff --git a/build/torch-cuda/quack/rounding.py b/build/torch-cuda/quack/rounding.py
new file mode 100644
index 0000000000000000000000000000000000000000..15f3541376a1d26920f4602031abae6a5cd27387
--- /dev/null
+++ b/build/torch-cuda/quack/rounding.py
@@ -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
diff --git a/build/torch-cuda/quack/sm90_utils.py b/build/torch-cuda/quack/sm90_utils.py
index ebabfe6509ffa2152aa90baedeb5982b76c95750..25a6966a51f1733745132fb1cd0bccba5121a1b1 100644
--- a/build/torch-cuda/quack/sm90_utils.py
+++ b/build/torch-cuda/quack/sm90_utils.py
@@ -1,6 +1,6 @@
 # Copyright (c) 2025, Tri Dao.
 
-from typing import Type, Union, Optional
+from typing import Literal, 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
 
 
+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,
diff --git a/build/torch-cuda/quack/utils.py b/build/torch-cuda/quack/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..7d2bb9c89b258194b860bd76cf387415b6303ed7
--- /dev/null
+++ b/build/torch-cuda/quack/utils.py
@@ -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)
diff --git a/build/torch-cuda/seqlen_info.py b/build/torch-cuda/seqlen_info.py
index 6d8c6feb279b53f0c93915abe44990f25719a4cb..334476d5afae7240db8dfc10e2d20a6fd1cdd5cc 100644
--- a/build/torch-cuda/seqlen_info.py
+++ b/build/torch-cuda/seqlen_info.py
@@ -5,6 +5,8 @@ import cutlass
 import cutlass.cute as cute
 from cutlass import Int32, const_expr
 
+from .quack import copy_utils
+
 """
 This consolidates all the info related to sequence length. This is so that we can do all
 the gmem reads once at the beginning of each tile, rather than having to repeat these reads
@@ -14,62 +16,95 @@ to compute various things like n_block_min, n_block_max, etc.
 
 @dataclass(frozen=True)
 class SeqlenInfo:
-    offset: cutlass.Int32
-    seqlen: cutlass.Int32
+    offset: Int32
+    offset_padded: Int32
+    seqlen: Int32
+    has_cu_seqlens: cutlass.Constexpr[bool] = False
 
     @staticmethod
     def create(
-        batch_idx: cutlass.Int32,
-        seqlen_static: cutlass.Int32,
+        batch_idx: Int32,
+        seqlen_static: Int32,
         cu_seqlens: Optional[cute.Tensor] = None,
         seqused: Optional[cute.Tensor] = None,
+        tile: cutlass.Constexpr[int] = 128,
     ):
         offset = 0 if const_expr(cu_seqlens is None) else cu_seqlens[batch_idx]
+        offset_padded = (
+            0
+            if const_expr(cu_seqlens is None)
+            # Add divby so that the compiler knows the alignment when moving by offset_padded
+            else cute.assume((offset + batch_idx * tile) // tile * tile, divby=tile)
+        )
         if const_expr(seqused is not None):
             seqlen = seqused[batch_idx]
         elif const_expr(cu_seqlens is not None):
             seqlen = cu_seqlens[batch_idx + 1] - cu_seqlens[batch_idx]
         else:
             seqlen = seqlen_static
-        return SeqlenInfo(offset, seqlen)
+        return SeqlenInfo(offset, offset_padded, seqlen, has_cu_seqlens=cu_seqlens is not None)
+
+    def offset_batch(
+        self,
+        mT: cute.Tensor,
+        batch_idx: Int32,
+        dim: int,
+        padded: cutlass.Constexpr[bool] = False,
+        multiple: int = 1,
+    ) -> cute.Tensor:
+        """Offset a tensor by batch index. batch dim is at position `dim`, seqlen is at dim=0."""
+        if const_expr(not self.has_cu_seqlens):
+            idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mT) - 1 - dim)
+            return mT[idx]
+        else:
+            off = multiple * (self.offset if const_expr(not padded) else self.offset_padded)
+            offset = off if const_expr(cute.rank(mT.shape[0]) == 1) else (0, off)
+            idx = (offset,) + (None,) * (cute.rank(mT) - 1)
+            return cute.domain_offset(idx, mT)
 
 
 @dataclass(frozen=True)
 class SeqlenInfoQK:
-    offset_q: cutlass.Int32
-    offset_k: cutlass.Int32
-    padded_offset_q: cutlass.Int32
-    padded_offset_k: cutlass.Int32
-    seqlen_q: cutlass.Int32
-    seqlen_k: cutlass.Int32
+    offset_q: Int32
+    offset_k: Int32
+    padded_offset_q: Int32
+    padded_offset_k: Int32
+    seqlen_q: Int32
+    seqlen_k: Int32
+    m_block_offset: Int32
+    block_idx_offset: Int32
+    num_n_blocks: Int32
     has_cu_seqlens_q: cutlass.Constexpr[bool]
     has_cu_seqlens_k: cutlass.Constexpr[bool]
     has_seqused_q: cutlass.Constexpr[bool]
     has_seqused_k: cutlass.Constexpr[bool]
+    has_cu_block_idx_offsets: cutlass.Constexpr[bool] = False
 
     @staticmethod
     def create(
-        batch_idx: cutlass.Int32,
-        seqlen_q_static: cutlass.Int32,
-        seqlen_k_static: cutlass.Int32,
+        batch_idx: Int32,
+        seqlen_q_static: Int32,
+        seqlen_k_static: Int32,
         mCuSeqlensQ: Optional[cute.Tensor] = None,
         mCuSeqlensK: Optional[cute.Tensor] = None,
         mSeqUsedQ: Optional[cute.Tensor] = None,
         mSeqUsedK: Optional[cute.Tensor] = None,
-        tile_m: cutlass.Constexpr[cutlass.Int32] = 128,
-        tile_n: cutlass.Constexpr[cutlass.Int32] = 128,
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
+        mCuBlockIdxOffsets: Optional[cute.Tensor] = None,
+        tile_m: cutlass.Constexpr[Int32] = 128,
+        tile_n: cutlass.Constexpr[Int32] = 128,
     ):
         offset_q = 0 if const_expr(mCuSeqlensQ is None) else mCuSeqlensQ[batch_idx]
         offset_k = 0 if const_expr(mCuSeqlensK is None) else mCuSeqlensK[batch_idx]
         padded_offset_q = (
             0
             if const_expr(mCuSeqlensQ is None)
-            else (offset_q + batch_idx * tile_m) // tile_m * tile_m
+            else cute.assume((offset_q + batch_idx * tile_m) // tile_m * tile_m, divby=tile_m)
         )
         padded_offset_k = (
             0
             if const_expr(mCuSeqlensK is None)
-            else (offset_k + batch_idx * tile_n) // tile_n * tile_n
+            else cute.assume((offset_k + batch_idx * tile_n) // tile_n * tile_n, divby=tile_n)
         )
         if const_expr(mSeqUsedQ is not None):
             seqlen_q = mSeqUsedQ[batch_idx]
@@ -87,10 +122,13 @@ class SeqlenInfoQK:
                 if const_expr(mCuSeqlensK is None)
                 else mCuSeqlensK[batch_idx + 1] - offset_k
             )
-        has_cu_seqlens_q: int = mCuSeqlensQ is not None
-        has_cu_seqlens_k: int = mCuSeqlensK is not None
-        has_seqused_q: int = mSeqUsedQ is not None
-        has_seqused_k: int = mSeqUsedK is not None
+        m_block_offset = 0 if const_expr(mCuTotalMBlocks is None) else mCuTotalMBlocks[batch_idx]
+        num_n_blocks = (seqlen_k + tile_n - 1) // tile_n
+        block_idx_offset = (
+            mCuBlockIdxOffsets[batch_idx]
+            if const_expr(mCuBlockIdxOffsets is not None)
+            else m_block_offset * num_n_blocks
+        )
         return SeqlenInfoQK(
             offset_q,
             offset_k,
@@ -98,10 +136,14 @@ class SeqlenInfoQK:
             padded_offset_k,
             seqlen_q,
             seqlen_k,
-            has_cu_seqlens_q,
-            has_cu_seqlens_k,
-            has_seqused_q,
-            has_seqused_k,
+            m_block_offset,
+            block_idx_offset,
+            num_n_blocks,
+            has_cu_seqlens_q=mCuSeqlensQ is not None,
+            has_cu_seqlens_k=mCuSeqlensK is not None,
+            has_seqused_q=mSeqUsedQ is not None,
+            has_seqused_k=mSeqUsedK is not None,
+            has_cu_block_idx_offsets=mCuBlockIdxOffsets is not None,
         )
 
     def offset_batch_Q(
@@ -110,16 +152,38 @@ class SeqlenInfoQK:
         batch_idx: Int32,
         dim: int,
         padded: cutlass.Constexpr[bool] = False,
+        ragged: cutlass.Constexpr[bool] = False,
     ) -> cute.Tensor:
         """Seqlen must be the first dimension of mQ"""
-        if const_expr(not self.has_cu_seqlens_q):
-            idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mQ) - 1 - dim)
-            return mQ[idx]
+        if const_expr(not ragged):
+            if const_expr(not self.has_cu_seqlens_q):
+                idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mQ) - 1 - dim)
+                return mQ[idx]
+            else:
+                offset_q = self.offset_q if const_expr(not padded) else self.padded_offset_q
+                offset_q = offset_q if const_expr(cute.rank(mQ.shape[0]) == 1) else (None, offset_q)
+                idx = (offset_q,) + (None,) * (cute.rank(mQ) - 1)
+                return cute.domain_offset(idx, mQ)
         else:
-            offset_q = self.offset_q if const_expr(not padded) else self.padded_offset_q
-            offset = offset_q if const_expr(cute.rank(mQ.shape[0]) == 1) else (0, offset_q)
-            idx = (offset,) + (0,) * (cute.rank(mQ) - 1)
-            return cute.domain_offset(idx, mQ)
+            if const_expr(not self.has_cu_seqlens_q):
+                offset_q = 0
+                idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mQ) - 1 - dim)
+                mQ = mQ[idx]
+            else:
+                offset_q = self.offset_q if const_expr(not padded) else self.padded_offset_q
+            if const_expr(cute.rank(mQ.shape[0]) == 1):
+                return copy_utils.offset_ragged_tensor(
+                    mQ, offset_q, self.seqlen_q, ragged_dim=0, ptr_shift=True
+                )
+            else:  # PackGQA
+                assert cute.rank(mQ.shape[0]) == 2
+                # Unpack before calling offset_ragged_tensor, then pack
+                idx = ((None, None),) + (None,) * (cute.rank(mQ) - 1)
+                mQ = mQ[idx]
+                mQ = copy_utils.offset_ragged_tensor(
+                    mQ, offset_q, self.seqlen_q, ragged_dim=1, ptr_shift=True
+                )
+                return cute.group_modes(mQ, 0, 2)
 
     def offset_batch_K(
         self,
@@ -127,12 +191,114 @@ class SeqlenInfoQK:
         batch_idx: Int32,
         dim: int,
         padded: cutlass.Constexpr[bool] = False,
+        ragged: cutlass.Constexpr[bool] = False,
+        multiple: int = 1,
     ) -> cute.Tensor:
         """Seqlen must be the first dimension of mK"""
-        if const_expr(not self.has_cu_seqlens_k):
-            idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mK) - 1 - dim)
-            return mK[idx]
+        if const_expr(not ragged):
+            if const_expr(not self.has_cu_seqlens_k):
+                idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mK) - 1 - dim)
+                return mK[idx]
+            else:
+                offset_k = self.offset_k if const_expr(not padded) else self.padded_offset_k
+                offset_k *= multiple
+                idx = (offset_k,) + (None,) * (cute.rank(mK) - 1)
+                return cute.domain_offset(idx, mK)
+        else:
+            if const_expr(not self.has_cu_seqlens_k):
+                offset_k = 0
+                idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mK) - 1 - dim)
+                mK = mK[idx]
+            else:
+                offset_k = self.offset_k if const_expr(not padded) else self.padded_offset_k
+                offset_k *= multiple
+            return copy_utils.offset_ragged_tensor(
+                mK, offset_k, self.seqlen_k, ragged_dim=0, ptr_shift=True
+            )
+
+
+@dataclass(frozen=True)
+class SeqlenInfoQKNewK:
+    """Sequence length info for append-KV with left-padding and new K support.
+
+    Extends SeqlenInfoQK with:
+    - leftpad_k: left padding for K (tokens to skip at the start of the KV cache)
+    - offset_k_new: offset into the new K tensor
+    - seqlen_k_og: original K length (before appending new K), excluding leftpad
+    - seqlen_k_new: length of new K to append
+    - seqlen_k: total K length (seqlen_k_og + seqlen_k_new)
+    - seqlen_rotary: position for rotary embedding computation
+    """
+
+    leftpad_k: Int32
+    offset_q: Int32
+    offset_k: Int32
+    offset_k_new: Int32
+    seqlen_q: Int32
+    seqlen_k_og: Int32
+    seqlen_k_new: Int32
+    seqlen_k: Int32
+    seqlen_rotary: Int32
+
+    @staticmethod
+    def create(
+        batch_idx: Int32,
+        seqlen_q_static: Int32,
+        seqlen_k_static: Int32,
+        shape_K_new_0: Int32,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        mCuSeqlensKNew: Optional[cute.Tensor] = None,
+        mSeqUsedQ: Optional[cute.Tensor] = None,
+        mSeqUsedK: Optional[cute.Tensor] = None,
+        mLeftpadK: Optional[cute.Tensor] = None,
+        mSeqlensRotary: Optional[cute.Tensor] = None,
+    ):
+        leftpad_k = 0 if const_expr(mLeftpadK is None) else mLeftpadK[batch_idx]
+        offset_q = 0 if const_expr(mCuSeqlensQ is None) else mCuSeqlensQ[batch_idx]
+        if const_expr(mCuSeqlensK is not None):
+            offset_k = mCuSeqlensK[batch_idx] + leftpad_k
+        else:
+            offset_k = leftpad_k if const_expr(mCuSeqlensQ is not None) else 0
+        offset_k_new = 0 if const_expr(mCuSeqlensKNew is None) else mCuSeqlensKNew[batch_idx]
+        # seqlen_q
+        if const_expr(mSeqUsedQ is not None):
+            seqlen_q = mSeqUsedQ[batch_idx]
+        elif const_expr(mCuSeqlensQ is not None):
+            seqlen_q = mCuSeqlensQ[batch_idx + 1] - mCuSeqlensQ[batch_idx]
         else:
-            offset_k = self.offset_k if const_expr(not padded) else self.padded_offset_k
-            idx = (offset_k,) + (0,) * (cute.rank(mK) - 1)
-            return cute.domain_offset(idx, mK)
+            seqlen_q = seqlen_q_static
+        # seqlen_k_og: original K length (excluding leftpad)
+        if const_expr(mSeqUsedK is not None):
+            seqlen_k_og = mSeqUsedK[batch_idx] - leftpad_k
+        elif const_expr(mCuSeqlensK is not None):
+            seqlen_k_og = mCuSeqlensK[batch_idx + 1] - mCuSeqlensK[batch_idx] - leftpad_k
+        else:
+            seqlen_k_og = (
+                seqlen_k_static - leftpad_k
+                if const_expr(mCuSeqlensQ is not None)
+                else seqlen_k_static
+            )
+        # seqlen_k_new
+        if const_expr(mCuSeqlensKNew is None):
+            seqlen_k_new = 0 if const_expr(mCuSeqlensQ is None) else shape_K_new_0
+        else:
+            seqlen_k_new = mCuSeqlensKNew[batch_idx + 1] - mCuSeqlensKNew[batch_idx]
+        seqlen_k = seqlen_k_og if const_expr(mCuSeqlensQ is None) else seqlen_k_og + seqlen_k_new
+
+        # seqlen_rotary: defaults to seqlen_k_og + leftpad_k unless explicitly provided
+        if const_expr(mSeqlensRotary is not None):
+            seqlen_rotary = mSeqlensRotary[batch_idx]
+        else:
+            seqlen_rotary = seqlen_k_og + leftpad_k
+        return SeqlenInfoQKNewK(
+            leftpad_k,
+            offset_q,
+            offset_k,
+            offset_k_new,
+            seqlen_q,
+            seqlen_k_og,
+            seqlen_k_new,
+            seqlen_k,
+            seqlen_rotary,
+        )
diff --git a/build/torch-cuda/sm100_hd256_2cta_fmha_backward.py b/build/torch-cuda/sm100_hd256_2cta_fmha_backward.py
new file mode 100644
index 0000000000000000000000000000000000000000..24f2a7307575973a925683a2751b3f5b47aa2ddf
--- /dev/null
+++ b/build/torch-cuda/sm100_hd256_2cta_fmha_backward.py
@@ -0,0 +1,225 @@
+# Copyright (c) 2025, Siyu Wang, Shengbin Di, Yuxi Chi, Johnsonms, Linfeng Zheng, Haoyan Huang, Lanbo Li, Yun Zhong, Man Yuan, Minmin Sun, Yong Li, Wei Lin.
+
+
+"""Fused multi-head attention (FMHA) backward for the SM100 architecture using CUTE DSL.
+
+Constraints:
+* Supported head dimensions: 256 only
+* mma_tiler_mn must be 64,64
+* Batch size must be the same for Q, K, and V tensors
+"""
+
+import cuda.bindings.driver as cuda
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.typing import Int32
+
+from .sm100_hd256_2cta_fmha_backward_dqkernel import (
+    BlackwellFusedMultiHeadAttentionBackwardDQKernel,
+)
+from .sm100_hd256_2cta_fmha_backward_dkdvkernel import (
+    BlackwellFusedMultiHeadAttentionBackwardDKDVKernel,
+)
+from .cute_dsl_utils import assume_tensor_aligned
+from .utils import AuxData, as_bshkrd_tensor, as_shhb_tensor
+
+
+class BlackwellFusedMultiHeadAttentionBackward:
+    """FMHA backward class for executing CuTeDSL kernel."""
+
+    def __init__(
+        self,
+        head_dim: int,
+        head_dim_v: int | None = None,
+        is_causal: bool = False,
+        is_local: bool = False,
+        qhead_per_kvhead: cutlass.Constexpr[int] = 1,
+        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,
+        q_subtile_factor: cutlass.Constexpr[int] = 1,
+        tile_m_dq: int = 128,
+        tile_n_dq: int = 128,
+        tile_m_dkdv: int = 128,
+        tile_n_dkdv: int = 64,
+        window_size_left: int | None = None,
+        window_size_right: int | None = None,
+        use_clc_scheduler: bool = False,
+    ):
+        """Initialization."""
+        head_dim_v = head_dim if head_dim_v is None else head_dim_v
+        assert head_dim == 256 and head_dim_v == 256, (
+            "SM100 dedicated backward kernel only supports (head_dim, head_dim_v) = (256, 256)"
+        )
+        assert not is_local, "SM100 backward with head_dim=256 does not support local attention"
+        assert tile_m_dq == 128 and tile_n_dq == 128, (
+            "SM100 dedicated backward kernel only supports tile_m_dq=128 and tile_n_dq=128"
+        )
+        assert tile_m_dkdv == 128 and tile_n_dkdv == 64, (
+            "SM100 dedicated backward kernel only supports tile_m_dkdv=128 and tile_n_dkdv=64"
+        )
+        assert score_mod is None and score_mod_bwd is None and mask_mod is None, (
+            "SM100 backward with head_dim=256 does not support score_mod/mask_mod"
+        )
+        assert not deterministic, (
+            "SM100 backward with head_dim=256 does not support deterministic mode"
+        )
+        assert not has_aux_tensors, "SM100 backward with head_dim=256 does not support aux_tensors"
+        assert cluster_size in (1, 2), (
+            "SM100 backward with head_dim=256 only supports cluster_size in {1, 2}"
+        )
+        assert use_2cta_instrs, "SM100 backward with head_dim=256 requires use_2cta_instrs=True"
+        # q_subtile_factor is accepted for interface parity with FlashAttentionBackwardSm100,
+        # but this dedicated kernel uses fixed internal behavior.
+
+        self.acc_dtype = cutlass.Float32
+        self.is_causal = is_causal
+        self.window_size_left = (
+            None if (window_size_left is None or window_size_left < 0) else window_size_left
+        )
+        self.window_size_right = (
+            None if (window_size_right is None or window_size_right < 0) else window_size_right
+        )
+        self.tile_m_dq = tile_m_dq
+        self.tile_n_dq = tile_n_dq
+        self.tile_m_dkdv = tile_m_dkdv
+        self.tile_n_dkdv = tile_n_dkdv
+        self.use_clc_scheduler = use_clc_scheduler
+
+        self.dq_kernel = BlackwellFusedMultiHeadAttentionBackwardDQKernel(
+            self.acc_dtype,
+            (self.tile_m_dq, self.tile_n_dq, 256),
+            self.is_causal,
+            self.window_size_left,
+            self.window_size_right,
+            False,  # is_persistent
+            False,  # split_head
+            use_clc_scheduler=self.use_clc_scheduler,
+        )
+        self.dkdv_kernel = BlackwellFusedMultiHeadAttentionBackwardDKDVKernel(
+            self.acc_dtype,
+            (self.tile_m_dkdv, self.tile_n_dkdv, 256),
+            self.is_causal,
+            self.window_size_left,
+            self.window_size_right,
+            use_clc_scheduler=self.use_clc_scheduler,
+        )
+
+    @cute.jit
+    def __call__(
+        self,
+        Q: cute.Tensor,
+        K: cute.Tensor,
+        V: cute.Tensor,
+        dO: cute.Tensor,
+        lse_log2: cute.Tensor,
+        dpsum: cute.Tensor,
+        dQ_accum: cute.Tensor | None,
+        dK: cute.Tensor,
+        dV: cute.Tensor,
+        scale_softmax: cutlass.Float32,
+        cumulative_s_q: cute.Tensor | None,
+        cumulative_s_k: cute.Tensor | None,
+        seqused_q: cute.Tensor | None = None,
+        seqused_k: cute.Tensor | None = None,
+        window_size_left: Int32 | None = None,
+        window_size_right: Int32 | None = None,
+        dQ_semaphore: cute.Tensor | None = None,
+        dK_semaphore: cute.Tensor | None = None,
+        dV_semaphore: cute.Tensor | None = None,
+        aux_data: AuxData = AuxData(),
+        block_sparse_tensors: cute.Tensor | None = None,
+        stream: cuda.CUstream = None,
+    ):
+        """Host function to launch CuTeDSL kernel."""
+        assert seqused_q is None and seqused_k is None, (
+            "SM100 backward with head_dim=256 does not support seqused_q/seqused_k"
+        )
+        assert window_size_left is None and window_size_right is None, (
+            "SM100 backward with head_dim=256 uses constructor-provided window sizes"
+        )
+        assert dQ_semaphore is None and dK_semaphore is None and dV_semaphore is None, (
+            "SM100 backward with head_dim=256 does not use semaphores"
+        )
+        assert block_sparse_tensors is None, (
+            "SM100 backward with head_dim=256 does not support block sparse tensors"
+        )
+        assert aux_data.tensors is None or len(aux_data.tensors) == 0, (
+            "SM100 backward with head_dim=256 does not support aux_tensors"
+        )
+        assert aux_data.scalars is None or len(aux_data.scalars) == 0, (
+            "SM100 backward with head_dim=256 does not support aux_scalars"
+        )
+        assert dQ_accum is not None, (
+            "SM100 backward with head_dim=256 expects dQ tensor at dQ_accum slot"
+        )
+        dQ = dQ_accum
+        varlen = cumulative_s_q is not None or cumulative_s_k is not None
+        q_rank = cute.rank(Q.layout)
+        k_rank = cute.rank(K.layout)
+        if cutlass.const_expr(q_rank == 5):
+            h_q = Q.shape[2] * Q.shape[3]
+        elif cutlass.const_expr(q_rank == 4):
+            h_q = Q.shape[2]
+        else:
+            h_q = Q.shape[1]
+        if cutlass.const_expr(k_rank == 5):
+            h_k = K.shape[2]
+        elif cutlass.const_expr(k_rank == 4):
+            h_k = K.shape[2]
+        else:
+            h_k = K.shape[1]
+        h_r = h_q // h_k
+        if cutlass.const_expr(cumulative_s_q is not None):
+            b = cumulative_s_q.shape[0] - 1
+        elif cutlass.const_expr(cumulative_s_k is not None):
+            b = cumulative_s_k.shape[0] - 1
+        else:
+            b = Q.shape[0]
+
+        Q, K, V, dQ, dK, dV, dO = [assume_tensor_aligned(t) for t in (Q, K, V, dQ, dK, dV, dO)]
+
+        Q = as_bshkrd_tensor(Q, h_k, h_r, varlen)
+        K = as_bshkrd_tensor(K, h_k, 1, varlen)
+        V = as_bshkrd_tensor(V, h_k, 1, varlen)
+        dQ = as_bshkrd_tensor(dQ, h_k, h_r, varlen)
+        dK = as_bshkrd_tensor(dK, h_k, 1, varlen)
+        dV = as_bshkrd_tensor(dV, h_k, 1, varlen)
+        dO = as_bshkrd_tensor(dO, h_k, h_r, varlen)
+        scaled_LSE = as_shhb_tensor(lse_log2, h_k, h_r, b, varlen)
+        sum_OdO = as_shhb_tensor(dpsum, h_k, h_r, b, varlen)
+
+        # Keep original order: dQ first, then dKdV.
+        self.dq_kernel(
+            Q,
+            K,
+            V,
+            dQ,
+            dO,
+            scaled_LSE,
+            sum_OdO,
+            cumulative_s_q,
+            cumulative_s_k,
+            scale_softmax,
+            stream,
+        )
+        self.dkdv_kernel(
+            Q,
+            K,
+            V,
+            dK,
+            dV,
+            dO,
+            scaled_LSE,
+            sum_OdO,
+            cumulative_s_q,
+            cumulative_s_k,
+            scale_softmax,
+            stream,
+        )
diff --git a/build/torch-cuda/sm100_hd256_2cta_fmha_backward_dkdvkernel.py b/build/torch-cuda/sm100_hd256_2cta_fmha_backward_dkdvkernel.py
new file mode 100644
index 0000000000000000000000000000000000000000..376ca94c90c717235503da9af629dd41426b3b7b
--- /dev/null
+++ b/build/torch-cuda/sm100_hd256_2cta_fmha_backward_dkdvkernel.py
@@ -0,0 +1,3411 @@
+# Copyright (c) 2025, Siyu Wang, Shengbin Di, Yuxi Chi, Johnsonms, Linfeng Zheng, Haoyan Huang, Lanbo Li, Yun Zhong, Man Yuan, Minmin Sun, Yong Li, Wei Lin.
+
+"""Fused multi-head attention (FMHA) backward for the SM100 architecture using CUTE DSL.
+
+Constraints:
+* Supported head dimensions: 256 only
+* cta_tiler_mn must be 64,128
+* Batch size must be the same for Q, K, and V tensors
+"""
+
+import enum
+import math
+
+import cuda.bindings.driver as cuda
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.nvgpu import cpasync, tcgen05
+import cutlass.utils as utils
+import cutlass.pipeline as pipeline
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass.cute.typing import Int32
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
+
+from cutlass.utils import ClcDynamicPersistentTileScheduler
+from .tile_scheduler import (
+    SchedulerState,
+    SM100_TMEM_CAPACITY_COLUMNS,
+    make_sm100_thread_cooperative_group as make_thread_cooperative_group,
+    Sm100FmhaClcDynamicTileSchedulerParams as FmhaClcDynamicTileSchedulerParams,
+    Sm100FmhaClcDynamicTileScheduler as FmhaClcDynamicTileScheduler,
+    Sm100FmhaStaticTileSchedulerParams as FmhaStaticTileSchedulerParams,
+)
+
+from . import copy_utils as fa_copy_utils
+
+LAYOUT_RANK_CONSTANT = 3
+
+
+@cute.jit
+def split_wg(
+    t: cute.Tensor,
+    num_warp_groups: Int32,
+    wg_idx: Int32,
+) -> cute.Tensor:
+    """Split warp group."""
+    # dishengbin, TODO：need to double check if more efficient to split in other dimensions
+    ret = None
+    if cutlass.const_expr(cute.rank(t.layout) == LAYOUT_RANK_CONSTANT):
+        p = cute.composition(
+            t,
+            cute.make_layout(
+                (
+                    t.shape[0],
+                    t.shape[1],
+                    (num_warp_groups, cute.size(t, mode=[2]) // num_warp_groups),
+                )
+            ),
+        )
+        ret = p[None, None, (wg_idx, None)]
+    else:
+        p = cute.composition(
+            t,
+            cute.make_layout(
+                (
+                    t.shape[0],
+                    t.shape[1],
+                    t.shape[2],
+                    (num_warp_groups, cute.size(t, mode=[3]) // num_warp_groups),
+                )
+            ),
+        )
+        ret = p[None, None, None, (wg_idx, None)]
+    return ret
+
+
+class MaskType(enum.Enum):
+    """Mask type used in FMHA backward."""
+
+    NO_MASK = enum.auto()
+    RESIDUAL_MASK_FOR_BACKWARD = enum.auto()
+    CAUSAL_MASK_FOR_BACKWARD = enum.auto()
+
+
+def Tmemory_offset(lane, col):
+    """Tensor memory offset."""
+    return (lane << 16) + col
+
+
+permute_order = (0, 1, 2, 3, 4)
+
+
+class BlackwellFusedMultiHeadAttentionBackwardDKDVKernel:
+    """FMHA backward class for executing CuTeDSL kernel."""
+
+    def __init__(
+        self,
+        acc_dtype: type[cutlass.Numeric],
+        cta_tiler: tuple[int, int, int],
+        is_causal: bool,
+        window_size_left: int | None,
+        window_size_right: int | None,
+        use_clc_scheduler: bool = False,
+    ):
+        """Initialization."""
+        self.acc_dtype = acc_dtype
+        self.cta_tiler = cta_tiler
+        self.use_clc_scheduler = use_clc_scheduler
+        self.sched_warp_id = 10 if use_clc_scheduler else None
+        # TODO: need check, not sure whether need to *2 if 2cta
+        self.tile_shape_Q = cta_tiler[0]
+        self.tile_shape_K = cta_tiler[1]
+        self.tile_shape_dQ_K = cta_tiler[2]
+        self.tile_shape_dV_dO = cta_tiler[2]
+        # For S
+        self.KQ_mma_tiler = (
+            cta_tiler[1] * 2,
+            cta_tiler[0],
+            cta_tiler[2],
+        )
+        # For dP
+        self.VdO_mma_tiler = (
+            cta_tiler[1] * 2,
+            cta_tiler[0],
+            cta_tiler[2],
+        )
+        # For dV
+        self.PdO_mma_tiler = (
+            cta_tiler[1] * 2,
+            cta_tiler[2],
+            cta_tiler[0],
+        )
+        # For dK
+        self.dSQ_mma_tiler = (
+            cta_tiler[1] * 2,
+            cta_tiler[2],
+            cta_tiler[0],
+        )
+        # For dQ, dishengbin, need to remove
+        self.dSK_mma_tiler = (
+            cta_tiler[0] * 2,
+            cta_tiler[2],
+            cta_tiler[1],
+        )
+        self.cluster_shape_mn = (2, 1)
+        self.is_causal = is_causal
+        self.window_size_left: int = -1 if window_size_left is None else window_size_left
+        self.window_size_right: int = -1 if window_size_right is None else window_size_right
+        self.has_sliding_window = False
+        if self.window_size_left > 0 or self.window_size_right > 0:
+            self.has_sliding_window = True
+        if self.is_causal:
+            self.window_size_right = 0
+
+        self.compute_warp_id = (0, 1, 2, 3, 4, 5, 6, 7)
+        self.mma_warp_id = 8
+        self.load_warp_id = 9
+        self.empty_warp_id = 10
+
+        self.num_compute_warps = 8
+
+        self.tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS
+
+        self.threads_per_warp = 32
+        self.threads_per_cta = self.threads_per_warp * (self.num_compute_warps + 4)
+
+        self.cta_sync_bar_id = 0
+        self.tmem_alloc_sync_bar_id = 1
+        self.compute_sync_bar_id = 2
+        self.epilogue_sync_bar_id = 3
+        self.reduce_sync_bar_id = 4
+
+        self.tmem_dK_offset = 0
+        self.tmem_dV_offset = Tmemory_offset(0, cta_tiler[2] // 2)
+        self.tmem_dP_offset = Tmemory_offset(0, cta_tiler[2] + cta_tiler[0] // 2)
+        self.tmem_S_offset = Tmemory_offset(0, cta_tiler[2])
+
+        self.num_regs_reduce = 152
+        self.num_regs_compute = 128
+        self.num_regs_mma = 128
+        self.num_regs_empty = 96
+        self.num_regs_load = 96
+
+        self.buffer_align_bytes = 128
+
+    def _setup_attributes(self):
+        """Settings for pipeline stage."""
+        self.load_mma_Q_stage = 1
+        self.load_mma_K_stage = 1
+        self.load_mma_V_stage = 1
+        self.load_mma_QT_stage = 1
+        self.load_mma_dO_stage = 1
+        self.load_compute_LSE_stage = 1
+        self.load_compute_sum_OdO_stage = 1
+        self.mma_compute_S_stage = 1
+        self.mma_compute_dP_stage = 1
+        self.compute_mma_P_stage = 1
+        self.compute_mma_dS_stage = 1
+        self.mma_compute_dKdV_stage = 2
+
+        if cutlass.const_expr(self.use_clc_scheduler):
+            self.num_clc_stage = 1
+            self.num_clc_response_bytes = 16
+
+    @cute.jit
+    def __call__(
+        self,
+        Q: cute.Tensor,
+        K: cute.Tensor,
+        V: cute.Tensor,
+        dK: cute.Tensor,
+        dV: cute.Tensor,
+        dO: cute.Tensor,
+        scaled_LSE: cute.Tensor,
+        sum_OdO: cute.Tensor,
+        cumulative_s_q: cute.Tensor | None,
+        cumulative_s_k: cute.Tensor | None,
+        scale_softmax: cutlass.Float32,
+        stream: cuda.CUstream,
+    ):
+        """Host function to launch CuTeDSL kernel."""
+        varlen = cumulative_s_q is not None or cumulative_s_k is not None
+        # Infer shape metadata from normalized 5D tensors (B, S, H_k, H_r, D).
+        h_r = Q.shape[3]
+        h_k = Q.shape[2]
+        if cutlass.const_expr(cumulative_s_q is not None):
+            b = cumulative_s_q.shape[0] - 1
+        elif cutlass.const_expr(cumulative_s_k is not None):
+            b = cumulative_s_k.shape[0] - 1
+        else:
+            b = Q.shape[0]
+        problem_shape = (
+            Q.shape[1],
+            K.shape[1],
+            Q.shape[4],
+            ((h_r, h_k), b),
+        )
+        hb = ((h_r, h_k), b)
+        # (b, s, h_k, h_r, d) -> (s, d, ((h_r, h_k), b))
+        Q = cute.make_tensor(
+            Q.iterator,
+            cute.make_layout(
+                (Q.shape[1], Q.shape[4], hb),
+                stride=(
+                    cute.assume(Q.stride[1], divby=64),
+                    Q.stride[4],
+                    (
+                        (Q.stride[3], Q.stride[2]),
+                        0 if cumulative_s_q is not None else cute.assume(Q.stride[0], divby=64),
+                    ),
+                ),
+            ),
+        )
+        # (b, s, h_k, 1, d) -> (s, d, ((1, h_k), b))
+        K = cute.make_tensor(
+            K.iterator,
+            cute.make_layout(
+                (K.shape[1], K.shape[4], hb),
+                stride=(
+                    cute.assume(K.stride[1], divby=64),
+                    K.stride[4],
+                    (
+                        (0, K.stride[2]),
+                        0 if cumulative_s_k is not None else cute.assume(K.stride[0], divby=64),
+                    ),
+                ),
+            ),
+        )
+        # (b, s, h_k, 1, d) -> (s, d, ((1, h_k), b))
+        V = cute.make_tensor(
+            V.iterator,
+            cute.make_layout(
+                (V.shape[1], V.shape[4], hb),
+                stride=(
+                    cute.assume(V.stride[1], divby=64),
+                    V.stride[4],
+                    (
+                        (0, V.stride[2]),
+                        0 if cumulative_s_k is not None else cute.assume(V.stride[0], divby=64),
+                    ),
+                ),
+            ),
+        )
+        # (s, d, ((h_r, h_k), b)) -> (d, s, ((h_r, h_k), b))
+        QT = cute.make_tensor(
+            Q.iterator,
+            cute.make_layout(
+                (Q.shape[1], Q.shape[0], Q.shape[2]),
+                stride=(
+                    Q.stride[1],
+                    Q.stride[0],
+                    Q.stride[2],
+                ),
+            ),
+        )
+        dK = cute.make_tensor(
+            dK.iterator,
+            cute.make_layout(
+                (dK.shape[1], dK.shape[4], hb),
+                stride=(
+                    cute.assume(dK.stride[1], divby=64),
+                    dK.stride[4],
+                    (
+                        (0, dK.stride[2]),
+                        0 if cumulative_s_k is not None else cute.assume(dK.stride[0], divby=64),
+                    ),
+                ),
+            ),
+        )
+        dV = cute.make_tensor(
+            dV.iterator,
+            cute.make_layout(
+                (dV.shape[1], dV.shape[4], hb),
+                stride=(
+                    cute.assume(dV.stride[1], divby=64),
+                    dV.stride[4],
+                    (
+                        (0, dV.stride[2]),
+                        0 if cumulative_s_k is not None else cute.assume(dV.stride[0], divby=64),
+                    ),
+                ),
+            ),
+        )
+        # (s, d, ((h_r, h_k), b))
+        dO = cute.make_tensor(
+            dO.iterator,
+            cute.make_layout(
+                (dO.shape[1], dO.shape[4], hb),
+                stride=(
+                    cute.assume(dO.stride[1], divby=64),
+                    dO.stride[4],
+                    (
+                        (dO.stride[3], dO.stride[2]),
+                        0 if cumulative_s_q is not None else cute.assume(dO.stride[0], divby=64),
+                    ),
+                ),
+            ),
+        )
+
+        # (s, d, ((h_r, h_k), b)) -> (d, s, ((h_r, h_k), b))
+        dOT = cute.make_tensor(
+            dO.iterator,
+            cute.make_layout(
+                (dO.shape[1], dO.shape[0], dO.shape[2]),
+                stride=(
+                    dO.stride[1],
+                    dO.stride[0],
+                    dO.stride[2],
+                ),
+            ),
+        )
+
+        self.Q_major_mode = utils.LayoutEnum.from_tensor(Q).mma_major_mode()
+        self.K_major_mode = utils.LayoutEnum.from_tensor(K).mma_major_mode()
+        self.dK_major_mode = utils.LayoutEnum.from_tensor(dK).mma_major_mode()
+        self.V_major_mode = utils.LayoutEnum.from_tensor(V).mma_major_mode()
+        self.dV_major_mode = utils.LayoutEnum.from_tensor(dV).mma_major_mode()
+        self.dO_major_mode = utils.LayoutEnum.from_tensor(dO).mma_major_mode()
+
+        if cutlass.const_expr(self.Q_major_mode != tcgen05.OperandMajorMode.K):
+            raise RuntimeError(f"The layout of q is not supported: {self.Q_major_mode}")
+        if cutlass.const_expr(self.K_major_mode != tcgen05.OperandMajorMode.K):
+            raise RuntimeError("The layout of k is not supported")
+        if cutlass.const_expr(self.dK_major_mode != tcgen05.OperandMajorMode.K):
+            raise RuntimeError("The layout of dk is not supported")
+        if cutlass.const_expr(self.V_major_mode != tcgen05.OperandMajorMode.K):
+            raise RuntimeError("The layout of v is not supported")
+        if cutlass.const_expr(self.dV_major_mode != tcgen05.OperandMajorMode.K):
+            raise RuntimeError("The layout of dv is not supported")
+
+        self._setup_attributes()
+
+        cta_group = tcgen05.CtaGroup.TWO
+        PT_source = tcgen05.OperandSource.SMEM
+
+        # compute S
+        KQ_tiled_mma = sm100_utils.make_trivial_tiled_mma(
+            K.element_type,
+            tcgen05.OperandMajorMode.K,
+            tcgen05.OperandMajorMode.K,
+            self.acc_dtype,
+            cta_group,
+            self.KQ_mma_tiler[:2],
+        )
+        # compute dP
+        VdO_tiled_mma = sm100_utils.make_trivial_tiled_mma(
+            V.element_type,
+            tcgen05.OperandMajorMode.K,
+            tcgen05.OperandMajorMode.K,
+            self.acc_dtype,
+            cta_group,
+            self.VdO_mma_tiler[:2],
+        )
+        # compute dV
+        PdO_tiled_mma = sm100_utils.make_trivial_tiled_mma(
+            dO.element_type,
+            tcgen05.OperandMajorMode.K,
+            tcgen05.OperandMajorMode.MN,
+            self.acc_dtype,
+            cta_group,
+            self.PdO_mma_tiler[:2],
+            PT_source,
+        )
+        # compute dK
+        dSQ_tiled_mma = sm100_utils.make_trivial_tiled_mma(
+            Q.element_type,
+            tcgen05.OperandMajorMode.K,
+            tcgen05.OperandMajorMode.MN,
+            self.acc_dtype,
+            cta_group,
+            self.dSQ_mma_tiler[:2],
+        )
+        # compute dQ
+        # dishengbin, need to remove, but used in dS_mem_layout_staged
+        dSK_tiled_mma = sm100_utils.make_trivial_tiled_mma(
+            K.element_type,
+            tcgen05.OperandMajorMode.MN,
+            tcgen05.OperandMajorMode.MN,
+            self.acc_dtype,
+            cta_group,
+            self.dSK_mma_tiler[:2],
+        )
+
+        atom_thr_size = cute.size(KQ_tiled_mma.thr_id.shape)
+        self.cluster_shape_mnk = (*self.cluster_shape_mn, 1)  # type: ignore[assignment]
+        self.cluster_layout_vmnk = cute.tiled_divide(
+            cute.make_layout(self.cluster_shape_mnk),
+            (atom_thr_size,),
+        )
+
+        K_smem_layout_staged = sm100_utils.make_smem_layout_a(
+            KQ_tiled_mma,
+            self.KQ_mma_tiler,
+            K.element_type,
+            1,
+        )
+        Q_smem_layout_staged = sm100_utils.make_smem_layout_b(
+            KQ_tiled_mma,
+            self.KQ_mma_tiler,
+            Q.element_type,
+            self.load_mma_Q_stage,
+        )
+        V_smem_layout_staged = sm100_utils.make_smem_layout_a(
+            VdO_tiled_mma,
+            self.VdO_mma_tiler,
+            V.element_type,
+            1,
+        )
+        dO_smem_layout_staged = sm100_utils.make_smem_layout_b(
+            VdO_tiled_mma,
+            self.VdO_mma_tiler,
+            dO.element_type,
+            self.load_mma_dO_stage,
+        )
+        # dishengbin, need to remove, but used for sPT, need to double check
+        dS_smem_layout_staged = sm100_utils.make_smem_layout_a(
+            dSK_tiled_mma,
+            self.dSK_mma_tiler,
+            Q.element_type,
+            self.compute_mma_dS_stage,
+        )
+        dST_smem_layout_staged = sm100_utils.make_smem_layout_a(
+            dSQ_tiled_mma,
+            self.dSQ_mma_tiler,
+            Q.element_type,
+            self.compute_mma_dS_stage,
+        )
+        tiled_mma = dSQ_tiled_mma
+        is_k_major = tiled_mma.op.a_major_mode == tcgen05.OperandMajorMode.K
+        a_major_mode = tcgen05.OperandMajorMode.K if is_k_major else tcgen05.OperandMajorMode.MN
+        tmp = cute.dice(self.dSQ_mma_tiler, (1, None, 1))
+        a_smem_shape = tiled_mma.partition_shape_A(
+            cute.dice(self.dSQ_mma_tiler, (1, None, 1)),
+        )
+        a_smem_shape_mn_k = (
+            cute.size(a_smem_shape[0][0]) * a_smem_shape[1],
+            cute.size(a_smem_shape[0][1]) * a_smem_shape[2],
+        )
+        smem_layout_atom_kind = sm100_utils.get_smem_layout_atom_ab(
+            a_major_mode,
+            K.element_type,
+            a_smem_shape_mn_k,
+        )
+        a_smem_layout_atom = sm100_utils.make_smem_layout_atom(
+            smem_layout_atom_kind,
+            K.element_type,
+        )
+
+        a_smem_shape = cute.append(
+            a_smem_shape,
+            self.compute_mma_dS_stage,
+        )
+        order = (2, 1, 3) if not is_k_major else (1, 2, 3)
+        dST_smem_layout_staged_tmp = sm100_utils.tile_to_mma_shape(
+            a_smem_layout_atom,
+            a_smem_shape,
+            order=order,
+        )
+        QT_smem_layout_staged = sm100_utils.make_smem_layout_b(
+            dSQ_tiled_mma,
+            self.dSQ_mma_tiler,
+            Q.element_type,
+            self.load_mma_QT_stage,
+        )
+        P_smem_layout_staged = sm100_utils.make_smem_layout_a(
+            PdO_tiled_mma,
+            self.PdO_mma_tiler,
+            Q.element_type,
+            self.compute_mma_P_stage,
+        )
+        dOT_smem_layout_staged = sm100_utils.make_smem_layout_b(
+            PdO_tiled_mma,
+            self.PdO_mma_tiler,
+            dO.element_type,
+            self.load_mma_dO_stage,
+        )
+        LSE_smem_layout = cute.make_layout((self.cta_tiler[0], self.load_compute_LSE_stage))
+        sum_OdO_smem_layout = cute.make_layout((self.cta_tiler[0], self.load_compute_sum_OdO_stage))
+
+        tma_load_op = cpasync.CopyBulkTensorTileG2SOp(cta_group)
+        tma_reduce_op = cpasync.CopyReduceBulkTensorTileS2GOp()
+
+        K_smem_layout = cute.select(K_smem_layout_staged, mode=[0, 1, 2])
+        tma_atom_K, tma_tensor_K = cute.nvgpu.make_tiled_tma_atom_A(
+            tma_load_op,
+            K,
+            K_smem_layout,
+            self.KQ_mma_tiler,
+            KQ_tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+
+        V_smem_layout = cute.select(V_smem_layout_staged, mode=[0, 1, 2])
+        tma_atom_V, tma_tensor_V = cute.nvgpu.make_tiled_tma_atom_A(
+            tma_load_op,
+            V,
+            V_smem_layout,
+            self.VdO_mma_tiler,
+            VdO_tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+
+        Q_smem_layout = cute.select(Q_smem_layout_staged, mode=[0, 1, 2])
+        tma_atom_Q, tma_tensor_Q = cute.nvgpu.make_tiled_tma_atom_B(
+            tma_load_op,
+            Q,
+            Q_smem_layout,
+            self.KQ_mma_tiler,
+            KQ_tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+        QT_smem_layout = cute.select(QT_smem_layout_staged, mode=[0, 1, 2])
+        tma_atom_QT, tma_tensor_QT = cute.nvgpu.make_tiled_tma_atom_B(
+            tma_load_op,
+            QT,
+            QT_smem_layout,
+            self.dSQ_mma_tiler,
+            dSQ_tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+
+        dO_smem_layout = cute.select(dO_smem_layout_staged, mode=[0, 1, 2])
+        tma_atom_dO, tma_tensor_dO = cute.nvgpu.make_tiled_tma_atom_B(
+            tma_load_op,
+            dO,
+            dO_smem_layout,
+            self.VdO_mma_tiler,
+            VdO_tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+        dOT_smem_layout = cute.select(dOT_smem_layout_staged, mode=[0, 1, 2])
+        tma_atom_dOT, tma_tensor_dOT = cute.nvgpu.make_tiled_tma_atom_B(
+            tma_load_op,
+            dOT,
+            dOT_smem_layout,
+            self.PdO_mma_tiler,
+            PdO_tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+
+        # for 2cta, tma_copy_QT_bytes is same as the tma_copy_Q_bytes
+        self.tma_copy_Q_bytes = cute.size_in_bytes(Q.element_type, Q_smem_layout) * atom_thr_size
+        self.tma_copy_K_bytes = cute.size_in_bytes(K.element_type, K_smem_layout) * atom_thr_size
+        self.tma_copy_V_bytes = cute.size_in_bytes(V.element_type, V_smem_layout) * atom_thr_size
+        self.tma_copy_dO_bytes = cute.size_in_bytes(dO.element_type, dO_smem_layout) * atom_thr_size
+
+        # Variant 3a epilogue: TMA store atoms (S2G) for dK / dV.
+        # Each compute warp group owns half the hd_v output via split_wg, so
+        # the per-WG epi tile is (cta_tiler[1], cta_tiler[2] / num_compute_wgs).
+        # SMEM staging will alias onto sP+sdST in subsequent commits; for now
+        # the atoms and layouts are built and threaded through but unused.
+        tma_store_op = cpasync.CopyBulkTensorTileS2GOp()
+        num_compute_wgs = self.num_compute_warps // 4
+        # Variant 3a Path 2: CTA-shared epilogue SMEM. epi_tile is (M, gcd(128B, ...))
+        # = (M, 64) for bf16 hd=256. Total stages = num_compute_wgs * num_epi_stages
+        # = 4 stages of (64, 64), virtually a per-CTA (64, 256) buffer aliased onto
+        # sP+sdST. Both warp-groups cooperatively populate this buffer; TMA fires
+        # one (64, 64) box per stage to the corresponding (64, 64) GMEM slice.
+        epi_cols_dKV = math.gcd(
+            128 // (dK.element_type.width // 8), self.cta_tiler[2] // num_compute_wgs
+        )
+        num_epi_stages_dKV = (self.cta_tiler[2] // num_compute_wgs) // epi_cols_dKV
+        epi_tile_dKV = (self.cta_tiler[1], epi_cols_dKV)
+        total_epi_stages = num_compute_wgs * num_epi_stages_dKV
+        dK_layout_enum = utils.LayoutEnum.from_tensor(dK)
+        dV_layout_enum = utils.LayoutEnum.from_tensor(dV)
+        sdK_epi_layout = sm100_utils.make_smem_layout_epi(
+            dK.element_type,
+            dK_layout_enum,
+            epi_tile_dKV,
+            total_epi_stages,
+        )
+        sdV_epi_layout = sm100_utils.make_smem_layout_epi(
+            dV.element_type,
+            dV_layout_enum,
+            epi_tile_dKV,
+            total_epi_stages,
+        )
+        tma_atom_dK, tma_tensor_dK = cpasync.make_tiled_tma_atom(
+            tma_store_op,
+            dK,
+            cute.select(sdK_epi_layout, mode=[0, 1]),
+            epi_tile_dKV,
+        )
+        tma_atom_dV, tma_tensor_dV = cpasync.make_tiled_tma_atom(
+            tma_store_op,
+            dV,
+            cute.select(sdV_epi_layout, mode=[0, 1]),
+            epi_tile_dKV,
+        )
+
+        @cute.struct
+        class SharedStorage:
+            # Pipeline barriers
+            load_mma_Q_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_mma_Q_stage * 2]
+            load_mma_K_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_mma_K_stage * 2]
+            load_mma_V_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_mma_V_stage * 2]
+            load_mma_QT_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_mma_QT_stage * 2]
+            load_mma_dO_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_mma_dO_stage * 2]
+            load_mma_dOT_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_mma_dO_stage * 2]
+            load_compute_lse_mbar_ptr: cute.struct.MemRange[
+                cutlass.Int64, self.load_compute_LSE_stage * 2
+            ]
+            load_compute_sum_OdO_mbar_ptr: cute.struct.MemRange[
+                cutlass.Int64, self.load_compute_sum_OdO_stage * 2
+            ]
+            mma_compute_S_mbar_ptr: cute.struct.MemRange[
+                cutlass.Int64, self.mma_compute_S_stage * 2
+            ]
+            mma_compute_dP_mbar_ptr: cute.struct.MemRange[
+                cutlass.Int64, self.mma_compute_dP_stage * 2
+            ]
+            compute_mma_P_mbar_ptr: cute.struct.MemRange[
+                cutlass.Int64, self.compute_mma_P_stage * 2
+            ]
+            compute_mma_dS_mbar_ptr: cute.struct.MemRange[
+                cutlass.Int64, self.compute_mma_dS_stage * 2
+            ]
+            mma_compute_dKdV_mbar_ptr: cute.struct.MemRange[
+                cutlass.Int64, self.mma_compute_dKdV_stage * 2
+            ]
+            tmem_holding_buf: cutlass.Int32
+            tmem_dealloc_mbar: cutlass.Int64
+            clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2]
+            clc_response: cute.struct.MemRange[Int32, 4]
+            # Smem tensors
+            sK: cute.struct.Align[
+                cute.struct.MemRange[K.element_type, cute.cosize(K_smem_layout_staged)],
+                self.buffer_align_bytes,
+            ]
+            # only used in 2cta
+            sV: cute.struct.Align[
+                cute.struct.MemRange[V.element_type, cute.cosize(V_smem_layout_staged)],
+                self.buffer_align_bytes,
+            ]
+            sQ: cute.struct.Align[
+                cute.struct.MemRange[Q.element_type, cute.cosize(Q_smem_layout_staged)],
+                self.buffer_align_bytes,
+            ]
+            sQT: cute.struct.Align[
+                cute.struct.MemRange[Q.element_type, cute.cosize(QT_smem_layout_staged)],
+                self.buffer_align_bytes,
+            ]
+            sdO: cute.struct.Align[
+                cute.struct.MemRange[dO.element_type, cute.cosize(dO_smem_layout_staged)],
+                self.buffer_align_bytes,
+            ]
+            sdOT: cute.struct.Align[
+                cute.struct.MemRange[dO.element_type, cute.cosize(dOT_smem_layout_staged)],
+                self.buffer_align_bytes,
+            ]
+            # only used in 2cta
+            # dishengbin checked whether we need sP
+            sP: cute.struct.Align[
+                cute.struct.MemRange[Q.element_type, cute.cosize(P_smem_layout_staged)],
+                self.buffer_align_bytes,
+            ]
+            sdST: cute.struct.Align[
+                cute.struct.MemRange[Q.element_type, cute.cosize(dST_smem_layout_staged)],
+                self.buffer_align_bytes,
+            ]
+
+            sLSE: cute.struct.Align[
+                cute.struct.MemRange[self.acc_dtype, cute.cosize(LSE_smem_layout)],
+                self.buffer_align_bytes,
+            ]
+            sSum_OdO: cute.struct.Align[
+                cute.struct.MemRange[self.acc_dtype, cute.cosize(sum_OdO_smem_layout)],
+                self.buffer_align_bytes,
+            ]
+
+        self.shared_storage = SharedStorage
+
+        # =============================== bwd ===============================
+        K_val = problem_shape[1]
+        _, H_K = problem_shape[3][0]
+        B = problem_shape[3][1]
+        problem_shape_mbh = (
+            cute.ceil_div(K_val, self.cta_tiler[1]),
+            cute.size(B),
+            cute.size(H_K),
+        )
+        if cutlass.const_expr(self.use_clc_scheduler):
+            self.tile_sched_params = FmhaClcDynamicTileSchedulerParams(
+                problem_shape_mbh,
+                (*self.cluster_shape_mn, 1),
+            )
+            bwd_grid = FmhaClcDynamicTileScheduler.get_grid_shape(self.tile_sched_params)
+        else:
+            self.tile_sched_params = FmhaStaticTileSchedulerParams(
+                is_persistent=False,
+                problem_shape_mbh=problem_shape_mbh,
+            )
+            bwd_grid = self._compute_bwd_grid(problem_shape, self.cta_tiler[1])
+            bwd_grid = cute.round_up(bwd_grid, self.cluster_shape_mnk)
+
+        self.dkdv_bwd(
+            KQ_tiled_mma,
+            VdO_tiled_mma,
+            PdO_tiled_mma,
+            dSQ_tiled_mma,
+            tma_atom_K,
+            tma_tensor_K,
+            K,
+            tma_atom_V,
+            tma_tensor_V,
+            tma_atom_Q,
+            tma_tensor_Q,
+            Q,
+            tma_atom_QT,
+            tma_tensor_QT,
+            tma_atom_dO,
+            tma_tensor_dO,
+            tma_atom_dOT,
+            tma_tensor_dOT,
+            dK,
+            dV,
+            tma_atom_dK,
+            tma_tensor_dK,
+            tma_atom_dV,
+            tma_tensor_dV,
+            scaled_LSE,
+            scale_softmax,
+            sum_OdO,
+            problem_shape,
+            cumulative_s_q,
+            cumulative_s_k,
+            self.cluster_layout_vmnk,
+            K_smem_layout_staged,
+            Q_smem_layout_staged,
+            V_smem_layout_staged,
+            dO_smem_layout_staged,
+            dS_smem_layout_staged,
+            dST_smem_layout_staged,
+            QT_smem_layout_staged,
+            dOT_smem_layout_staged,
+            P_smem_layout_staged,
+            LSE_smem_layout,
+            sum_OdO_smem_layout,
+            sdK_epi_layout,
+            sdV_epi_layout,
+            self.tile_sched_params,
+        ).launch(
+            grid=bwd_grid,
+            block=[self.threads_per_cta, 1, 1],
+            cluster=self.cluster_shape_mnk,
+            smem=self.shared_storage.size_in_bytes(),  # type: ignore [attr-defined]
+            stream=stream,
+            min_blocks_per_mp=1,
+        )
+
+    @cute.kernel
+    def dkdv_bwd(
+        self,
+        KQ_tiled_mma: cute.TiledMma,
+        VdO_tiled_mma: cute.TiledMma,
+        PdO_tiled_mma: cute.TiledMma,
+        dSQ_tiled_mma: cute.TiledMma,
+        tma_atom_K: cute.CopyAtom,
+        K_in: cute.Tensor,
+        K_ref: cute.Tensor,
+        tma_atom_V: cute.CopyAtom,
+        V_in: cute.Tensor,
+        tma_atom_Q: cute.CopyAtom,
+        Q_in: cute.Tensor,
+        Q_ref: cute.Tensor,
+        tma_atom_QT: cute.CopyAtom,
+        QT_in: cute.Tensor,
+        tma_atom_dO: cute.CopyAtom,
+        dO_in: cute.Tensor,
+        tma_atom_dOT: cute.CopyAtom,
+        dOT_in: cute.Tensor,
+        dK: cute.Tensor,
+        dV: cute.Tensor,
+        tma_atom_dK: cute.CopyAtom,
+        dK_tma: cute.Tensor,
+        tma_atom_dV: cute.CopyAtom,
+        dV_tma: cute.Tensor,
+        LSE: cute.Tensor,
+        scale_softmax: cutlass.Float32,
+        sum_OdO: cute.Tensor,
+        problem_shape: tuple[Int32, Int32, Int32, tuple[tuple[Int32, Int32], Int32]],
+        cumulative_s_q: cute.Tensor | None,
+        cumulative_s_k: cute.Tensor | None,
+        cluster_layout_vmnk: cute.Layout,
+        K_smem_layout_staged: cute.ComposedLayout,
+        Q_smem_layout_staged: cute.ComposedLayout,
+        V_smem_layout_staged: cute.ComposedLayout,
+        dO_smem_layout_staged: cute.ComposedLayout,
+        dS_smem_layout_staged: cute.ComposedLayout,
+        dST_smem_layout_staged: cute.ComposedLayout,
+        QT_smem_layout_staged: cute.ComposedLayout,
+        dOT_smem_layout_staged: cute.ComposedLayout,
+        P_smem_layout_staged: cute.ComposedLayout,
+        LSE_smem_layout: cute.Layout,
+        sum_OdO_smem_layout: cute.Layout,
+        sdK_epi_layout: cute.ComposedLayout,
+        sdV_epi_layout: cute.ComposedLayout,
+        tile_sched_params: FmhaStaticTileSchedulerParams | FmhaClcDynamicTileSchedulerParams,
+    ):
+        """Core CuTeDSL backward kernel."""
+        bidx, bidy, bidz = cute.arch.block_idx()
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+        varlen = cumulative_s_q is not None or cumulative_s_k is not None
+
+        mma_tile_coord_v = bidx % cute.size(KQ_tiled_mma.thr_id.shape)
+
+        if warp_idx == self.load_warp_id:
+            cpasync.prefetch_descriptor(tma_atom_K)
+            cpasync.prefetch_descriptor(tma_atom_Q)
+            cpasync.prefetch_descriptor(tma_atom_QT)
+            cpasync.prefetch_descriptor(tma_atom_V)
+            cpasync.prefetch_descriptor(tma_atom_dO)
+            cpasync.prefetch_descriptor(tma_atom_dOT)
+
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(self.shared_storage)
+
+        load_mma_Q_producer, load_mma_Q_consumer = pipeline.PipelineTmaUmma.create(
+            num_stages=self.load_mma_Q_stage,
+            producer_group=make_thread_cooperative_group(len([self.load_warp_id])),
+            consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            tx_count=self.tma_copy_Q_bytes,
+            barrier_storage=storage.load_mma_Q_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        load_mma_K_producer, load_mma_K_consumer = pipeline.PipelineTmaUmma.create(
+            num_stages=self.load_mma_K_stage,
+            producer_group=make_thread_cooperative_group(len([self.load_warp_id])),
+            consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            tx_count=self.tma_copy_K_bytes,
+            barrier_storage=storage.load_mma_K_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        load_mma_V_producer, load_mma_V_consumer = pipeline.PipelineTmaUmma.create(
+            num_stages=self.load_mma_V_stage,
+            producer_group=make_thread_cooperative_group(len([self.load_warp_id])),
+            consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            tx_count=self.tma_copy_V_bytes,
+            barrier_storage=storage.load_mma_V_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        load_mma_QT_producer, load_mma_QT_consumer = pipeline.PipelineTmaUmma.create(
+            num_stages=self.load_mma_QT_stage,
+            producer_group=make_thread_cooperative_group(len([self.load_warp_id])),
+            consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            tx_count=self.tma_copy_Q_bytes,
+            barrier_storage=storage.load_mma_QT_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        load_mma_dO_producer, load_mma_dO_consumer = pipeline.PipelineTmaUmma.create(
+            num_stages=self.load_mma_dO_stage,
+            producer_group=make_thread_cooperative_group(len([self.load_warp_id])),
+            consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            tx_count=self.tma_copy_dO_bytes,
+            barrier_storage=storage.load_mma_dO_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        load_mma_dOT_producer, load_mma_dOT_consumer = pipeline.PipelineTmaUmma.create(
+            num_stages=self.load_mma_dO_stage,
+            producer_group=make_thread_cooperative_group(len([self.load_warp_id])),
+            consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            tx_count=self.tma_copy_dO_bytes,
+            barrier_storage=storage.load_mma_dOT_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        load_compute_LSE_producer, load_compute_LSE_consumer = pipeline.PipelineCpAsync.create(
+            num_stages=self.load_compute_LSE_stage,
+            producer_group=make_thread_cooperative_group(self.threads_per_warp),
+            consumer_group=make_thread_cooperative_group(
+                self.threads_per_warp * self.num_compute_warps
+            ),
+            barrier_storage=storage.load_compute_lse_mbar_ptr.data_ptr(),
+        ).make_participants()
+        load_compute_sum_OdO_producer, load_compute_sum_OdO_consumer = (
+            pipeline.PipelineCpAsync.create(
+                num_stages=self.load_compute_sum_OdO_stage,
+                producer_group=make_thread_cooperative_group(self.threads_per_warp),
+                consumer_group=make_thread_cooperative_group(
+                    self.threads_per_warp * self.num_compute_warps
+                ),
+                barrier_storage=storage.load_compute_sum_OdO_mbar_ptr.data_ptr(),
+            ).make_participants()
+        )
+        mma_compute_S_producer, mma_compute_S_consumer = pipeline.PipelineUmmaAsync.create(
+            num_stages=self.mma_compute_S_stage,
+            producer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            consumer_group=make_thread_cooperative_group(
+                self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0]
+            ),
+            barrier_storage=storage.mma_compute_S_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        mma_compute_dP_producer, mma_compute_dP_consumer = pipeline.PipelineUmmaAsync.create(
+            num_stages=self.mma_compute_dP_stage,
+            producer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            consumer_group=make_thread_cooperative_group(
+                self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0]
+            ),
+            barrier_storage=storage.mma_compute_dP_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        compute_mma_P_producer, compute_mma_P_consumer = pipeline.PipelineAsyncUmma.create(
+            num_stages=self.compute_mma_P_stage,
+            producer_group=make_thread_cooperative_group(
+                self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0]
+            ),
+            consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            barrier_storage=storage.compute_mma_P_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        compute_mma_dS_producer, compute_mma_dS_consumer = pipeline.PipelineAsyncUmma.create(
+            num_stages=self.compute_mma_dS_stage,
+            producer_group=make_thread_cooperative_group(
+                self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0]
+            ),
+            consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            barrier_storage=storage.compute_mma_dS_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        mma_compute_dKdV_producer, mma_compute_dKdV_consumer = pipeline.PipelineUmmaAsync.create(
+            num_stages=self.mma_compute_dKdV_stage,
+            producer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            consumer_group=make_thread_cooperative_group(
+                self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0]
+            ),
+            barrier_storage=storage.mma_compute_dKdV_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+
+        cute.arch.barrier(barrier_id=self.cta_sync_bar_id, number_of_threads=self.threads_per_cta)
+
+        # setup mma
+        sQ = storage.sQ.get_tensor(Q_smem_layout_staged.outer, swizzle=Q_smem_layout_staged.inner)
+        sK = storage.sK.get_tensor(K_smem_layout_staged.outer, swizzle=K_smem_layout_staged.inner)
+        sV = storage.sV.get_tensor(V_smem_layout_staged.outer, swizzle=V_smem_layout_staged.inner)
+        sdO = storage.sdO.get_tensor(
+            dO_smem_layout_staged.outer, swizzle=dO_smem_layout_staged.inner
+        )
+        sLSE = storage.sLSE.get_tensor(LSE_smem_layout)
+        sSum_OdO = storage.sSum_OdO.get_tensor(sum_OdO_smem_layout)
+        tmem_holding_buf = storage.tmem_holding_buf
+        # for 2cta, QT use different mem from Q
+
+        sQT = storage.sQT.get_tensor(
+            QT_smem_layout_staged.outer, swizzle=QT_smem_layout_staged.inner
+        )
+        sdST = storage.sdST.get_tensor(
+            dST_smem_layout_staged.outer, swizzle=dST_smem_layout_staged.inner
+        )
+        tP_fake_ptr = cute.make_ptr(sQ.element_type, 0, cute.AddressSpace.tmem)
+        tP = cute.make_tensor(tP_fake_ptr, P_smem_layout_staged.outer)
+
+        sP = storage.sP.get_tensor(P_smem_layout_staged.outer, swizzle=P_smem_layout_staged.inner)
+
+        sdOT = storage.sdOT.get_tensor(
+            dOT_smem_layout_staged.outer, swizzle=dOT_smem_layout_staged.inner
+        )
+
+        # tSTrK shape : (MMA, MMA_M, MMA_K, STAGE)
+        tSTrK = KQ_tiled_mma.make_fragment_A(sK)
+        # tSTrQ shape : (MMA, MMA_N, MMA_K, STAGE)
+        tSTrQ = KQ_tiled_mma.make_fragment_B(sQ)
+
+        # tdPTrV shape : (MMA, MMA_M, MMA_K, STAGE)
+        tdPTrV = VdO_tiled_mma.make_fragment_A(sV)
+        # tdPTrdO shape : (MMA, MMA_N, MMA_K, STAGE)
+        tdPTrdO = VdO_tiled_mma.make_fragment_B(sdO)
+
+        # tdKrdST shape: (MMA, MMA_M, MMA_K, STAGE)
+        tdKrdST = dSQ_tiled_mma.make_fragment_A(sdST)
+        # tdKrQT shape : (MMA, MMA_N, MMA_K, STAGE)
+        tdKrQT = dSQ_tiled_mma.make_fragment_B(sQT)
+
+        tmem_alloc_barrier = pipeline.NamedBarrier(
+            barrier_id=self.tmem_alloc_sync_bar_id,
+            num_threads=self.threads_per_cta,
+        )
+
+        tmem = utils.TmemAllocator(
+            storage.tmem_holding_buf.ptr,
+            barrier_for_retrieve=tmem_alloc_barrier,
+            allocator_warp_id=self.load_warp_id,
+            is_two_cta=True,
+            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
+        )
+
+        tmem.allocate(self.tmem_alloc_cols)
+
+        # wait for tmem allocation and retrieve the pointer
+        tmem.wait_for_alloc()
+        tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+
+        # Cluster arrive after barrier init
+        # is_relaxed=False has memory consistency guarantee
+        pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=False)
+
+        if cutlass.const_expr(self.use_clc_scheduler):
+            clc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+            cluster_size = cute.size(self.cluster_shape_mnk)
+            num_clc_consumer_threads = self.threads_per_warp * (
+                1  # sched_warp (CTA 0 only)
+                + cluster_size
+                * (
+                    len(self.compute_warp_id)
+                    + 1  # mma_warp
+                    + 1  # load_warp
+                )
+            )
+            clc_pipeline_consumer_group = pipeline.CooperativeGroup(
+                pipeline.Agent.Thread, num_clc_consumer_threads
+            )
+            clc_response_ptr = storage.clc_response.data_ptr()
+            clc = SchedulerState.create_clc(
+                hw_scheduler=ClcDynamicPersistentTileScheduler.create(
+                    self.tile_sched_params.clc_hw_params(),
+                    cute.arch.block_idx(),
+                    cute.arch.grid_dim(),
+                    clc_response_ptr,
+                ),
+                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,
+                ),
+                consumer_state=pipeline.make_pipeline_state(
+                    pipeline.PipelineUserType.Consumer, self.num_clc_stage
+                ),
+                producer_state=pipeline.make_pipeline_state(
+                    pipeline.PipelineUserType.Producer, self.num_clc_stage
+                ),
+            )
+            tile_sched = FmhaClcDynamicTileScheduler.create(
+                tile_sched_params,
+                cute.arch.block_idx(),
+                cute.arch.grid_dim(),
+                clc_response_ptr,
+                clc,
+            )
+            work_tile = tile_sched.initial_work_tile_info()
+        else:
+            clc = None
+            clc_response_ptr = None
+
+        tSTtST_shape = KQ_tiled_mma.partition_shape_C(cute.select(self.KQ_mma_tiler, mode=[0, 1]))
+        tSTtST = KQ_tiled_mma.make_fragment_C(tSTtST_shape)
+        # tSTtST shape : (MMA, MMA_M, MMA_N)
+        tSTtST = cute.make_tensor(tmem_ptr + self.tmem_S_offset, tSTtST.layout)
+
+        # tdVrP shape : (MMA, MMA_M, MMA_K, STAGE)
+        tdVrP = PdO_tiled_mma.make_fragment_A(sP)
+        # tdVrdOT shape : (MMA, MMA_N, MMA_K, STAGE)
+        tdVrdOT = PdO_tiled_mma.make_fragment_B(sdOT)
+
+        tdPTtdPT_shape = VdO_tiled_mma.partition_shape_C(
+            cute.select(self.VdO_mma_tiler, mode=[0, 1])
+        )
+        tdPTtdPT = VdO_tiled_mma.make_fragment_C(tdPTtdPT_shape)
+        # tdPTtdPT shape : (MMA, MMA_M, MMA_N)
+        tdPTtdPT = cute.make_tensor(tmem_ptr + self.tmem_dP_offset, tdPTtdPT.layout)
+
+        tdKtdK_shape = dSQ_tiled_mma.partition_shape_C(cute.select(self.dSQ_mma_tiler, mode=[0, 1]))
+        tdKtdK = dSQ_tiled_mma.make_fragment_C(tdKtdK_shape)
+        # tdKtdK shape : (MMA, MMA_M, MMA_N)
+        tdKtdK = cute.make_tensor(tmem_ptr + self.tmem_dK_offset, tdKtdK.layout)
+
+        tdVtdV_shape = PdO_tiled_mma.partition_shape_C(cute.select(self.PdO_mma_tiler, mode=[0, 1]))
+        tdVtdV = PdO_tiled_mma.make_fragment_C(tdVtdV_shape)
+        # tdVtdV shape : (MMA, MMA_M, MMA_N)
+        tdVtdV = cute.make_tensor(tmem_ptr + self.tmem_dV_offset, tdVtdV.layout)
+
+        # get the current batch problem shape
+
+        if cutlass.const_expr(self.use_clc_scheduler):
+            # ===== CLC PERSISTENT PATH: per-warp while loops =====
+            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
+            is_sched_warp = warp_idx == self.sched_warp_id and is_first_cta_in_cluster
+
+            # Register allocation ONCE (before any loop)
+            if warp_idx == self.load_warp_id:
+                cute.arch.warpgroup_reg_dealloc(self.num_regs_load)
+            elif warp_idx == self.mma_warp_id:
+                cute.arch.warpgroup_reg_alloc(self.num_regs_mma)
+            elif warp_idx >= self.compute_warp_id[0] and warp_idx <= self.compute_warp_id[-1]:
+                cute.arch.warpgroup_reg_alloc(self.num_regs_compute)
+            else:
+                cute.arch.warpgroup_reg_dealloc(self.num_regs_empty)
+
+            # Cluster wait
+            pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk)
+
+            # ===== SCHEDULER WARP =====
+            if is_sched_warp:
+                while work_tile.is_valid_tile:
+                    tile_sched.prefetch_next_work()
+                    work_tile = tile_sched.advance_to_next_work()
+                tile_sched.producer_tail()
+
+            # ===== LOAD WARP =====
+            elif warp_idx == self.load_warp_id:
+                while work_tile.is_valid_tile:
+                    # Decode coordinates from CLC work tile
+                    blk_coord_k = work_tile.tile_idx[0]
+                    blk_coord_b = work_tile.tile_idx[2][0]
+                    blk_coord_h_k = work_tile.tile_idx[2][1]
+                    blk_coord = (
+                        Int32(0),
+                        blk_coord_k,
+                        Int32(0),
+                        ((Int32(0), blk_coord_h_k), blk_coord_b),
+                    )
+                    seqlen_q_cur_batch = Q_ref.shape[0]
+                    seqlen_k_cur_batch = K_ref.shape[0]
+                    blk_offset = (Int32(0), Int32(0), Int32(0), ((Int32(0), Int32(0)), Int32(0)))
+                    if cutlass.const_expr(varlen):
+                        assert isinstance(cumulative_s_q, cute.Tensor)
+                        assert isinstance(cumulative_s_k, cute.Tensor)
+                        seqlen_q_cur_batch = (
+                            cumulative_s_q[blk_coord_b + 1] - cumulative_s_q[blk_coord_b]
+                        )
+                        seqlen_k_cur_batch = (
+                            cumulative_s_k[blk_coord_b + 1] - cumulative_s_k[blk_coord_b]
+                        )
+                        blk_offset = (
+                            cumulative_s_q[blk_coord_b],
+                            cumulative_s_k[blk_coord_b],
+                            Int32(0),
+                            ((Int32(0), Int32(0)), Int32(0)),
+                        )
+                    iter_start, iter_end = self.get_Q_block_min_max(
+                        seqlen_q_cur_batch,
+                        seqlen_k_cur_batch,
+                        blk_coord_k,
+                        is_2cta=True,
+                    )
+                    iter_count = (iter_end - iter_start) * problem_shape[3][0][0]
+                    if iter_count <= 0:
+                        if blk_coord_k * self.tile_shape_K < seqlen_k_cur_batch:
+                            problem_shape_cur_batch = (
+                                seqlen_q_cur_batch,
+                                seqlen_k_cur_batch,
+                                problem_shape[2],
+                                problem_shape[3],
+                            )
+                            self.epilogue_clear(
+                                blk_coord, blk_offset, problem_shape_cur_batch, dK, dV
+                            )
+                    else:
+                        problem_shape_cur_batch = (
+                            seqlen_q_cur_batch,
+                            seqlen_k_cur_batch,
+                            problem_shape[2],
+                            problem_shape[3],
+                        )
+                        (
+                            load_mma_Q_producer,
+                            load_mma_K_producer,
+                            load_mma_V_producer,
+                            load_compute_LSE_producer,
+                            load_mma_dO_producer,
+                            load_mma_dOT_producer,
+                            load_compute_sum_OdO_producer,
+                            load_mma_QT_producer,
+                        ) = self.load(
+                            K_in,
+                            V_in,
+                            Q_in,
+                            QT_in,
+                            dO_in,
+                            dOT_in,
+                            LSE,
+                            sum_OdO,
+                            sK,
+                            sQ,
+                            sQT,
+                            sV,
+                            sdO,
+                            sdOT,
+                            sLSE,
+                            sSum_OdO,
+                            KQ_tiled_mma,
+                            VdO_tiled_mma,
+                            PdO_tiled_mma,
+                            dSQ_tiled_mma,
+                            tma_atom_K,
+                            tma_atom_Q,
+                            tma_atom_QT,
+                            tma_atom_V,
+                            tma_atom_dO,
+                            tma_atom_dOT,
+                            blk_offset,
+                            problem_shape_cur_batch,
+                            varlen,
+                            iter_count,
+                            iter_start,
+                            iter_end,
+                            load_mma_Q_producer,
+                            load_mma_Q_consumer,
+                            load_mma_K_producer,
+                            load_mma_K_consumer,
+                            load_mma_V_producer,
+                            load_mma_V_consumer,
+                            load_compute_LSE_producer,
+                            load_compute_LSE_consumer,
+                            load_mma_dO_producer,
+                            load_mma_dO_consumer,
+                            load_mma_dOT_producer,
+                            load_mma_dOT_consumer,
+                            load_compute_sum_OdO_producer,
+                            load_compute_sum_OdO_consumer,
+                            load_mma_QT_producer,
+                            load_mma_QT_consumer,
+                            blk_coord_k,
+                            blk_coord_h_k,
+                            blk_coord_b,
+                        )
+                    # CLC advance
+                    work_tile = tile_sched.advance_to_next_work()
+                # producer_tail after loop
+                load_mma_K_producer.tail()
+                load_mma_V_producer.tail()
+                load_mma_Q_producer.tail()
+                load_compute_LSE_producer.tail()
+                load_mma_dO_producer.tail()
+                load_mma_dOT_producer.tail()
+                load_compute_sum_OdO_producer.tail()
+                load_mma_QT_producer.tail()
+
+            # ===== MMA WARP =====
+            elif warp_idx == self.mma_warp_id:
+                while work_tile.is_valid_tile:
+                    blk_coord_k = work_tile.tile_idx[0]
+                    blk_coord_b = work_tile.tile_idx[2][0]
+                    blk_coord_h_k = work_tile.tile_idx[2][1]
+                    blk_coord = (
+                        Int32(0),
+                        blk_coord_k,
+                        Int32(0),
+                        ((Int32(0), blk_coord_h_k), blk_coord_b),
+                    )
+                    seqlen_q_cur_batch = Q_ref.shape[0]
+                    seqlen_k_cur_batch = K_ref.shape[0]
+                    blk_offset = (Int32(0), Int32(0), Int32(0), ((Int32(0), Int32(0)), Int32(0)))
+                    if cutlass.const_expr(varlen):
+                        assert isinstance(cumulative_s_q, cute.Tensor)
+                        assert isinstance(cumulative_s_k, cute.Tensor)
+                        seqlen_q_cur_batch = (
+                            cumulative_s_q[blk_coord_b + 1] - cumulative_s_q[blk_coord_b]
+                        )
+                        seqlen_k_cur_batch = (
+                            cumulative_s_k[blk_coord_b + 1] - cumulative_s_k[blk_coord_b]
+                        )
+                        blk_offset = (
+                            cumulative_s_q[blk_coord_b],
+                            cumulative_s_k[blk_coord_b],
+                            Int32(0),
+                            ((Int32(0), Int32(0)), Int32(0)),
+                        )
+                    iter_start, iter_end = self.get_Q_block_min_max(
+                        seqlen_q_cur_batch,
+                        seqlen_k_cur_batch,
+                        blk_coord_k,
+                        is_2cta=True,
+                    )
+                    iter_count = (iter_end - iter_start) * problem_shape[3][0][0]
+                    if iter_count <= 0:
+                        if blk_coord_k * self.tile_shape_K < seqlen_k_cur_batch:
+                            problem_shape_cur_batch = (
+                                seqlen_q_cur_batch,
+                                seqlen_k_cur_batch,
+                                problem_shape[2],
+                                problem_shape[3],
+                            )
+                            self.epilogue_clear(
+                                blk_coord, blk_offset, problem_shape_cur_batch, dK, dV
+                            )
+                    else:
+                        (
+                            mma_compute_S_producer,
+                            mma_compute_dP_producer,
+                            mma_compute_dKdV_producer,
+                            load_mma_Q_consumer,
+                            load_mma_K_consumer,
+                            load_mma_V_consumer,
+                            load_mma_dO_consumer,
+                            load_mma_dOT_consumer,
+                            compute_mma_P_consumer,
+                            compute_mma_dS_consumer,
+                            load_mma_QT_consumer,
+                        ) = self.mma_2cta(
+                            KQ_tiled_mma,
+                            VdO_tiled_mma,
+                            PdO_tiled_mma,
+                            dSQ_tiled_mma,
+                            tSTtST,
+                            tSTrQ,
+                            tSTrK,
+                            tdPTtdPT,
+                            tdPTrV,
+                            tdPTrdO,
+                            tdVtdV,
+                            tdVrP,
+                            tdVrdOT,
+                            tdKrdST,
+                            tdKtdK,
+                            tdKrQT,
+                            iter_count,
+                            load_mma_Q_consumer,
+                            load_mma_K_consumer,
+                            load_mma_V_consumer,
+                            mma_compute_S_producer,
+                            load_mma_dO_consumer,
+                            mma_compute_dP_producer,
+                            load_mma_dOT_consumer,
+                            compute_mma_P_consumer,
+                            compute_mma_dS_consumer,
+                            load_mma_QT_consumer,
+                            mma_compute_dKdV_producer,
+                        )
+                    # CLC advance
+                    work_tile = tile_sched.advance_to_next_work()
+                # producer_tail after loop
+                mma_compute_S_producer.tail()
+                mma_compute_dP_producer.tail()
+                mma_compute_dKdV_producer.tail()
+
+            # ===== COMPUTE WARPS =====
+            elif warp_idx >= self.compute_warp_id[0] and warp_idx <= self.compute_warp_id[-1]:
+                while work_tile.is_valid_tile:
+                    blk_coord_k = work_tile.tile_idx[0]
+                    blk_coord_b = work_tile.tile_idx[2][0]
+                    blk_coord_h_k = work_tile.tile_idx[2][1]
+                    blk_coord = (
+                        Int32(0),
+                        blk_coord_k,
+                        Int32(0),
+                        ((Int32(0), blk_coord_h_k), blk_coord_b),
+                    )
+                    seqlen_q_cur_batch = Q_ref.shape[0]
+                    seqlen_k_cur_batch = K_ref.shape[0]
+                    blk_offset = (Int32(0), Int32(0), Int32(0), ((Int32(0), Int32(0)), Int32(0)))
+                    if cutlass.const_expr(varlen):
+                        assert isinstance(cumulative_s_q, cute.Tensor)
+                        assert isinstance(cumulative_s_k, cute.Tensor)
+                        seqlen_q_cur_batch = (
+                            cumulative_s_q[blk_coord_b + 1] - cumulative_s_q[blk_coord_b]
+                        )
+                        seqlen_k_cur_batch = (
+                            cumulative_s_k[blk_coord_b + 1] - cumulative_s_k[blk_coord_b]
+                        )
+                        blk_offset = (
+                            cumulative_s_q[blk_coord_b],
+                            cumulative_s_k[blk_coord_b],
+                            Int32(0),
+                            ((Int32(0), Int32(0)), Int32(0)),
+                        )
+                    iter_start, iter_end = self.get_Q_block_min_max(
+                        seqlen_q_cur_batch,
+                        seqlen_k_cur_batch,
+                        blk_coord_k,
+                        is_2cta=True,
+                    )
+                    iter_count = (iter_end - iter_start) * problem_shape[3][0][0]
+                    if iter_count <= 0:
+                        if blk_coord_k * self.tile_shape_K < seqlen_k_cur_batch:
+                            problem_shape_cur_batch = (
+                                seqlen_q_cur_batch,
+                                seqlen_k_cur_batch,
+                                problem_shape[2],
+                                problem_shape[3],
+                            )
+                            self.epilogue_clear(
+                                blk_coord, blk_offset, problem_shape_cur_batch, dK, dV
+                            )
+                    else:
+                        problem_shape_cur_batch = (
+                            seqlen_q_cur_batch,
+                            seqlen_k_cur_batch,
+                            problem_shape[2],
+                            problem_shape[3],
+                        )
+                        (
+                            compute_mma_P_producer,
+                            compute_mma_dS_producer,
+                            mma_compute_S_consumer,
+                            compute_mma_P_consumer,
+                            load_compute_LSE_consumer,
+                            load_compute_sum_OdO_consumer,
+                            mma_compute_dP_consumer,
+                            compute_mma_dS_consumer,
+                            mma_compute_dKdV_consumer,
+                        ) = self.compute(
+                            tSTtST,
+                            tdPTtdPT,
+                            tdVrP,
+                            sP,
+                            sLSE,
+                            sdST,
+                            sdOT,
+                            sSum_OdO,
+                            dK,
+                            dV,
+                            tdKtdK,
+                            tdVtdV,
+                            PdO_tiled_mma,
+                            dSQ_tiled_mma,
+                            blk_coord,
+                            blk_offset,
+                            problem_shape_cur_batch,
+                            iter_count,
+                            iter_start,
+                            iter_end,
+                            scale_softmax,
+                            mma_compute_S_producer,
+                            mma_compute_S_consumer,
+                            compute_mma_P_producer,
+                            compute_mma_P_consumer,
+                            load_compute_LSE_producer,
+                            load_compute_LSE_consumer,
+                            load_compute_sum_OdO_producer,
+                            load_compute_sum_OdO_consumer,
+                            mma_compute_dP_producer,
+                            mma_compute_dP_consumer,
+                            compute_mma_dS_producer,
+                            compute_mma_dS_consumer,
+                            mma_compute_dKdV_producer,
+                            mma_compute_dKdV_consumer,
+                            varlen,
+                            sK,
+                            seqlen_k_cur_batch,
+                            tma_atom_dK,
+                            dK_tma,
+                            tma_atom_dV,
+                            dV_tma,
+                            sdK_epi_layout,
+                            sdV_epi_layout,
+                        )
+                        cute.arch.barrier(
+                            barrier_id=self.epilogue_sync_bar_id,
+                            number_of_threads=self.num_compute_warps * self.threads_per_warp,
+                        )
+                    # CLC advance
+                    work_tile = tile_sched.advance_to_next_work()
+                # producer_tail after loop
+                compute_mma_P_producer.tail()
+                compute_mma_dS_producer.tail()
+
+        else:
+            # ===== STATIC PATH: original non-persistent code =====
+            blk_coord = (Int32(0), bidx, Int32(0), ((Int32(0), bidy), bidz))
+            seqlen_q_cur_batch = Q_ref.shape[0]
+            seqlen_k_cur_batch = K_ref.shape[0]
+            blk_offset = (Int32(0), Int32(0), Int32(0), ((Int32(0), Int32(0)), Int32(0)))
+            if cutlass.const_expr(varlen):
+                assert isinstance(cumulative_s_q, cute.Tensor)
+                assert isinstance(cumulative_s_k, cute.Tensor)
+                seqlen_q_cur_batch = cumulative_s_q[bidz + 1] - cumulative_s_q[bidz]
+                seqlen_k_cur_batch = cumulative_s_k[bidz + 1] - cumulative_s_k[bidz]
+                blk_offset = (
+                    cumulative_s_q[bidz],
+                    cumulative_s_k[bidz],
+                    Int32(0),
+                    ((Int32(0), Int32(0)), Int32(0)),
+                )
+
+            iter_start, iter_end = self.get_Q_block_min_max(
+                seqlen_q_cur_batch,
+                seqlen_k_cur_batch,
+                blk_coord[1],
+                is_2cta=True,
+            )
+
+            # Cluster wait
+            pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk)
+
+            iter_count = (iter_end - iter_start) * problem_shape[3][0][0]
+            problem_shape_cur_batch = (
+                seqlen_q_cur_batch,
+                seqlen_k_cur_batch,
+                problem_shape[2],
+                problem_shape[3],
+            )
+            if iter_count <= 0:
+                if bidx * self.tile_shape_K < seqlen_k_cur_batch:
+                    self.epilogue_clear(
+                        blk_coord,
+                        blk_offset,
+                        problem_shape_cur_batch,
+                        dK,
+                        dV,
+                    )
+            # ///////////////////////////////////////////////////////////////////////////////
+            #  LOAD
+            # ///////////////////////////////////////////////////////////////////////////////
+            elif warp_idx == self.load_warp_id:
+                cute.arch.warpgroup_reg_dealloc(self.num_regs_load)
+
+                self.load(
+                    K_in,
+                    V_in,
+                    Q_in,
+                    QT_in,
+                    dO_in,
+                    dOT_in,
+                    LSE,
+                    sum_OdO,
+                    sK,
+                    sQ,
+                    sQT,
+                    sV,
+                    sdO,
+                    sdOT,
+                    sLSE,
+                    sSum_OdO,
+                    KQ_tiled_mma,
+                    VdO_tiled_mma,
+                    PdO_tiled_mma,
+                    dSQ_tiled_mma,
+                    tma_atom_K,
+                    tma_atom_Q,
+                    tma_atom_QT,
+                    tma_atom_V,
+                    tma_atom_dO,
+                    tma_atom_dOT,
+                    blk_offset,
+                    problem_shape_cur_batch,
+                    varlen,
+                    iter_count,
+                    iter_start,
+                    iter_end,
+                    load_mma_Q_producer,
+                    load_mma_Q_consumer,
+                    load_mma_K_producer,
+                    load_mma_K_consumer,
+                    load_mma_V_producer,
+                    load_mma_V_consumer,
+                    load_compute_LSE_producer,
+                    load_compute_LSE_consumer,
+                    load_mma_dO_producer,
+                    load_mma_dO_consumer,
+                    load_mma_dOT_producer,
+                    load_mma_dOT_consumer,
+                    load_compute_sum_OdO_producer,
+                    load_compute_sum_OdO_consumer,
+                    load_mma_QT_producer,
+                    load_mma_QT_consumer,
+                )
+
+            # ///////////////////////////////////////////////////////////////////////////////
+            #  MMA
+            # ///////////////////////////////////////////////////////////////////////////////
+            elif warp_idx == self.mma_warp_id:
+                cute.arch.warpgroup_reg_alloc(self.num_regs_mma)
+
+                self.mma_2cta(
+                    KQ_tiled_mma,
+                    VdO_tiled_mma,
+                    PdO_tiled_mma,
+                    dSQ_tiled_mma,
+                    tSTtST,
+                    tSTrQ,
+                    tSTrK,
+                    tdPTtdPT,
+                    tdPTrV,
+                    tdPTrdO,
+                    tdVtdV,
+                    tdVrP,
+                    tdVrdOT,
+                    tdKrdST,
+                    tdKtdK,
+                    tdKrQT,
+                    iter_count,
+                    load_mma_Q_consumer,
+                    load_mma_K_consumer,
+                    load_mma_V_consumer,
+                    mma_compute_S_producer,
+                    load_mma_dO_consumer,
+                    mma_compute_dP_producer,
+                    load_mma_dOT_consumer,
+                    compute_mma_P_consumer,
+                    compute_mma_dS_consumer,
+                    load_mma_QT_consumer,
+                    mma_compute_dKdV_producer,
+                )
+
+            # ///////////////////////////////////////////////////////////////////////////////
+            #  Compute
+            # ///////////////////////////////////////////////////////////////////////////////
+            elif warp_idx >= self.compute_warp_id[0] and warp_idx <= self.compute_warp_id[-1]:
+                cute.arch.warpgroup_reg_alloc(self.num_regs_compute)
+
+                self.compute(
+                    tSTtST,
+                    tdPTtdPT,
+                    tdVrP,
+                    sP,
+                    sLSE,
+                    sdST,
+                    sdOT,
+                    sSum_OdO,
+                    dK,
+                    dV,
+                    tdKtdK,
+                    tdVtdV,
+                    PdO_tiled_mma,
+                    dSQ_tiled_mma,
+                    blk_coord,
+                    blk_offset,
+                    problem_shape_cur_batch,
+                    iter_count,
+                    iter_start,
+                    iter_end,
+                    scale_softmax,
+                    mma_compute_S_producer,
+                    mma_compute_S_consumer,
+                    compute_mma_P_producer,
+                    compute_mma_P_consumer,
+                    load_compute_LSE_producer,
+                    load_compute_LSE_consumer,
+                    load_compute_sum_OdO_producer,
+                    load_compute_sum_OdO_consumer,
+                    mma_compute_dP_producer,
+                    mma_compute_dP_consumer,
+                    compute_mma_dS_producer,
+                    compute_mma_dS_consumer,
+                    mma_compute_dKdV_producer,
+                    mma_compute_dKdV_consumer,
+                    varlen,
+                    sK,
+                    seqlen_k_cur_batch,
+                    tma_atom_dK,
+                    dK_tma,
+                    tma_atom_dV,
+                    dV_tma,
+                    sdK_epi_layout,
+                    sdV_epi_layout,
+                )
+
+                cute.arch.barrier(
+                    barrier_id=self.epilogue_sync_bar_id,
+                    number_of_threads=self.num_compute_warps * self.threads_per_warp,
+                )
+
+            else:
+                cute.arch.warpgroup_reg_dealloc(self.num_regs_empty)
+
+        cute.arch.cluster_arrive()
+        cute.arch.cluster_wait()
+        # dishengbin Deallocate tmem for early exit
+        # Dealloc the tensor memory
+        tmem.relinquish_alloc_permit()
+        tmem.free(tmem_ptr)
+
+    @cute.jit
+    def get_Q_block_min_max(
+        self,
+        seq_Q: Int32,
+        seq_K: Int32,
+        blk_coord_k: Int32,
+        is_2cta: bool,
+    ):
+        """Get Q tiles range."""
+        Q_block_max = cute.ceil_div(seq_Q, self.tile_shape_Q)
+        Q_block_min = cutlass.Int32(0)
+        if cutlass.const_expr(self.has_sliding_window):
+            # For 2cta, use the last K block in the cluster so both CTAs get the same Q_block_max
+            blk_coord_k_for_max = (blk_coord_k // 2) * 2 + 1
+            Q_block_max_tmp = cute.ceil_div(
+                (blk_coord_k_for_max + 1) * self.tile_shape_K
+                + seq_Q
+                - seq_K
+                + self.window_size_left,
+                self.tile_shape_Q,
+            )
+            Q_block_max = min(Q_block_max, Q_block_max_tmp)
+        if cutlass.const_expr(self.is_causal or self.has_sliding_window):
+            # For 2cta, use the first K block in the cluster so both CTAs get the same Q_block_min.
+            # This ensures both CTAs in a cluster run the same number of pipeline iterations,
+            # avoiding hang from mismatched producer_commit / consumer_wait counts.
+            blk_coord_k_for_min = (blk_coord_k // 2) * 2
+            Q_block_min_tmp = (
+                blk_coord_k_for_min * self.tile_shape_K + seq_Q - seq_K - self.window_size_right
+            ) // self.tile_shape_Q
+            # Consider the case of 2cta, we need to ensure the K block is aligned to 2
+            Q_block_min_tmp = Q_block_min_tmp - Q_block_min_tmp % 2
+            Q_block_min = max(Q_block_min_tmp, Q_block_min)
+        return Q_block_min, Q_block_max
+
+    @cute.jit
+    def load(
+        self,
+        K_in: cute.Tensor,
+        V_in: cute.Tensor,
+        Q_in: cute.Tensor,
+        QT_in: cute.Tensor,
+        dO_in: cute.Tensor,
+        dOT_in: cute.Tensor,
+        LSE_in: cute.Tensor,
+        sum_OdO_in: cute.Tensor,
+        sK: cute.Tensor,
+        sQ: cute.Tensor,
+        sQT: cute.Tensor,
+        sV: cute.Tensor,
+        sdO: cute.Tensor,
+        sdOT: cute.Tensor,
+        sLSE: cute.Tensor,
+        sSum_OdO: cute.Tensor,
+        KQ_tiled_mma: cute.TiledMma,
+        VdO_tiled_mma: cute.TiledMma,
+        PdO_tiled_mma: cute.TiledMma,
+        dSQ_tiled_mma: cute.TiledMma,
+        tma_atom_K: cute.CopyAtom,
+        tma_atom_Q: cute.CopyAtom,
+        tma_atom_QT: cute.CopyAtom,
+        tma_atom_V: cute.CopyAtom,
+        tma_atom_dO: cute.CopyAtom,
+        tma_atom_dOT: cute.CopyAtom,
+        blk_offset: cute.Shape,
+        problem_shape: tuple[Int32, Int32, Int32, tuple[tuple[Int32, Int32], Int32]],
+        varlen: bool,
+        iter_count: Int32,
+        iter_start: Int32,
+        iter_end: Int32,
+        load_mma_Q_producer,
+        load_mma_Q_consumer,
+        load_mma_K_producer,
+        load_mma_K_consumer,
+        load_mma_V_producer,
+        load_mma_V_consumer,
+        load_compute_LSE_producer,
+        load_compute_LSE_consumer,
+        load_mma_dO_producer,
+        load_mma_dO_consumer,
+        load_mma_dOT_producer,
+        load_mma_dOT_consumer,
+        load_compute_sum_OdO_producer,
+        load_compute_sum_OdO_consumer,
+        load_mma_QT_producer,
+        load_mma_QT_consumer,
+        blk_coord_k_override: Int32 = Int32(-1),
+        blk_coord_h_k_override: Int32 = Int32(-1),
+        blk_coord_b_override: Int32 = Int32(-1),
+    ):
+        """TMA load."""
+        tidx, _, _ = cute.arch.thread_idx()
+        if cutlass.const_expr(self.use_clc_scheduler):
+            blk_coord_k = blk_coord_k_override
+            blk_coord_h_k = blk_coord_h_k_override
+            blk_coord_b = blk_coord_b_override
+        else:
+            blk_coord_k, blk_coord_h_k, blk_coord_b = cute.arch.block_idx()
+        blk_coord_h_r = Int32(0)
+        blk_coord_h = (blk_coord_h_r, blk_coord_h_k)
+        iter_index = iter_start
+        mma_tile_coord_v = blk_coord_k % cute.size(KQ_tiled_mma.thr_id.shape)
+        mma_tile_coord_m = blk_coord_k // cute.size(KQ_tiled_mma.thr_id.shape)
+
+        K = cute.domain_offset(cute.select(blk_offset, mode=[1, 2, 3]), K_in)
+        V = cute.domain_offset(cute.select(blk_offset, mode=[1, 2, 3]), V_in)
+        Q = cute.domain_offset(cute.select(blk_offset, mode=[0, 2, 3]), Q_in)
+        QT = cute.domain_offset(cute.select(blk_offset, mode=[2, 0, 3]), QT_in)
+        dO = cute.domain_offset(cute.select(blk_offset, mode=[0, 2, 3]), dO_in)
+        dOT = cute.domain_offset(cute.select(blk_offset, mode=[2, 0, 3]), dOT_in)
+        blk_offset_stats = blk_offset
+        if cutlass.const_expr(varlen):
+            cuseqlen_q_stats = cute.assume(
+                (blk_offset[0] + blk_coord_b * self.tile_shape_Q)
+                // self.tile_shape_Q
+                * self.tile_shape_Q,
+                divby=self.tile_shape_Q,
+            )
+            blk_offset_stats = (
+                cuseqlen_q_stats,
+                blk_offset[1],
+                blk_offset[2],
+                blk_offset[3],
+            )
+        LSE = cute.domain_offset(cute.select(blk_offset_stats, mode=[0, 3]), LSE_in)
+        sum_OdO = cute.domain_offset(cute.select(blk_offset_stats, mode=[0, 3]), sum_OdO_in)
+
+        gK = cute.local_tile(K, cute.select(self.KQ_mma_tiler, mode=[0, 2]), (None, None, None))
+        gQ = cute.local_tile(Q, cute.select(self.KQ_mma_tiler, mode=[1, 2]), (None, None, None))
+        gQT = cute.local_tile(QT, cute.select(self.dSQ_mma_tiler, mode=[1, 2]), (None, None, None))
+        gV = cute.local_tile(V, cute.select(self.VdO_mma_tiler, mode=[0, 2]), (None, None, None))
+        gdO = cute.local_tile(dO, cute.select(self.VdO_mma_tiler, mode=[1, 2]), (None, None, None))
+        gdOT = cute.local_tile(
+            dOT, cute.select(self.PdO_mma_tiler, mode=[1, 2]), (None, None, None)
+        )
+
+        KQ_thr_mma = KQ_tiled_mma.get_slice(mma_tile_coord_v)
+        VdO_thr_mma = VdO_tiled_mma.get_slice(mma_tile_coord_v)
+        PdO_thr_mma = PdO_tiled_mma.get_slice(mma_tile_coord_v)
+        dSQ_thr_mma = dSQ_tiled_mma.get_slice(mma_tile_coord_v)
+
+        tSTgK = KQ_thr_mma.partition_A(gK)
+        tSTgQ = KQ_thr_mma.partition_B(gQ)
+        tdKgQT = dSQ_thr_mma.partition_B(gQT)
+        tdPTgV = VdO_thr_mma.partition_A(gV)
+        tdPTgdO = VdO_thr_mma.partition_B(gdO)
+        tdVgdOT = PdO_thr_mma.partition_B(gdOT)
+
+        cta_layout_mnk = cute.make_layout(self.cluster_shape_mnk)
+        cta_layout_vmnk = cute.tiled_divide(cta_layout_mnk, (KQ_tiled_mma.thr_id,))
+        cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cute.arch.block_idx_in_cluster())
+
+        tKsK, tKgK_mkl = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_K,
+            cta_in_cluster_coord_vmnk[2],
+            cute.make_layout(cute.size(cta_layout_vmnk, mode=[2])),
+            cute.group_modes(sK, 0, 3),
+            cute.group_modes(tSTgK, 0, 3),
+        )
+        tQsQ, tQgQ_mkl = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_Q,
+            cta_in_cluster_coord_vmnk[1],
+            cute.make_layout(cute.size(cta_layout_vmnk, mode=[1])),
+            cute.group_modes(sQ, 0, 3),
+            cute.group_modes(tSTgQ, 0, 3),
+        )
+        tQTsQT, tQTgQT_mkl = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_QT,
+            cta_in_cluster_coord_vmnk[1],
+            cute.make_layout(cute.size(cta_layout_vmnk, mode=[1])),
+            cute.group_modes(sQT, 0, 3),
+            cute.group_modes(tdKgQT, 0, 3),
+        )
+        tVsV, tVgV_mkl = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_V,
+            cta_in_cluster_coord_vmnk[2],
+            cute.make_layout(cute.size(cta_layout_vmnk, mode=[2])),
+            cute.group_modes(sV, 0, 3),
+            cute.group_modes(tdPTgV, 0, 3),
+        )
+        tdOsdO, tdOgdO_mkl = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_dO,
+            cta_in_cluster_coord_vmnk[1],
+            cute.make_layout(cute.size(cta_layout_vmnk, mode=[1])),
+            cute.group_modes(sdO, 0, 3),
+            cute.group_modes(tdPTgdO, 0, 3),
+        )
+        tdOTsdOT, tdOTgdOT_mkl = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_dOT,
+            cta_in_cluster_coord_vmnk[1],
+            cute.make_layout(cute.size(cta_layout_vmnk, mode=[1])),
+            cute.group_modes(sdOT, 0, 3),
+            cute.group_modes(tdVgdOT, 0, 3),
+        )
+
+        k_handle = load_mma_K_producer.acquire_and_advance()
+        cute.copy(
+            tma_atom_K,
+            tKgK_mkl[(None, mma_tile_coord_m, 0, (blk_coord_h, blk_coord_b))],
+            tKsK[None, 0],
+            tma_bar_ptr=k_handle.barrier,
+        )
+
+        q_handle = load_mma_Q_producer.acquire_and_advance()
+        cute.copy(
+            tma_atom_Q,
+            tQgQ_mkl[(None, iter_index, 0, (blk_coord_h, blk_coord_b))],
+            tQsQ[None, q_handle.index],
+            tma_bar_ptr=q_handle.barrier,
+        )
+
+        lse_handle = load_compute_LSE_producer.acquire_and_advance()
+        thread_idx = tidx % self.threads_per_warp
+        async_copy_num_elts = sLSE.shape[0] // self.threads_per_warp
+        atom_async_copy = cute.make_copy_atom(
+            cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.ALWAYS),
+            self.acc_dtype,
+            num_bits_per_copy=self.acc_dtype.width,
+        )
+        sLSE_for_copy = cute.flat_divide(sLSE, (1,))
+        LSE_for_copy = cute.flat_divide(LSE, (1,))
+        # Warp-coalesced: at each i, lane T accesses index `T + i*W` (stride-1
+        # across the warp) instead of `T*N + i` (stride-N across the warp).
+        for i in cutlass.range_constexpr(async_copy_num_elts):
+            LSE_idx = self.tile_shape_Q * iter_index + thread_idx + i * self.threads_per_warp
+            sLSE_idx = thread_idx + i * self.threads_per_warp
+            if cute.elem_less(LSE_idx, problem_shape[0]):
+                cute.copy(
+                    atom_async_copy,
+                    LSE_for_copy[None, LSE_idx, (blk_coord_h, blk_coord_b)],
+                    sLSE_for_copy[None, sLSE_idx, lse_handle.index],
+                )
+            else:
+                sLSE_for_copy[None, sLSE_idx, lse_handle.index].fill(0.0)
+        lse_handle.commit()
+
+        v_handle = load_mma_V_producer.acquire_and_advance()
+        cute.copy(
+            tma_atom_V,
+            tVgV_mkl[(None, mma_tile_coord_m, 0, (blk_coord_h, blk_coord_b))],
+            tVsV[(None, 0)],
+            tma_bar_ptr=v_handle.barrier,
+        )
+
+        do_handle = load_mma_dO_producer.acquire_and_advance()
+        cute.copy(
+            tma_atom_dO,
+            tdOgdO_mkl[(None, iter_index, 0, (blk_coord_h, blk_coord_b))],
+            tdOsdO[(None, do_handle.index)],
+            tma_bar_ptr=do_handle.barrier,
+        )
+
+        sum_odo_handle = load_compute_sum_OdO_producer.acquire_and_advance()
+        sSum_OdO_for_copy = cute.flat_divide(sSum_OdO, (1,))
+        sum_OdO_for_copy = cute.flat_divide(sum_OdO, (1,))
+        for i in cutlass.range_constexpr(async_copy_num_elts):
+            sum_OdO_idx = self.tile_shape_Q * iter_index + thread_idx + i * self.threads_per_warp
+            sSum_OdO_idx = thread_idx + i * self.threads_per_warp
+            if cute.elem_less(sum_OdO_idx, problem_shape[0]):
+                cute.copy(
+                    atom_async_copy,
+                    sum_OdO_for_copy[None, sum_OdO_idx, (blk_coord_h, blk_coord_b)],
+                    sSum_OdO_for_copy[None, sSum_OdO_idx, sum_odo_handle.index],
+                )
+            else:
+                sSum_OdO_for_copy[None, sSum_OdO_idx, sum_odo_handle.index].fill(0.0)
+        sum_odo_handle.commit()
+
+        dot_handle = load_mma_dOT_producer.acquire_and_advance()
+        cute.copy(
+            tma_atom_dOT,
+            tdOTgdOT_mkl[(None, 0, iter_index, (blk_coord_h, blk_coord_b))],
+            tdOTsdOT[None, dot_handle.index],
+            tma_bar_ptr=dot_handle.barrier,
+        )
+
+        qt_handle = load_mma_QT_producer.acquire_and_advance()
+        cute.copy(
+            tma_atom_QT,
+            tQTgQT_mkl[(None, 0, iter_index, (blk_coord_h, blk_coord_b))],
+            tQTsQT[None, qt_handle.index],
+            tma_bar_ptr=qt_handle.barrier,
+        )
+
+        iter_count -= 1
+        iter_index += 1
+
+        while iter_count > 0:
+            if iter_index == iter_end:
+                iter_index = iter_start
+                blk_coord_h_r += 1
+                blk_coord_h = (blk_coord_h_r, blk_coord_h_k)
+
+            q_handle = load_mma_Q_producer.acquire_and_advance()
+            cute.copy(
+                tma_atom_Q,
+                tQgQ_mkl[(None, iter_index, 0, (blk_coord_h, blk_coord_b))],
+                tQsQ[None, q_handle.index],
+                tma_bar_ptr=q_handle.barrier,
+            )
+
+            lse_handle = load_compute_LSE_producer.acquire_and_advance()
+            sLSE_for_copy = cute.flat_divide(sLSE, (1,))
+            LSE_for_copy = cute.flat_divide(LSE, (1,))
+            for i in cutlass.range_constexpr(async_copy_num_elts):
+                LSE_idx = self.tile_shape_Q * iter_index + thread_idx + i * self.threads_per_warp
+                sLSE_idx = thread_idx + i * self.threads_per_warp
+                if cute.elem_less(LSE_idx, problem_shape[0]):
+                    cute.copy(
+                        atom_async_copy,
+                        LSE_for_copy[None, LSE_idx, (blk_coord_h, blk_coord_b)],
+                        sLSE_for_copy[None, sLSE_idx, lse_handle.index],
+                    )
+                else:
+                    sLSE_for_copy[None, sLSE_idx, lse_handle.index].fill(0.0)
+            lse_handle.commit()
+
+            do_handle = load_mma_dO_producer.acquire_and_advance()
+            cute.copy(
+                tma_atom_dO,
+                tdOgdO_mkl[(None, iter_index, 0, (blk_coord_h, blk_coord_b))],
+                tdOsdO[None, do_handle.index],
+                tma_bar_ptr=do_handle.barrier,
+            )
+
+            sum_odo_handle = load_compute_sum_OdO_producer.acquire_and_advance()
+            sSum_OdO_for_copy = cute.flat_divide(sSum_OdO, (1,))
+            sum_OdO_for_copy = cute.flat_divide(sum_OdO, (1,))
+            for i in cutlass.range_constexpr(async_copy_num_elts):
+                sum_OdO_idx = (
+                    self.tile_shape_Q * iter_index + thread_idx + i * self.threads_per_warp
+                )
+                sSum_OdO_idx = thread_idx + i * self.threads_per_warp
+                if cute.elem_less(sum_OdO_idx, problem_shape[0]):
+                    cute.copy(
+                        atom_async_copy,
+                        sum_OdO_for_copy[None, sum_OdO_idx, (blk_coord_h, blk_coord_b)],
+                        sSum_OdO_for_copy[None, sSum_OdO_idx, sum_odo_handle.index],
+                    )
+                else:
+                    sSum_OdO_for_copy[None, sSum_OdO_idx, sum_odo_handle.index].fill(0.0)
+            sum_odo_handle.commit()
+
+            dot_handle = load_mma_dOT_producer.acquire_and_advance()
+            cute.copy(
+                tma_atom_dOT,
+                tdOTgdOT_mkl[(None, 0, iter_index, (blk_coord_h, blk_coord_b))],
+                tdOTsdOT[None, dot_handle.index],
+                tma_bar_ptr=dot_handle.barrier,
+            )
+
+            qt_handle = load_mma_QT_producer.acquire_and_advance()
+            cute.copy(
+                tma_atom_QT,
+                tQTgQT_mkl[(None, 0, iter_index, (blk_coord_h, blk_coord_b))],
+                tQTsQT[None, qt_handle.index],
+                tma_bar_ptr=qt_handle.barrier,
+            )
+
+            iter_count -= 1
+            iter_index += 1
+
+        if not cutlass.const_expr(self.use_clc_scheduler):
+            load_mma_K_producer.tail()
+            load_mma_V_producer.tail()
+            load_mma_Q_producer.tail()
+            load_compute_LSE_producer.tail()
+            load_mma_dO_producer.tail()
+            load_mma_dOT_producer.tail()
+            load_compute_sum_OdO_producer.tail()
+            load_mma_QT_producer.tail()
+
+        return (
+            load_mma_Q_producer,
+            load_mma_K_producer,
+            load_mma_V_producer,
+            load_compute_LSE_producer,
+            load_mma_dO_producer,
+            load_mma_dOT_producer,
+            load_compute_sum_OdO_producer,
+            load_mma_QT_producer,
+        )
+
+    @cute.jit
+    def mma_2cta(
+        self,
+        KQ_tiled_mma: cute.TiledMma,
+        VdO_tiled_mma: cute.TiledMma,
+        PdO_tiled_mma: cute.TiledMma,
+        dSQ_tiled_mma: cute.TiledMma,
+        tSTtST: cute.Tensor,
+        tSTrQ: cute.Tensor,
+        tSTrK: cute.Tensor,
+        tdPTtdPT: cute.Tensor,
+        tdPTrV: cute.Tensor,
+        tdPTrdO: cute.Tensor,
+        tdVtdV: cute.Tensor,
+        tdVrP: cute.Tensor,
+        tdVrdOT: cute.Tensor,
+        tdKrdST: cute.Tensor,
+        tdKtdK: cute.Tensor,
+        tdKrQT: cute.Tensor,
+        iter_count: Int32,
+        load_mma_Q_consumer,
+        load_mma_K_consumer,
+        load_mma_V_consumer,
+        mma_compute_S_producer,
+        load_mma_dO_consumer,
+        mma_compute_dP_producer,
+        load_mma_dOT_consumer,
+        compute_mma_P_consumer,
+        compute_mma_dS_consumer,
+        load_mma_QT_consumer,
+        mma_compute_dKdV_producer,
+    ):
+        """CuTeDSL kernel for mma pipeline."""
+        load_mma_Q_releaser = load_mma_Q_consumer.clone()
+        load_mma_K_releaser = load_mma_K_consumer.clone()
+        load_mma_V_releaser = load_mma_V_consumer.clone()
+
+        cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+        is_leader_cta = cta_rank_in_cluster % 2 == 0
+
+        if is_leader_cta:
+            s_handle = mma_compute_S_producer.acquire_and_advance()
+            k_handle = load_mma_K_consumer.wait_and_advance()
+            q_handle = load_mma_Q_consumer.wait_and_advance()
+
+            # Compute S = K * Q
+            for k_block in cutlass.range(0, cute.size(tSTrQ, mode=[2]), unroll_full=True):
+                KQ_tiled_mma.set(tcgen05.Field.ACCUMULATE, k_block != 0)
+                cute.gemm(
+                    KQ_tiled_mma,
+                    tSTtST,
+                    tSTrK[None, None, k_block, 0],
+                    tSTrQ[None, None, k_block, q_handle.index],
+                    tSTtST,
+                )
+            q_handle.release()
+
+            cute.arch.fence_view_async_tmem_store()
+            s_handle.commit()
+
+            do_handle = load_mma_dO_consumer.wait_and_advance()
+            v_handle = load_mma_V_consumer.wait_and_advance()
+
+            dp_handle = mma_compute_dP_producer.acquire_and_advance()
+
+            # Compute dP = V * dO
+            VdO_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+            for k_block in cutlass.range(0, cute.size(tdPTrV, mode=[2]), unroll_full=True):
+                cute.gemm(
+                    VdO_tiled_mma,
+                    tdPTtdPT,
+                    tdPTrV[None, None, k_block, 0],
+                    tdPTrdO[None, None, k_block, do_handle.index],
+                    tdPTtdPT,
+                )
+                VdO_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+
+            dp_handle.commit()
+            do_handle.release()
+            # V only produced once by load(); hold v_handle until end, release there via releaser
+
+            p_handle = compute_mma_P_consumer.wait_and_advance()
+            dot_handle = load_mma_dOT_consumer.wait_and_advance()
+
+            # Compute dV = P * dO (First iteration)
+            PdO_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+            for k_block in cutlass.range(0, cute.size(tdVrP, mode=[2]), unroll_full=True):
+                cute.gemm(
+                    PdO_tiled_mma,
+                    tdVtdV,
+                    tdVrP[None, None, k_block, 0],
+                    tdVrdOT[None, None, k_block, dot_handle.index],
+                    tdVtdV,
+                )
+                PdO_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+
+            dot_handle.release()
+            p_handle.release()
+
+        iter_count -= 1
+
+        dSQ_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+        while iter_count > 0:
+            if is_leader_cta:
+                q_handle = load_mma_Q_consumer.wait_and_advance()
+                s_handle = mma_compute_S_producer.acquire_and_advance()
+
+                # Compute S = K * Q
+                for k_block in cutlass.range(0, cute.size(tSTrQ, mode=[2]), unroll_full=True):
+                    KQ_tiled_mma.set(tcgen05.Field.ACCUMULATE, k_block != 0)
+                    cute.gemm(
+                        KQ_tiled_mma,
+                        tSTtST,
+                        tSTrK[None, None, k_block, 0],
+                        tSTrQ[None, None, k_block, q_handle.index],
+                        tSTtST,
+                    )
+                q_handle.release()
+                s_handle.commit()
+
+            if is_leader_cta:
+                qt_handle = load_mma_QT_consumer.wait_and_advance()
+                ds_handle = compute_mma_dS_consumer.wait_and_advance()
+
+                # Compute dK = dS * QT
+                for k_block in cutlass.range(0, cute.size(tdKrdST, mode=[2]), unroll_full=True):
+                    cute.gemm(
+                        dSQ_tiled_mma,
+                        tdKtdK,
+                        tdKrdST[
+                            None,
+                            None,
+                            k_block,
+                            ds_handle.index,
+                        ],
+                        tdKrQT[None, None, k_block, qt_handle.index],
+                        tdKtdK,
+                    )
+                    dSQ_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                qt_handle.release()
+                ds_handle.release()
+
+            if is_leader_cta:
+                dp_handle = mma_compute_dP_producer.acquire_and_advance()
+                do_handle = load_mma_dO_consumer.wait_and_advance()
+                # V only produced once by load(); reuse same V (index 0) for all loop iterations
+                VdO_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                for k_block in cutlass.range(0, cute.size(tdPTrV, mode=[2]), unroll_full=True):
+                    cute.gemm(
+                        VdO_tiled_mma,
+                        tdPTtdPT,
+                        tdPTrV[None, None, k_block, 0],
+                        tdPTrdO[None, None, k_block, do_handle.index],
+                        tdPTtdPT,
+                    )
+                    VdO_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+
+                dp_handle.commit()
+                do_handle.release()
+
+            if is_leader_cta:
+                p_handle = compute_mma_P_consumer.wait_and_advance()
+                dot_handle = load_mma_dOT_consumer.wait_and_advance()
+
+                # Compute dV = P * dO (Loop iterations)
+                for k_block in cutlass.range(0, cute.size(tdVrP, mode=[2]), unroll_full=True):
+                    cute.gemm(
+                        PdO_tiled_mma,
+                        tdVtdV,
+                        tdVrP[None, None, k_block, 0],
+                        tdVrdOT[None, None, k_block, dot_handle.index],
+                        tdVtdV,
+                    )
+                    PdO_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+
+                p_handle.release()
+                dot_handle.release()
+
+            iter_count -= 1
+
+        if is_leader_cta:
+            dkdv_handle = mma_compute_dKdV_producer.acquire_and_advance()
+            dkdv_handle.commit()
+
+            load_mma_K_releaser.release()
+            load_mma_K_releaser.advance()
+            load_mma_V_releaser.release()
+            load_mma_V_releaser.advance()
+
+        if is_leader_cta:
+            dkdv_handle = mma_compute_dKdV_producer.acquire_and_advance()
+
+            ds_handle = compute_mma_dS_consumer.wait_and_advance()
+            qt_handle = load_mma_QT_consumer.wait_and_advance()
+
+            # Compute dK = dS * Q
+            for k_block in cutlass.range(0, cute.size(tdKrdST, mode=[2]), unroll_full=True):
+                cute.gemm(
+                    dSQ_tiled_mma,
+                    tdKtdK,
+                    tdKrdST[None, None, k_block, ds_handle.index],
+                    tdKrQT[None, None, k_block, qt_handle.index],
+                    tdKtdK,
+                )
+                dSQ_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+
+            dkdv_handle.commit()
+            qt_handle.release()
+            ds_handle.release()
+
+        if not cutlass.const_expr(self.use_clc_scheduler):
+            mma_compute_S_producer.tail()
+            mma_compute_dP_producer.tail()
+            mma_compute_dKdV_producer.tail()
+
+        return (
+            mma_compute_S_producer,
+            mma_compute_dP_producer,
+            mma_compute_dKdV_producer,
+            load_mma_Q_consumer,
+            load_mma_K_consumer,
+            load_mma_V_consumer,
+            load_mma_dO_consumer,
+            load_mma_dOT_consumer,
+            compute_mma_P_consumer,
+            compute_mma_dS_consumer,
+            load_mma_QT_consumer,
+        )
+
+    @cute.jit
+    def reg_to_smem_mma64x64(
+        self,
+        regs: cute.Tensor,
+        smem: cute.Tensor,
+        index: Int32,
+        tiler_mn: tuple[Int32, Int32],
+        dp_idx: Int32,
+        wg_idx: Int32,
+    ):
+        smem_slice = smem[None, None, None, index]
+        # TODO: double check the layout of the data in reg.
+        # TODO: this may introduce additional smem transpose.
+        thread_layout = cute.make_ordered_layout(
+            tiler_mn,
+            (0, 1),  # TODO: (0,1) or (1,0) ???
+        )
+        smem_slice_tmp = cute.composition(smem_slice, thread_layout)
+
+        # TODO: temporary code for tile 64 x 64.
+        tmp_shape = ((8, 2, 4), (2, 4, 2, 2, 2))
+        tmp_stride = ((64, 512, 1024), (1, 2, 8, 16, 32))
+        smem_copy = cute.composition(smem_slice_tmp, cute.make_layout(tmp_shape, stride=tmp_stride))
+
+        # TODO: the following code is only for tile 64 x 64.
+        # TODO: need to modify the code for other tile sizes.
+        lane_idx = dp_idx % 32
+        reg_shape = regs.shape
+        atom_loops = reg_shape[0][0][2]
+        block_loops = reg_shape[2]
+        # | 00 ~ 07 | 08 ~ 15 | 16 ~ 23 | 24 ~ 31 | 32 ~ 39 | 40 ~ 47 | 48 ~ 55 | 56 ~ 63 |
+        # |---- atom size ----|---- atom size ----|---- atom size ----|---- atom size ----|
+        # |----     wg0   ----|----     wg1   ----|----     wg0   ----|----     wg1   ----|
+        for ia in cutlass.range(atom_loops):
+            for ib in cutlass.range(block_loops):
+                # the lower 8 lines
+                regs_copy = regs[((None, 0, ia), 0), 0, ib]  # two elements
+                smem_copy_slice = smem_copy[
+                    (lane_idx // 4, 0, dp_idx // 32),
+                    (None, lane_idx % 4, ia, wg_idx, ib),
+                ]
+                cute.autovec_copy(regs_copy, smem_copy_slice)
+                # the upper 8 lines
+                regs_copy = regs[((None, 1, ia), 0), 0, ib]
+                smem_copy_slice = smem_copy[
+                    (lane_idx // 4, 1, dp_idx // 32),
+                    (None, lane_idx % 4, ia, wg_idx, ib),
+                ]
+                cute.autovec_copy(regs_copy, smem_copy_slice)
+
+    @cute.jit
+    def reg_to_smem_mma128x128_2cta(
+        self,
+        regs: cute.Tensor,
+        smem: cute.Tensor,
+        index: Int32,
+        tiler_mn: tuple[Int32, Int32],
+        dp_idx: Int32,
+        wg_idx: Int32,
+        smem_RowMajor: bool = True,
+    ):
+        smem_slice = smem[None, None, None, index]
+        # K>> smem_slice:  tensor<ptr<f16, smem, align<1024>, S<3,4,3>> o ((64,16),1,(4,2)):((64,1),0,(16,4096))>
+        thread_layout = cute.make_ordered_layout(
+            # (tileN, tileM)
+            tiler_mn,
+            (0, 1),
+        )
+        # K>> thread_layout:  (64,128):(128,1)
+        smem_slice_tmp = cute.composition(smem_slice, thread_layout)
+
+        # NOTE: hardcode for tcgen05.ld.32x32b.x8 & mma128x64+2cta
+        # tmp_shape = ((32, 2), (8, 2, 2, 2)) # for 64x64 tile
+        # tmp_stride = ((64, 32*64), (1, 8, 16, 32))
+        # NOTE: hardcode for tcgen05.ld.32x32b.x16 & mma128x64+2cta
+        tmp_shape = ((32, 2), (16, 2, 2, 2))  # for 128x64 tile
+        tmp_stride = ((64, 32 * 64), (1, 16, 32, 64 * 64))
+        # smem_copy = cute.composition(smem_slice_tmp, cute.make_layout(tmp_shape, stride=tmp_stride))
+        smem_copy = cute.make_tensor(
+            smem_slice_tmp.iterator, cute.make_layout(tmp_shape, stride=tmp_stride)
+        )
+
+        warp_idx = dp_idx // 32
+        warp_row_idx = warp_idx % 2
+        warp_col_idx = warp_idx // 2  # corresponding to the second 64 cols in smem
+        lane_idx = dp_idx % 32
+        reg_shape = (
+            regs.shape
+        )  # ((8,1),1,2):((1,0),0,8) for 64x64, ((16,1),1,2):((1,0),0,16) for 128x64
+        block_loops = reg_shape[2]
+
+        # TODO: maybe can use cp.async for optimization
+        for ib in cutlass.range(block_loops):
+            regs_copy = regs[(None, 0), 0, ib]
+            smem_copy_slice = smem_copy[(lane_idx, warp_row_idx), (None, wg_idx, ib, warp_col_idx)]
+            cute.autovec_copy(regs_copy, smem_copy_slice)
+
+    @cute.jit
+    def reg_to_smem_mma128x64_2cta(
+        self,
+        regs: cute.Tensor,
+        smem: cute.Tensor,
+        index: Int32,
+        tiler_mn: tuple[Int32, Int32],
+        dp_idx: Int32,
+        wg_idx: Int32,
+        smem_RowMajor: bool = True,
+    ):
+        smem_slice = smem[None, None, None, index]
+        thread_layout = cute.make_ordered_layout(
+            # (tileN, tileM)
+            tiler_mn,
+            (1, 0) if smem_RowMajor else (0, 1),
+        )
+        smem_slice_tmp = cute.composition(smem_slice, thread_layout)
+        # NOTE: hardcode for tcgen05.ld.32x32b.x8 & mma128x64+2cta
+        tmp_shape = ((32, 2), (8, 2, 2, 2))
+        tmp_stride = ((64, 32 * 64), (1, 8, 16, 32))
+        smem_copy = cute.composition(smem_slice_tmp, cute.make_layout(tmp_shape, stride=tmp_stride))
+
+        warp_idx = dp_idx // 32
+        warp_row_idx = warp_idx % 2
+        warp_col_idx = warp_idx // 2
+        lane_idx = dp_idx % 32
+        reg_shape = regs.shape  # ((8,1),1,2):((1,0),0,8)
+        block_loops = reg_shape[2]
+
+        # TODO: maybe can use cp.async for optimization
+        for ib in cutlass.range(block_loops):
+            regs_copy = regs[(None, 0), 0, ib]
+            smem_copy_slice = smem_copy[(lane_idx, warp_row_idx), (None, wg_idx, ib, warp_col_idx)]
+            cute.autovec_copy(regs_copy, smem_copy_slice)
+
+    @cute.jit
+    def compute(
+        self,
+        tSTtST: cute.Tensor,
+        tdPTtdPT: cute.Tensor,
+        tdVrP: cute.Tensor,
+        sP: cute.Tensor,
+        sLSE: cute.Tensor,
+        # sdS: cute.Tensor,
+        sdST: cute.Tensor,
+        sdOT: cute.Tensor,
+        sSum_OdO: cute.Tensor,
+        dK: cute.Tensor,
+        dV: cute.Tensor,
+        tdKtdK: cute.Tensor,
+        tdVtdV: cute.Tensor,
+        PdO_tiled_mma: cute.TiledMma,
+        dSQ_tiled_mma: cute.TiledMma,
+        blk_coord: cute.Coord,
+        blk_offset: cute.Shape,
+        problem_shape: tuple[Int32, Int32, Int32, tuple[tuple[Int32, Int32], Int32]],
+        iter_count: Int32,
+        iter_start: Int32,
+        iter_end: Int32,
+        scale_softmax: cutlass.Float32,
+        mma_compute_S_producer,
+        mma_compute_S_consumer,
+        compute_mma_P_producer,
+        compute_mma_P_consumer,
+        load_compute_LSE_producer,
+        load_compute_LSE_consumer,
+        load_compute_sum_OdO_producer,
+        load_compute_sum_OdO_consumer,
+        mma_compute_dP_producer,
+        mma_compute_dP_consumer,
+        compute_mma_dS_producer,
+        compute_mma_dS_consumer,
+        mma_compute_dKdV_producer,
+        mma_compute_dKdV_consumer,
+        varlen: bool,
+        sK: cute.Tensor,
+        problem_shape_k_cur_batch: Int32,
+        tma_atom_dK: cute.CopyAtom,
+        dK_tma: cute.Tensor,
+        tma_atom_dV: cute.CopyAtom,
+        dV_tma: cute.Tensor,
+        sdK_epi_layout: cute.ComposedLayout,
+        sdV_epi_layout: cute.ComposedLayout,
+    ):
+        """CuTeDSL kernel for recomputing softmax and producing dk and dv."""
+        tidx, _, _ = cute.arch.thread_idx()
+        Q, K, _, _ = problem_shape
+        _, blk_coord_k, _, _ = blk_coord
+
+        iter_index = iter_start
+
+        # adi: TMEM_ST, TMEM_DPT
+        tmem_load_op = tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(16))
+        tmem_load_atom = cute.make_copy_atom(
+            tmem_load_op,
+            self.acc_dtype,
+        )
+
+        tSTtST = tSTtST[(None, None), 0, 0]
+        tdPTtdPT = tdPTtdPT[(None, None), 0, 0]
+
+        cST = cute.make_identity_tensor(cute.select(self.cta_tiler, mode=[1, 0]))
+        cdPT = cute.make_identity_tensor(cute.select(self.cta_tiler, mode=[1, 0]))
+
+        num_warp_groups = self.num_compute_warps // 4
+        dp_idx = tidx % 128
+        wg_idx = (tidx % (self.num_compute_warps * self.threads_per_warp)) // 128
+        tiled_t2r = tcgen05.make_tmem_copy(tmem_load_atom, tSTtST)
+        thr_t2r = tiled_t2r.get_slice(dp_idx)
+
+        tTR_cST = thr_t2r.partition_D(cST)
+        tTR_cST = split_wg(tTR_cST, num_warp_groups, wg_idx)
+        tTR_rST = cute.make_rmem_tensor(tTR_cST.shape, self.acc_dtype)
+
+        tTR_tST = thr_t2r.partition_S(tSTtST)
+        tTR_tST = split_wg(tTR_tST, num_warp_groups, wg_idx)
+
+        tTR_cdPT_p = thr_t2r.partition_D(cdPT)
+        tTR_cdPT = split_wg(tTR_cdPT_p, num_warp_groups, wg_idx)
+        tTR_rdPT = cute.make_rmem_tensor(tTR_cdPT.shape, self.acc_dtype)
+
+        tTR_tdPT = thr_t2r.partition_S(tdPTtdPT)
+        tTR_tdPT = split_wg(tTR_tdPT, num_warp_groups, wg_idx)
+
+        tdVcST = PdO_tiled_mma.get_slice(0).partition_A(cST)
+
+        is_residual_k = blk_coord_k * self.tile_shape_K + self.tile_shape_K > K
+        last_iter = iter_end - 1
+
+        while iter_count > 0:
+            s_handle = mma_compute_S_consumer.wait_and_advance()
+            p_handle = compute_mma_P_producer.acquire_and_advance()
+            lse_handle = load_compute_LSE_consumer.wait_and_advance()
+
+            leading_causal_masking = cutlass.Boolean(False)
+            if cutlass.const_expr(self.is_causal):
+                # TODO: could be optimized by specify an exact iter_index
+                #
+                # NOTE (causal + 2CTA correctness):
+                # `iter_start` can be rounded down to an even index for 2CTA. When the true
+                # causal boundary Q tile is odd, it becomes (iter_start + 1). In that case,
+                # we must treat both (iter_start, iter_start + 1) as "masked tiles" so the
+                # per-element causal mask is applied on the boundary tile too.
+                leading_causal_masking = iter_index == iter_start
+                q_block_min_unaligned = (
+                    blk_coord_k * self.tile_shape_K + Q - K - self.window_size_right
+                ) // self.tile_shape_Q
+                boundary_in_second = (q_block_min_unaligned % 2) == 1
+                leading_causal_masking = leading_causal_masking or (
+                    boundary_in_second and (iter_index == iter_start + 1)
+                )
+                offset_partial_tile = (K - Q) % self.tile_shape_K
+                need_additional_mask = offset_partial_tile and iter_index == iter_start + 1
+                leading_causal_masking = leading_causal_masking or need_additional_mask
+                leading_causal_masking = cute.arch.shuffle_sync(leading_causal_masking, 0)
+
+            trailing_residual_masking = cutlass.Boolean(False)
+            trailing_residual_masking = iter_index == last_iter or is_residual_k
+            trailing_residual_masking = cute.arch.shuffle_sync(trailing_residual_masking, 0)
+
+            # For causal, every tile may contain (q,k) with k > q; we must apply per-element mask for all Q tiles.
+            is_masked_tile = (
+                leading_causal_masking
+                or trailing_residual_masking
+                or self.has_sliding_window
+                or cutlass.const_expr(self.is_causal)
+            )
+
+            # Compute P = softmax(S, LSE)
+            cute.copy(tiled_t2r, tTR_tST, tTR_rST)
+
+            if is_masked_tile:
+                for i in cutlass.range(cute.size(tTR_rST), unroll_full=True):
+                    c_transpose = tTR_cST[i]
+                    pos = (
+                        cute.get(c_transpose, mode=[1]) + iter_index * self.tile_shape_Q,
+                        cute.get(c_transpose, mode=[0]) + blk_coord_k * self.tile_shape_K,
+                    )
+                    if cutlass.const_expr(self.has_sliding_window):
+                        if cutlass.const_expr(self.window_size_left < 0):
+                            tTR_rST[i] = (
+                                -cutlass.Float32.inf
+                                if pos[1] > pos[0] + K - Q + self.window_size_right
+                                else tTR_rST[i]
+                            )
+                        else:
+                            max_K_index = min(pos[0] + K - Q + self.window_size_right, K)
+                            min_K_index = max(0, pos[0] + K - Q - self.window_size_left)
+                            tTR_rST[i] = (
+                                -cutlass.Float32.inf
+                                if pos[1] > max_K_index or pos[1] < min_K_index
+                                else tTR_rST[i]
+                            )
+                    if cutlass.const_expr(self.is_causal) and (
+                        pos[0] + K - Q < pos[1] or not cute.elem_less(pos, (Q, K))
+                    ):
+                        tTR_rST[i] = -cutlass.Float32.inf
+                    if not cute.elem_less(pos, (Q, K)):
+                        tTR_rST[i] = -cutlass.Float32.inf
+
+            log2_e = cutlass.Float32(math.log2(math.e))
+            softmax_scale_log2_e = scale_softmax * log2_e
+
+            for i in cutlass.range(0, cute.size(tTR_rST), 2, unroll_full=True):
+                lse = (
+                    -sLSE[
+                        cute.get(tTR_cST[i], mode=[1]),
+                        lse_handle.index,
+                    ],
+                    -sLSE[
+                        cute.get(tTR_cST[i + 1], mode=[1]),
+                        lse_handle.index,
+                    ],
+                )
+                tTR_rST[i], tTR_rST[i + 1] = cute.arch.fma_packed_f32x2(
+                    (tTR_rST[i], tTR_rST[i + 1]),
+                    (softmax_scale_log2_e, softmax_scale_log2_e),
+                    lse,
+                )
+                tTR_rST[i] = cute.math.exp2(tTR_rST[i], fastmath=True)
+                tTR_rST[i + 1] = cute.math.exp2(tTR_rST[i + 1], fastmath=True)
+
+            # convert fp32 P to fp16 P which will be used in the PdO
+            tTR_rPT = self.quantize(tTR_rST, dV.element_type)  # tTR_rST is ST in fp32 in RF.
+            self.reg_to_smem_mma128x128_2cta(
+                tTR_rPT,
+                sP,
+                p_handle.index,
+                (self.tile_shape_K, self.tile_shape_Q),
+                dp_idx,
+                wg_idx,
+            )
+            cute.arch.fence_view_async_shared()
+            cute.arch.barrier(
+                barrier_id=self.compute_sync_bar_id,
+                number_of_threads=self.num_compute_warps * self.threads_per_warp,
+            )
+
+            p_handle.commit()
+
+            s_handle.release()
+            lse_handle.release()
+
+            sum_odo_handle = load_compute_sum_OdO_consumer.wait_and_advance()
+            dp_handle = mma_compute_dP_consumer.wait_and_advance()
+            ds_handle = compute_mma_dS_producer.acquire_and_advance()
+
+            # Compute dS = dsoftmax(P, dP, sum_OdO)
+            cute.copy(tiled_t2r, tTR_tdPT, tTR_rdPT)
+
+            for i in cutlass.range(0, cute.size(tTR_rdPT), 2, unroll_full=True):
+                dpsum_0 = -sSum_OdO[
+                    cute.get(tTR_cdPT[i], mode=[1]),
+                    sum_odo_handle.index,
+                ]
+                dpsum_1 = -sSum_OdO[
+                    cute.get(tTR_cdPT[i + 1], mode=[1]),
+                    sum_odo_handle.index,
+                ]
+                if cutlass.const_expr(varlen):
+                    if not cute.elem_less(cute.get(tTR_cdPT[i], mode=[1]), Q):
+                        dpsum_0 = 0.0
+                    if not cute.elem_less(cute.get(tTR_cdPT[i + 1], mode=[1]), Q):
+                        dpsum_1 = 0.0
+                tTR_rdPT[i], tTR_rdPT[i + 1] = cute.arch.add_packed_f32x2(
+                    (tTR_rdPT[i], tTR_rdPT[i + 1]),
+                    (dpsum_0, dpsum_1),
+                )
+                tTR_rdPT[i], tTR_rdPT[i + 1] = cute.arch.mul_packed_f32x2(
+                    (tTR_rdPT[i], tTR_rdPT[i + 1]), (tTR_rST[i], tTR_rST[i + 1])
+                )
+            # For causal, force dS to zero at masked (q,k) so dK/dV accumulation is correct
+            if cutlass.const_expr(self.is_causal):
+                for i in cutlass.range(cute.size(tTR_rdPT), unroll_full=True):
+                    c_transpose = tTR_cdPT[i]
+                    pos = (
+                        cute.get(c_transpose, mode=[1]) + iter_index * self.tile_shape_Q,
+                        cute.get(c_transpose, mode=[0]) + blk_coord_k * self.tile_shape_K,
+                    )
+                    if pos[0] + K - Q < pos[1] or not cute.elem_less(pos, (Q, K)):
+                        tTR_rdPT[i] = cutlass.Float32(0.0)
+            # convert fp32 dS to fp16 dS which will be used in the computation of dK and DQ
+            tTR_rdST = self.quantize(tTR_rdPT, dV.element_type)
+
+            cute.arch.fence_view_async_tmem_load()
+            dp_handle.release()
+
+            self.reg_to_smem_mma128x128_2cta(
+                tTR_rdST,
+                sdST,
+                ds_handle.index,
+                (self.tile_shape_K, self.tile_shape_Q),
+                dp_idx,
+                wg_idx,
+            )
+            cute.arch.fence_view_async_shared()
+            cute.arch.barrier(
+                barrier_id=self.compute_sync_bar_id,
+                number_of_threads=self.num_compute_warps * self.threads_per_warp,
+            )
+
+            ds_handle.commit()
+            sum_odo_handle.release()
+
+            iter_count -= 1
+            iter_index += 1
+            if iter_index == iter_end:
+                iter_index = iter_start
+
+        # Epilogue
+        mma_compute_dKdV_consumer = self.epilogue(
+            blk_coord,
+            blk_offset,
+            problem_shape,
+            dK,
+            dV,
+            tdKtdK,
+            tdVtdV,
+            scale_softmax,
+            mma_compute_dKdV_producer,
+            mma_compute_dKdV_consumer,
+            problem_shape_k_cur_batch,
+            tma_atom_dK,
+            dK_tma,
+            tma_atom_dV,
+            dV_tma,
+            sdK_epi_layout,
+            sdV_epi_layout,
+            varlen,
+            sdOT,
+            sP,
+        )
+
+        if not cutlass.const_expr(self.use_clc_scheduler):
+            compute_mma_P_producer.tail()
+            compute_mma_dS_producer.tail()
+
+        return (
+            compute_mma_P_producer,
+            compute_mma_dS_producer,
+            mma_compute_S_consumer,
+            compute_mma_P_consumer,
+            load_compute_LSE_consumer,
+            load_compute_sum_OdO_consumer,
+            mma_compute_dP_consumer,
+            compute_mma_dS_consumer,
+            mma_compute_dKdV_consumer,
+        )
+
+    @cute.jit
+    def quantize(
+        self,
+        input_t: cute.Tensor,
+        element_dtype: type[cutlass.Numeric],
+    ) -> cute.Tensor:
+        """Convert Float32 to element dtype."""
+        output = cute.make_rmem_tensor(input_t.shape, element_dtype)
+        output.store(input_t.load().to(element_dtype))
+        return output
+
+    @cute.jit
+    def store(
+        self,
+        gmem: cute.Tensor,
+        regs: cute.Tensor,
+        coord: cute.Tensor,
+        tensor_shape: cute.Shape,
+    ):
+        for i in cutlass.range(cute.size(coord, mode=[2]), unroll_full=True):
+            coord_i = coord[None, 0, i]
+            gmem_i = gmem[None, 0, i]
+            regs_i = regs[None, 0, i]
+            if cute.elem_less(coord_i[0], tensor_shape):
+                gmem_i.store(regs_i.load().to(gmem.element_type))
+
+    @cute.jit
+    def epilogue_clear(
+        self,
+        blk_coord: cute.Coord,
+        blk_offset: cute.Shape,
+        problem_shape: tuple[Int32, Int32, Int32, tuple[tuple[Int32, Int32], Int32]],
+        dK: cute.Tensor,
+        dV: cute.Tensor,
+    ):
+        """Early stopping needs to clear dK and dV."""
+        tidx, _, _ = cute.arch.thread_idx()
+        block_dim_x, _, _ = cute.arch.block_dim()
+        _, K, _, HB = problem_shape
+        _, blk_coord_k, _, blk_coord_batch = blk_coord
+
+        mdK_offset = cute.assume(blk_offset[1] * dK.stride[0], divby=64)
+        mdK = cute.make_tensor(
+            dK.iterator + mdK_offset,
+            cute.make_layout((K, self.tile_shape_dQ_K, HB), stride=dK.stride),
+        )
+        gdK = cute.local_tile(mdK, (self.cta_tiler[1], self.cta_tiler[2]), (None, None, None))
+        gdK = gdK[None, None, blk_coord_k, 0, blk_coord_batch]
+        cdK = cute.domain_offset(
+            (blk_coord_k * self.tile_shape_K, 0),
+            cute.make_identity_tensor((self.cta_tiler[1], self.cta_tiler[2])),
+        )
+
+        mdV_offset = cute.assume(blk_offset[1] * dV.stride[0], divby=64)
+        mdV = cute.make_tensor(
+            dV.iterator + mdV_offset,
+            cute.make_layout((K, self.tile_shape_dV_dO, HB), stride=dV.stride),
+        )
+        gdV = cute.local_tile(mdV, (self.cta_tiler[1], self.cta_tiler[2]), (None, None, None))
+        gdV = gdV[None, None, blk_coord_k, 0, blk_coord_batch]
+        cdV = cute.domain_offset(
+            (blk_coord_k * self.tile_shape_K, 0),
+            cute.make_identity_tensor((self.cta_tiler[1], self.cta_tiler[2])),
+        )
+
+        num_zero_epi_threads = 256
+
+        tiled_copy_r2g = fa_copy_utils.tiled_copy_2d(
+            dK.element_type, self.cta_tiler[2], num_zero_epi_threads
+        )
+
+        thr_copy_r2g = tiled_copy_r2g.get_slice(tidx)
+
+        tRG_gdK = thr_copy_r2g.partition_D(gdK)
+        tRG_cdK = thr_copy_r2g.partition_D(cdK)
+        tRG_gdV = thr_copy_r2g.partition_D(gdV)
+        tRG_cdV = thr_copy_r2g.partition_D(cdV)
+
+        zero_frg = cute.make_rmem_tensor_like(tRG_gdK[None, 0, None])
+        zero_frg.fill(dK.element_type(0.0))
+
+        # check we don't need zero fragment duplication
+        V_frg_size = cute.size(tRG_gdV[None, 0, None])
+        assert cute.size(zero_frg) == V_frg_size
+
+        if tidx < num_zero_epi_threads:
+            for n in cutlass.range(cute.size(tRG_gdK.shape[1]), unroll_full=True):
+                if cute.elem_less(tRG_cdK[0, n, 0][0], problem_shape[1]):
+                    cute.copy(tiled_copy_r2g, zero_frg, tRG_gdK[None, n, None])
+
+            for n in cutlass.range(cute.size(tRG_gdV.shape[1]), unroll_full=True):
+                if cute.elem_less(tRG_cdV[0, n, 0][0], problem_shape[1]):
+                    cute.copy(tiled_copy_r2g, zero_frg, tRG_gdV[None, n, None])
+
+    @cute.jit
+    def epilogue(
+        self,
+        blk_coord: cute.Coord,
+        blk_offset: cute.Shape,
+        problem_shape: tuple[Int32, Int32, Int32, tuple[tuple[Int32, Int32], Int32]],
+        dK: cute.Tensor,
+        dV: cute.Tensor,
+        tdKtdK: cute.Tensor,
+        tdVtdV: cute.Tensor,
+        scale_softmax: cutlass.Float32,
+        mma_compute_dKdV_producer,
+        mma_compute_dKdV_consumer,
+        problem_shape_k_cur_batch: Int32,
+        tma_atom_dK: cute.CopyAtom,
+        dK_tma: cute.Tensor,
+        tma_atom_dV: cute.CopyAtom,
+        dV_tma: cute.Tensor,
+        sdK_epi_layout: cute.ComposedLayout,
+        sdV_epi_layout: cute.ComposedLayout,
+        varlen: bool,
+        sdOT: cute.Tensor,
+        sP: cute.Tensor,
+    ):
+        """Variant 3a (5/5) Path 2: CTA-shared SMEM with cooperative WG writes + TMA bulk store.
+
+        Both warp-groups cooperatively populate a per-CTA (64, 256) virtual SMEM
+        buffer (4 stages of (64, 64) aliased onto sP+sdST). Per-thread t2r N
+        coverage is interleaved across the full hd=256, so per-WG TMA is not
+        viable — instead we treat SMEM as one shared per-CTA buffer and let
+        each thread's `self.store`-equivalent write into it via a (64, 256)
+        virtual tensor whose N axis maps (n%64, n//64) → (N_within, stage).
+        After an inter-WG barrier (256 threads), the leader warp fires 4 TMA
+        bulk stores, one per stage, to the corresponding (64, 64) GMEM slice.
+        Varlen falls back to per-thread self.store as in flash_bwd_sm100.py.
+        """
+        tidx, _, _ = cute.arch.thread_idx()
+        _, K, D, HB = problem_shape
+        _, blk_coord_k, _, blk_coord_batch = blk_coord
+
+        # adi: TMEM_DK, TMEM_DV
+        tmem_copy_op = tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32))
+        load_op = cute.make_copy_atom(
+            tmem_copy_op,
+            self.acc_dtype,
+        )
+
+        tdKtdK = tdKtdK[(None, None), 0, 0]
+        mdK_offset = cute.assume(blk_offset[1] * dK.stride[0], divby=64)
+        mdK = cute.make_tensor(
+            dK.iterator + mdK_offset,
+            cute.make_layout((K, self.tile_shape_dQ_K, HB), stride=dK.stride),
+        )
+        gdK = cute.local_tile(mdK, (self.cta_tiler[1], self.cta_tiler[2]), (None, None, None))
+        gdK = gdK[None, None, blk_coord_k, 0, blk_coord_batch]
+        cdK = cute.domain_offset(
+            (blk_coord_k * self.tile_shape_K, 0),
+            cute.make_identity_tensor((self.cta_tiler[1], self.cta_tiler[2])),
+        )
+
+        num_warp_groups = self.num_compute_warps // 4
+        dp_idx = tidx % 128
+        wg_idx = (tidx % (self.num_compute_warps * self.threads_per_warp)) // 128
+        leader_warp = (cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4) == 0
+
+        # Path 2 SMEM staging. dV stages through sdOT (already-consumed by the
+        # dV MMA before the dV epilogue begins). dK stages through sP+sdST
+        # (dead after dK MMA completes, before dK epilogue runs).
+        s_epi_dK = cute.make_tensor(
+            cute.recast_ptr(sP.iterator, sdK_epi_layout.inner, dK.element_type),
+            sdK_epi_layout.outer,
+        )
+        s_epi_dV = cute.make_tensor(
+            cute.recast_ptr(sdOT.iterator, sdV_epi_layout.inner, dV.element_type),
+            sdV_epi_layout.outer,
+        )
+
+        # Compile-time: stage tile shape and number of stages.
+        epi_cols_dKV = math.gcd(
+            128 // (dV.element_type.width // 8), self.cta_tiler[2] // num_warp_groups
+        )
+        num_epi_stages_dKV = (self.cta_tiler[2] // num_warp_groups) // epi_cols_dKV
+        total_epi_stages = num_warp_groups * num_epi_stages_dKV
+        epi_tile_dKV = (self.cta_tiler[1], epi_cols_dKV)
+
+        # Local (M, N) coord tensor for SMEM indexing (no global domain offset
+        # — cdK/cdV are domain-offset by blk_coord_k * tile_shape_K to match
+        # the GMEM destination, but the SMEM indexing must be per-CTA-local).
+        cdV_local = cute.make_identity_tensor((self.cta_tiler[1], self.cta_tiler[2]))
+        cdK_local = cdV_local
+
+        tiled_t2r_dK = tcgen05.make_tmem_copy(load_op, tdKtdK)
+        thread_t2r_dK = tiled_t2r_dK.get_slice(dp_idx)
+
+        tTR_cdK = thread_t2r_dK.partition_D(cdK)
+        tTR_cdK = split_wg(tTR_cdK, num_warp_groups, wg_idx)
+        tTR_cdK_local = thread_t2r_dK.partition_D(cdK_local)
+        tTR_cdK_local = split_wg(tTR_cdK_local, num_warp_groups, wg_idx)
+        tTR_gdK = thread_t2r_dK.partition_D(gdK)
+        tTR_gdK = split_wg(tTR_gdK, num_warp_groups, wg_idx)
+        tTR_rdK = cute.make_rmem_tensor(tTR_cdK.shape, self.acc_dtype)
+        tTR_tdK = thread_t2r_dK.partition_S(tdKtdK)
+        tTR_tdK = split_wg(tTR_tdK, num_warp_groups, wg_idx)
+
+        mdV_in = cute.make_tensor(
+            dV.iterator, cute.make_layout((K, self.cta_tiler[2], HB), stride=dV.stride)
+        )
+        offset_mdV = cute.assume(blk_offset[1] * mdV_in.stride[0], divby=64)
+        mdV = cute.make_tensor(mdV_in.iterator + offset_mdV, mdV_in.layout)
+        gdV = cute.local_tile(mdV, (self.cta_tiler[1], self.cta_tiler[2]), (None, None, None))
+        gdV = gdV[None, None, blk_coord_k, 0, blk_coord_batch]
+
+        cdV = cute.domain_offset(
+            (blk_coord_k * self.cta_tiler[1], 0),
+            cute.make_identity_tensor((self.cta_tiler[1], self.cta_tiler[2])),
+        )
+
+        tdVtdV = tdVtdV[(None, None), 0, 0]
+
+        tiled_t2r_dV = tcgen05.make_tmem_copy(load_op, tdVtdV)
+        thread_t2r_dV = tiled_t2r_dV.get_slice(dp_idx)
+
+        tTR_cdV = thread_t2r_dV.partition_D(cdV)
+        tTR_cdV = split_wg(tTR_cdV, num_warp_groups, wg_idx)
+        tTR_cdV_local = thread_t2r_dV.partition_D(cdV_local)
+        tTR_cdV_local = split_wg(tTR_cdV_local, num_warp_groups, wg_idx)
+        tTR_gdV = thread_t2r_dV.partition_D(gdV)
+        tTR_gdV = split_wg(tTR_gdV, num_warp_groups, wg_idx)
+        tTR_rdV = cute.make_rmem_tensor(tTR_cdV.shape, self.acc_dtype)
+        tTR_tdV = thread_t2r_dV.partition_S(tdVtdV)
+        tTR_tdV = split_wg(tTR_tdV, num_warp_groups, wg_idx)
+
+        # GMEM destinations for the multi-stage TMA path (gated on not-varlen).
+        if cutlass.const_expr(not varlen):
+            mdV_tma_3d = cute.make_tensor(
+                dV_tma.iterator,
+                cute.make_layout((K, self.cta_tiler[2], HB), stride=dV_tma.stride),
+            )
+            mdV_tma_cur = mdV_tma_3d[None, None, blk_coord_batch]
+            gdV_tma = cute.local_tile(
+                mdV_tma_cur, (self.cta_tiler[1], self.cta_tiler[2]), (blk_coord_k, 0)
+            )
+            gdV_tma_epi = cute.local_tile(gdV_tma, epi_tile_dKV, (0, None))
+
+            mdK_tma_3d = cute.make_tensor(
+                dK_tma.iterator,
+                cute.make_layout((K, self.cta_tiler[2], HB), stride=dK_tma.stride),
+            )
+            mdK_tma_cur = mdK_tma_3d[None, None, blk_coord_batch]
+            gdK_tma = cute.local_tile(
+                mdK_tma_cur, (self.cta_tiler[1], self.cta_tiler[2]), (blk_coord_k, 0)
+            )
+            gdK_tma_epi = cute.local_tile(gdK_tma, epi_tile_dKV, (0, None))
+
+        cta_threads = self.num_compute_warps * self.threads_per_warp
+
+        dkdv_handle = mma_compute_dKdV_consumer.wait_and_advance()
+
+        if blk_coord_k * self.tile_shape_K < problem_shape_k_cur_batch:
+            cute.copy(tiled_t2r_dV, tTR_tdV, tTR_rdV)
+            tTR_rdV_cast = cute.make_rmem_tensor(tTR_rdV.shape, dV.element_type)
+            tTR_rdV_cast.store(tTR_rdV.load().to(dV.element_type))
+
+            if cutlass.const_expr(not varlen):
+                # reg -> SMEM via per-element indexed stores using tTR_cdV's
+                # per-thread (M, N) coords. (M, N) is per-CTA cdV space (M=0..63,
+                # N=0..255). We map N=(n%epi_cols, n//epi_cols) → (N_within, stage)
+                # of the 3D s_epi_dV tensor.
+                for _i in cutlass.range_constexpr(cute.size(tTR_cdV_local, mode=[2])):
+                    for _j in cutlass.range_constexpr(cute.size(tTR_cdV_local[None, 0, _i])):
+                        c = tTR_cdV_local[None, 0, _i][_j]
+                        m_pos = c[0]
+                        n_pos = c[1]
+                        stage_pos = n_pos // epi_cols_dKV
+                        n_within_pos = n_pos % epi_cols_dKV
+                        v = tTR_rdV_cast[None, 0, _i][_j]
+                        s_epi_dV[m_pos, n_within_pos, stage_pos] = v
+                cute.arch.fence_view_async_shared()
+                # Inter-WG barrier — both warp-groups must finish their writes
+                # before the leader warp reads SMEM via TMA.
+                cute.arch.barrier(barrier_id=5, number_of_threads=cta_threads)
+                # TMA bulk store, one (64, 64) box per stage.
+                if leader_warp and wg_idx == 0:
+                    for _stage in cutlass.range_constexpr(total_epi_stages):
+                        sdV_stage = s_epi_dV[None, None, _stage]
+                        gdV_stage = gdV_tma_epi[None, None, _stage]
+                        td_sdV, td_gdV = cpasync.tma_partition(
+                            tma_atom_dV,
+                            0,
+                            cute.make_layout(1),
+                            cute.group_modes(sdV_stage, 0, 2),
+                            cute.group_modes(gdV_stage, 0, 2),
+                        )
+                        cute.copy(tma_atom_dV, td_sdV, td_gdV)
+                        cute.arch.cp_async_bulk_commit_group()
+                cute.arch.cp_async_bulk_wait_group(0, read=True)
+            else:
+                self.store(tTR_gdV, tTR_rdV, tTR_cdV, (K, D))
+
+        cute.arch.fence_view_async_tmem_load()
+        dkdv_handle.release()
+
+        dkdv_handle = mma_compute_dKdV_consumer.wait_and_advance()
+
+        if blk_coord_k * self.tile_shape_K < problem_shape_k_cur_batch:
+            cute.copy(tiled_t2r_dK, tTR_tdK, tTR_rdK)
+
+            for i in cutlass.range(cute.size(tTR_rdK), unroll_full=True):
+                tTR_rdK[i] = scale_softmax * tTR_rdK[i]
+
+            tTR_rdK_cast = cute.make_rmem_tensor(tTR_rdK.shape, dK.element_type)
+            tTR_rdK_cast.store(tTR_rdK.load().to(dK.element_type))
+
+            if cutlass.const_expr(not varlen):
+                for _i in cutlass.range_constexpr(cute.size(tTR_cdK_local, mode=[2])):
+                    for _j in cutlass.range_constexpr(cute.size(tTR_cdK_local[None, 0, _i])):
+                        c = tTR_cdK_local[None, 0, _i][_j]
+                        m_pos = c[0]
+                        n_pos = c[1]
+                        stage_pos = n_pos // epi_cols_dKV
+                        n_within_pos = n_pos % epi_cols_dKV
+                        v = tTR_rdK_cast[None, 0, _i][_j]
+                        s_epi_dK[m_pos, n_within_pos, stage_pos] = v
+                cute.arch.fence_view_async_shared()
+                cute.arch.barrier(barrier_id=6, number_of_threads=cta_threads)
+                if leader_warp and wg_idx == 0:
+                    for _stage in cutlass.range_constexpr(total_epi_stages):
+                        sdK_stage = s_epi_dK[None, None, _stage]
+                        gdK_stage = gdK_tma_epi[None, None, _stage]
+                        td_sdK, td_gdK = cpasync.tma_partition(
+                            tma_atom_dK,
+                            0,
+                            cute.make_layout(1),
+                            cute.group_modes(sdK_stage, 0, 2),
+                            cute.group_modes(gdK_stage, 0, 2),
+                        )
+                        cute.copy(tma_atom_dK, td_sdK, td_gdK)
+                        cute.arch.cp_async_bulk_commit_group()
+                cute.arch.cp_async_bulk_wait_group(0, read=True)
+            else:
+                self.store(tTR_gdK, tTR_rdK, tTR_cdK, (K, D))
+
+        cute.arch.fence_view_async_tmem_load()
+        dkdv_handle.release()
+
+        return mma_compute_dKdV_consumer
+
+    def get_workspace_tensor(
+        self,
+        problem_shape: tuple[Int32, Int32, Int32, tuple[tuple[Int32, Int32], Int32]],
+        workspace: cute.Tensor,
+        acc_dtype: type[cutlass.Numeric],
+        varlen: bool,
+    ) -> tuple[cute.Tensor, cute.Tensor, cute.Tensor]:
+        """Get workspace tensor."""
+        D = problem_shape[2]
+        H, B = cute.size(problem_shape[3][0]), cute.size(problem_shape[3][1])
+        H_r, H_k = problem_shape[3][0]
+        D = cute.round_up(D, 8)
+
+        # b = 1 for varlen, else batch_size
+        b = workspace.shape[0]
+        # s_q_sum for varlen, else s_q_max, already rounded to 8
+        S_Q = workspace.shape[1]
+
+        acc_bytes = acc_dtype.width // 8
+        sum_OdO_bytes = cute.assume(b * H * S_Q * acc_bytes, divby=acc_bytes)
+        scaled_lse_bytes = cute.assume(b * H * S_Q * acc_bytes, divby=acc_bytes)
+
+        sum_OdO_iter = workspace.iterator
+        scaled_lse_iter = sum_OdO_iter + sum_OdO_bytes
+
+        sum_OdO_iter = cute.recast_ptr(sum_OdO_iter, dtype=self.acc_dtype)
+        scaled_lse_iter = cute.recast_ptr(scaled_lse_iter, dtype=self.acc_dtype)
+
+        sum_OdO = cute.make_tensor(
+            sum_OdO_iter,
+            cute.make_layout(
+                (S_Q, ((H_r, H_k), B)),
+                stride=(1, ((S_Q, S_Q * H_r), 0 if varlen else S_Q * H)),
+            ),
+        )
+        scaled_lse = cute.make_tensor(
+            scaled_lse_iter,
+            cute.make_layout(
+                (S_Q, ((H_r, H_k), B)),
+                stride=(1, ((S_Q, S_Q * H_r), 0 if varlen else S_Q * H)),
+            ),
+        )
+
+        return sum_OdO, scaled_lse
+
+    @staticmethod
+    def _compute_sum_OdO_grid(
+        problem_shape: tuple[Int32, Int32, Int32, tuple[tuple[Int32, Int32], Int32]],
+        block_q: int,
+    ) -> tuple[Int32, Int32, Int32]:
+        """Compute grid shape for sum_OdO kernel."""
+        return (
+            cute.ceil_div(cute.size(problem_shape[0]), block_q),
+            cute.size(problem_shape[3][0]),  # H
+            cute.size(problem_shape[3][1]),  # B
+        )
+
+    @staticmethod
+    def _compute_bwd_grid(
+        problem_shape: tuple[Int32, Int32, Int32, tuple[tuple[Int32, Int32], Int32]],
+        block_k: int,
+    ) -> tuple[Int32, Int32, Int32]:
+        """Compute grid shape for bwd kernel."""
+        K = problem_shape[1]
+        _, H_K = problem_shape[3][0]
+        B = problem_shape[3][1]
+        return (cute.ceil_div(K, block_k), cute.size(H_K), cute.size(B))
+
+    @staticmethod
+    def get_workspace_size(s_q: int, d: int, h: int, b: int, acc_dtype: type[cutlass.Numeric]):
+        """Get workspace size."""
+        d = (d + 7) // 8 * 8  # round up to 8
+        s_q = (s_q + 7) // 8 * 8  # round up to 8
+        workspace_bytes = 0
+        # OdO vector
+        workspace_bytes += acc_dtype.width // 8
+        # scaled LSE vector
+        workspace_bytes += acc_dtype.width // 8
+        # FP32 versions of outputs that are churned (start off with Q only)
+        workspace_bytes += d * acc_dtype.width // 8
+        return (b, s_q, h, workspace_bytes)
+
+    def make_and_init_load_mma_K_pipeline(self, load_mma_K_mbar_ptr, cluster_layout_vmnk):
+        """Create and initialize barrier for load mma Q."""
+        load_mma_K_producer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, len([self.load_warp_id])
+        )
+        load_mma_K_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, len([self.mma_warp_id])
+        )
+        return pipeline.PipelineTmaUmma.create(
+            barrier_storage=load_mma_K_mbar_ptr,
+            num_stages=self.load_mma_K_stage,
+            producer_group=load_mma_K_producer_group,
+            consumer_group=load_mma_K_consumer_group,
+            tx_count=self.tma_copy_K_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+        )
+
+    def make_and_init_load_mma_V_pipeline(self, load_mma_V_mbar_ptr, cluster_layout_vmnk):
+        """Create and initialize barrier for load mma Q."""
+        load_mma_V_producer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, len([self.load_warp_id])
+        )
+        load_mma_V_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, len([self.mma_warp_id])
+        )
+        return pipeline.PipelineTmaUmma.create(
+            barrier_storage=load_mma_V_mbar_ptr,
+            num_stages=self.load_mma_V_stage,
+            producer_group=load_mma_V_producer_group,
+            consumer_group=load_mma_V_consumer_group,
+            tx_count=self.tma_copy_V_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+        )
+
+    def make_and_init_load_mma_Q_pipeline(self, load_mma_Q_mbar_ptr, cluster_layout_vmnk):
+        """Create and initialize barrier for load mma Q."""
+        load_mma_Q_producer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, len([self.load_warp_id])
+        )
+        load_mma_Q_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, len([self.mma_warp_id])
+        )
+        return pipeline.PipelineTmaUmma.create(
+            barrier_storage=load_mma_Q_mbar_ptr,
+            num_stages=self.load_mma_Q_stage,
+            producer_group=load_mma_Q_producer_group,
+            consumer_group=load_mma_Q_consumer_group,
+            tx_count=self.tma_copy_Q_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+        )
+
+    def make_and_init_load_mma_QT_pipeline(self, load_mma_QT_mbar_ptr, cluster_layout_vmnk):
+        """Create and initialize barrier for load mma QT."""
+        load_mma_QT_producer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, len([self.load_warp_id])
+        )
+        load_mma_QT_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, len([self.mma_warp_id])
+        )
+        return pipeline.PipelineTmaUmma.create(
+            barrier_storage=load_mma_QT_mbar_ptr,
+            num_stages=self.load_mma_QT_stage,
+            producer_group=load_mma_QT_producer_group,
+            consumer_group=load_mma_QT_consumer_group,
+            tx_count=self.tma_copy_Q_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+        )
+
+    def make_and_init_load_mma_dO_pipeline(self, load_mma_dO_mbar_ptr, cluster_layout_vmnk):
+        """Create and initialize barrier for load mma dO."""
+        load_mma_dO_producer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, len([self.load_warp_id])
+        )
+        load_mma_dO_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, len([self.mma_warp_id])
+        )
+        return pipeline.PipelineTmaUmma.create(
+            barrier_storage=load_mma_dO_mbar_ptr,
+            num_stages=self.load_mma_dO_stage,
+            producer_group=load_mma_dO_producer_group,
+            consumer_group=load_mma_dO_consumer_group,
+            tx_count=self.tma_copy_dO_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+        )
+
+    def make_and_init_load_compute_LSE_pipeline(self, load_compute_lse_mbar_ptr):
+        """Create and initialize barrier for load compute lse."""
+        load_compute_lse_producer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            self.threads_per_warp,
+            # self.threads_per_warp,
+        )
+        load_compute_lse_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            self.threads_per_warp * self.num_compute_warps,
+            # self.threads_per_warp * self.num_compute_warps,
+        )
+        return pipeline.PipelineCpAsync.create(
+            barrier_storage=load_compute_lse_mbar_ptr,
+            num_stages=self.load_compute_LSE_stage,
+            producer_group=load_compute_lse_producer_group,
+            consumer_group=load_compute_lse_consumer_group,
+        )
+
+    def make_and_init_load_compute_sum_OdO_pipeline(self, load_compute_sum_OdO_mbar_ptr):
+        """Create and initialize barrier for load sum OdO."""
+        load_compute_sum_OdO_producer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            self.threads_per_warp,
+            # self.threads_per_warp,
+        )
+        load_compute_sum_OdO_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            self.threads_per_warp * self.num_compute_warps,
+            # self.threads_per_warp * self.num_compute_warps,
+        )
+        return pipeline.PipelineCpAsync.create(
+            barrier_storage=load_compute_sum_OdO_mbar_ptr,
+            num_stages=self.load_compute_sum_OdO_stage,
+            producer_group=load_compute_sum_OdO_producer_group,
+            consumer_group=load_compute_sum_OdO_consumer_group,
+        )
+
+    def make_and_init_mma_compute_S_pipeline(self, mma_compute_S_mbar_ptr, cluster_layout_vmnk):
+        """Create and initialize barrier for mma S."""
+        mma_compute_S_producer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            len([self.mma_warp_id]),
+        )
+        mma_compute_S_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0],
+        )
+        return pipeline.PipelineUmmaAsync.create(
+            barrier_storage=mma_compute_S_mbar_ptr,
+            num_stages=self.mma_compute_S_stage,
+            producer_group=mma_compute_S_producer_group,
+            consumer_group=mma_compute_S_consumer_group,
+            cta_layout_vmnk=cluster_layout_vmnk,
+        )
+
+    #  Barrier to between dP = v * dO and consume of dP in compute()
+    def make_and_init_mma_compute_dP_pipeline(self, mma_compute_dP_mbar_ptr, cluster_layout_vmnk):
+        """Create and initialize barrier for mma Q."""
+        mma_compute_dP_producer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            len([self.mma_warp_id]),
+        )
+        mma_compute_dP_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0],
+        )
+        return pipeline.PipelineUmmaAsync.create(
+            barrier_storage=mma_compute_dP_mbar_ptr,
+            num_stages=self.mma_compute_dP_stage,
+            producer_group=mma_compute_dP_producer_group,
+            consumer_group=mma_compute_dP_consumer_group,
+            cta_layout_vmnk=cluster_layout_vmnk,
+        )
+
+    def make_and_init_compute_mma_P_pipeline(self, compute_mma_P_mbar_ptr, cluster_layout_vmnk):
+        """Create and initialize barrier for mma P."""
+        compute_mma_P_producer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0],
+            # self.num_compute_warps * self.threads_per_warp,
+        )
+        compute_mma_P_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            len([self.mma_warp_id]),
+        )
+        return pipeline.PipelineAsyncUmma.create(
+            barrier_storage=compute_mma_P_mbar_ptr,
+            num_stages=self.compute_mma_P_stage,
+            producer_group=compute_mma_P_producer_group,
+            consumer_group=compute_mma_P_consumer_group,
+            cta_layout_vmnk=cluster_layout_vmnk,
+        )
+
+    def make_and_init_compute_mma_dS_pipeline(self, compute_mma_dS_mbar_ptr, cluster_layout_vmnk):
+        """Create and initialize barrier for mma dS."""
+
+        compute_mma_dS_producer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0],
+            # self.num_compute_warps * self.threads_per_warp,
+        )
+        compute_mma_dS_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            len([self.mma_warp_id]),
+        )
+
+        return pipeline.PipelineAsyncUmma.create(
+            barrier_storage=compute_mma_dS_mbar_ptr,
+            num_stages=self.compute_mma_dS_stage,
+            producer_group=compute_mma_dS_producer_group,
+            consumer_group=compute_mma_dS_consumer_group,
+            cta_layout_vmnk=cluster_layout_vmnk,
+        )
+
+    def make_and_init_mma_compute_dKdV_pipeline(
+        self, mma_compute_dKdV_mbar_ptr, cluster_layout_vmnk
+    ):
+        """Create and initialize barrier for mma dKdV."""
+        mma_compute_dKdV_producer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            len([self.mma_warp_id]),
+        )
+        mma_compute_dKdV_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0],
+            # self.num_compute_warps * self.threads_per_warp,
+        )
+        return pipeline.PipelineUmmaAsync.create(
+            barrier_storage=mma_compute_dKdV_mbar_ptr,
+            num_stages=self.mma_compute_dKdV_stage,
+            producer_group=mma_compute_dKdV_producer_group,
+            consumer_group=mma_compute_dKdV_consumer_group,
+            cta_layout_vmnk=cluster_layout_vmnk,
+        )
diff --git a/build/torch-cuda/sm100_hd256_2cta_fmha_backward_dqkernel.py b/build/torch-cuda/sm100_hd256_2cta_fmha_backward_dqkernel.py
new file mode 100644
index 0000000000000000000000000000000000000000..7ffcf783ace91ddfd0ac167b7d999329a9bdcbbc
--- /dev/null
+++ b/build/torch-cuda/sm100_hd256_2cta_fmha_backward_dqkernel.py
@@ -0,0 +1,2317 @@
+# Copyright (c) 2025, Siyu Wang, Shengbin Di, Yuxi Chi, Johnsonms, Linfeng Zheng, Haoyan Huang, Lanbo Li, Yun Zhong, Man Yuan, Minmin Sun, Yong Li, Wei Lin.
+
+from typing import Type, Tuple, Optional
+
+import cuda.bindings.driver as cuda
+
+import math
+import cutlass
+import cutlass.cute as cute
+import cutlass.cute.nvgpu.tcgen05 as tcgen05
+from cutlass.cute.nvgpu import cpasync
+import cutlass.utils as utils
+import cutlass.pipeline as pipeline
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass.cute.typing import Int32, Int64, Float32
+
+from cutlass.utils import ClcDynamicPersistentTileScheduler
+from .tile_scheduler import (
+    SchedulerState,
+    compute_sm100_fmha_grid as compute_grid,
+    compute_sm100_fmha_grid_clc as compute_grid_clc,
+    make_sm100_thread_cooperative_group as make_thread_cooperative_group,
+    Sm100FmhaStaticTileScheduler as FmhaStaticTileScheduler,
+    Sm100FmhaStaticTileSchedulerParams as FmhaStaticTileSchedulerParams,
+    Sm100FmhaClcDynamicTileScheduler as FmhaClcDynamicTileScheduler,
+    Sm100FmhaClcDynamicTileSchedulerParams as FmhaClcDynamicTileSchedulerParams,
+)
+from .mask import (
+    Sm100FusedMask as FusedMask,
+)
+from .tile_scheduler import SM100_TMEM_CAPACITY_COLUMNS
+from . import copy_utils as fa_copy_utils
+
+
+class BlackwellFusedMultiHeadAttentionBackwardDQKernel:
+    def __init__(
+        self,
+        acc_dtype: Type[cutlass.Numeric],
+        mma_tiler: Tuple[int, int, int],
+        is_causal: bool,
+        window_size_left: int | None,
+        window_size_right: int | None,
+        is_persistent: bool,
+        split_head: bool,
+        use_clc_scheduler: bool = False,
+    ):
+        self.acc_dtype = acc_dtype
+        self.mma_tiler = mma_tiler
+        self.is_causal = is_causal
+        self.window_size_left = window_size_left
+        # Keep original behavior (known-good in this repo)
+        window_size_left = (
+            None
+            if (window_size_left is None or window_size_left < 0)
+            else cutlass.Int32(window_size_left)
+        )
+        window_size_right = (
+            None
+            if (window_size_right is None or window_size_right < 0)
+            else cutlass.Int32(window_size_right)
+        )
+        self.window_size_left = None if self.is_causal else window_size_left
+        self.window_size_right = cutlass.Int32(0) if self.is_causal else window_size_right
+        self.is_local = (not self.is_causal) and (
+            self.window_size_left is not None or self.window_size_right is not None
+        )
+        assert mma_tiler[0] == 128 and mma_tiler[1] == 128, "Only 128x128 tile impl is supported"
+        assert mma_tiler[2] == 256, "Only 256 is supported for 128x128 tile impl"
+        self.cta_tiler = (
+            mma_tiler[0],
+            mma_tiler[1],
+            mma_tiler[2],
+        )
+        self.qk_mma_tiler = (
+            2 * mma_tiler[0],
+            mma_tiler[1],
+            min(self.cta_tiler[2], 128) if split_head else self.cta_tiler[2],
+        )
+        self.dov_mma_tiler = self.qk_mma_tiler
+        self.dsk_mma_tiler = (
+            2 * mma_tiler[0],
+            min(self.cta_tiler[2], 128) if split_head else self.cta_tiler[2],
+            mma_tiler[1],
+        )
+
+        self.dsk_block_tiler = (
+            self.dsk_mma_tiler[0] // 2,
+            self.dsk_mma_tiler[1],
+            self.dsk_mma_tiler[2],
+        )
+        self.iterations_qk = self.cta_tiler[2] // self.qk_mma_tiler[2]
+        self.iterations_dov = self.cta_tiler[2] // self.dov_mma_tiler[2]
+        self.iterations_dsk = self.cta_tiler[2] // self.dsk_mma_tiler[1]
+        self.cluster_shape_mn = (2, 1)
+        self.tmem_warp_shape_mn = (4, 1)
+        self.is_persistent = is_persistent
+        self.use_clc_scheduler = use_clc_scheduler
+        self.use_semantic_trip_range = self.is_causal or self.is_local
+
+        self.compute_warp_ids = (0, 1, 2, 3)
+        self.epilogue_warp_ids = (4, 5, 6, 7)
+        self.mma_warp_id = 8
+        self.load_warp_id = 9
+        self.empty_warp_id = (10, 11)
+        self.sched_warp_id = self.empty_warp_id[0] if use_clc_scheduler else None
+        self.tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS
+        self.num_compute_warps = len(self.compute_warp_ids)
+
+        self.cta_sync_bar_id = 0
+        self.tmem_alloc_sync_bar_id = 1
+        self.compute_sync_bar_id = 2
+
+        self.threads_per_warp = 32
+        self.threads_per_cta = self.threads_per_warp * len(
+            (
+                *self.compute_warp_ids,  # this is to get a round num threads
+                *self.epilogue_warp_ids,
+                self.mma_warp_id,
+                self.load_warp_id,
+                *self.empty_warp_id,
+            )
+        )
+
+        self.tmem_alloc_barrier = pipeline.NamedBarrier(
+            barrier_id=1,
+            num_threads=self.threads_per_cta,
+        )
+
+        self.tmem_s_offset = 0
+        self.tmem_dp_offset = 128
+        self.tmem_dq_offset = 256
+
+        self.num_regs_compute = 256
+        self.num_regs_epilogue = 160
+        self.num_regs_other = 32
+
+        self.buffer_align_bytes = 1024
+
+    def _setup_attributes(self):
+        self.q_stage = self.iterations_qk
+        self.k_stage = self.iterations_qk
+        self.do_stage = self.iterations_dov
+        self.v_stage = self.iterations_dov
+        self.kt_stage = 1
+        self.qk_acc_stage = 1
+        self.dov_acc_stage = 1
+        self.dsk_acc_stage = 1
+        self.epi_stage = 1
+        self.load_compute_LSE_stage = 1
+        self.load_compute_sum_OdO_stage = 1
+        if cutlass.const_expr(self.use_clc_scheduler):
+            self.num_clc_stage = 1
+            self.num_clc_response_bytes = 16
+
+    @cute.jit
+    def __call__(
+        self,
+        q_tensor: cute.Tensor,
+        k_tensor: cute.Tensor,
+        v_tensor: cute.Tensor,
+        dq_tensor: cute.Tensor,
+        do_tensor: cute.Tensor,
+        lse_tensor: cute.Tensor,
+        sum_odo_tensor: cute.Tensor,
+        cum_seqlen_q: Optional[cute.Tensor],
+        cum_seqlen_k: Optional[cute.Tensor],
+        scale_softmax: cutlass.Float32,
+        stream: cuda.CUstream,
+    ):
+        varlen = cum_seqlen_q is not None or cum_seqlen_k is not None
+        # Infer shape metadata from normalized 5D tensors (B, S, H_k, H_r, D),
+        # similar to the dedicated hd256 forward path.
+        s_q = q_tensor.shape[1]
+        s_k = k_tensor.shape[1]
+        d = q_tensor.shape[4]
+        h_k = q_tensor.shape[2]
+        h_r = q_tensor.shape[3]
+        if cutlass.const_expr(cum_seqlen_q is not None):
+            b = cum_seqlen_q.shape[0] - 1
+        elif cutlass.const_expr(cum_seqlen_k is not None):
+            b = cum_seqlen_k.shape[0] - 1
+        else:
+            b = q_tensor.shape[0]
+        # `lse_tensor` / `sum_odo_tensor` are preallocated FP32 buffers (lse_log2 / dpsum)
+        # whose sequence dimension is padded (rounded up) by the caller.
+        # Use their leading dimension as the LSE length to ensure correct batch strides.
+        s_lse = lse_tensor.shape[0]
+        s_q64 = Int64(s_q)
+        s_k64 = Int64(s_k)
+        s_lse64 = Int64(s_lse)
+        h_r64 = Int64(h_r)
+        h_k64 = Int64(h_k)
+        b64 = Int64(b)
+        # Packed-varlen representation uses batch-dim = 1 and sequence-dim = total_{q,k}.
+        # Keep the *physical* sequence extent in the tensor layouts so that applying
+        # `cuseqlen_*` offsets stays within the tensor domain.
+        s_q_total = q_tensor.shape[1] if cum_seqlen_q is not None else s_q64
+        s_k_total = k_tensor.shape[1] if cum_seqlen_k is not None else s_k64
+        b_lse = b64 if cum_seqlen_q is None else 1
+        stride_b_lse = h_r64 * h_k64 * s_lse64 if cum_seqlen_q is None else 0
+
+        # (s, d, ((h_r, h_k), b))
+        q_layout = cute.make_layout(
+            (s_q_total, d, ((h_r, h_k), b)),
+            stride=(
+                cute.assume(q_tensor.stride[1], divby=64),
+                q_tensor.stride[4],
+                (
+                    (q_tensor.stride[3], q_tensor.stride[2]),
+                    0 if cum_seqlen_q is not None else cute.assume(q_tensor.stride[0], divby=64),
+                ),
+            ),
+        )
+        q = cute.make_tensor(q_tensor.iterator, q_layout)
+        # (s, d, ((h_r, h_k), b))
+        do_layout = cute.make_layout(
+            (s_q_total, d, ((h_r, h_k), b)),
+            stride=(
+                cute.assume(do_tensor.stride[1], divby=64),
+                do_tensor.stride[4],
+                (
+                    (do_tensor.stride[3], do_tensor.stride[2]),
+                    0 if cum_seqlen_q is not None else cute.assume(do_tensor.stride[0], divby=64),
+                ),
+            ),
+        )
+        do = cute.make_tensor(do_tensor.iterator, do_layout)
+        # (s, d, ((h_r, h_k), b)), 0-stride for h_r to broadcast
+        k_layout = cute.make_layout(
+            (s_k_total, d, ((h_r, h_k), b)),
+            stride=(
+                cute.assume(k_tensor.stride[1], divby=64),
+                k_tensor.stride[4],
+                (
+                    (0, k_tensor.stride[2]),
+                    0 if cum_seqlen_k is not None else cute.assume(k_tensor.stride[0], divby=64),
+                ),
+            ),
+        )
+        k = cute.make_tensor(k_tensor.iterator, k_layout)
+        # (d, s, ((h_r, h_k), b)), 0-stride for h_r to broadcast
+        kt_layout = cute.make_layout(
+            (d, s_k_total, ((h_r, h_k), b)),
+            stride=(
+                k_tensor.stride[4],
+                cute.assume(k_tensor.stride[1], divby=64),
+                (
+                    (0, k_tensor.stride[2]),
+                    0 if cum_seqlen_k is not None else cute.assume(k_tensor.stride[0], divby=64),
+                ),
+            ),
+        )
+        kt = cute.make_tensor(k_tensor.iterator, kt_layout)
+        # (s, d, ((h_r, h_k), b)), 0-stride for h_r to broadcast
+        v_layout = cute.make_layout(
+            (s_k_total, d, ((h_r, h_k), b)),
+            stride=(
+                cute.assume(v_tensor.stride[1], divby=64),
+                v_tensor.stride[4],
+                (
+                    (0, v_tensor.stride[2]),
+                    0 if cum_seqlen_k is not None else cute.assume(v_tensor.stride[0], divby=64),
+                ),
+            ),
+        )
+        v = cute.make_tensor(v_tensor.iterator, v_layout)
+        # (s, ((h_r, h_k), b))
+        lse = cute.make_tensor(lse_tensor.iterator, lse_tensor.layout)
+        # (s, ((h_r, h_k), b))
+        sum_odo_layout = cute.make_layout(
+            (s_lse64, ((h_r, h_k), b_lse)),
+            stride=(1, ((s_lse64, h_r64 * s_lse64), stride_b_lse)),
+        )
+        sum_odo = cute.make_tensor(sum_odo_tensor.iterator, sum_odo_layout)
+        # (s, d, ((h_r, h_k), b))
+        dq_layout = cute.make_layout(
+            (s_q_total, d, ((h_r, h_k), b)),
+            stride=(
+                cute.assume(dq_tensor.stride[1], divby=64),
+                dq_tensor.stride[4],
+                (
+                    (dq_tensor.stride[3], dq_tensor.stride[2]),
+                    0 if cum_seqlen_q is not None else cute.assume(dq_tensor.stride[0], divby=64),
+                ),
+            ),
+        )
+        dq = cute.make_tensor(dq_tensor.iterator, dq_layout)
+
+        # setup static attributes before smem/grid/tma computation
+        self.q_dtype = q.element_type
+        self.k_dtype = k.element_type
+        self.v_dtype = v.element_type
+        self.do_dtype = do.element_type
+        self.dq_dtype = dq.element_type
+        self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.q_dtype.width
+
+        if cutlass.const_expr(self.use_clc_scheduler):
+            self.tile_sched_params, grid = compute_grid_clc(
+                (s_q, dq.shape[1], dq.shape[2]) if cum_seqlen_q is not None else dq.shape,
+                self.cta_tiler,
+                (*self.cluster_shape_mn, 1),
+            )
+        else:
+            self.tile_sched_params, grid = compute_grid(
+                (s_q, dq.shape[1], dq.shape[2]) if cum_seqlen_q is not None else dq.shape,
+                self.cta_tiler,
+                self.is_persistent,
+            )
+
+        self.q_major_mode = utils.LayoutEnum.from_tensor(q).mma_major_mode()
+        self.do_major_mode = utils.LayoutEnum.from_tensor(do).mma_major_mode()
+        self.k_major_mode = utils.LayoutEnum.from_tensor(k).mma_major_mode()
+        self.v_major_mode = utils.LayoutEnum.from_tensor(v).mma_major_mode()
+        self.dq_layout = utils.LayoutEnum.from_tensor(dq)
+
+        if cutlass.const_expr(self.q_major_mode != tcgen05.OperandMajorMode.K):
+            raise RuntimeError("The layout of q is not supported")
+        if cutlass.const_expr(self.k_major_mode != tcgen05.OperandMajorMode.K):
+            raise RuntimeError("The layout of k is not supported")
+        if cutlass.const_expr(self.v_major_mode != tcgen05.OperandMajorMode.K):
+            raise RuntimeError("The layout of v is not supported")
+        if cutlass.const_expr(self.do_major_mode != tcgen05.OperandMajorMode.K):
+            raise RuntimeError("The layout of v is not supported")
+
+        # check type consistency
+        if cutlass.const_expr(self.q_dtype != self.k_dtype):
+            raise TypeError(f"Type mismatch: {self.q_dtype} != {self.k_dtype}")
+        if cutlass.const_expr(self.q_dtype != self.v_dtype):
+            raise TypeError(f"Type mismatch: {self.q_dtype} != {self.v_dtype}")
+        if cutlass.const_expr(self.q_dtype != self.do_dtype):
+            raise TypeError(f"Type mismatch: {self.q_dtype} != {self.do_dtype}")
+
+        self._setup_attributes()
+
+        cta_group = tcgen05.CtaGroup.TWO
+        # the intermediate tensor p is from tmem & k-major
+        ds_source = tcgen05.OperandSource.TMEM
+        ds_major_mode = tcgen05.OperandMajorMode.K
+        k_trans_major_mode = tcgen05.OperandMajorMode.MN
+        qk_tiled_mma = sm100_utils.make_trivial_tiled_mma(
+            self.q_dtype,
+            self.q_major_mode,
+            self.k_major_mode,
+            self.acc_dtype,
+            cta_group,
+            self.qk_mma_tiler[:2],
+        )
+        dov_tiled_mma = sm100_utils.make_trivial_tiled_mma(
+            self.do_dtype,
+            self.do_major_mode,
+            self.v_major_mode,
+            self.acc_dtype,
+            cta_group,
+            self.dov_mma_tiler[:2],
+        )
+        dsk_tiled_mma = sm100_utils.make_trivial_tiled_mma(
+            self.q_dtype,
+            ds_major_mode,
+            k_trans_major_mode,
+            self.acc_dtype,
+            cta_group,
+            self.dsk_mma_tiler[:2],
+            ds_source,
+        )
+
+        self.cluster_shape_mnk = (*self.cluster_shape_mn, 1)
+        self.cluster_layout_vmnk = cute.tiled_divide(
+            cute.make_layout(self.cluster_shape_mnk),
+            (qk_tiled_mma.thr_id.shape,),
+        )
+
+        self.epi_tile = self.dsk_block_tiler[:2]
+
+        q_smem_layout_staged = sm100_utils.make_smem_layout_a(
+            qk_tiled_mma,
+            self.qk_mma_tiler,
+            self.q_dtype,
+            self.q_stage,
+        )
+        k_smem_layout_staged = sm100_utils.make_smem_layout_b(
+            qk_tiled_mma,
+            self.qk_mma_tiler,
+            self.k_dtype,
+            self.k_stage,
+        )
+        do_smem_layout_staged = sm100_utils.make_smem_layout_a(
+            dov_tiled_mma,
+            self.dov_mma_tiler,
+            self.do_dtype,
+            self.do_stage,
+        )
+        v_smem_layout_staged = sm100_utils.make_smem_layout_b(
+            dov_tiled_mma,
+            self.dov_mma_tiler,
+            self.v_dtype,
+            self.v_stage,
+        )
+        ds_tmem_layout_staged = sm100_utils.make_smem_layout_a(
+            dsk_tiled_mma,
+            self.dsk_mma_tiler,
+            self.q_dtype,
+            self.qk_acc_stage,
+        )
+        ds_tmem_layout = cute.select(ds_tmem_layout_staged, mode=[0, 1, 2])
+        kt_smem_layout_staged = sm100_utils.make_smem_layout_b(
+            dsk_tiled_mma,
+            self.dsk_mma_tiler,
+            self.k_dtype,
+            self.dsk_acc_stage,
+        )
+
+        # TMA load for Q
+        tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group)
+
+        q_smem_layout = cute.select(q_smem_layout_staged, mode=[0, 1, 2])
+        tma_atom_q, tma_tensor_q = cute.nvgpu.make_tiled_tma_atom_A(
+            tma_load_op,
+            q,
+            q_smem_layout,
+            self.qk_mma_tiler,
+            qk_tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+        # TMA load for K
+        k_smem_layout = cute.select(k_smem_layout_staged, mode=[0, 1, 2])
+        tma_atom_k, tma_tensor_k = cute.nvgpu.make_tiled_tma_atom_B(
+            tma_load_op,
+            k,
+            k_smem_layout,
+            self.qk_mma_tiler,
+            qk_tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+        # TMA load for dO
+        do_smem_layout = cute.select(do_smem_layout_staged, mode=[0, 1, 2])
+        tma_atom_do, tma_tensor_do = cute.nvgpu.make_tiled_tma_atom_A(
+            tma_load_op,
+            do,
+            do_smem_layout,
+            self.dov_mma_tiler,
+            dov_tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+        # TMA load for V
+        v_smem_layout = cute.select(v_smem_layout_staged, mode=[0, 1, 2])
+        tma_atom_v, tma_tensor_v = cute.nvgpu.make_tiled_tma_atom_B(
+            tma_load_op,
+            v,
+            v_smem_layout,
+            self.dov_mma_tiler,
+            dov_tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+        # TMA load for KT
+        kt_smem_layout = cute.select(kt_smem_layout_staged, mode=[0, 1, 2])
+        tma_atom_kt, tma_tensor_kt = cute.nvgpu.make_tiled_tma_atom_B(
+            tma_load_op,
+            kt,
+            kt_smem_layout,
+            self.dsk_mma_tiler,
+            dsk_tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+        lse_smem_layout = cute.make_layout((self.cta_tiler[0], self.load_compute_LSE_stage))
+        sum_odo_smem_layout = cute.make_layout((self.cta_tiler[0], self.load_compute_sum_OdO_stage))
+
+        q_copy_size = cute.size_in_bytes(self.q_dtype, q_smem_layout)
+        k_copy_size = cute.size_in_bytes(self.k_dtype, k_smem_layout)
+        do_copy_size = cute.size_in_bytes(self.do_dtype, do_smem_layout)
+        v_copy_size = cute.size_in_bytes(self.v_dtype, v_smem_layout)
+        kt_copy_size = cute.size_in_bytes(self.k_dtype, kt_smem_layout)
+        self.tma_copy_q_bytes = q_copy_size * cute.size(qk_tiled_mma.thr_id.shape)
+        self.tma_copy_k_bytes = k_copy_size * cute.size(qk_tiled_mma.thr_id.shape)
+        self.tma_copy_do_bytes = do_copy_size * cute.size(qk_tiled_mma.thr_id.shape)
+        self.tma_copy_v_bytes = v_copy_size * cute.size(qk_tiled_mma.thr_id.shape)
+        self.tma_copy_kt_bytes = kt_copy_size * cute.size(qk_tiled_mma.thr_id.shape)
+
+        # dQ epilogue TMA store. Use a single (M, 64) SMEM stage and rotate it
+        # across the hd=256 N dimension to stay within the SMEM budget.
+        tma_store_op = cpasync.CopyBulkTensorTileS2GOp()
+        epi_cols_dQ = math.gcd(128 // (dq.element_type.width // 8), self.epi_tile[1])
+        epi_tile_dQ = (self.epi_tile[0], epi_cols_dQ)
+        sdQ_epi_layout = sm100_utils.make_smem_layout_epi(
+            dq.element_type,
+            self.dq_layout,
+            epi_tile_dQ,
+            1,
+        )
+        tma_atom_dQ, tma_tensor_dQ = cpasync.make_tiled_tma_atom(
+            tma_store_op,
+            dq,
+            cute.select(sdQ_epi_layout, mode=[0, 1]),
+            epi_tile_dQ,
+        )
+
+        @cute.struct
+        class SharedStorage:
+            # TMA G2S load barriers: LOAD warp (producer) -> MMA warp (consumer)
+            load_q_mbar_ptr: cute.struct.MemRange[
+                Int64, self.q_stage * 2
+            ]  # load_q_{producer,consumer}
+            load_do_mbar_ptr: cute.struct.MemRange[
+                Int64, self.do_stage * 2
+            ]  # load_do_{producer,consumer}
+            load_k_mbar_ptr: cute.struct.MemRange[
+                Int64, self.k_stage * 2
+            ]  # load_k_{producer,consumer}
+            load_kt_mbar_ptr: cute.struct.MemRange[
+                Int64, self.kt_stage * 2
+            ]  # load_kt_{producer,consumer}
+            load_v_mbar_ptr: cute.struct.MemRange[
+                Int64, self.v_stage * 2
+            ]  # load_v_{producer,consumer}
+            mma_s_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2]
+            mma_dp_mbar_ptr: cute.struct.MemRange[Int64, self.dov_acc_stage * 2]
+            mma_dq_mbar_ptr: cute.struct.MemRange[Int64, self.epi_stage * 2]
+            ds_mma_mbar_ptr: cute.struct.MemRange[Int64, self.dsk_acc_stage * 2]
+            lse_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_compute_LSE_stage * 2]
+            sum_odo_mbar_ptr: cute.struct.MemRange[
+                cutlass.Int64, self.load_compute_sum_OdO_stage * 2
+            ]
+            # A CTA-wide "TMEM lifetime" barrier used to safely deallocate TMEM after all users finish.
+            tmem_dealloc_mbar: Int64
+            # Tmem holding buffer
+            tmem_holding_buf: Int32
+            # CLC pipeline barriers and response buffer
+            clc_mbar_ptr: cute.struct.MemRange[Int64, 2]
+            clc_response: cute.struct.MemRange[Int32, 4]
+
+        self.shared_storage = SharedStorage
+
+        grid = cute.round_up(grid, self.cluster_shape_mnk)
+        # Launch the kernel synchronously
+        self.kernel(
+            qk_tiled_mma,
+            dov_tiled_mma,
+            dsk_tiled_mma,
+            tma_atom_q,
+            tma_tensor_q,
+            tma_atom_k,
+            tma_tensor_k,
+            tma_atom_v,
+            tma_tensor_v,
+            tma_atom_do,
+            tma_tensor_do,
+            tma_atom_kt,
+            tma_tensor_kt,
+            tma_atom_dQ,
+            tma_tensor_dQ,
+            lse,
+            sum_odo,
+            dq,
+            cum_seqlen_q,
+            cum_seqlen_k,
+            scale_softmax,
+            self.window_size_left,
+            self.window_size_right,
+            self.cluster_layout_vmnk,
+            q_smem_layout_staged,
+            k_smem_layout_staged,
+            v_smem_layout_staged,
+            do_smem_layout_staged,
+            kt_smem_layout_staged,
+            ds_tmem_layout,
+            sdQ_epi_layout,
+            lse_smem_layout,
+            sum_odo_smem_layout,
+            self.tile_sched_params,
+        ).launch(
+            grid=grid,
+            block=[self.threads_per_cta, 1, 1],
+            cluster=self.cluster_shape_mnk,
+            stream=stream,
+            min_blocks_per_mp=1,
+        )
+
+    #  GPU device kernel
+    @cute.kernel
+    def kernel(
+        self,
+        qk_tiled_mma: cute.TiledMma,
+        dov_tiled_mma: cute.TiledMma,
+        dsk_tiled_mma: cute.TiledMma,
+        tma_atom_q: cute.CopyAtom,
+        mQ_qdl: cute.Tensor,
+        tma_atom_k: cute.CopyAtom,
+        mK_kdl: cute.Tensor,
+        tma_atom_v: cute.CopyAtom,
+        mV_dkl: cute.Tensor,
+        tma_atom_do: cute.CopyAtom,
+        mdO_qdl: cute.Tensor,
+        tma_atom_kt: cute.CopyAtom,
+        mK_dkl: cute.Tensor,
+        tma_atom_dQ: cute.CopyAtom,
+        mdQ_tma: cute.Tensor,
+        mLSE: cute.Tensor,
+        mSum_OdO: cute.Tensor,
+        mdQ_qdl: cute.Tensor,
+        cum_seqlen_q: Optional[cute.Tensor],
+        cum_seqlen_k: Optional[cute.Tensor],
+        scale_softmax: Float32,
+        window_size_left: Optional[Int32],
+        window_size_right: Optional[Int32],
+        cluster_layout_vmnk: cute.Layout,
+        q_smem_layout_staged: cute.ComposedLayout,
+        k_smem_layout_staged: cute.ComposedLayout,
+        v_smem_layout_staged: cute.ComposedLayout,
+        do_smem_layout_staged: cute.ComposedLayout,
+        kt_smem_layout_staged: cute.ComposedLayout,
+        ds_tmem_layout_staged: cute.ComposedLayout,
+        sdQ_epi_layout: cute.ComposedLayout,
+        lse_smem_layout: cute.Layout,
+        sum_odo_smem_layout: cute.Layout,
+        tile_sched_params: FmhaStaticTileSchedulerParams | FmhaClcDynamicTileSchedulerParams,
+    ):
+        # llvm.inline_asm(
+        #     None,
+        #     [],
+        #     '.pragma "global knob CommonIntoMultiBlockLoop=1";',
+        #     "",
+        #     has_side_effects=True,
+        #     is_align_stack=False,
+        #     asm_dialect=llvm.AsmDialect.AD_ATT,
+        # )
+
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+        #
+        # Prefetch tma desc
+        #
+        if warp_idx == self.load_warp_id:
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_q)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_k)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_v)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_do)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_kt)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_dQ)
+
+        bidx, bidy, bidz = cute.arch.block_idx()
+        tidx, _, _ = cute.arch.thread_idx()
+        varlen = cum_seqlen_q is not None or cum_seqlen_k is not None
+        mma_tile_coord_v = bidx % cute.size(qk_tiled_mma.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)
+
+        # Alloc
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(self.shared_storage)
+
+        load_q_producer, load_q_consumer = pipeline.PipelineTmaUmma.create(
+            num_stages=self.q_stage,
+            producer_group=make_thread_cooperative_group(len([self.load_warp_id])),
+            consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            tx_count=self.tma_copy_q_bytes,
+            barrier_storage=storage.load_q_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        load_k_producer, load_k_consumer = pipeline.PipelineTmaUmma.create(
+            num_stages=self.k_stage,
+            producer_group=make_thread_cooperative_group(len([self.load_warp_id])),
+            consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            tx_count=self.tma_copy_k_bytes,
+            barrier_storage=storage.load_k_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        load_v_producer, load_v_consumer = pipeline.PipelineTmaUmma.create(
+            num_stages=self.v_stage,
+            producer_group=make_thread_cooperative_group(len([self.load_warp_id])),
+            consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            tx_count=self.tma_copy_v_bytes,
+            barrier_storage=storage.load_v_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        load_do_producer, load_do_consumer = pipeline.PipelineTmaUmma.create(
+            num_stages=self.do_stage,
+            producer_group=make_thread_cooperative_group(len([self.load_warp_id])),
+            consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            tx_count=self.tma_copy_do_bytes,
+            barrier_storage=storage.load_do_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        load_kt_producer, load_kt_consumer = pipeline.PipelineTmaUmma.create(
+            num_stages=self.kt_stage,
+            producer_group=make_thread_cooperative_group(len([self.load_warp_id])),
+            consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            tx_count=self.tma_copy_kt_bytes,
+            barrier_storage=storage.load_kt_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        mma_s_producer, mma_s_consumer = pipeline.PipelineUmmaAsync.create(
+            num_stages=self.qk_acc_stage,
+            producer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            consumer_group=make_thread_cooperative_group(
+                len(self.compute_warp_ids) * self.threads_per_warp * self.cluster_shape_mnk[0],
+            ),
+            barrier_storage=storage.mma_s_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        mma_dp_producer, mma_dp_consumer = pipeline.PipelineUmmaAsync.create(
+            num_stages=self.dov_acc_stage,
+            producer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            consumer_group=make_thread_cooperative_group(
+                len(self.compute_warp_ids) * self.threads_per_warp * self.cluster_shape_mnk[0],
+            ),
+            barrier_storage=storage.mma_dp_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        ds_mma_producer, ds_mma_consumer = pipeline.PipelineAsyncUmma.create(
+            num_stages=self.dsk_acc_stage,
+            producer_group=make_thread_cooperative_group(
+                len(self.compute_warp_ids) * self.threads_per_warp * self.cluster_shape_mnk[0],
+            ),
+            consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            barrier_storage=storage.ds_mma_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        mma_dq_producer, mma_dq_consumer = pipeline.PipelineUmmaAsync.create(
+            num_stages=self.epi_stage,
+            producer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            consumer_group=make_thread_cooperative_group(
+                len(self.epilogue_warp_ids) * self.threads_per_warp * self.cluster_shape_mnk[0],
+            ),
+            barrier_storage=storage.mma_dq_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+
+        load_lse_producer, load_lse_consumer = pipeline.PipelineCpAsync.create(
+            num_stages=self.load_compute_LSE_stage,
+            producer_group=make_thread_cooperative_group(self.threads_per_warp),
+            consumer_group=make_thread_cooperative_group(
+                self.threads_per_warp * self.num_compute_warps
+            ),
+            barrier_storage=storage.lse_mbar_ptr.data_ptr(),
+        ).make_participants()
+        load_sum_odo_producer, load_sum_odo_consumer = pipeline.PipelineCpAsync.create(
+            num_stages=self.load_compute_sum_OdO_stage,
+            producer_group=make_thread_cooperative_group(self.threads_per_warp),
+            consumer_group=make_thread_cooperative_group(
+                self.threads_per_warp * self.num_compute_warps
+            ),
+            barrier_storage=storage.sum_odo_mbar_ptr.data_ptr(),
+        ).make_participants()
+
+        # Tensor memory dealloc barrier init
+        tmem = utils.TmemAllocator(
+            storage.tmem_holding_buf.ptr,
+            barrier_for_retrieve=self.tmem_alloc_barrier,
+            allocator_warp_id=self.epilogue_warp_ids[0],
+            is_two_cta=True,
+            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
+        )
+        tmem.allocate(self.tmem_alloc_cols)
+        tmem.wait_for_alloc()
+        tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
+
+        # Initialize CLC state if using dynamic scheduler
+        if cutlass.const_expr(self.use_clc_scheduler):
+            clc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+            cluster_size = cute.size(self.cluster_shape_mnk)
+            num_clc_consumer_threads = self.threads_per_warp * (
+                1  # sched_warp (CTA 0 only)
+                + cluster_size
+                * (
+                    len(self.compute_warp_ids)
+                    + len(self.epilogue_warp_ids)
+                    + 1  # mma_warp
+                    + 1  # load_warp
+                )
+            )
+            clc_pipeline_consumer_group = pipeline.CooperativeGroup(
+                pipeline.Agent.Thread, num_clc_consumer_threads
+            )
+            clc_response_ptr = storage.clc_response.data_ptr()
+            clc = SchedulerState.create_clc(
+                hw_scheduler=ClcDynamicPersistentTileScheduler.create(
+                    self.tile_sched_params.clc_hw_params(),
+                    cute.arch.block_idx(),
+                    cute.arch.grid_dim(),
+                    clc_response_ptr,
+                ),
+                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,
+                ),
+                consumer_state=pipeline.make_pipeline_state(
+                    pipeline.PipelineUserType.Consumer, self.num_clc_stage
+                ),
+                producer_state=pipeline.make_pipeline_state(
+                    pipeline.PipelineUserType.Producer, self.num_clc_stage
+                ),
+            )
+        else:
+            clc = None
+            clc_response_ptr = None
+
+        # Cluster arrive after barrier init
+        pipeline.pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True)
+
+        sQ = smem.allocate_tensor(
+            element_type=self.q_dtype,
+            layout=q_smem_layout_staged.outer,
+            swizzle=q_smem_layout_staged.inner,
+            byte_alignment=128,
+        )
+        sK = smem.allocate_tensor(
+            element_type=self.k_dtype,
+            layout=k_smem_layout_staged.outer,
+            swizzle=k_smem_layout_staged.inner,
+            byte_alignment=128,
+        )
+        # K and V now use separate memory since we removed the transform stage
+        sV = smem.allocate_tensor(
+            element_type=self.v_dtype,
+            layout=v_smem_layout_staged.outer,
+            swizzle=v_smem_layout_staged.inner,
+            byte_alignment=128,
+        )
+        sdO = smem.allocate_tensor(
+            element_type=self.do_dtype,
+            layout=do_smem_layout_staged.outer,
+            swizzle=do_smem_layout_staged.inner,
+            byte_alignment=128,
+        )
+        sKT = smem.allocate_tensor(
+            element_type=self.k_dtype,
+            layout=kt_smem_layout_staged.outer,
+            swizzle=kt_smem_layout_staged.inner,
+            byte_alignment=128,
+        )
+        sLSE = smem.allocate_tensor(
+            element_type=self.acc_dtype,
+            layout=lse_smem_layout,
+            byte_alignment=128,
+        )
+        sSum_OdO = smem.allocate_tensor(
+            element_type=self.acc_dtype,
+            layout=sum_odo_smem_layout,
+            byte_alignment=128,
+        )
+        # Alias the dQ TMA epilogue staging buffer onto sdO.  A standalone
+        # (128, 64) bf16 allocation exceeds B200's per-block SMEM limit for
+        # this kernel, while sdO is no longer needed by the epilogue warp once
+        # the dQ accumulator is ready to store.
+        s_epi_dQ = cute.make_tensor(
+            cute.recast_ptr(sdO.iterator, sdQ_epi_layout.inner, self.dq_dtype),
+            sdQ_epi_layout.outer,
+        )
+        qk_thr_mma = qk_tiled_mma.get_slice(mma_tile_coord_v)  # default 1sm
+        dov_thr_mma = dov_tiled_mma.get_slice(mma_tile_coord_v)  # default 1sm
+        dsk_thr_mma = dsk_tiled_mma.get_slice(mma_tile_coord_v)  # default 1sm
+        tSrQ = qk_thr_mma.make_fragment_A(sQ)
+        tSrK = qk_thr_mma.make_fragment_B(sK)
+        tdPrdO = dov_thr_mma.make_fragment_A(sdO)
+        tdPrV = dov_thr_mma.make_fragment_B(sV)
+        tdQrKT = dsk_thr_mma.make_fragment_B(sKT)
+        qk_acc_shape = qk_thr_mma.partition_shape_C((self.qk_mma_tiler[0], self.qk_mma_tiler[1]))
+        tStS = qk_thr_mma.make_fragment_C(cute.append(qk_acc_shape, self.qk_acc_stage))
+        dov_acc_shape = dov_thr_mma.partition_shape_C(
+            (self.dov_mma_tiler[0], self.dov_mma_tiler[1])
+        )
+        tdPtdP = dov_thr_mma.make_fragment_C(cute.append(dov_acc_shape, self.dov_acc_stage))
+        dsk_acc_shape = dsk_thr_mma.partition_shape_C(
+            (self.dsk_mma_tiler[0], self.dsk_mma_tiler[1])
+        )
+        tdQtdQ = dsk_thr_mma.make_fragment_C(dsk_acc_shape)
+        tdQtdQ_layout = cute.append(
+            tdQtdQ.layout,
+            cute.make_layout(
+                self.iterations_dsk,
+                stride=self.dsk_mma_tiler[1] // self.tmem_warp_shape_mn[1],
+            ),
+        )
+        tStS = cute.make_tensor(tStS.iterator + self.tmem_s_offset, tStS.layout)
+        tdPtdP = cute.make_tensor(tdPtdP.iterator + self.tmem_dp_offset, tdPtdP.layout)
+        tdQtdQ_staged = cute.make_tensor(tdQtdQ.iterator + self.tmem_dq_offset, tdQtdQ_layout)
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        #  EMPTY
+        # ///////////////////////////////////////////////////////////////////////////////
+        for _i in cutlass.range_constexpr(len(self.empty_warp_id)):
+            if warp_idx == self.empty_warp_id[_i]:
+                cute.arch.warpgroup_reg_dealloc(self.num_regs_other)
+
+        if cutlass.const_expr(self.use_clc_scheduler):
+            tile_sched = FmhaClcDynamicTileScheduler.create(
+                tile_sched_params,
+                cute.arch.block_idx(),
+                cute.arch.grid_dim(),
+                clc_response_ptr,
+                clc,
+            )
+        else:
+            blk_idx = cute.arch.block_idx()
+            tile_sched = FmhaStaticTileScheduler(
+                tile_sched_params, blk_idx[0], blk_idx, cute.arch.grid_dim()
+            )
+        work_tile = tile_sched.initial_work_tile_info()
+
+        # Cluster wait
+        pipeline.pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk)
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        #  LOAD
+        # ///////////////////////////////////////////////////////////////////////////////
+        if warp_idx == self.load_warp_id:
+            cute.arch.warpgroup_reg_dealloc(self.num_regs_other)
+
+            while work_tile.is_valid_tile:
+                curr_block_coord = work_tile.tile_idx
+                mma_block_coord = (
+                    curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape),
+                    curr_block_coord[1],
+                    curr_block_coord[2],
+                )
+                batch_coord = curr_block_coord[2][1]
+                seqlen_q = mQ_qdl.shape[0]
+                seqlen_k = mK_kdl.shape[0]
+                cuseqlen_q = Int32(0)
+                cuseqlen_k = Int32(0)
+                is_valid_q = True
+                if cutlass.const_expr(cum_seqlen_q is not None):
+                    cuseqlen_q = cum_seqlen_q[batch_coord]
+                    seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q
+                    is_valid_q = FmhaStaticTileScheduler.check_valid_work_for_seqlen_q(
+                        self.qk_mma_tiler[0],
+                        mma_block_coord[0],
+                        seqlen_q,
+                    )
+                if cutlass.const_expr(cum_seqlen_k is not None):
+                    cuseqlen_k = cum_seqlen_k[batch_coord]
+                    seqlen_k = cum_seqlen_k[batch_coord + 1] - cuseqlen_k
+                seqlen_kv_loop_start, seqlen_kv_loop_steps = (
+                    FusedMask.get_trip_start_count_via_block_info(
+                        mma_block_coord,
+                        self.qk_mma_tiler,
+                        seqlen_q,
+                        seqlen_k,
+                        self.is_causal,
+                        self.is_local,
+                        window_size_left,
+                        window_size_right,
+                    )
+                )
+                is_valid_k = seqlen_kv_loop_steps > 0
+                has_work = is_valid_q and is_valid_k
+
+                if has_work:
+                    block_offset = (
+                        cuseqlen_q,
+                        cuseqlen_k,
+                        Int32(0),
+                        ((Int32(0), Int32(0)), Int32(0)),
+                    )
+                    mQ_qdl_ = cute.domain_offset(cute.select(block_offset, mode=[0, 2, 3]), mQ_qdl)
+                    mK_kdl_ = cute.domain_offset(cute.select(block_offset, mode=[1, 2, 3]), mK_kdl)
+                    mdO_qdl_ = cute.domain_offset(
+                        cute.select(block_offset, mode=[0, 2, 3]), mdO_qdl
+                    )
+                    mV_dkl_ = cute.domain_offset(cute.select(block_offset, mode=[1, 2, 3]), mV_dkl)
+                    mK_dkl_ = cute.domain_offset(cute.select(block_offset, mode=[2, 1, 3]), mK_dkl)
+                    block_offset_stats = block_offset
+                    if cutlass.const_expr(cum_seqlen_q is not None):
+                        cuseqlen_q_stats = cute.assume(
+                            (cuseqlen_q + batch_coord * self.cta_tiler[0])
+                            // self.cta_tiler[0]
+                            * self.cta_tiler[0],
+                            divby=self.cta_tiler[0],
+                        )
+                        block_offset_stats = (
+                            cuseqlen_q_stats,
+                            block_offset[1],
+                            block_offset[2],
+                            block_offset[3],
+                        )
+                    LSE = cute.domain_offset(cute.select(block_offset_stats, mode=[0, 3]), mLSE)
+                    sum_OdO = cute.domain_offset(
+                        cute.select(block_offset_stats, mode=[0, 3]), mSum_OdO
+                    )
+
+                    # Local tile partition global tensors
+                    q_cta_layout = cute.make_layout(
+                        cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape
+                    )
+                    # (bM, bK, loopM, loopK, loopL)
+                    gQ_qdl = cute.flat_divide(mQ_qdl_, cute.select(self.qk_mma_tiler, mode=[0, 2]))
+                    tSgQ_qdl = qk_thr_mma.partition_A(gQ_qdl)
+                    tQsQ, tQgQ_qdl = cute.nvgpu.cpasync.tma_partition(
+                        tma_atom_q,
+                        block_in_cluster_coord_vmnk[2],
+                        q_cta_layout,
+                        cute.group_modes(sQ, 0, 3),
+                        cute.group_modes(tSgQ_qdl, 0, 3),
+                    )
+                    k_cta_layout = cute.make_layout(
+                        cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape
+                    )
+                    gK_kdl = cute.flat_divide(mK_kdl_, cute.select(self.qk_mma_tiler, mode=[1, 2]))
+                    tSgK_kdl = qk_thr_mma.partition_B(gK_kdl)
+                    tKsK, tKgK_kdl = cute.nvgpu.cpasync.tma_partition(
+                        tma_atom_k,
+                        block_in_cluster_coord_vmnk[1],
+                        k_cta_layout,
+                        cute.group_modes(sK, 0, 3),
+                        cute.group_modes(tSgK_kdl, 0, 3),
+                    )
+                    do_cta_layout = cute.make_layout(
+                        cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape
+                    )
+                    # (bM, bK, loopM, loopK, loopL)
+                    gdO_qdl = cute.flat_divide(
+                        mdO_qdl_, cute.select(self.dov_mma_tiler, mode=[0, 2])
+                    )
+                    tdPgdO_qdl = dov_thr_mma.partition_A(gdO_qdl)
+                    tdOsdO, tdOgdO_qdl = cute.nvgpu.cpasync.tma_partition(
+                        tma_atom_do,
+                        block_in_cluster_coord_vmnk[2],
+                        do_cta_layout,
+                        cute.group_modes(sdO, 0, 3),
+                        cute.group_modes(tdPgdO_qdl, 0, 3),
+                    )
+                    v_cta_layout = cute.make_layout(
+                        cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape
+                    )
+                    gV_dkl = cute.flat_divide(mV_dkl_, cute.select(self.dov_mma_tiler, mode=[1, 2]))
+                    tSgV_dkl = dov_thr_mma.partition_B(gV_dkl)
+                    tVsV, tVgV_dkl = cute.nvgpu.cpasync.tma_partition(
+                        tma_atom_v,
+                        block_in_cluster_coord_vmnk[1],
+                        v_cta_layout,
+                        cute.group_modes(sV, 0, 3),
+                        cute.group_modes(tSgV_dkl, 0, 3),
+                    )
+                    # kt layout
+                    kt_cta_layout = cute.make_layout(
+                        cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape
+                    )
+                    gK_dkl = cute.flat_divide(mK_dkl_, cute.select(self.dsk_mma_tiler, mode=[1, 2]))
+                    tdQgK_dkl = dsk_thr_mma.partition_B(gK_dkl)
+                    tKTsKT, tKgK_dkl = cute.nvgpu.cpasync.tma_partition(
+                        tma_atom_kt,
+                        block_in_cluster_coord_vmnk[1],
+                        kt_cta_layout,
+                        cute.group_modes(sKT, 0, 3),
+                        cute.group_modes(tdQgK_dkl, 0, 3),
+                    )
+                    # ((atom_v, rest_v), RestK)
+                    tQgQ = tQgQ_qdl[None, mma_block_coord[0], None, mma_block_coord[2]]
+                    # ((atom_v, rest_v), RestK)
+                    tdOgdO = tdOgdO_qdl[None, mma_block_coord[0], None, mma_block_coord[2]]
+                    # ((atom_v, rest_v), RestN, RestK)
+                    tKgK = tKgK_kdl[None, None, None, mma_block_coord[2]]
+                    # ((atom_v, rest_v), RestN, RestK)
+                    tVgV = tVgV_dkl[None, None, None, mma_block_coord[2]]
+                    # ((atom_v, rest_v), RestN, RestK)
+                    tKTgKT = tKgK_dkl[None, None, None, mma_block_coord[2]]
+
+                    # LSE
+                    lse_handle = load_lse_producer.acquire_and_advance()
+                    # 32 threads loading 128 values of 32b each
+                    # so 4*32b = 128b
+                    thread_idx = tidx % self.threads_per_warp
+                    async_copy_num_elts = sLSE.shape[0] // self.threads_per_warp
+                    atom_async_copy = cute.make_copy_atom(
+                        cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.ALWAYS),
+                        self.acc_dtype,
+                        num_bits_per_copy=self.acc_dtype.width,
+                    )
+                    sLSE_for_copy = cute.flat_divide(sLSE, (1,))
+                    LSE_for_copy = cute.flat_divide(LSE, (1,))
+                    for i in cutlass.range_constexpr(async_copy_num_elts):
+                        LSE_idx = (
+                            self.cta_tiler[0] * curr_block_coord[0]
+                            + thread_idx * async_copy_num_elts
+                        )
+                        if cute.elem_less(LSE_idx + i, seqlen_q):
+                            cute.copy(
+                                atom_async_copy,
+                                LSE_for_copy[None, LSE_idx + i, curr_block_coord[2]],
+                                sLSE_for_copy[
+                                    None,
+                                    thread_idx * async_copy_num_elts + i,
+                                    lse_handle.index,
+                                ],
+                            )
+                        else:
+                            sLSE_for_copy[
+                                None,
+                                thread_idx * async_copy_num_elts + i,
+                                lse_handle.index,
+                            ].fill(0.0)
+                    lse_handle.commit()
+
+                    sum_odo_handle = load_sum_odo_producer.acquire_and_advance()
+                    sSum_OdO_for_copy = cute.flat_divide(sSum_OdO, (1,))
+                    sum_OdO_for_copy = cute.flat_divide(sum_OdO, (1,))
+                    for i in cutlass.range_constexpr(async_copy_num_elts):
+                        sum_OdO_idx = (
+                            self.cta_tiler[0] * curr_block_coord[0]
+                            + thread_idx * async_copy_num_elts
+                        )
+                        if cute.elem_less(sum_OdO_idx + i, seqlen_q):
+                            cute.copy(
+                                atom_async_copy,
+                                sum_OdO_for_copy[None, sum_OdO_idx + i, curr_block_coord[2]],
+                                sSum_OdO_for_copy[
+                                    None,
+                                    thread_idx * async_copy_num_elts + i,
+                                    sum_odo_handle.index,
+                                ],
+                            )
+                        else:
+                            sSum_OdO_for_copy[
+                                None,
+                                thread_idx * async_copy_num_elts + i,
+                                sum_odo_handle.index,
+                            ].fill(0.0)
+                    sum_odo_handle.commit()
+
+                    # Q
+                    for iter in cutlass.range(self.iterations_qk, unroll=1):
+                        q_handle = load_q_producer.acquire_and_advance()
+                        cute.copy(
+                            tma_atom_q,
+                            tQgQ[None, iter],
+                            tQsQ[None, q_handle.index],
+                            tma_bar_ptr=q_handle.barrier,
+                        )
+                    # dO
+                    for iter in cutlass.range(self.iterations_dov, unroll=1):
+                        do_handle = load_do_producer.acquire_and_advance()
+                        cute.copy(
+                            tma_atom_do,
+                            tdOgdO[None, iter],
+                            tdOsdO[None, do_handle.index],
+                            tma_bar_ptr=do_handle.barrier,
+                        )
+
+                    kv_coord = seqlen_kv_loop_start
+                    for i in cutlass.range(0, seqlen_kv_loop_steps, 1, unroll=1):
+                        # Ki
+                        for iter in cutlass.range(self.iterations_qk, unroll=1):
+                            k_handle = load_k_producer.acquire_and_advance()
+                            cute.copy(
+                                tma_atom_k,
+                                tKgK[None, kv_coord, iter],
+                                tKsK[None, k_handle.index],
+                                tma_bar_ptr=k_handle.barrier,
+                            )
+                        # Vi
+                        for iter in cutlass.range(self.iterations_dov, unroll=1):
+                            v_handle = load_v_producer.acquire_and_advance()
+                            cute.copy(
+                                tma_atom_v,
+                                tVgV[None, kv_coord, iter],
+                                tVsV[None, v_handle.index],
+                                tma_bar_ptr=v_handle.barrier,
+                            )
+                        # KTi
+                        for iter in cutlass.range(self.iterations_dsk, unroll=1):
+                            kt_handle = load_kt_producer.acquire_and_advance()
+                            cute.copy(
+                                tma_atom_kt,
+                                tKTgKT[None, iter, kv_coord],
+                                tKTsKT[None, kt_handle.index],
+                                tma_bar_ptr=kt_handle.barrier,
+                            )
+                        kv_coord += 1
+
+                work_tile = tile_sched.advance_to_next_work()
+                # End of persistent scheduler loop
+            load_k_producer.tail()
+            load_v_producer.tail()
+            load_kt_producer.tail()
+            load_q_producer.tail()
+            load_do_producer.tail()
+            load_lse_producer.tail()
+            load_sum_odo_producer.tail()
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        #  MMA
+        # ///////////////////////////////////////////////////////////////////////////////
+        if warp_idx == self.mma_warp_id:
+            cute.arch.warpgroup_reg_dealloc(self.num_regs_other)
+
+            cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+            is_leader_cta = cta_rank_in_cluster % 2 == 0
+
+            while work_tile.is_valid_tile:
+                curr_block_coord = work_tile.tile_idx
+                mma_block_coord = (
+                    curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape),
+                    curr_block_coord[1],
+                    curr_block_coord[2],
+                )
+                seqlen_q = mQ_qdl.shape[0]
+                seqlen_k = mK_kdl.shape[0]
+                batch_coord = curr_block_coord[2][1]
+                is_valid_q = True
+                if cutlass.const_expr(cum_seqlen_q is not None):
+                    cuseqlen_q = cum_seqlen_q[batch_coord]
+                    seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q
+                    is_valid_q = FmhaStaticTileScheduler.check_valid_work_for_seqlen_q(
+                        self.qk_mma_tiler[0],
+                        mma_block_coord[0],
+                        seqlen_q,
+                    )
+                if cutlass.const_expr(cum_seqlen_k is not None):
+                    cuseqlen_k = cum_seqlen_k[batch_coord]
+                    seqlen_k = cum_seqlen_k[batch_coord + 1] - cuseqlen_k
+                seqlen_kv_loop_start, seqlen_kv_loop_steps = (
+                    FusedMask.get_trip_start_count_via_block_info(
+                        mma_block_coord,
+                        self.qk_mma_tiler,
+                        seqlen_q,
+                        seqlen_k,
+                        self.is_causal,
+                        self.is_local,
+                        window_size_left,
+                        window_size_right,
+                    )
+                )
+                is_valid_k = seqlen_kv_loop_steps > 0
+                has_work = is_valid_q and is_valid_k
+
+                if has_work:
+                    # dq_handle = mma_dq_producer.acquire_and_advance()
+                    load_q_releaser = load_q_consumer.clone()
+                    load_do_releaser = load_do_consumer.clone()
+                    dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+
+                    num_innerloop = 8
+
+                    if is_leader_cta:
+                        dq_handle = mma_dq_producer.acquire_and_advance()
+                        if seqlen_kv_loop_steps > 1:
+                            # QK0
+                            s_handle = mma_s_producer.acquire_and_advance()
+                            tStS_slice = tStS[None, None, None, s_handle.index]
+                            qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                            for iter in cutlass.range(self.iterations_qk, unroll=1):
+                                load_q_consumer.wait_and_advance()
+                                tSrQ_slice = tSrQ[None, None, None, iter]
+
+                                k_handle = load_k_consumer.wait_and_advance()
+                                tSrK_trans_slice = tSrK[None, None, None, k_handle.index]
+                                num_kphases = cute.size(tSrQ_slice, mode=[2])
+                                if cutlass.const_expr(num_kphases % num_innerloop == 0):
+                                    num_outer_iter = num_kphases // num_innerloop
+                                    for outer_iter in cutlass.range(num_outer_iter, unroll=1):
+                                        for kphase_idx in cutlass.range(
+                                            num_innerloop, unroll_full=True
+                                        ):
+                                            kphase_coord = (
+                                                None,
+                                                None,
+                                                outer_iter * num_innerloop + kphase_idx,
+                                            )
+                                            cute.gemm(
+                                                qk_tiled_mma,
+                                                tStS_slice,
+                                                tSrQ_slice[kphase_coord],
+                                                tSrK_trans_slice[kphase_coord],
+                                                tStS_slice,
+                                            )
+                                            qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                else:
+                                    for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
+                                        kphase_coord = (None, None, kphase_idx)
+                                        cute.gemm(
+                                            qk_tiled_mma,
+                                            tStS_slice,
+                                            tSrQ_slice[kphase_coord],
+                                            tSrK_trans_slice[kphase_coord],
+                                            tStS_slice,
+                                        )
+                                        qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                k_handle.release()
+                            cute.arch.fence_view_async_tmem_store()
+                            s_handle.commit()
+
+                            # dOV0
+                            dp_handle = mma_dp_producer.acquire_and_advance()
+                            tdPtdP_slice = tdPtdP[None, None, None, dp_handle.index]
+                            dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                            for iter in cutlass.range(self.iterations_dov, unroll=1):
+                                load_do_consumer.wait_and_advance()
+                                tdPrdO_slice = tdPrdO[None, None, None, iter]
+                                v_handle = load_v_consumer.wait_and_advance()
+                                tdPrV_trans_slice = tdPrV[None, None, None, v_handle.index]
+                                num_kphases = cute.size(tdPrdO_slice, mode=[2])
+                                if cutlass.const_expr(num_kphases % num_innerloop == 0):
+                                    num_outer_iter = num_kphases // num_innerloop
+                                    for outer_iter in cutlass.range(num_outer_iter, unroll=1):
+                                        for kphase_idx in cutlass.range(
+                                            num_innerloop, unroll_full=True
+                                        ):
+                                            kphase_coord = (
+                                                None,
+                                                None,
+                                                outer_iter * num_innerloop + kphase_idx,
+                                            )
+                                            cute.gemm(
+                                                dov_tiled_mma,
+                                                tdPtdP_slice,
+                                                tdPrdO_slice[kphase_coord],
+                                                tdPrV_trans_slice[kphase_coord],
+                                                tdPtdP_slice,
+                                            )
+                                            dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                else:
+                                    for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
+                                        kphase_coord = (None, None, kphase_idx)
+                                        cute.gemm(
+                                            dov_tiled_mma,
+                                            tdPtdP_slice,
+                                            tdPrdO_slice[kphase_coord],
+                                            tdPrV_trans_slice[kphase_coord],
+                                            tdPtdP_slice,
+                                        )
+                                        dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                v_handle.release()
+                            cute.arch.fence_view_async_tmem_store()
+                            dp_handle.commit()
+
+                            for i in cutlass.range(1, seqlen_kv_loop_steps - 1, 1, unroll=1):
+                                # QKi
+                                s_handle = mma_s_producer.acquire_and_advance()
+
+                                tStS_slice = tStS[None, None, None, s_handle.index]
+                                qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                                for iter in cutlass.range(self.iterations_qk, unroll=1):
+                                    tSrQ_slice = tSrQ[None, None, None, iter]
+                                    k_handle = load_k_consumer.wait_and_advance()
+                                    tSrK_trans_slice = tSrK[None, None, None, k_handle.index]
+                                    num_kphases = cute.size(tSrQ_slice, mode=[2])
+                                    if cutlass.const_expr(num_kphases % num_innerloop == 0):
+                                        num_outer_iter = num_kphases // num_innerloop
+                                        for outer_iter in cutlass.range(num_outer_iter, unroll=1):
+                                            for kphase_idx in cutlass.range(
+                                                num_innerloop, unroll_full=True
+                                            ):
+                                                kphase_coord = (
+                                                    None,
+                                                    None,
+                                                    outer_iter * num_innerloop + kphase_idx,
+                                                )
+                                                cute.gemm(
+                                                    qk_tiled_mma,
+                                                    tStS_slice,
+                                                    tSrQ_slice[kphase_coord],
+                                                    tSrK_trans_slice[kphase_coord],
+                                                    tStS_slice,
+                                                )
+                                                qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                    else:
+                                        for kphase_idx in cutlass.range(
+                                            num_kphases, unroll_full=True
+                                        ):
+                                            kphase_coord = (None, None, kphase_idx)
+                                            cute.gemm(
+                                                qk_tiled_mma,
+                                                tStS_slice,
+                                                tSrQ_slice[kphase_coord],
+                                                tSrK_trans_slice[kphase_coord],
+                                                tStS_slice,
+                                            )
+                                            qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                    k_handle.release()
+                                s_handle.commit()
+
+                                # dSKTi
+                                ds_handle = ds_mma_consumer.wait_and_advance()
+                                dsk_whether_acc = dsk_tiled_mma.get(tcgen05.Field.ACCUMULATE)
+                                for iter in cutlass.range(self.iterations_dsk, unroll=1):
+                                    kt_handle = load_kt_consumer.wait_and_advance()
+                                    dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, dsk_whether_acc)
+                                    tdQtdQ_slice = tdQtdQ_staged[None, None, None, iter]
+                                    tdStdS_slice = tdPtdP[None, None, None, ds_handle.index]
+                                    tdS = cute.make_tensor(
+                                        tdStdS_slice.iterator, ds_tmem_layout_staged.outer
+                                    )
+                                    tdQrdS = dsk_thr_mma.make_fragment_A(tdS)
+                                    tdQrdS_slice = cute.make_tensor(
+                                        cute.recast_ptr(tdStdS_slice.iterator, dtype=self.q_dtype),
+                                        tdQrdS.layout,
+                                    )
+
+                                    tdQrKT_slice = tdQrKT[None, None, None, kt_handle.index]
+                                    num_kphases = cute.size(tdQrKT_slice, mode=[2])
+                                    if cutlass.const_expr(num_kphases % num_innerloop == 0):
+                                        num_outer_iter = num_kphases // num_innerloop
+                                        for outer_iter in cutlass.range(num_outer_iter, unroll=1):
+                                            for kphase_idx in cutlass.range(
+                                                num_innerloop, unroll_full=True
+                                            ):
+                                                kphase_coord = (
+                                                    None,
+                                                    None,
+                                                    outer_iter * num_innerloop + kphase_idx,
+                                                )
+                                                cute.gemm(
+                                                    dsk_tiled_mma,
+                                                    tdQtdQ_slice,
+                                                    tdQrdS_slice[kphase_coord],
+                                                    tdQrKT_slice[kphase_coord],
+                                                    tdQtdQ_slice,
+                                                )
+                                                dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                    else:
+                                        for kphase_idx in cutlass.range(
+                                            num_kphases, unroll_full=True
+                                        ):
+                                            kphase_coord = (None, None, kphase_idx)
+                                            cute.gemm(
+                                                dsk_tiled_mma,
+                                                tdQtdQ_slice,
+                                                tdQrdS_slice[kphase_coord],
+                                                tdQrKT_slice[kphase_coord],
+                                                tdQtdQ_slice,
+                                            )
+                                            dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                    kt_handle.release()
+                                ds_handle.release()
+
+                                # dOVi
+                                dp_handle = mma_dp_producer.acquire_and_advance()
+                                tdPtdP_slice = tdPtdP[None, None, None, dp_handle.index]
+                                dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                                for iter in cutlass.range(self.iterations_dov, unroll=1):
+                                    tdPrdO_slice = tdPrdO[None, None, None, iter]
+                                    v_handle = load_v_consumer.wait_and_advance()
+                                    tdPrV_trans_slice = tdPrV[None, None, None, v_handle.index]
+                                    num_kphases = cute.size(tdPrdO_slice, mode=[2])
+                                    if cutlass.const_expr(num_kphases % num_innerloop == 0):
+                                        num_outer_iter = num_kphases // num_innerloop
+                                        for outer_iter in cutlass.range(num_outer_iter, unroll=1):
+                                            for kphase_idx in cutlass.range(
+                                                num_innerloop, unroll_full=True
+                                            ):
+                                                kphase_coord = (
+                                                    None,
+                                                    None,
+                                                    outer_iter * num_innerloop + kphase_idx,
+                                                )
+                                                cute.gemm(
+                                                    dov_tiled_mma,
+                                                    tdPtdP_slice,
+                                                    tdPrdO_slice[kphase_coord],
+                                                    tdPrV_trans_slice[kphase_coord],
+                                                    tdPtdP_slice,
+                                                )
+                                                dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                    else:
+                                        for kphase_idx in cutlass.range(
+                                            num_kphases, unroll_full=True
+                                        ):
+                                            kphase_coord = (None, None, kphase_idx)
+                                            cute.gemm(
+                                                dov_tiled_mma,
+                                                tdPtdP_slice,
+                                                tdPrdO_slice[kphase_coord],
+                                                tdPrV_trans_slice[kphase_coord],
+                                                tdPtdP_slice,
+                                            )
+                                            dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                    v_handle.release()
+                                dp_handle.commit()
+
+                            # QKend
+                            s_handle = mma_s_producer.acquire_and_advance()
+                            tStS_slice = tStS[None, None, None, s_handle.index]
+                            qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                            for iter in cutlass.range(self.iterations_qk, unroll=1):
+                                tSrQ_slice = tSrQ[None, None, None, iter]
+                                k_handle = load_k_consumer.wait_and_advance()
+
+                                tSrK_trans_slice = tSrK[None, None, None, k_handle.index]
+
+                                num_kphases = cute.size(tSrQ_slice, mode=[2])
+                                if cutlass.const_expr(num_kphases % num_innerloop == 0):
+                                    num_outer_iter = num_kphases // num_innerloop
+                                    for outer_iter in cutlass.range(num_outer_iter, unroll=1):
+                                        for kphase_idx in cutlass.range(
+                                            num_innerloop, unroll_full=True
+                                        ):
+                                            kphase_coord = (
+                                                None,
+                                                None,
+                                                outer_iter * num_innerloop + kphase_idx,
+                                            )
+                                            cute.gemm(
+                                                qk_tiled_mma,
+                                                tStS_slice,
+                                                tSrQ_slice[kphase_coord],
+                                                tSrK_trans_slice[kphase_coord],
+                                                tStS_slice,
+                                            )
+                                            qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                else:
+                                    for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
+                                        kphase_coord = (None, None, kphase_idx)
+                                        cute.gemm(
+                                            qk_tiled_mma,
+                                            tStS_slice,
+                                            tSrQ_slice[kphase_coord],
+                                            tSrK_trans_slice[kphase_coord],
+                                            tStS_slice,
+                                        )
+                                        qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                k_handle.release()
+                                load_q_releaser.release()
+                                load_q_releaser.advance()
+                            s_handle.commit()
+
+                            # dSKTend - 1
+                            ds_handle = ds_mma_consumer.wait_and_advance()
+                            dsk_whether_acc = dsk_tiled_mma.get(tcgen05.Field.ACCUMULATE)
+                            for iter in cutlass.range(self.iterations_dsk, unroll=1):
+                                kt_handle = load_kt_consumer.wait_and_advance()
+                                dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, dsk_whether_acc)
+                                tdQtdQ_slice = tdQtdQ_staged[None, None, None, iter]
+                                tdStdS_slice = tdPtdP[None, None, None, ds_handle.index]
+                                tdS = cute.make_tensor(
+                                    tdStdS_slice.iterator, ds_tmem_layout_staged.outer
+                                )
+                                tdQrdS = dsk_thr_mma.make_fragment_A(tdS)
+                                tdQrdS_slice = cute.make_tensor(
+                                    cute.recast_ptr(tdStdS_slice.iterator, dtype=self.q_dtype),
+                                    tdQrdS.layout,
+                                )
+
+                                tdQrKT_slice = tdQrKT[None, None, None, kt_handle.index]
+                                num_kphases = cute.size(tdQrKT_slice, mode=[2])
+                                if cutlass.const_expr(num_kphases % num_innerloop == 0):
+                                    num_outer_iter = num_kphases // num_innerloop
+                                    for outer_iter in cutlass.range(num_outer_iter, unroll=1):
+                                        for kphase_idx in cutlass.range(
+                                            num_innerloop, unroll_full=True
+                                        ):
+                                            kphase_coord = (
+                                                None,
+                                                None,
+                                                outer_iter * num_innerloop + kphase_idx,
+                                            )
+                                            cute.gemm(
+                                                dsk_tiled_mma,
+                                                tdQtdQ_slice,
+                                                tdQrdS_slice[kphase_coord],
+                                                tdQrKT_slice[kphase_coord],
+                                                tdQtdQ_slice,
+                                            )
+                                            dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                else:
+                                    for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
+                                        kphase_coord = (None, None, kphase_idx)
+                                        cute.gemm(
+                                            dsk_tiled_mma,
+                                            tdQtdQ_slice,
+                                            tdQrdS_slice[kphase_coord],
+                                            tdQrKT_slice[kphase_coord],
+                                            tdQtdQ_slice,
+                                        )
+                                        dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                kt_handle.release()
+                            ds_handle.release()
+
+                            # dOVend
+                            dp_handle = mma_dp_producer.acquire_and_advance()
+                            tdPtdP_slice = tdPtdP[None, None, None, dp_handle.index]
+                            dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                            for iter in cutlass.range(self.iterations_dov, unroll=1):
+                                tdPrdO_slice = tdPrdO[None, None, None, iter]
+                                v_handle = load_v_consumer.wait_and_advance()
+                                tdPrV_trans_slice = tdPrV[None, None, None, v_handle.index]
+                                num_kphases = cute.size(tdPrdO_slice, mode=[2])
+                                if cutlass.const_expr(num_kphases % num_innerloop == 0):
+                                    num_outer_iter = num_kphases // num_innerloop
+                                    for outer_iter in cutlass.range(num_outer_iter, unroll=1):
+                                        for kphase_idx in cutlass.range(
+                                            num_innerloop, unroll_full=True
+                                        ):
+                                            kphase_coord = (
+                                                None,
+                                                None,
+                                                outer_iter * num_innerloop + kphase_idx,
+                                            )
+                                            cute.gemm(
+                                                dov_tiled_mma,
+                                                tdPtdP_slice,
+                                                tdPrdO_slice[kphase_coord],
+                                                tdPrV_trans_slice[kphase_coord],
+                                                tdPtdP_slice,
+                                            )
+                                            dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                else:
+                                    for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
+                                        kphase_coord = (None, None, kphase_idx)
+                                        cute.gemm(
+                                            dov_tiled_mma,
+                                            tdPtdP_slice,
+                                            tdPrdO_slice[kphase_coord],
+                                            tdPrV_trans_slice[kphase_coord],
+                                            tdPtdP_slice,
+                                        )
+                                        dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                v_handle.release()
+                                load_do_releaser.release()
+                                load_do_releaser.advance()
+                            dp_handle.commit()
+                            # dSKTend
+                            ds_handle = ds_mma_consumer.wait_and_advance()
+                            dsk_whether_acc = dsk_tiled_mma.get(tcgen05.Field.ACCUMULATE)
+                            for iter in cutlass.range(self.iterations_dsk, unroll=1):
+                                kt_handle = load_kt_consumer.wait_and_advance()
+                                dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, dsk_whether_acc)
+                                tdQtdQ_slice = tdQtdQ_staged[None, None, None, iter]
+                                tdStdS_slice = tdPtdP[None, None, None, ds_handle.index]
+                                tdS = cute.make_tensor(
+                                    tdStdS_slice.iterator, ds_tmem_layout_staged.outer
+                                )
+                                tdQrdS = dsk_thr_mma.make_fragment_A(tdS)
+                                tdQrdS_slice = cute.make_tensor(
+                                    cute.recast_ptr(tdStdS_slice.iterator, dtype=self.q_dtype),
+                                    tdQrdS.layout,
+                                )
+
+                                tdQrKT_slice = tdQrKT[None, None, None, kt_handle.index]
+                                num_kphases = cute.size(tdQrKT_slice, mode=[2])
+                                if cutlass.const_expr(num_kphases % num_innerloop == 0):
+                                    num_outer_iter = num_kphases // num_innerloop
+                                    for outer_iter in cutlass.range(num_outer_iter, unroll=1):
+                                        for kphase_idx in cutlass.range(
+                                            num_innerloop, unroll_full=True
+                                        ):
+                                            kphase_coord = (
+                                                None,
+                                                None,
+                                                outer_iter * num_innerloop + kphase_idx,
+                                            )
+                                            cute.gemm(
+                                                dsk_tiled_mma,
+                                                tdQtdQ_slice,
+                                                tdQrdS_slice[kphase_coord],
+                                                tdQrKT_slice[kphase_coord],
+                                                tdQtdQ_slice,
+                                            )
+                                            dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                else:
+                                    for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
+                                        kphase_coord = (None, None, kphase_idx)
+                                        cute.gemm(
+                                            dsk_tiled_mma,
+                                            tdQtdQ_slice,
+                                            tdQrdS_slice[kphase_coord],
+                                            tdQrKT_slice[kphase_coord],
+                                            tdQtdQ_slice,
+                                        )
+                                        dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                kt_handle.release()
+                            ds_handle.release()
+                        else:
+                            # QK0
+                            s_handle = mma_s_producer.acquire_and_advance()
+                            tStS_slice = tStS[None, None, None, s_handle.index]
+                            qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+
+                            for iter in cutlass.range(self.iterations_qk, unroll=1):
+                                load_q_consumer.wait_and_advance()
+                                tSrQ_slice = tSrQ[None, None, None, iter]
+                                k_handle = load_k_consumer.wait_and_advance()
+                                tSrK_trans_slice = tSrK[None, None, None, k_handle.index]
+                                num_kphases = cute.size(tSrQ_slice, mode=[2])
+                                if cutlass.const_expr(num_kphases % num_innerloop == 0):
+                                    num_outer_iter = num_kphases // num_innerloop
+                                    for outer_iter in cutlass.range(num_outer_iter, unroll=1):
+                                        for kphase_idx in cutlass.range(
+                                            num_innerloop, unroll_full=True
+                                        ):
+                                            kphase_coord = (
+                                                None,
+                                                None,
+                                                outer_iter * num_innerloop + kphase_idx,
+                                            )
+                                            cute.gemm(
+                                                qk_tiled_mma,
+                                                tStS_slice,
+                                                tSrQ_slice[kphase_coord],
+                                                tSrK_trans_slice[kphase_coord],
+                                                tStS_slice,
+                                            )
+                                            qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                else:
+                                    for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
+                                        kphase_coord = (None, None, kphase_idx)
+                                        cute.gemm(
+                                            qk_tiled_mma,
+                                            tStS_slice,
+                                            tSrQ_slice[kphase_coord],
+                                            tSrK_trans_slice[kphase_coord],
+                                            tStS_slice,
+                                        )
+                                        qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                k_handle.release()
+                                load_q_releaser.release()
+                                load_q_releaser.advance()
+                            s_handle.commit()
+
+                            # dOV0
+                            dp_handle = mma_dp_producer.acquire_and_advance()
+                            tdPtdP_slice = tdPtdP[None, None, None, dp_handle.index]
+                            dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                            for iter in cutlass.range(self.iterations_dov, unroll=1):
+                                load_do_consumer.wait_and_advance()
+                                tdPrdO_slice = tdPrdO[None, None, None, iter]
+                                v_handle = load_v_consumer.wait_and_advance()
+                                tdPrV_trans_slice = tdPrV[None, None, None, v_handle.index]
+                                num_kphases = cute.size(tdPrdO_slice, mode=[2])
+                                if cutlass.const_expr(num_kphases % num_innerloop == 0):
+                                    num_outer_iter = num_kphases // num_innerloop
+                                    for outer_iter in cutlass.range(num_outer_iter, unroll=1):
+                                        for kphase_idx in cutlass.range(
+                                            num_innerloop, unroll_full=True
+                                        ):
+                                            kphase_coord = (
+                                                None,
+                                                None,
+                                                outer_iter * num_innerloop + kphase_idx,
+                                            )
+                                            cute.gemm(
+                                                dov_tiled_mma,
+                                                tdPtdP_slice,
+                                                tdPrdO_slice[kphase_coord],
+                                                tdPrV_trans_slice[kphase_coord],
+                                                tdPtdP_slice,
+                                            )
+                                            dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                else:
+                                    for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
+                                        kphase_coord = (None, None, kphase_idx)
+                                        cute.gemm(
+                                            dov_tiled_mma,
+                                            tdPtdP_slice,
+                                            tdPrdO_slice[kphase_coord],
+                                            tdPrV_trans_slice[kphase_coord],
+                                            tdPtdP_slice,
+                                        )
+                                        dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                v_handle.release()
+                                load_do_releaser.release()
+                                load_do_releaser.advance()
+                            dp_handle.commit()
+
+                            # dSKT0
+                            ds_handle = ds_mma_consumer.wait_and_advance()
+                            dsk_whether_acc = dsk_tiled_mma.get(tcgen05.Field.ACCUMULATE)
+                            for iter in cutlass.range(self.iterations_dsk, unroll=1):
+                                kt_handle = load_kt_consumer.wait_and_advance()
+                                dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, dsk_whether_acc)
+                                tdQtdQ_slice = tdQtdQ_staged[None, None, None, iter]
+                                tdStdS_slice = tdPtdP[None, None, None, ds_handle.index]
+                                tdS = cute.make_tensor(
+                                    tdStdS_slice.iterator, ds_tmem_layout_staged.outer
+                                )
+                                tdQrdS = dsk_thr_mma.make_fragment_A(tdS)
+                                tdQrdS_slice = cute.make_tensor(
+                                    cute.recast_ptr(tdStdS_slice.iterator, dtype=self.q_dtype),
+                                    tdQrdS.layout,
+                                )
+
+                                tdQrKT_slice = tdQrKT[None, None, None, kt_handle.index]
+                                num_kphases = cute.size(tdQrKT_slice, mode=[2])
+                                if cutlass.const_expr(num_kphases % num_innerloop == 0):
+                                    num_outer_iter = num_kphases // num_innerloop
+                                    for outer_iter in cutlass.range(num_outer_iter, unroll=1):
+                                        for kphase_idx in cutlass.range(
+                                            num_innerloop, unroll_full=True
+                                        ):
+                                            kphase_coord = (
+                                                None,
+                                                None,
+                                                outer_iter * num_innerloop + kphase_idx,
+                                            )
+                                            cute.gemm(
+                                                dsk_tiled_mma,
+                                                tdQtdQ_slice,
+                                                tdQrdS_slice[kphase_coord],
+                                                tdQrKT_slice[kphase_coord],
+                                                tdQtdQ_slice,
+                                            )
+                                            dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                else:
+                                    for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
+                                        kphase_coord = (None, None, kphase_idx)
+                                        cute.gemm(
+                                            dsk_tiled_mma,
+                                            tdQtdQ_slice,
+                                            tdQrdS_slice[kphase_coord],
+                                            tdQrKT_slice[kphase_coord],
+                                            tdQtdQ_slice,
+                                        )
+                                        dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                kt_handle.release()
+                            ds_handle.release()
+                        dq_handle.commit()
+                work_tile = tile_sched.advance_to_next_work()
+            # End of persistent scheduler loop
+            mma_s_producer.tail()
+            mma_dp_producer.tail()
+            mma_dq_producer.tail()
+
+        # Softmax and dSoftmax warp
+        if warp_idx >= self.compute_warp_ids[0] and warp_idx <= self.compute_warp_ids[-1]:
+            # increase register after decreasing
+            cute.arch.warpgroup_reg_alloc(self.num_regs_compute)
+            while work_tile.is_valid_tile:
+                curr_block_coord = work_tile.tile_idx
+                mma_block_coord = (
+                    curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape),
+                    curr_block_coord[1],
+                    curr_block_coord[2],
+                )
+                batch_coord = curr_block_coord[2][1]
+                seqlen_q = mQ_qdl.shape[0]
+                seqlen_k = mK_kdl.shape[0]
+                cuseqlen_q = Int32(0)
+                is_valid_q = True
+                if cutlass.const_expr(cum_seqlen_q is not None):
+                    cuseqlen_q = cum_seqlen_q[batch_coord]
+                    seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q
+                    is_valid_q = FmhaStaticTileScheduler.check_valid_work_for_seqlen_q(
+                        self.qk_mma_tiler[0],
+                        mma_block_coord[0],
+                        seqlen_q,
+                    )
+                if cutlass.const_expr(cum_seqlen_k is not None):
+                    cuseqlen_k = cum_seqlen_k[batch_coord]
+                    seqlen_k = cum_seqlen_k[batch_coord + 1] - cuseqlen_k
+                start_count, trip_count = FusedMask.get_trip_start_count_via_block_info(
+                    mma_block_coord,
+                    self.qk_mma_tiler,
+                    seqlen_q,
+                    seqlen_k,
+                    self.is_causal,
+                    self.is_local,
+                    window_size_left,
+                    window_size_right,
+                )
+                is_valid_k = trip_count > 0
+                has_work = is_valid_q and is_valid_k
+
+                if has_work:
+                    end_count = start_count + trip_count
+                    if cutlass.const_expr(self.use_semantic_trip_range):
+                        n_block_min_causal_local_mask, n_block_min_before_local_mask = (
+                            FusedMask.get_trip_mask_bounds_via_block_info(
+                                mma_block_coord,
+                                self.qk_mma_tiler,
+                                seqlen_q,
+                                seqlen_k,
+                                self.is_causal,
+                                self.is_local,
+                                window_size_left,
+                                window_size_right,
+                            )
+                        )
+
+                    cS_base = cute.make_identity_tensor(
+                        (self.qk_mma_tiler[0], self.qk_mma_tiler[1])
+                    )
+                    cS = cute.domain_offset((mma_block_coord[0] * self.qk_mma_tiler[0], 0), cS_base)
+                    tScS = qk_thr_mma.partition_C(cS)
+
+                    cdP_base = cute.make_identity_tensor(
+                        (self.dov_mma_tiler[0], self.dov_mma_tiler[1])
+                    )
+                    cdP = cute.domain_offset(
+                        (mma_block_coord[0] * self.dov_mma_tiler[0], 0), cdP_base
+                    )
+                    tdPcdP = dov_thr_mma.partition_C(cdP)
+
+                    lse_handle = load_lse_consumer.wait_and_advance()
+                    sum_odo_handle = load_sum_odo_consumer.wait_and_advance()
+                    for step in cutlass.range(start_count, end_count, 1, unroll=1):
+                        cS_iter = cute.domain_offset((0, step * self.qk_mma_tiler[1]), cS)
+                        tScS_iter = qk_thr_mma.partition_C(cS_iter)
+
+                        cdP_iter = cute.domain_offset((0, step * self.dov_mma_tiler[1]), cdP)
+
+                        tdPcdP_iter = dov_thr_mma.partition_C(cdP_iter)
+
+                        # Si, dPi -> dSi
+                        if cutlass.const_expr(self.use_semantic_trip_range):
+                            need_apply_mask = (
+                                step >= n_block_min_causal_local_mask
+                                or step < n_block_min_before_local_mask
+                            )
+                        else:
+                            need_apply_mask = step == end_count - 1
+                        mma_s_consumer, mma_dp_consumer, ds_mma_producer = self.compute_step(
+                            (need_apply_mask, window_size_left, window_size_right),
+                            (
+                                seqlen_q,
+                                seqlen_k,
+                                scale_softmax,
+                                batch_coord,
+                                curr_block_coord[0],
+                                varlen,
+                            ),
+                            (tStS, tScS_iter, tdPtdP, tdPcdP_iter, sLSE, sSum_OdO),
+                            (
+                                mma_s_consumer,
+                                mma_dp_consumer,
+                                ds_mma_producer,
+                                lse_handle,
+                                sum_odo_handle,
+                            ),
+                            step,
+                        )
+                    lse_handle.release()
+                    sum_odo_handle.release()
+
+                work_tile = tile_sched.advance_to_next_work()
+            ds_mma_producer.tail()
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        #  Epilogue
+        # ///////////////////////////////////////////////////////////////////////////////
+        if warp_idx >= self.epilogue_warp_ids[0] and warp_idx <= self.epilogue_warp_ids[-1]:
+            cute.arch.warpgroup_reg_dealloc(self.num_regs_epilogue)
+
+            while work_tile.is_valid_tile:
+                curr_block_coord = work_tile.tile_idx
+                mma_block_coord = (
+                    curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape),
+                    curr_block_coord[1],
+                    curr_block_coord[2],
+                )
+                batch_coord = curr_block_coord[2][1]
+                # cute.printf("batch_coord={}", batch_coord)
+                seqlen_q = mQ_qdl.shape[0]
+                seqlen_k = mK_kdl.shape[0]
+                cuseqlen_q = Int32(0)
+                is_valid_q = True
+                if cutlass.const_expr(cum_seqlen_q is not None):
+                    cuseqlen_q = cum_seqlen_q[batch_coord]
+                    seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q
+                    is_valid_q = FmhaStaticTileScheduler.check_valid_work_for_seqlen_q(
+                        self.qk_mma_tiler[0],
+                        mma_block_coord[0],
+                        seqlen_q,
+                    )
+                if cutlass.const_expr(cum_seqlen_k is not None):
+                    cuseqlen_k = cum_seqlen_k[batch_coord]
+                    seqlen_k = cum_seqlen_k[batch_coord + 1] - cuseqlen_k
+                seqlen_kv_loop_start, seqlen_kv_loop_steps = (
+                    FusedMask.get_trip_start_count_via_block_info(
+                        mma_block_coord,
+                        self.qk_mma_tiler,
+                        seqlen_q,
+                        seqlen_k,
+                        self.is_causal,
+                        self.is_local,
+                        window_size_left,
+                        window_size_right,
+                    )
+                )
+                is_valid_k = seqlen_kv_loop_steps > 0
+                has_work = is_valid_q and is_valid_k
+
+                mdQ_qdl_eff = mdQ_qdl
+                if cutlass.const_expr(cum_seqlen_q is not None):
+                    block_offset_dQ = (cuseqlen_q,) + (None,) * 2
+                    mdQ_qdl_eff = cute.domain_offset(block_offset_dQ, mdQ_qdl)
+
+                # (bM, bN, loopM, loopN, loopL)
+                gdQ_qdl = cute.flat_divide(
+                    mdQ_qdl_eff, cute.select(self.dsk_block_tiler, mode=[0, 1])
+                )
+                cdQ_qdl = cute.flat_divide(
+                    cute.make_identity_tensor(mdQ_qdl_eff.shape),
+                    cute.select(self.dsk_block_tiler, mode=[0, 1]),
+                )
+
+                gdQ_staged = gdQ_qdl[None, None, curr_block_coord[0], None, curr_block_coord[2]]
+                cdQ_staged = cdQ_qdl[None, None, curr_block_coord[0], None, curr_block_coord[2]]
+                gdQ_tma_staged = gdQ_staged
+
+                if cutlass.const_expr(not varlen):
+                    gdQ_tma_qdl = cute.flat_divide(
+                        mdQ_tma, cute.select(self.dsk_block_tiler, mode=[0, 1])
+                    )
+                    gdQ_tma_staged = gdQ_tma_qdl[
+                        None, None, curr_block_coord[0], None, curr_block_coord[2]
+                    ]
+
+                if has_work:
+                    # dQ TMEM to GMEM
+                    mma_dq_consumer = self.dQ_epilogue(
+                        seqlen_q,
+                        (mma_dq_consumer, gdQ_staged, cdQ_staged, tdQtdQ_staged),
+                        self.epi_tile,
+                        (tma_atom_dQ, gdQ_tma_staged, s_epi_dQ, varlen),
+                    )
+                else:
+                    self.dQ_epilogue_write_zero(
+                        seqlen_q,
+                        gdQ_staged,
+                        cdQ_staged,
+                    )
+
+                work_tile = tile_sched.advance_to_next_work()
+            # NOTE: tmem.free() moved to kernel end to enable cluster-wide sync
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        #  Scheduler Warp (only for CLC dynamic scheduler)
+        # ///////////////////////////////////////////////////////////////////////////////
+        if cutlass.const_expr(self.use_clc_scheduler):
+            is_first_cta_in_cluster = cta_rank_in_cluster == 0
+
+            if warp_idx == self.sched_warp_id and is_first_cta_in_cluster:
+                cute.arch.warpgroup_reg_dealloc(self.num_regs_other)
+                while work_tile.is_valid_tile:
+                    tile_sched.prefetch_next_work()
+                    work_tile = tile_sched.advance_to_next_work()
+                tile_sched.producer_tail()
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        #  Empty warps reg dealloc
+        # ///////////////////////////////////////////////////////////////////////////////
+        if cutlass.const_expr(self.use_clc_scheduler):
+            if warp_idx > self.load_warp_id:
+                if not (warp_idx == self.sched_warp_id and is_first_cta_in_cluster):
+                    cute.arch.warpgroup_reg_dealloc(self.num_regs_other)
+        else:
+            if warp_idx > self.load_warp_id:
+                cute.arch.warpgroup_reg_dealloc(self.num_regs_other)
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        #  Cooperative TMEM Deallocation (2CTA)
+        # ///////////////////////////////////////////////////////////////////////////////
+        # All warps (including scheduler) have finished by this point.
+        # Cluster-wide sync ensures both CTAs reach here before dealloc.
+        cute.arch.cluster_arrive()
+        cute.arch.cluster_wait()
+        tmem.relinquish_alloc_permit()
+        tmem.free(tmem_ptr)
+
+        return
+
+    @cute.jit
+    def compute_step(
+        self,
+        mask_args: Tuple,
+        value_args: Tuple,
+        tensor_args: Tuple,
+        pipeline_args: Tuple,
+        step: Int32,
+    ) -> Tuple[Float32, Float32, pipeline.PipelineConsumer, pipeline.PipelineProducer]:
+        need_apply_mask, window_size_left, window_size_right = mask_args
+        seqlen_q, seqlen_k, scale_softmax, batch_coord, block_m_idx, varlen = value_args
+        tStS, tScS, tdPtdP, tdPcdP, sLSE, sSum_OdO = tensor_args
+        mma_s_consumer, mma_dp_consumer, ds_mma_producer, lse_handle, sum_odo_handle = pipeline_args
+
+        bidx = block_m_idx
+        tidx, _, _ = cute.arch.thread_idx()
+        thread_idx = tidx % (self.threads_per_warp * len(self.compute_warp_ids))
+        s_handle = mma_s_consumer.wait_and_advance()
+        tStS_slice = tStS[(None, None), 0, 0, s_handle.index]
+        tScS_slice = tScS[(None, None), 0, 0]
+        tmem_load_atom = cute.make_copy_atom(
+            tcgen05.Ld32x32bOp(tcgen05.Repetition(16)), self.acc_dtype
+        )
+        tmem_tiled_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS_slice)
+        thr_load = tmem_tiled_load.get_slice(thread_idx)
+        tTMEM_LOADtS = thr_load.partition_S(tStS_slice)
+        tTMEM_LOADcS = thr_load.partition_D(tScS_slice)
+        tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.acc_dtype)
+        cute.copy(tmem_tiled_load, tTMEM_LOADtS, tTMEM_LOADrS)
+        cute.arch.fence_view_async_tmem_load()
+        s_handle.release()
+        if need_apply_mask:
+            FusedMask.apply_mask_via_causal_local(
+                tTMEM_LOADrS,
+                tTMEM_LOADcS,
+                seqlen_q,
+                seqlen_k,
+                self.use_semantic_trip_range,
+                self.is_causal,
+                self.is_local,
+                window_size_left,
+                window_size_right,
+            )
+
+        log2_e = cutlass.Float32(math.log2(math.e))
+        softmax_scale_log2_e = scale_softmax * log2_e
+        tTMEM_STORErP = cute.make_rmem_tensor(tTMEM_LOADrS.shape, self.q_dtype)
+        for k in cutlass.range(0, cute.size(tTMEM_LOADrS), 2, unroll_full=True):
+            lse = (
+                -sLSE[
+                    cute.get(tTMEM_LOADcS[k], mode=[0]) - bidx * self.cta_tiler[0],
+                    lse_handle.index,
+                ],
+                -sLSE[
+                    cute.get(tTMEM_LOADcS[k + 1], mode=[0]) - bidx * self.cta_tiler[0],
+                    lse_handle.index,
+                ],
+            )
+            tTMEM_LOADrS[k], tTMEM_LOADrS[k + 1] = cute.arch.fma_packed_f32x2(
+                (tTMEM_LOADrS[k], tTMEM_LOADrS[k + 1]),
+                (softmax_scale_log2_e, softmax_scale_log2_e),
+                lse,
+            )
+            tTMEM_LOADrS[k] = cute.math.exp2(tTMEM_LOADrS[k], fastmath=True)
+            tTMEM_LOADrS[k + 1] = cute.math.exp2(tTMEM_LOADrS[k + 1], fastmath=True)
+
+        dp_handle = mma_dp_consumer.wait_and_advance()
+        tdPtdP_slice = tdPtdP[(None, None), 0, 0, dp_handle.index]
+        tdPcdP_slice = tdPcdP[(None, None), 0, 0]
+        thr_load = tmem_tiled_load.get_slice(thread_idx)
+        tTMEM_LOADtdP = thr_load.partition_S(tdPtdP_slice)
+        tTMEM_LOADcdP = thr_load.partition_D(tdPcdP_slice)
+        tTMEM_LOADrdP = cute.make_rmem_tensor(tTMEM_LOADcdP.shape, self.acc_dtype)
+        cute.copy(tmem_tiled_load, tTMEM_LOADtdP, tTMEM_LOADrdP)
+        cute.arch.fence_view_async_tmem_load()
+        dp_handle.release()
+        tTMEM_STORErdP = cute.make_rmem_tensor(tTMEM_LOADrdP.shape, self.q_dtype)
+
+        for k in cutlass.range(0, cute.size(tTMEM_LOADrdP), 2, unroll_full=True):
+            dpsum_0 = -sSum_OdO[
+                cute.get(tTMEM_LOADcdP[k], mode=[0]) - bidx * self.cta_tiler[0],
+                sum_odo_handle.index,
+            ]
+            dpsum_1 = -sSum_OdO[
+                cute.get(tTMEM_LOADcdP[k + 1], mode=[0]) - bidx * self.cta_tiler[0],
+                sum_odo_handle.index,
+            ]
+            if cutlass.const_expr(varlen):
+                if not cute.elem_less(cute.get(tTMEM_LOADcdP[k], mode=[0]), seqlen_q):
+                    dpsum_0 = 0.0
+                if not cute.elem_less(cute.get(tTMEM_LOADcdP[k + 1], mode=[0]), seqlen_q):
+                    dpsum_1 = 0.0
+            tTMEM_LOADrdP[k], tTMEM_LOADrdP[k + 1] = cute.arch.add_packed_f32x2(
+                (tTMEM_LOADrdP[k], tTMEM_LOADrdP[k + 1]),
+                (dpsum_0, dpsum_1),
+            )
+            tTMEM_LOADrdP[k], tTMEM_LOADrdP[k + 1] = cute.arch.mul_packed_f32x2(
+                (tTMEM_LOADrdP[k], tTMEM_LOADrdP[k + 1]), (tTMEM_LOADrS[k], tTMEM_LOADrS[k + 1])
+            )
+            tTMEM_LOADrdP[k], tTMEM_LOADrdP[k + 1] = cute.arch.mul_packed_f32x2(
+                (tTMEM_LOADrdP[k], tTMEM_LOADrdP[k + 1]), (scale_softmax, scale_softmax)
+            )
+        dp_vec = tTMEM_LOADrdP.load()
+        tTMEM_STORErdP.store(dp_vec.to(self.q_dtype))
+
+        ds_handle = ds_mma_producer.acquire_and_advance()
+        tmem_store_atom = cute.make_copy_atom(
+            tcgen05.St32x32bOp(tcgen05.Repetition(32)), self.acc_dtype
+        )
+        tilePlikeFP32 = tdPtdP_slice.shape[1] // Float32.width * self.q_dtype.width
+        tdPtdP_dS_layout = cute.composition(
+            tdPtdP_slice.layout, cute.make_layout((tdPtdP_slice.shape[0], tilePlikeFP32))
+        )
+        tdPtdP_dS = cute.make_tensor(tdPtdP_slice.iterator, tdPtdP_dS_layout)
+        tdPcdP_dS_layout = cute.composition(
+            tdPcdP_slice.layout, cute.make_layout((tdPcdP_slice.shape[0], tilePlikeFP32))
+        )
+        tdPcdP_dS = cute.make_tensor(tdPcdP_slice.iterator, tdPcdP_dS_layout)
+        tmem_tiled_store = tcgen05.make_tmem_copy(tmem_store_atom, tdPtdP_dS)
+
+        thr_store = tmem_tiled_store.get_slice(thread_idx)
+        tTMEM_STOREtdS = thr_store.partition_D(tdPtdP_dS)
+        tTMEM_STOREcdP = thr_store.partition_S(tdPcdP_dS)
+        tTMEM_STORErdS_ = cute.make_tensor(
+            cute.recast_ptr(tTMEM_STORErdP.iterator, dtype=self.acc_dtype),
+            tTMEM_STOREcdP.shape,
+        )
+        cute.copy(tmem_tiled_store, tTMEM_STORErdS_, tTMEM_STOREtdS)
+        cute.arch.fence_view_async_tmem_store()
+        ds_handle.commit()
+        return mma_s_consumer, mma_dp_consumer, ds_mma_producer
+
+    @cute.jit
+    def dQ_epilogue(
+        self,
+        seqlen_q: int,
+        dq_args: Tuple,
+        epi_tile: cute.Tile,
+        tma_args: Tuple,
+    ) -> Tuple[pipeline.PipelineConsumer, pipeline.PipelineProducer]:
+        (mma_dq_consumer, gdQ_staged, cdQ_staged, tdQtdQ_staged) = dq_args
+        tma_atom_dQ, gdQ_tma_staged, s_epi_dQ, varlen = tma_args
+        dq_handle = mma_dq_consumer.wait_and_advance()
+        tidx, _, _ = cute.arch.thread_idx()
+        bidx, bidy, bidz = cute.arch.block_idx()
+        cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+        cute.arch.fence_view_async_shared()
+
+        epi_cols_dQ = math.gcd(128 // (self.dq_dtype.width // 8), epi_tile[1])
+        num_epi_stages_dQ = epi_tile[1] // epi_cols_dQ
+        epi_tile_dQ = (epi_tile[0], epi_cols_dQ)
+        leader_warp = (cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4) == 0
+
+        for iter in cutlass.range(self.iterations_dsk):
+            gdQ = gdQ_staged[None, None, iter]
+            cdQ = cdQ_staged[None, None, iter]
+            tdQtdQ = tdQtdQ_staged[(None, None), 0, 0, iter]
+            tdQtdQ_epi = cute.zipped_divide(tdQtdQ, epi_tile)
+            cdQ_epi = cute.zipped_divide(cdQ, epi_tile)
+            gdQ_epi = cute.zipped_divide(gdQ, epi_tile)
+            cdQ_local = cute.make_identity_tensor(epi_tile)
+            cdQ_local_epi = cute.zipped_divide(cdQ_local, epi_tile)
+            tidx, _, _ = cute.arch.thread_idx()
+            thread_idx = tidx % (self.threads_per_warp * len(self.epilogue_warp_ids))
+            tmem_copy_atom = cute.make_copy_atom(
+                tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype
+            )
+            tiled_tmem_load = tcgen05.make_tmem_copy(tmem_copy_atom, tdQtdQ_epi)
+            thr_tmem_load = tiled_tmem_load.get_slice(thread_idx)
+            tTMEM_LOADtdQ = thr_tmem_load.partition_S(tdQtdQ_epi)
+            tTMEM_LOADgdQ = thr_tmem_load.partition_D(gdQ_epi)
+            tTMEM_LOADcdQ = thr_tmem_load.partition_D(cdQ_epi)
+            tTMEM_LOADcdQ_local = thr_tmem_load.partition_D(cdQ_local_epi)
+
+            if cutlass.const_expr(not varlen):
+                gdQ_tma = gdQ_tma_staged[None, None, iter]
+                gdQ_tma_epi = cute.local_tile(gdQ_tma, epi_tile_dQ, (0, None))
+                sdQ_stage = s_epi_dQ[None, None, 0]
+
+                for stage_k in cutlass.range_constexpr(num_epi_stages_dQ):
+                    for i in cutlass.range(cute.size(tTMEM_LOADtdQ, mode=[1]), unroll_full=True):
+                        tTMEM_LOADtdQ_i = tTMEM_LOADtdQ[None, i, 0]
+                        tTMEM_LOADcdQ_i_local = tTMEM_LOADcdQ_local[None, i, 0]
+                        tTMrdQ = cute.make_rmem_tensor(tTMEM_LOADcdQ_i_local.shape, self.acc_dtype)
+                        cute.copy(tiled_tmem_load, tTMEM_LOADtdQ_i, tTMrdQ)
+                        tSMrdQ = cute.make_rmem_tensor(tTMrdQ.shape, self.q_dtype)
+                        dq_vec = tTMrdQ.load()
+                        tSMrdQ.store(dq_vec.to(self.q_dtype))
+                        for j in cutlass.range_constexpr(cute.size(tTMEM_LOADcdQ_i_local)):
+                            c = tTMEM_LOADcdQ_i_local[j]
+                            m_pos = c[0]
+                            n_pos = c[1]
+                            if n_pos // epi_cols_dQ == stage_k:
+                                s_epi_dQ[m_pos, n_pos % epi_cols_dQ, 0] = tSMrdQ[j]
+
+                    cute.arch.fence_view_async_shared()
+                    cute.arch.barrier(barrier_id=3, number_of_threads=128)
+
+                    if leader_warp:
+                        gdQ_stage = gdQ_tma_epi[None, None, stage_k]
+                        td_sdQ, td_gdQ = cpasync.tma_partition(
+                            tma_atom_dQ,
+                            0,
+                            cute.make_layout(1),
+                            cute.group_modes(sdQ_stage, 0, 2),
+                            cute.group_modes(gdQ_stage, 0, 2),
+                        )
+                        cute.copy(tma_atom_dQ, td_sdQ, td_gdQ)
+                        cute.arch.cp_async_bulk_commit_group()
+                        cute.arch.cp_async_bulk_wait_group(0, read=True)
+                    # Non-issuing threads must not rotate the SMEM buffer
+                    # until the leader warp's TMA read has completed.
+                    cute.arch.barrier(barrier_id=3, number_of_threads=128)
+            else:
+                for i in cutlass.range(cute.size(tTMEM_LOADtdQ, mode=[1]), unroll_full=True):
+                    tTMEM_LOADtdQ_i = tTMEM_LOADtdQ[None, i, 0]
+                    tTMEM_LOADgdQ_i = tTMEM_LOADgdQ[None, i, 0]
+                    tTMEM_LOADcdQ_i = tTMEM_LOADcdQ[None, i, 0]
+                    tTMrdQ = cute.make_rmem_tensor(tTMEM_LOADcdQ[None, 0, i].shape, self.acc_dtype)
+                    cute.copy(tiled_tmem_load, tTMEM_LOADtdQ_i, tTMrdQ)
+                    tSMrdQ = cute.make_rmem_tensor(tTMrdQ.shape, self.q_dtype)
+                    dq_vec = tTMrdQ.load()
+                    tSMrdQ.store(dq_vec.to(self.q_dtype))
+                    if cute.elem_less(tTMEM_LOADcdQ_i[0][0], seqlen_q):
+                        cute.autovec_copy(tSMrdQ, tTMEM_LOADgdQ_i)
+        dq_handle.release()
+        return mma_dq_consumer
+
+    @cute.jit
+    def dQ_epilogue_write_zero(
+        self,
+        seqlen_q,
+        gdQ_staged,
+        cdQ_staged,
+    ):
+        num_epi_threads = self.threads_per_warp * len(self.epilogue_warp_ids)
+        tidx = cute.arch.thread_idx()[0] % num_epi_threads
+
+        tiled_copy_r2g = fa_copy_utils.tiled_copy_2d(
+            self.dq_dtype, cute.size(gdQ_staged.shape[1]), num_epi_threads
+        )
+
+        thr_copy_r2g = tiled_copy_r2g.get_slice(tidx)
+        tdQgdQ_staged = thr_copy_r2g.partition_D(gdQ_staged)
+        tdQcdQ_staged = thr_copy_r2g.partition_D(cdQ_staged)
+
+        tdQrdQ = cute.make_rmem_tensor_like(tdQgdQ_staged[None, 0, None, 0])
+        tdQrdQ.fill(self.dq_dtype(0.0))
+
+        for iter in cutlass.range(self.iterations_dsk, unroll_full=True):
+            tdQgdQ = tdQgdQ_staged[None, None, None, iter]
+            tdQcdQ = tdQcdQ_staged[None, None, None, iter]
+            for m in cutlass.range(cute.size(tdQgdQ.shape[1]), unroll_full=True):
+                if cute.elem_less(tdQcdQ[0, m, 0][0], seqlen_q):
+                    cute.copy(tiled_copy_r2g, tdQrdQ, tdQgdQ[None, m, None])
diff --git a/build/torch-cuda/sm100_hd256_2cta_fmha_forward.py b/build/torch-cuda/sm100_hd256_2cta_fmha_forward.py
new file mode 100644
index 0000000000000000000000000000000000000000..eb122a913e2201f856fd62ff56370bbdbc286582
--- /dev/null
+++ b/build/torch-cuda/sm100_hd256_2cta_fmha_forward.py
@@ -0,0 +1,1920 @@
+# Copyright (c) 2025, Siyu Wang, Shengbin Di, Yuxi Chi, Johnsonms, Linfeng Zheng, Haoyan Huang, Lanbo Li, Yun Zhong, Man Yuan, Minmin Sun, Yong Li, Wei Lin.
+
+import math
+from typing import Tuple, Optional
+
+import cuda.bindings.driver as cuda
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.cute.nvgpu.tcgen05 as tcgen05
+import cutlass.utils as utils
+import cutlass.pipeline as pipeline
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass.cute.typing import Int32, Int64, Float32
+
+from cutlass.utils import ClcDynamicPersistentTileScheduler
+from .tile_scheduler import (
+    SchedulerState,
+    compute_sm100_fmha_grid as compute_grid,
+    compute_sm100_fmha_grid_clc as compute_grid_clc,
+    make_sm100_thread_cooperative_group as make_thread_cooperative_group,
+    Sm100FmhaStaticTileScheduler as FmhaStaticTileScheduler,
+    Sm100FmhaStaticTileSchedulerParams as FmhaStaticTileSchedulerParams,
+    Sm100FmhaClcDynamicTileScheduler as FmhaClcDynamicTileScheduler,
+    Sm100FmhaClcDynamicTileSchedulerParams as FmhaClcDynamicTileSchedulerParams,
+)
+from .mask import (
+    Sm100FusedMask as FusedMask,
+)
+from .tile_scheduler import SM100_TMEM_CAPACITY_COLUMNS
+from .flash_fwd_sm100 import DescaleTensors, _TUNING_CONFIG
+from .utils import ex2_emulation_2, as_bshkrd_tensor, AuxData
+
+
+class BlackwellFusedMultiHeadAttentionForward:
+    def __init__(
+        self,
+        head_dim: int,
+        head_dim_v: Optional[int] = None,
+        qhead_per_kvhead: int = 1,
+        is_causal: bool = False,
+        is_local: bool = False,
+        is_split_kv: bool = False,
+        pack_gqa: bool = False,
+        q_subtile_factor: int = 1,
+        kv_subtile_factor: int = 1,
+        m_block_size: int = 128,
+        n_block_size: int = 128,
+        q_stage: int = 2,
+        is_static_persistent: bool = True,
+        score_mod=None,
+        mask_mod=None,
+        has_aux_tensors: bool = False,
+        paged_kv_non_tma: bool = False,
+        is_varlen_q: bool = False,
+        use_2cta_instrs: bool = False,
+        use_clc_scheduler: bool = False,
+        has_tile_count_semaphore: bool = False,
+        seqlen_k_per_split: Optional[int] = None,
+    ):
+        head_dim_v = head_dim if head_dim_v is None else head_dim_v
+        assert head_dim == 256 and head_dim_v == 256, (
+            "SM100 dedicated kernel only supports (head_dim, head_dim_v) = (256, 256)"
+        )
+        assert score_mod is None, "SM100 forward with head_dim=256 does not support score_mod"
+        assert mask_mod is None, "SM100 forward with head_dim=256 does not support mask_mod"
+        assert not has_aux_tensors, "SM100 forward with head_dim=256 does not support aux tensors"
+        assert not paged_kv_non_tma, (
+            "SM100 hd256 2CTA supports TMA paged KV only (page_size must equal tile_n=128)"
+        )
+        assert not pack_gqa, "SM100 forward with head_dim=256 does not support pack_gqa"
+        assert not is_split_kv, "SM100 forward with head_dim=256 does not support SplitKV"
+        assert q_subtile_factor == 1, (
+            "SM100 forward with head_dim=256 does not support q_subtile_factor"
+        )
+        assert kv_subtile_factor == 1, (
+            "SM100 forward with head_dim=256 does not support kv_subtile_factor"
+        )
+        assert m_block_size == 128 and n_block_size == 128, (
+            "SM100 dedicated kernel only supports tile_m=128 and tile_n=128"
+        )
+        # q_stage / persistence / scheduler knobs are accepted for interface parity,
+        # but this dedicated kernel uses fixed internal settings.
+
+        qk_acc_dtype = cutlass.Float32
+        pv_acc_dtype = cutlass.Float32
+        mma_tiler = (128, 128, head_dim)
+        self.qk_acc_dtype = qk_acc_dtype
+        self.pv_acc_dtype = pv_acc_dtype
+        self.qhead_per_kvhead = qhead_per_kvhead
+        self.mma_tiler = mma_tiler
+        assert mma_tiler[0] == 128 and mma_tiler[1] == 128, "Only 128x128 tile impl is supported"
+        assert mma_tiler[2] == 256, "Only 256 is supported for 128x128 tile impl"
+        self.cta_tiler = (
+            mma_tiler[0],
+            mma_tiler[1],
+            mma_tiler[2],
+        )
+        self.qk_mma_tiler = (
+            2 * mma_tiler[0],
+            mma_tiler[1],
+            min(self.cta_tiler[2], 128),
+        )
+        self.pv_mma_tiler = self.qk_mma_tiler
+        self.pv_block_tiler = (
+            self.pv_mma_tiler[0] // 2,
+            self.pv_mma_tiler[1],
+            self.pv_mma_tiler[2],
+        )
+        self.iterations_qk = self.cta_tiler[2] // self.qk_mma_tiler[2]
+        self.iterations_pv = self.cta_tiler[2] // self.pv_mma_tiler[1]
+        self.cluster_shape_mn = (2, 1)
+        self.tmem_warp_shape_mn = (4, 1)
+        # Dedicated hd256 kernel uses fixed scheduling policy.
+        self.is_persistent = False
+        self.is_causal = is_causal
+        self.is_local = is_local
+        self.use_semantic_trip_range = is_causal or is_local
+        self.use_clc_scheduler = False
+
+        self.softmax_warp_ids = (0, 1, 2, 3)
+        self.correction_warp_ids = (4, 5, 6, 7)
+        self.mma_warp_id = 8
+        self.load_warp_id = 9
+        self.empty_warp_id = (10, 11)
+        self.sched_warp_id = self.empty_warp_id[0] if use_clc_scheduler else None
+        self.tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS
+
+        self.threads_per_warp = 32
+        self.threads_per_cta = self.threads_per_warp * len(
+            (
+                *self.softmax_warp_ids,  # this is to get a round num threads
+                *self.correction_warp_ids,
+                self.mma_warp_id,
+                self.load_warp_id,
+                *self.empty_warp_id,
+            )
+        )
+
+        self.tmem_alloc_barrier = pipeline.NamedBarrier(
+            barrier_id=1,
+            num_threads=self.threads_per_cta,
+        )
+
+        self.tmem_s_offset = 0
+        self.tmem_o_offset = 256
+        self.tmem_p_offset = self.tmem_s_offset
+
+        _tune_key = (True, is_causal, 256, False)  # hd256: always 2cta, no sm103 variant
+        _tune = _TUNING_CONFIG.get(_tune_key, {})
+        self.num_regs_softmax = _tune.get("num_regs_softmax", 256)
+        self.num_regs_correction = _tune.get("num_regs_correction", 160)
+        self.num_regs_other = 32  # fixed for hd256; not derived from 512 budget like other kernels
+        self.ex2_emu_freq = _tune.get("ex2_emu_freq", 4)
+        self.ex2_emu_res = _tune.get("ex2_emu_res", 3)
+        self.ex2_emu_start_frg = _tune.get("ex2_emu_start_frg", 0)
+
+        self.buffer_align_bytes = 1024
+
+    def _setup_attributes(self):
+        self.q_stage = self.iterations_qk
+        self.kv_stage = 4
+        self.qk_acc_stage = 2
+        self.mma_corr_stage = 1
+        if cutlass.const_expr(self.use_clc_scheduler):
+            self.num_clc_stage = 1
+            self.num_clc_response_bytes = 16
+
+    @cute.jit
+    def __call__(
+        self,
+        mQ: cute.Tensor,
+        mK: cute.Tensor,
+        mV: cute.Tensor,
+        mO: cute.Tensor,
+        mLSE: Optional[cute.Tensor],
+        softmax_scale: Float32,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        mSeqUsedQ: Optional[cute.Tensor] = None,
+        mSeqUsedK: Optional[cute.Tensor] = None,
+        mPageTable: Optional[cute.Tensor] = None,
+        window_size_left: Int32 | int | None = None,
+        window_size_right: Int32 | int | None = None,
+        learnable_sink: Optional[cute.Tensor] = None,
+        descale_tensors: Optional[DescaleTensors] = None,
+        blocksparse_tensors: Optional[cute.Tensor] = None,
+        aux_data: AuxData = AuxData(),
+        stream: cuda.CUstream = None,
+    ):
+        # Keep parity with FlashAttentionForwardSm100.__call__ interface.
+        # (TODO@wangsiyu) Implement these features.
+        assert mSeqUsedQ is None and mSeqUsedK is None, (
+            "SM100 forward with head_dim=256 does not support seqused_q/seqused_k"
+        )
+        assert learnable_sink is None, (
+            "SM100 forward with head_dim=256 does not support learnable_sink"
+        )
+        assert blocksparse_tensors is None, (
+            "SM100 forward with head_dim=256 does not support block sparsity"
+        )
+        assert aux_data.tensors is None, (
+            "SM100 forward with head_dim=256 does not support aux_tensors"
+        )
+        assert aux_data.scalars is None, (
+            "SM100 forward with head_dim=256 does not support aux_scalars"
+        )
+        assert not self.is_local, (
+            "SM100 forward with head_dim=256 does not support local attention yet"
+        )
+        assert window_size_left is None and window_size_right is None, (
+            "SM100 forward with head_dim=256 does not support runtime window_size overrides"
+        )
+        assert descale_tensors is None, (
+            "SM100 forward with head_dim=256 does not support descale_tensors"
+        )
+
+        q_tensor, k_tensor, v_tensor, o_tensor = mQ, mK, mV, mO
+        lse_tensor = mLSE
+        cum_seqlen_q = mCuSeqlensQ
+        cum_seqlen_k = mCuSeqlensK
+
+        q_rank = len(mQ.shape)
+        k_rank = len(mK.shape)
+        if cutlass.const_expr(cum_seqlen_q is not None):
+            # Varlen path accepts either legacy 5D tensors or standard 3D tensors.
+            if cutlass.const_expr(q_rank == 5):
+                s_q = mQ.shape[1]
+                h_q = mQ.shape[2] * mQ.shape[3]
+                d = mQ.shape[4]
+            elif cutlass.const_expr(q_rank == 3):
+                s_q = mQ.shape[0]
+                h_q = mQ.shape[1]
+                d = mQ.shape[2]
+            else:
+                raise RuntimeError(f"hd256 forward varlen expects q rank 3 or 5, got rank {q_rank}")
+        else:
+            # Non-varlen path accepts either legacy 5D tensors or standard 4D tensors.
+            if cutlass.const_expr(q_rank == 5):
+                s_q = mQ.shape[1]
+                h_q = mQ.shape[2] * mQ.shape[3]
+                d = mQ.shape[4]
+            elif cutlass.const_expr(q_rank == 4):
+                s_q = mQ.shape[1]
+                h_q = mQ.shape[2]
+                d = mQ.shape[3]
+            else:
+                raise RuntimeError(
+                    f"hd256 forward non-varlen expects q rank 4 or 5, got rank {q_rank}"
+                )
+
+        if cutlass.const_expr(cum_seqlen_k is not None):
+            if cutlass.const_expr(k_rank == 5):
+                s_k = mK.shape[1]
+                h_k = mK.shape[2]
+            elif cutlass.const_expr(k_rank == 3):
+                s_k = mK.shape[0]
+                h_k = mK.shape[1]
+            else:
+                raise RuntimeError(f"hd256 forward varlen expects k rank 3 or 5, got rank {k_rank}")
+        else:
+            if cutlass.const_expr(k_rank == 5):
+                s_k = mK.shape[1]
+                h_k = mK.shape[2]
+            elif cutlass.const_expr(k_rank == 4):
+                s_k = mK.shape[1]
+                h_k = mK.shape[2]
+            else:
+                raise RuntimeError(
+                    f"hd256 forward non-varlen expects k rank 4 or 5, got rank {k_rank}"
+                )
+        if cutlass.const_expr(cum_seqlen_q is not None):
+            b = mCuSeqlensQ.shape[0] - 1
+        elif cutlass.const_expr(cum_seqlen_k is not None):
+            b = mCuSeqlensK.shape[0] - 1
+        else:
+            b = mQ.shape[0]
+
+        scale_softmax = softmax_scale
+        scale_softmax_log2 = softmax_scale * math.log2(math.exp(1.0))
+        scale_output = 1.0
+        s_lse = s_q
+        h_r = h_q // h_k
+        s_q64 = Int64(s_q)
+        s_k64 = Int64(s_k)
+        s_lse64 = Int64(s_lse)
+        h_r64 = Int64(h_r)
+        h_k64 = Int64(h_k)
+        b64 = Int64(b)
+        s_q_total = (
+            q_tensor.shape[1]
+            if cum_seqlen_q is not None and q_rank == 5
+            else (q_tensor.shape[0] if cum_seqlen_q is not None else s_q64)
+        )
+        s_k_total = (
+            k_tensor.shape[1]
+            if cum_seqlen_k is not None and k_rank == 5
+            else (k_tensor.shape[0] if cum_seqlen_k is not None else s_k64)
+        )
+        b_lse = b64 if cum_seqlen_q is None else 1
+        stride_b_lse = h_r64 * h_k64 * s_lse64 if cum_seqlen_q is None else 0
+
+        varlen_q = cum_seqlen_q is not None
+        varlen_k = cum_seqlen_k is not None
+        q_norm = as_bshkrd_tensor(q_tensor, h_k, h_r, varlen_q)
+        o_norm = as_bshkrd_tensor(o_tensor, h_k, h_r, varlen_q)
+
+        # Forward layout: (s, d, ((h_r, h_k), b)). Stride picks from canonical
+        # positions 1=S, 4=D, 3=H_r, 2=H_k, 0=B.
+        q = cute.make_tensor(
+            q_norm.iterator,
+            cute.make_layout(
+                (s_q_total, d, ((h_r, h_k), b)),
+                stride=(
+                    q_norm.stride[1],
+                    q_norm.stride[4],
+                    ((q_norm.stride[3], q_norm.stride[2]), q_norm.stride[0]),
+                ),
+            ),
+        )
+        if cutlass.const_expr(mPageTable is not None):
+            # Paged: input k/v are rank-4 (num_pages, page_size, h_k, d); the kernel
+            # consumes K as (page_size, d, h_k, num_pages) and V as
+            # (d, page_size, h_k, num_pages).
+            # cute.select reorders modes while preserving input strides
+            page_size = k_tensor.shape[1]
+            max_seqlen_k_paged = Int32(mPageTable.shape[1] * page_size)
+            k = cute.make_tensor(k_tensor.iterator, cute.select(k_tensor.layout, mode=[1, 3, 2, 0]))
+            v = cute.make_tensor(v_tensor.iterator, cute.select(v_tensor.layout, mode=[3, 1, 2, 0]))
+            page_table = cute.make_tensor(
+                mPageTable.iterator,
+                cute.make_layout(
+                    (b, mPageTable.shape[1]),
+                    stride=(mPageTable.stride[0], mPageTable.stride[1]),
+                ),
+            )
+        else:
+            # K/V have no h_r dim; pass h_r=1 to the normalizer and override the
+            # h_r stride to 0 below to broadcast across the query-grouped heads.
+            k_norm = as_bshkrd_tensor(k_tensor, h_k, 1, varlen_k)
+            v_norm = as_bshkrd_tensor(v_tensor, h_k, 1, varlen_k)
+            # (s, d, ((h_r, h_k), b)), 0-stride for h_r to broadcast
+            k = cute.make_tensor(
+                k_norm.iterator,
+                cute.make_layout(
+                    (s_k_total, d, ((h_r, h_k), b)),
+                    stride=(
+                        k_norm.stride[1],
+                        k_norm.stride[4],
+                        ((0, k_norm.stride[2]), k_norm.stride[0]),
+                    ),
+                ),
+            )
+            # (d, s, ((h_r, h_k), b)), 0-stride for h_r to broadcast
+            v = cute.make_tensor(
+                v_norm.iterator,
+                cute.make_layout(
+                    (d, s_k_total, ((h_r, h_k), b)),
+                    stride=(
+                        v_norm.stride[4],
+                        v_norm.stride[1],
+                        ((0, v_norm.stride[2]), v_norm.stride[0]),
+                    ),
+                ),
+            )
+            page_table = None
+            max_seqlen_k_paged = None
+        # (s, d, ((h_r, h_k), b))
+        o = cute.make_tensor(
+            o_norm.iterator,
+            cute.make_layout(
+                (s_q_total, d, ((h_r, h_k), b)),
+                stride=(
+                    o_norm.stride[1],
+                    o_norm.stride[4],
+                    ((o_norm.stride[3], o_norm.stride[2]), o_norm.stride[0]),
+                ),
+            ),
+        )
+        if cutlass.const_expr(lse_tensor is not None):
+            # (s, ((h_r, h_k), b))
+            lse_layout = cute.make_layout(
+                (s_lse64, ((h_r, h_k), b_lse)),
+                stride=(1, ((s_lse64, h_r64 * s_lse64), stride_b_lse)),
+            )
+            lse = cute.make_tensor(lse_tensor.iterator, lse_layout)
+        else:
+            lse = None
+
+        # setup static attributes before smem/grid/tma computation
+        self.q_dtype = q.element_type
+        self.k_dtype = k.element_type
+        self.v_dtype = v.element_type
+        self.o_dtype = o.element_type
+        self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.q_dtype.width
+
+        if cutlass.const_expr(self.use_clc_scheduler):
+            self.tile_sched_params, grid = compute_grid_clc(
+                (s_q, o.shape[1], o.shape[2]) if cum_seqlen_q is not None else o.shape,
+                self.cta_tiler,
+                (*self.cluster_shape_mn, 1),
+            )
+        else:
+            self.tile_sched_params, grid = compute_grid(
+                (s_q, o.shape[1], o.shape[2]) if cum_seqlen_q is not None else o.shape,
+                self.cta_tiler,
+                self.is_persistent,
+            )
+
+        self.q_major_mode = utils.LayoutEnum.from_tensor(q).mma_major_mode()
+        self.k_major_mode = utils.LayoutEnum.from_tensor(k).mma_major_mode()
+        self.v_major_mode = utils.LayoutEnum.from_tensor(v).mma_major_mode()
+        self.o_layout = utils.LayoutEnum.from_tensor(o)
+
+        if cutlass.const_expr(self.q_major_mode != tcgen05.OperandMajorMode.K):
+            raise RuntimeError("The layout of q is not supported")
+        if cutlass.const_expr(self.k_major_mode != tcgen05.OperandMajorMode.K):
+            raise RuntimeError("The layout of k is not supported")
+        if cutlass.const_expr(self.v_major_mode != tcgen05.OperandMajorMode.MN):
+            raise RuntimeError("The layout of v is not supported")
+
+        # check type consistency
+        if cutlass.const_expr(self.q_dtype != self.k_dtype):
+            raise TypeError(f"Type mismatch: {self.q_dtype} != {self.k_dtype}")
+        if cutlass.const_expr(self.q_dtype != self.v_dtype):
+            raise TypeError(f"Type mismatch: {self.q_dtype} != {self.v_dtype}")
+        self._setup_attributes()
+
+        cta_group = tcgen05.CtaGroup.TWO
+        # the intermediate tensor p is from tmem & k-major
+        p_source = tcgen05.OperandSource.TMEM
+        p_major_mode = tcgen05.OperandMajorMode.K
+        qk_tiled_mma = sm100_utils.make_trivial_tiled_mma(
+            self.q_dtype,
+            self.q_major_mode,
+            self.k_major_mode,
+            self.qk_acc_dtype,
+            cta_group,
+            self.qk_mma_tiler[:2],
+        )
+        pv_tiled_mma = sm100_utils.make_trivial_tiled_mma(
+            self.v_dtype,
+            p_major_mode,
+            self.v_major_mode,
+            self.pv_acc_dtype,
+            cta_group,
+            self.pv_mma_tiler[:2],
+            p_source,
+        )
+
+        self.cluster_shape_mnk = (*self.cluster_shape_mn, 1)
+        self.cluster_layout_vmnk = cute.tiled_divide(
+            cute.make_layout(self.cluster_shape_mnk),
+            (qk_tiled_mma.thr_id.shape,),
+        )
+
+        self.epi_tile = self.pv_block_tiler[:2]
+
+        q_smem_layout_staged = sm100_utils.make_smem_layout_a(
+            qk_tiled_mma,
+            self.qk_mma_tiler,
+            self.q_dtype,
+            self.q_stage,
+        )
+        k_smem_layout_staged = sm100_utils.make_smem_layout_b(
+            qk_tiled_mma,
+            self.qk_mma_tiler,
+            self.k_dtype,
+            self.kv_stage,
+        )
+        p_tmem_layout_staged = sm100_utils.make_smem_layout_a(
+            pv_tiled_mma,
+            self.pv_mma_tiler,
+            self.q_dtype,
+            self.qk_acc_stage,
+        )
+        p_tmem_layout = cute.select(p_tmem_layout_staged, mode=[0, 1, 2])
+        v_smem_layout_staged = sm100_utils.make_smem_layout_b(
+            pv_tiled_mma,
+            self.pv_mma_tiler,
+            self.v_dtype,
+            self.kv_stage,
+        )
+        # TMA load for Q
+        tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group)
+
+        q_smem_layout = cute.select(q_smem_layout_staged, mode=[0, 1, 2])
+        tma_atom_q, tma_tensor_q = cute.nvgpu.make_tiled_tma_atom_A(
+            tma_load_op,
+            q,
+            q_smem_layout,
+            self.qk_mma_tiler,
+            qk_tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+
+        # TMA load for K
+        k_smem_layout = cute.select(k_smem_layout_staged, mode=[0, 1, 2])
+        tma_atom_k, tma_tensor_k = cute.nvgpu.make_tiled_tma_atom_B(
+            tma_load_op,
+            k,
+            k_smem_layout,
+            self.qk_mma_tiler,
+            qk_tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+        # TMA load for V
+        v_smem_layout = cute.select(v_smem_layout_staged, mode=[0, 1, 2])
+        tma_atom_v, tma_tensor_v = cute.nvgpu.make_tiled_tma_atom_B(
+            tma_load_op,
+            v,
+            v_smem_layout,
+            self.pv_mma_tiler,
+            pv_tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+
+        q_copy_size = cute.size_in_bytes(self.q_dtype, q_smem_layout)
+        k_copy_size = cute.size_in_bytes(self.k_dtype, k_smem_layout)
+        self.tma_copy_q_bytes = q_copy_size * cute.size(qk_tiled_mma.thr_id.shape)
+        self.tma_copy_kv_bytes = k_copy_size * cute.size(qk_tiled_mma.thr_id.shape)
+
+        @cute.struct
+        class SharedStorage:
+            # TMA G2S load barriers: LOAD warp (producer) -> MMA warp (consumer)
+            load_q_mbar_ptr: cute.struct.MemRange[
+                Int64, self.q_stage * 2
+            ]  # load_q_{producer,consumer}
+            load_kv_mbar_ptr: cute.struct.MemRange[
+                Int64, self.kv_stage * 2
+            ]  # load_kv_{producer,consumer}
+            mma_s_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2]
+            p_mma_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2]
+            # Softmax -> Correction signaling barriers (row_max/row_sum vec ready)
+            s_corr_mbar_ptr: cute.struct.MemRange[
+                Int64, self.qk_acc_stage * 2
+            ]  # s_corr_{producer,consumer}
+            sum_mbar_ptr: cute.struct.MemRange[Int64, 2]
+            # MMA -> Correction ownership barriers for O_partial tokens (online rescale/finalize)
+            mma_corr_mbar_ptr: cute.struct.MemRange[
+                Int64, self.mma_corr_stage * 2
+            ]  # mma_corr_{producer,consumer}
+            # A CTA-wide "TMEM lifetime" barrier used to safely deallocate TMEM after all users finish.
+            tmem_dealloc_mbar: Int64
+            # Tmem holding buffer
+            tmem_holding_buf: Int32
+            # CLC pipeline barriers and response buffer
+            clc_mbar_ptr: cute.struct.MemRange[Int64, 2]
+            clc_response: cute.struct.MemRange[Int32, 4]
+
+        self.shared_storage = SharedStorage
+
+        grid = cute.round_up(grid, self.cluster_shape_mnk)
+        # Launch the kernel synchronously
+        self.kernel(
+            qk_tiled_mma,
+            pv_tiled_mma,
+            tma_atom_q,
+            tma_tensor_q,
+            tma_atom_k,
+            tma_tensor_k,
+            tma_atom_v,
+            tma_tensor_v,
+            o,
+            cum_seqlen_q,
+            cum_seqlen_k,
+            lse,
+            scale_softmax_log2,
+            scale_softmax,
+            scale_output,
+            page_table,
+            max_seqlen_k_paged,
+            window_size_left,
+            window_size_right,
+            self.cluster_layout_vmnk,
+            q_smem_layout_staged,
+            k_smem_layout_staged,
+            p_tmem_layout,
+            v_smem_layout_staged,
+            self.tile_sched_params,
+        ).launch(
+            grid=grid,
+            block=[self.threads_per_cta, 1, 1],
+            cluster=self.cluster_shape_mnk,
+            stream=stream,
+            min_blocks_per_mp=1,
+        )
+
+    #  GPU device kernel
+    @cute.kernel
+    def kernel(
+        self,
+        qk_tiled_mma: cute.TiledMma,
+        pv_tiled_mma: cute.TiledMma,
+        tma_atom_q: cute.CopyAtom,
+        mQ_qdl: cute.Tensor,
+        tma_atom_k: cute.CopyAtom,
+        mK_kdl: cute.Tensor,
+        tma_atom_v: cute.CopyAtom,
+        mV_dkl: cute.Tensor,
+        mO_qdl: cute.Tensor,
+        cum_seqlen_q: Optional[cute.Tensor],
+        cum_seqlen_k: Optional[cute.Tensor],
+        mLSE: Optional[cute.Tensor],
+        scale_softmax_log2: Float32,
+        scale_softmax: Float32,
+        scale_output: Float32,
+        mPageTable: Optional[cute.Tensor],
+        max_seqlen_k: Optional[Int32],
+        window_size_left: Optional[Int32],
+        window_size_right: Optional[Int32],
+        cluster_layout_vmnk: cute.Layout,
+        q_smem_layout_staged: cute.ComposedLayout,
+        k_smem_layout_staged: cute.ComposedLayout,
+        p_tmem_layout_staged: cute.ComposedLayout,
+        v_smem_layout_staged: cute.ComposedLayout,
+        tile_sched_params: FmhaStaticTileSchedulerParams | FmhaClcDynamicTileSchedulerParams,
+    ):
+        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+
+        #
+        # Prefetch tma desc
+        #
+        if warp_idx == self.load_warp_id:
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_q)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_k)
+            cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_v)
+
+        bidx, _, _ = cute.arch.block_idx()
+        mma_tile_coord_v = bidx % cute.size(qk_tiled_mma.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)
+
+        # Alloc
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(self.shared_storage)
+
+        load_q_producer, load_q_consumer = pipeline.PipelineTmaUmma.create(
+            num_stages=self.q_stage,
+            producer_group=make_thread_cooperative_group(len([self.load_warp_id])),
+            consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            tx_count=self.tma_copy_q_bytes,
+            barrier_storage=storage.load_q_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        load_kv_producer, load_kv_consumer = pipeline.PipelineTmaUmma.create(
+            num_stages=self.kv_stage,
+            producer_group=make_thread_cooperative_group(len([self.load_warp_id])),
+            consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            tx_count=self.tma_copy_kv_bytes,
+            barrier_storage=storage.load_kv_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        mma_s_producer, mma_s_consumer = pipeline.PipelineUmmaAsync.create(
+            num_stages=self.qk_acc_stage,
+            producer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            consumer_group=make_thread_cooperative_group(
+                len(self.softmax_warp_ids) * self.threads_per_warp * self.cluster_shape_mnk[0],
+            ),
+            barrier_storage=storage.mma_s_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        p_mma_producer, p_mma_consumer = pipeline.PipelineAsyncUmma.create(
+            num_stages=self.qk_acc_stage,
+            producer_group=make_thread_cooperative_group(
+                len(self.softmax_warp_ids) * self.threads_per_warp * self.cluster_shape_mnk[0],
+            ),
+            consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            barrier_storage=storage.p_mma_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        s_corr_producer, s_corr_consumer = pipeline.PipelineAsync.create(
+            num_stages=self.qk_acc_stage,
+            producer_group=make_thread_cooperative_group(
+                self.threads_per_warp * len(self.softmax_warp_ids)
+            ),
+            consumer_group=make_thread_cooperative_group(
+                self.threads_per_warp * len(self.correction_warp_ids)
+            ),
+            barrier_storage=storage.s_corr_mbar_ptr.data_ptr(),
+            defer_sync=True,
+        ).make_participants()
+        sum_producer, sum_consumer = pipeline.PipelineAsync.create(
+            num_stages=1,
+            producer_group=make_thread_cooperative_group(
+                self.threads_per_warp * len(self.softmax_warp_ids)
+            ),
+            consumer_group=make_thread_cooperative_group(
+                self.threads_per_warp * len(self.correction_warp_ids)
+            ),
+            barrier_storage=storage.sum_mbar_ptr.data_ptr(),
+            defer_sync=True,
+        ).make_participants()
+        mma_corr_producer, mma_corr_consumer = pipeline.PipelineUmmaAsync.create(
+            num_stages=self.mma_corr_stage,
+            producer_group=make_thread_cooperative_group(len([self.mma_warp_id])),
+            consumer_group=make_thread_cooperative_group(
+                len(self.correction_warp_ids) * self.threads_per_warp * self.cluster_shape_mnk[0],
+            ),
+            barrier_storage=storage.mma_corr_mbar_ptr.data_ptr(),
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        ).make_participants()
+        # Tensor memory dealloc barrier init
+        tmem = utils.TmemAllocator(
+            storage.tmem_holding_buf.ptr,
+            barrier_for_retrieve=self.tmem_alloc_barrier,
+            allocator_warp_id=self.correction_warp_ids[0],
+            is_two_cta=True,
+            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
+        )
+        tmem.allocate(self.tmem_alloc_cols)
+        tmem.wait_for_alloc()
+        tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype)
+        # Initialize CLC state if using dynamic scheduler
+        if cutlass.const_expr(self.use_clc_scheduler):
+            clc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+            cluster_size = cute.size(self.cluster_shape_mnk)
+            num_clc_consumer_threads = self.threads_per_warp * (
+                1  # sched_warp (CTA 0 only)
+                + cluster_size
+                * (
+                    len(self.softmax_warp_ids)
+                    + len(self.correction_warp_ids)
+                    + 1  # mma_warp
+                    + 1  # load_warp
+                )
+            )
+            clc_pipeline_consumer_group = pipeline.CooperativeGroup(
+                pipeline.Agent.Thread, num_clc_consumer_threads
+            )
+            clc_response_ptr = storage.clc_response.data_ptr()
+            clc = SchedulerState.create_clc(
+                hw_scheduler=ClcDynamicPersistentTileScheduler.create(
+                    self.tile_sched_params.clc_hw_params(),
+                    cute.arch.block_idx(),
+                    cute.arch.grid_dim(),
+                    clc_response_ptr,
+                ),
+                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,
+                ),
+                consumer_state=pipeline.make_pipeline_state(
+                    pipeline.PipelineUserType.Consumer, self.num_clc_stage
+                ),
+                producer_state=pipeline.make_pipeline_state(
+                    pipeline.PipelineUserType.Producer, self.num_clc_stage
+                ),
+            )
+        else:
+            clc = None
+            clc_response_ptr = None
+
+        # Cluster arrive after barrier init
+        pipeline.pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True)
+
+        sQ = smem.allocate_tensor(
+            element_type=self.q_dtype,
+            layout=q_smem_layout_staged.outer,
+            swizzle=q_smem_layout_staged.inner,
+            byte_alignment=128,
+        )
+        sK = smem.allocate_tensor(
+            element_type=self.k_dtype,
+            layout=k_smem_layout_staged.outer,
+            swizzle=k_smem_layout_staged.inner,
+            byte_alignment=128,
+        )
+        # K and V now use separate memory since we removed the transform stage
+        sV = smem.allocate_tensor(
+            element_type=self.v_dtype,
+            layout=v_smem_layout_staged.outer,
+            swizzle=v_smem_layout_staged.inner,
+            byte_alignment=128,
+        )
+
+        sSum = smem.allocate_tensor(
+            element_type=self.qk_acc_dtype,
+            layout=cute.make_layout(len(self.softmax_warp_ids) * self.threads_per_warp),
+            byte_alignment=128,
+        )
+        qk_thr_mma = qk_tiled_mma.get_slice(mma_tile_coord_v)  # default 1sm
+        pv_thr_mma = pv_tiled_mma.get_slice(mma_tile_coord_v)  # default 1sm
+        tSrQ = qk_thr_mma.make_fragment_A(sQ)
+        tSrK = qk_thr_mma.make_fragment_B(sK)
+        tOrV = pv_thr_mma.make_fragment_B(sV)
+        qk_acc_shape = qk_thr_mma.partition_shape_C((self.qk_mma_tiler[0], self.qk_mma_tiler[1]))
+        tStS = qk_thr_mma.make_fragment_C(cute.append(qk_acc_shape, self.qk_acc_stage))
+        pv_acc_shape = pv_thr_mma.partition_shape_C((self.pv_mma_tiler[0], self.pv_mma_tiler[1]))
+        tOtO = pv_thr_mma.make_fragment_C(pv_acc_shape)
+        tOtO_layout = cute.append(
+            tOtO.layout,
+            cute.make_layout(
+                self.iterations_pv,
+                stride=self.pv_mma_tiler[1] // self.tmem_warp_shape_mn[1],
+            ),
+        )
+        tStS = cute.make_tensor(tStS.iterator + self.tmem_s_offset, tStS.layout)
+        tOtO_staged = cute.make_tensor(tOtO.iterator + self.tmem_o_offset, tOtO_layout)
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        #  EMPTY
+        # ///////////////////////////////////////////////////////////////////////////////
+        for _i in cutlass.range_constexpr(len(self.empty_warp_id)):
+            if warp_idx == self.empty_warp_id[_i]:
+                cute.arch.warpgroup_reg_dealloc(self.num_regs_other)
+
+        if cutlass.const_expr(self.use_clc_scheduler):
+            tile_sched = FmhaClcDynamicTileScheduler.create(
+                tile_sched_params,
+                cute.arch.block_idx(),
+                cute.arch.grid_dim(),
+                clc_response_ptr,
+                clc,
+            )
+        else:
+            blk_idx = cute.arch.block_idx()
+            tile_sched = FmhaStaticTileScheduler(
+                tile_sched_params, blk_idx[0], blk_idx, cute.arch.grid_dim()
+            )
+        work_tile = tile_sched.initial_work_tile_info()
+
+        # Cluster wait
+        pipeline.pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk)
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        #  LOAD
+        # ///////////////////////////////////////////////////////////////////////////////
+        if warp_idx == self.load_warp_id:
+            cute.arch.warpgroup_reg_dealloc(self.num_regs_other)
+            while work_tile.is_valid_tile:
+                curr_block_coord = work_tile.tile_idx  # (q_tile_idx, 0, (head_idx, batch_idx))
+                mma_block_coord = (
+                    curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape),
+                    curr_block_coord[1],
+                    curr_block_coord[2],
+                )
+                continue_cond = False
+                batch_coord = curr_block_coord[2][1]
+                seqlen_q = mQ_qdl.shape[0]
+                seqlen_k = (
+                    mK_kdl.shape[0] if cutlass.const_expr(mPageTable is None) else max_seqlen_k
+                )
+                cuseqlen_q = Int32(0)
+                cuseqlen_k = Int32(0)
+                block_offset = (
+                    Int32(0),
+                    Int32(0),
+                    Int32(0),
+                    ((Int32(0), Int32(0)), Int32(0)),
+                )
+                if cutlass.const_expr(cum_seqlen_q is not None):
+                    cuseqlen_q = cum_seqlen_q[batch_coord]
+                    seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q
+                    if cutlass.const_expr(cum_seqlen_k is not None):
+                        cuseqlen_k = cum_seqlen_k[batch_coord]
+                        seqlen_k = cum_seqlen_k[batch_coord + 1] - cuseqlen_k
+                    block_offset = (
+                        cuseqlen_q,
+                        cuseqlen_k,
+                        Int32(0),
+                        ((Int32(0), Int32(0)), Int32(0)),
+                    )
+                    continue_cond = not FmhaStaticTileScheduler.check_valid_work_for_seqlen_q(
+                        self.qk_mma_tiler[0],
+                        mma_block_coord[0],
+                        seqlen_q,
+                    )
+                if not continue_cond:
+                    mQ_qdl_ = cute.domain_offset(cute.select(block_offset, mode=[0, 2, 3]), mQ_qdl)
+                    # Local tile partition global tensors
+                    q_cta_layout = cute.make_layout(
+                        cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape
+                    )
+                    # (bM, bK, loopM, loopK, loopL)
+                    gQ_qdl = cute.flat_divide(mQ_qdl_, cute.select(self.qk_mma_tiler, mode=[0, 2]))
+                    tSgQ_qdl = qk_thr_mma.partition_A(gQ_qdl)
+                    tQsQ, tQgQ_qdl = cute.nvgpu.cpasync.tma_partition(
+                        tma_atom_q,
+                        block_in_cluster_coord_vmnk[2],
+                        q_cta_layout,
+                        cute.group_modes(sQ, 0, 3),
+                        cute.group_modes(tSgQ_qdl, 0, 3),
+                    )
+                    kv_cta_layout = cute.make_layout(
+                        cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape
+                    )
+                    if cutlass.const_expr(mPageTable is None):
+                        # Dense path: domain_offset K/V by batch block, select batch via mma_block_coord[2].
+                        mK_kdl_ = cute.domain_offset(
+                            cute.select(block_offset, mode=[1, 2, 3]), mK_kdl
+                        )
+                        mV_dkl_ = cute.domain_offset(
+                            cute.select(block_offset, mode=[2, 1, 3]), mV_dkl
+                        )
+                        gK_kdl = cute.flat_divide(
+                            mK_kdl_, cute.select(self.qk_mma_tiler, mode=[1, 2])
+                        )
+                        tSgK_kdl = qk_thr_mma.partition_B(gK_kdl)
+                        tKsK, tKgK_kdl = cute.nvgpu.cpasync.tma_partition(
+                            tma_atom_k,
+                            block_in_cluster_coord_vmnk[1],
+                            kv_cta_layout,
+                            cute.group_modes(sK, 0, 3),
+                            cute.group_modes(tSgK_kdl, 0, 3),
+                        )
+                        gV_dkl = cute.flat_divide(
+                            mV_dkl_, cute.select(self.pv_mma_tiler, mode=[1, 2])
+                        )
+                        tSgV_dkl = pv_thr_mma.partition_B(gV_dkl)
+                        tVsV, tVgV_dkl = cute.nvgpu.cpasync.tma_partition(
+                            tma_atom_v,
+                            block_in_cluster_coord_vmnk[1],
+                            kv_cta_layout,
+                            cute.group_modes(sV, 0, 3),
+                            cute.group_modes(tSgV_dkl, 0, 3),
+                        )
+                        # ((atom_v, rest_v), RestN, RestK)
+                        tKgK = tKgK_kdl[None, None, None, mma_block_coord[2]]
+                        tVgV = tVgV_dkl[None, None, None, mma_block_coord[2]]
+                    else:
+                        # Paged path: slice K/V by KV head, keep num_pages dim for page_idx-based TMA.
+                        head_kv_coord = curr_block_coord[2][0] // self.qhead_per_kvhead
+                        mK_kdl_ = mK_kdl[None, None, head_kv_coord, None]
+                        mV_dkl_ = mV_dkl[None, None, head_kv_coord, None]
+                        gK_kdl = cute.flat_divide(
+                            mK_kdl_, cute.select(self.qk_mma_tiler, mode=[1, 2])
+                        )
+                        tSgK_kdl = qk_thr_mma.partition_B(gK_kdl)
+                        tKsK, tKgK_kdl = cute.nvgpu.cpasync.tma_partition(
+                            tma_atom_k,
+                            block_in_cluster_coord_vmnk[1],
+                            kv_cta_layout,
+                            cute.group_modes(sK, 0, 3),
+                            cute.group_modes(tSgK_kdl, 0, 3),
+                        )
+                        gV_dkl = cute.flat_divide(
+                            mV_dkl_, cute.select(self.pv_mma_tiler, mode=[1, 2])
+                        )
+                        tSgV_dkl = pv_thr_mma.partition_B(gV_dkl)
+                        tVsV, tVgV_dkl = cute.nvgpu.cpasync.tma_partition(
+                            tma_atom_v,
+                            block_in_cluster_coord_vmnk[1],
+                            kv_cta_layout,
+                            cute.group_modes(sV, 0, 3),
+                            cute.group_modes(tSgV_dkl, 0, 3),
+                        )
+                        tKgK = tKgK_kdl
+                        tVgV = tVgV_dkl
+                    # ((atom_v, rest_v), RestK)
+                    tQgQ = tQgQ_qdl[None, mma_block_coord[0], None, mma_block_coord[2]]
+
+                    seqlen_kv_loop_start, seqlen_kv_loop_steps = (
+                        FusedMask.get_trip_start_count_via_block_info(
+                            mma_block_coord,
+                            self.qk_mma_tiler,
+                            seqlen_q,
+                            seqlen_k,
+                            self.is_causal,
+                            self.is_local,
+                            window_size_left,
+                            window_size_right,
+                        )
+                    )
+                    seqlen_kv_loop_end = seqlen_kv_loop_start + seqlen_kv_loop_steps
+                    # Q
+                    for iter in cutlass.range(self.iterations_qk, unroll=1):
+                        q_handle = load_q_producer.acquire_and_advance()
+                        cute.copy(
+                            tma_atom_q,
+                            tQgQ[None, iter],
+                            tQsQ[None, q_handle.index],
+                            tma_bar_ptr=q_handle.barrier,
+                        )
+
+                    # K0
+                    kv_coord = seqlen_kv_loop_start
+                    k_page_idx = (
+                        mPageTable[batch_coord, kv_coord]
+                        if cutlass.const_expr(mPageTable is not None)
+                        else None
+                    )
+                    for iter in cutlass.range(self.iterations_qk, unroll=1):
+                        k_handle = load_kv_producer.acquire_and_advance()
+                        cute.copy(
+                            tma_atom_k,
+                            tKgK[None, kv_coord, iter]
+                            if cutlass.const_expr(mPageTable is None)
+                            else tKgK[None, 0, iter, k_page_idx],
+                            tKsK[None, k_handle.index],
+                            tma_bar_ptr=k_handle.barrier,
+                        )
+                    kv_coord += 1
+                    # v_page_idx_prev carries K[i-1]'s page index for use as V[i-1]'s page
+                    # (K and V for the same KV block share the same physical page).
+                    # Also serves as the Vend page index when seqlen_kv_loop_steps == 1.
+                    v_page_idx_prev = (
+                        k_page_idx if cutlass.const_expr(mPageTable is not None) else None
+                    )
+                    # Prefetch K1 page after K0 TMA dispatch to hide L2 latency.
+                    if cutlass.const_expr(mPageTable is not None):
+                        if seqlen_kv_loop_steps > 1:
+                            k_page_idx = mPageTable[batch_coord, kv_coord]
+
+                    for i in cutlass.range(1, seqlen_kv_loop_steps, 1, unroll=1):
+                        # Ki: k_page_idx was prefetched at end of previous iteration
+                        # (or in the prologue for i==1); L2 latency already hidden.
+                        for iter in cutlass.range(self.iterations_qk, unroll=1):
+                            k_handle = load_kv_producer.acquire_and_advance()
+                            cute.copy(
+                                tma_atom_k,
+                                tKgK[None, kv_coord, iter]
+                                if cutlass.const_expr(mPageTable is None)
+                                else tKgK[None, 0, iter, k_page_idx],
+                                tKsK[None, k_handle.index],
+                                tma_bar_ptr=k_handle.barrier,
+                            )
+                        # Vi-1: reuse v_page_idx_prev (= K[i-1]'s page), no extra GMEM read.
+                        for iter in cutlass.range(self.iterations_pv, unroll=1):
+                            v_handle = load_kv_producer.acquire_and_advance()
+                            cute.copy(
+                                tma_atom_v,
+                                tVgV[None, iter, kv_coord - 1]
+                                if cutlass.const_expr(mPageTable is None)
+                                else tVgV[None, iter, 0, v_page_idx_prev],
+                                tVsV[None, v_handle.index],
+                                tma_bar_ptr=v_handle.barrier,
+                            )
+                        v_page_idx_prev = (
+                            k_page_idx if cutlass.const_expr(mPageTable is not None) else None
+                        )
+                        kv_coord += 1
+                        # Prefetch next K page while V TMA is in flight.
+                        if cutlass.const_expr(mPageTable is not None):
+                            if kv_coord < seqlen_kv_loop_end:
+                                k_page_idx = mPageTable[batch_coord, kv_coord]
+                    # Vend: reuse v_page_idx_prev (= K[end-1]'s page), no extra GMEM read.
+                    for iter in cutlass.range(self.iterations_pv, unroll=1):
+                        v_handle = load_kv_producer.acquire_and_advance()
+                        cute.copy(
+                            tma_atom_v,
+                            tVgV[None, iter, seqlen_kv_loop_end - 1]
+                            if cutlass.const_expr(mPageTable is None)
+                            else tVgV[None, iter, 0, v_page_idx_prev],
+                            tVsV[None, v_handle.index],
+                            tma_bar_ptr=v_handle.barrier,
+                        )
+
+                work_tile = tile_sched.advance_to_next_work()
+                # End of persistent scheduler loop
+            load_kv_producer.tail()
+            load_q_producer.tail()
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        #  MMA
+        # ///////////////////////////////////////////////////////////////////////////////
+        if warp_idx == self.mma_warp_id:
+            cute.arch.warpgroup_reg_dealloc(self.num_regs_other)
+
+            cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+            is_leader_cta = cta_rank_in_cluster % 2 == 0
+
+            while work_tile.is_valid_tile:
+                curr_block_coord = work_tile.tile_idx
+                mma_block_coord = (
+                    curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape),
+                    curr_block_coord[1],
+                    curr_block_coord[2],
+                )
+                continue_cond = False
+                seqlen_q = mQ_qdl.shape[0]
+                seqlen_k = (
+                    mK_kdl.shape[0] if cutlass.const_expr(mPageTable is None) else max_seqlen_k
+                )
+                batch_coord = curr_block_coord[2][1]
+                if cutlass.const_expr(cum_seqlen_q is not None):
+                    cuseqlen_q = cum_seqlen_q[batch_coord]
+                    seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q
+                    continue_cond = not FmhaStaticTileScheduler.check_valid_work_for_seqlen_q(
+                        self.qk_mma_tiler[0],
+                        mma_block_coord[0],
+                        seqlen_q,
+                    )
+
+                if not continue_cond:
+                    if cutlass.const_expr(cum_seqlen_k is not None):
+                        cuseqlen_k = cum_seqlen_k[batch_coord]
+                        seqlen_k = cum_seqlen_k[batch_coord + 1] - cuseqlen_k
+
+                    seqlen_kv_loop_start, seqlen_kv_loop_steps = (
+                        FusedMask.get_trip_start_count_via_block_info(
+                            mma_block_coord,
+                            self.qk_mma_tiler,
+                            seqlen_q,
+                            seqlen_k,
+                            self.is_causal,
+                            self.is_local,
+                            window_size_left,
+                            window_size_right,
+                        )
+                    )
+                    seqlen_kv_loop_end = seqlen_kv_loop_start + seqlen_kv_loop_steps
+
+                    load_q_releaser = load_q_consumer.clone()
+                    pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                    if seqlen_kv_loop_steps > 1:
+                        # QK0
+                        if is_leader_cta:
+                            s_handle = mma_s_producer.acquire_and_advance()
+                            tStS_slice = tStS[None, None, None, s_handle.index]
+                            qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                            for iter in cutlass.range(self.iterations_qk, unroll=1):
+                                load_q_consumer.wait_and_advance()
+                                tSrQ_slice = tSrQ[None, None, None, iter]
+                                k_handle = load_kv_consumer.wait_and_advance()
+                                tSrK_trans_slice = tSrK[None, None, None, k_handle.index]
+                                num_kphases = cute.size(tSrQ_slice, mode=[2])
+                                for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
+                                    kphase_coord = (None, None, kphase_idx)
+                                    cute.gemm(
+                                        qk_tiled_mma,
+                                        tStS_slice,
+                                        tSrQ_slice[kphase_coord],
+                                        tSrK_trans_slice[kphase_coord],
+                                        tStS_slice,
+                                    )
+                                    qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                k_handle.release()
+                            s_handle.commit()
+                        for i in cutlass.range(1, seqlen_kv_loop_steps - 1, 1, unroll=1):
+                            # QKi
+                            if is_leader_cta:
+                                s_handle = mma_s_producer.acquire_and_advance()
+                                tStS_slice = tStS[None, None, None, s_handle.index]
+                                qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                                for iter in cutlass.range(self.iterations_qk, unroll=1):
+                                    tSrQ_slice = tSrQ[None, None, None, iter]
+                                    k_handle = load_kv_consumer.wait_and_advance()
+                                    tSrK_trans_slice = tSrK[None, None, None, k_handle.index]
+                                    num_kphases = cute.size(tSrQ_slice, mode=[2])
+                                    for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
+                                        kphase_coord = (None, None, kphase_idx)
+                                        cute.gemm(
+                                            qk_tiled_mma,
+                                            tStS_slice,
+                                            tSrQ_slice[kphase_coord],
+                                            tSrK_trans_slice[kphase_coord],
+                                            tStS_slice,
+                                        )
+                                        qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                    k_handle.release()
+                                s_handle.commit()
+
+                                # PVi-1
+                                p_handle = p_mma_consumer.wait_and_advance()
+                                o_handle = mma_corr_producer.acquire_and_advance()
+                                pv_whether_acc = pv_tiled_mma.get(tcgen05.Field.ACCUMULATE)
+                                for iter in cutlass.range(self.iterations_pv, unroll=1):
+                                    v_handle = load_kv_consumer.wait_and_advance()
+                                    pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, pv_whether_acc)
+                                    tOtO_slice = tOtO_staged[None, None, None, iter]
+                                    tStS_slice = tStS[None, None, None, p_handle.index]
+                                    tP = cute.make_tensor(
+                                        tStS_slice.iterator, p_tmem_layout_staged.outer
+                                    )
+                                    tOrP = pv_thr_mma.make_fragment_A(tP)
+                                    tOrP_slice = cute.make_tensor(
+                                        cute.recast_ptr(tStS_slice.iterator, dtype=self.q_dtype),
+                                        tOrP.layout,
+                                    )
+                                    tOrV_slice = tOrV[None, None, None, v_handle.index]
+                                    num_kphases = cute.size(tOrV_slice, mode=[2])
+                                    for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
+                                        kphase_coord = (None, None, kphase_idx)
+                                        cute.gemm(
+                                            pv_tiled_mma,
+                                            tOtO_slice,
+                                            tOrP_slice[kphase_coord],
+                                            tOrV_slice[kphase_coord],
+                                            tOtO_slice,
+                                        )
+                                        pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                    v_handle.release()
+                                o_handle.commit()
+                                p_handle.release()
+                        if is_leader_cta:
+                            # QKend
+                            s_handle = mma_s_producer.acquire_and_advance()
+                            tStS_slice = tStS[None, None, None, s_handle.index]
+                            qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                            for iter in cutlass.range(self.iterations_qk, unroll=1):
+                                tSrQ_slice = tSrQ[None, None, None, iter]
+                                k_handle = load_kv_consumer.wait_and_advance()
+                                tSrK_trans_slice = tSrK[None, None, None, k_handle.index]
+                                num_kphases = cute.size(tSrQ_slice, mode=[2])
+                                for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
+                                    kphase_coord = (None, None, kphase_idx)
+                                    cute.gemm(
+                                        qk_tiled_mma,
+                                        tStS_slice,
+                                        tSrQ_slice[kphase_coord],
+                                        tSrK_trans_slice[kphase_coord],
+                                        tStS_slice,
+                                    )
+                                    qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                k_handle.release()
+                                load_q_releaser.release()
+                                load_q_releaser.advance()
+                            s_handle.commit()
+
+                            # PVend-1
+                            p_handle = p_mma_consumer.wait_and_advance()
+                            o_handle = mma_corr_producer.acquire_and_advance()
+                            pv_whether_acc = pv_tiled_mma.get(tcgen05.Field.ACCUMULATE)
+                            for iter in cutlass.range(self.iterations_pv, unroll=1):
+                                v_handle = load_kv_consumer.wait_and_advance()
+                                pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, pv_whether_acc)
+                                tOtO_slice = tOtO_staged[None, None, None, iter]
+                                tStS_slice = tStS[None, None, None, p_handle.index]
+                                tP = cute.make_tensor(
+                                    tStS_slice.iterator, p_tmem_layout_staged.outer
+                                )
+                                tOrP = pv_thr_mma.make_fragment_A(tP)
+                                tOrP_slice = cute.make_tensor(
+                                    cute.recast_ptr(tStS_slice.iterator, dtype=self.q_dtype),
+                                    tOrP.layout,
+                                )
+                                tOrV_slice = tOrV[None, None, None, v_handle.index]
+                                num_kphases = cute.size(tOrV_slice, mode=[2])
+                                for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
+                                    kphase_coord = (None, None, kphase_idx)
+                                    cute.gemm(
+                                        pv_tiled_mma,
+                                        tOtO_slice,
+                                        tOrP_slice[kphase_coord],
+                                        tOrV_slice[kphase_coord],
+                                        tOtO_slice,
+                                    )
+                                    pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                v_handle.release()
+                            o_handle.commit()
+                            p_handle.release()
+                    else:
+                        if is_leader_cta:
+                            # QK0
+                            s_handle = mma_s_producer.acquire_and_advance()
+                            tStS_slice = tStS[None, None, None, s_handle.index]
+                            qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                            for iter in cutlass.range(self.iterations_qk, unroll=1):
+                                load_q_consumer.wait_and_advance()
+                                tSrQ_slice = tSrQ[None, None, None, iter]
+                                k_handle = load_kv_consumer.wait_and_advance()
+                                tSrK_trans_slice = tSrK[None, None, None, k_handle.index]
+                                num_kphases = cute.size(tSrQ_slice, mode=[2])
+                                for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
+                                    kphase_coord = (None, None, kphase_idx)
+                                    cute.gemm(
+                                        qk_tiled_mma,
+                                        tStS_slice,
+                                        tSrQ_slice[kphase_coord],
+                                        tSrK_trans_slice[kphase_coord],
+                                        tStS_slice,
+                                    )
+                                    qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                                k_handle.release()
+                                load_q_releaser.release()
+                                load_q_releaser.advance()
+                            s_handle.commit()
+
+                    if is_leader_cta:
+                        # PVend
+                        p_handle = p_mma_consumer.wait_and_advance()
+                        o_handle = mma_corr_producer.acquire_and_advance()
+                        pv_whether_acc = pv_tiled_mma.get(tcgen05.Field.ACCUMULATE)
+                        for iter in cutlass.range(self.iterations_pv, unroll=1):
+                            v_handle = load_kv_consumer.wait_and_advance()
+                            pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, pv_whether_acc)
+                            tOtO_slice = tOtO_staged[None, None, None, iter]
+                            tStS_slice = tStS[None, None, None, p_handle.index]
+                            tP = cute.make_tensor(tStS_slice.iterator, p_tmem_layout_staged.outer)
+                            tOrP = pv_thr_mma.make_fragment_A(tP)
+                            tOrP_slice = cute.make_tensor(
+                                cute.recast_ptr(tStS_slice.iterator, dtype=self.q_dtype),
+                                tOrP.layout,
+                            )
+                            tOrV_slice = tOrV[None, None, None, v_handle.index]
+                            num_kphases = cute.size(tOrV_slice, mode=[2])
+                            for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
+                                kphase_coord = (None, None, kphase_idx)
+                                cute.gemm(
+                                    pv_tiled_mma,
+                                    tOtO_slice,
+                                    tOrP_slice[kphase_coord],
+                                    tOrV_slice[kphase_coord],
+                                    tOtO_slice,
+                                )
+                                pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                            v_handle.release()
+                        o_handle.commit()
+                        p_handle.release()
+                work_tile = tile_sched.advance_to_next_work()
+            # End of persistent scheduler loop
+            mma_s_producer.tail()
+            mma_corr_producer.tail()
+
+        if warp_idx < self.correction_warp_ids[0] and warp_idx >= self.softmax_warp_ids[0]:
+            # increase register after decreasing
+            cute.arch.warpgroup_reg_alloc(self.num_regs_softmax)
+
+            while work_tile.is_valid_tile:
+                curr_block_coord = work_tile.tile_idx
+                mma_block_coord = (
+                    curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape),
+                    curr_block_coord[1],
+                    curr_block_coord[2],
+                )
+                batch_coord = curr_block_coord[2][1]
+                continue_cond = False
+                seqlen_q = mQ_qdl.shape[0]
+                seqlen_k = (
+                    mK_kdl.shape[0] if cutlass.const_expr(mPageTable is None) else max_seqlen_k
+                )
+                cuseqlen_q = Int32(0)
+                if cutlass.const_expr(cum_seqlen_q is not None):
+                    cuseqlen_q = cum_seqlen_q[batch_coord]
+                    seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q
+                    continue_cond = not FmhaStaticTileScheduler.check_valid_work_for_seqlen_q(
+                        self.qk_mma_tiler[0],
+                        mma_block_coord[0],
+                        seqlen_q,
+                    )
+                if not continue_cond:
+                    if cutlass.const_expr(cum_seqlen_k is not None):
+                        cuseqlen_k = cum_seqlen_k[batch_coord]
+                        seqlen_k = cum_seqlen_k[batch_coord + 1] - cuseqlen_k
+
+                    row_max = -Float32.inf
+                    row_max_prev = -Float32.inf
+                    row_sum = 0.0
+
+                    start_count, trip_count = FusedMask.get_trip_start_count_via_block_info(
+                        mma_block_coord,
+                        self.qk_mma_tiler,
+                        seqlen_q,
+                        seqlen_k,
+                        self.is_causal,
+                        self.is_local,
+                        window_size_left,
+                        window_size_right,
+                    )
+                    end_count = start_count + trip_count
+                    # require at least one softmax iteration for zero trip_count case;
+                    # rely on masking this iteration for correctness
+                    if end_count <= start_count:
+                        start_count = 0
+                        end_count = 1
+                    if cutlass.const_expr(self.use_semantic_trip_range):
+                        n_block_min_causal_local_mask, n_block_min_before_local_mask = (
+                            FusedMask.get_trip_mask_bounds_via_block_info(
+                                mma_block_coord,
+                                self.qk_mma_tiler,
+                                seqlen_q,
+                                seqlen_k,
+                                self.is_causal,
+                                self.is_local,
+                                window_size_left,
+                                window_size_right,
+                            )
+                        )
+                    cS_base = cute.make_identity_tensor(
+                        (self.qk_mma_tiler[0], self.qk_mma_tiler[1])
+                    )
+                    cS = cute.domain_offset((mma_block_coord[0] * self.qk_mma_tiler[0], 0), cS_base)
+                    tScS = qk_thr_mma.partition_C(cS)
+
+                    for step in cutlass.range(start_count, end_count, 1, unroll=1):
+                        cS_iter = cute.domain_offset((0, step * self.qk_mma_tiler[1]), cS)
+                        tScS_iter = qk_thr_mma.partition_C(cS_iter)
+                        if cutlass.const_expr(self.use_semantic_trip_range):
+                            need_apply_mask = (
+                                step >= n_block_min_causal_local_mask
+                                or step < n_block_min_before_local_mask
+                                or step == end_count - 1
+                            )
+                        else:
+                            # Residual path only needs seqlen masking on the last K tile.
+                            need_apply_mask = step == end_count - 1
+                        # Si -> Pi
+                        (
+                            row_max,
+                            row_sum,
+                            mma_s_consumer,
+                            p_mma_producer,
+                            s_corr_producer,
+                        ) = self.softmax_step(
+                            (need_apply_mask, window_size_left, window_size_right),
+                            (
+                                row_max_prev,
+                                row_sum,
+                                seqlen_q,
+                                seqlen_k,
+                                scale_softmax_log2,
+                            ),
+                            (tStS, tScS_iter),
+                            (mma_s_consumer, p_mma_producer, s_corr_producer),
+                        )
+                        row_max_prev = row_max
+                    sum_producer = self.store_sum_max(
+                        row_max,
+                        mLSE,
+                        row_sum,
+                        sSum,
+                        sum_producer,
+                        curr_block_coord,
+                        seqlen_q,
+                        cum_seqlen_q,
+                        cuseqlen_q,
+                        scale_softmax,
+                    )
+                work_tile = tile_sched.advance_to_next_work()
+            p_mma_producer.tail()
+            s_corr_producer.tail()
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        #  Correction
+        # ///////////////////////////////////////////////////////////////////////////////
+        if warp_idx >= self.correction_warp_ids[0] and warp_idx < self.mma_warp_id:
+            cute.arch.warpgroup_reg_dealloc(self.num_regs_correction)
+
+            while work_tile.is_valid_tile:
+                curr_block_coord = work_tile.tile_idx
+                mma_block_coord = (
+                    curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape),
+                    curr_block_coord[1],
+                    curr_block_coord[2],
+                )
+                batch_coord = curr_block_coord[2][1]
+                seqlen_q = mQ_qdl.shape[0]
+                seqlen_k = (
+                    mK_kdl.shape[0] if cutlass.const_expr(mPageTable is None) else max_seqlen_k
+                )
+                continue_cond = False
+                cuseqlen_q = Int32(0)
+                if cutlass.const_expr(cum_seqlen_q is not None):
+                    cuseqlen_q = cum_seqlen_q[batch_coord]
+                    seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q
+                    continue_cond = not FmhaStaticTileScheduler.check_valid_work_for_seqlen_q(
+                        self.qk_mma_tiler[0],
+                        mma_block_coord[0],
+                        seqlen_q,
+                    )
+
+                if not continue_cond:
+                    if cutlass.const_expr(cum_seqlen_k is not None):
+                        cuseqlen_k = cum_seqlen_k[batch_coord]
+                        seqlen_k = cum_seqlen_k[batch_coord + 1] - cuseqlen_k
+
+                    mO_qdl_eff = mO_qdl
+                    if cutlass.const_expr(cum_seqlen_q is not None):
+                        block_offset_o = (
+                            cuseqlen_q,
+                            Int32(0),
+                            Int32(0),
+                            ((Int32(0), Int32(0)), Int32(0)),
+                        )
+                        mO_qdl_eff = cute.domain_offset(
+                            cute.select(block_offset_o, mode=[0, 2, 3]), mO_qdl
+                        )
+
+                    # (bM, bN, loopM, loopN, loopL)
+                    gO_qdl = cute.flat_divide(
+                        mO_qdl_eff, cute.select(self.pv_block_tiler, mode=[0, 1])
+                    )
+                    cO_qdl = cute.flat_divide(
+                        cute.make_identity_tensor(mO_qdl_eff.shape),
+                        cute.select(self.pv_block_tiler, mode=[0, 1]),
+                    )
+
+                    _, seqlen_kv_loop_steps = FusedMask.get_trip_start_count_via_block_info(
+                        mma_block_coord,
+                        self.qk_mma_tiler,
+                        seqlen_q,
+                        seqlen_k,
+                        self.is_causal,
+                        self.is_local,
+                        window_size_left,
+                        window_size_right,
+                    )
+                    gO_staged = gO_qdl[None, None, curr_block_coord[0], None, curr_block_coord[2]]
+                    cO_staged = cO_qdl[None, None, curr_block_coord[0], None, curr_block_coord[2]]
+                    cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1]))
+                    tScS = qk_thr_mma.partition_C(cS)
+
+                    # Empty step as the first step is no need for correction
+                    stats_handle = s_corr_consumer.wait_and_advance()
+                    stats_handle.release()
+                    for step in cutlass.range(1, seqlen_kv_loop_steps, 1, unroll=1):
+                        # Oi-1 -> Oi
+                        mma_corr_consumer, s_corr_consumer = self.correction_rescale(
+                            scale_softmax_log2,
+                            (s_corr_consumer, tStS, tScS),
+                            (mma_corr_consumer, tOtO_staged, cO_staged),
+                            self.epi_tile,
+                        )
+                    # O_partial -> O_final
+                    mma_corr_consumer, sum_consumer = self.correction_epilog(
+                        (seqlen_q, scale_output),
+                        (sum_consumer, sSum),
+                        (mma_corr_consumer, gO_staged, cO_staged, tOtO_staged),
+                        self.epi_tile,
+                    )
+                work_tile = tile_sched.advance_to_next_work()
+            # NOTE: tmem.free() moved to kernel end to enable cluster-wide sync
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        #  Scheduler Warp (only for CLC dynamic scheduler)
+        # ///////////////////////////////////////////////////////////////////////////////
+        if cutlass.const_expr(self.use_clc_scheduler):
+            is_first_cta_in_cluster = cta_rank_in_cluster == 0
+
+            if warp_idx == self.sched_warp_id and is_first_cta_in_cluster:
+                cute.arch.warpgroup_reg_dealloc(self.num_regs_other)
+                while work_tile.is_valid_tile:
+                    tile_sched.prefetch_next_work()
+                    work_tile = tile_sched.advance_to_next_work()
+                tile_sched.producer_tail()
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        #  Empty warps reg dealloc
+        # ///////////////////////////////////////////////////////////////////////////////
+        if cutlass.const_expr(self.use_clc_scheduler):
+            if warp_idx > self.load_warp_id:
+                if not (warp_idx == self.sched_warp_id and is_first_cta_in_cluster):
+                    cute.arch.warpgroup_reg_dealloc(self.num_regs_other)
+        else:
+            if warp_idx > self.load_warp_id:
+                cute.arch.warpgroup_reg_dealloc(self.num_regs_other)
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        #  Cooperative TMEM Deallocation (2CTA)
+        # ///////////////////////////////////////////////////////////////////////////////
+        # All warps (including scheduler) have finished by this point.
+        # Cluster-wide sync ensures both CTAs reach here before dealloc.
+        cute.arch.cluster_arrive()
+        cute.arch.cluster_wait()
+        tmem.relinquish_alloc_permit()
+        tmem.free(tmem_ptr)
+
+        return
+
+    @cute.jit
+    def softmax_step(
+        self,
+        mask_args: Tuple,
+        value_args: Tuple,
+        tensor_args: Tuple,
+        pipeline_args: Tuple,
+    ) -> Tuple[Float32, Float32, pipeline.PipelineConsumer, pipeline.PipelineProducer]:
+        need_apply_mask, window_size_left, window_size_right = mask_args
+        row_max, row_sum, seqlen_q, seqlen_k, scale_softmax_log2 = value_args
+        tStS, tScS = tensor_args
+        mma_s_consumer, p_mma_producer, s_corr_producer = pipeline_args
+        tidx, _, _ = cute.arch.thread_idx()
+        thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids))
+        s_handle = mma_s_consumer.wait_and_advance()
+        tStS_slice = tStS[(None, None), 0, 0, s_handle.index]
+        tScS_slice = tScS[(None, None), 0, 0]
+        tmem_load_atom = cute.make_copy_atom(
+            tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), self.qk_acc_dtype
+        )
+        tmem_tiled_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS_slice)
+        thr_load = tmem_tiled_load.get_slice(thread_idx)
+        tTMEM_LOADtS = thr_load.partition_S(tStS_slice)
+        tTMEM_LOADcS = thr_load.partition_D(tScS_slice)
+        tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype)
+        cute.copy(tmem_tiled_load, tTMEM_LOADtS, tTMEM_LOADrS)
+
+        cute.arch.fence_view_async_tmem_load()
+        s_handle.release()
+        if need_apply_mask:
+            FusedMask.apply_mask_via_causal_local(
+                tTMEM_LOADrS,
+                tTMEM_LOADcS,
+                seqlen_q,
+                seqlen_k,
+                self.use_semantic_trip_range,
+                self.is_causal,
+                self.is_local,
+                window_size_left,
+                window_size_right,
+            )
+        old_row_max = row_max
+        row_max = tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, row_max, 0)
+        row_max_safe = row_max
+        if row_max == -cutlass.Float32.inf:
+            row_max_safe = 0.0
+
+        stats_handle = s_corr_producer.acquire_and_advance()
+        stats_layout = cute.composition(
+            tStS_slice.layout, cute.make_layout((tStS_slice.shape[0], 2))
+        )
+        stats_c_layout = cute.composition(
+            tScS_slice.layout, cute.make_layout((tScS_slice.shape[0], 2))
+        )
+        tOtStats = cute.make_tensor(tStS_slice.iterator + self.tilePlikeFP32, stats_layout)
+        tOcStats = cute.make_tensor(tScS_slice.iterator, stats_c_layout)
+        tmem_store_stats_atom = cute.make_copy_atom(
+            tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(2)),
+            self.qk_acc_dtype,
+        )
+        tiled_tmem_store_stats = tcgen05.make_tmem_copy(tmem_store_stats_atom, tOtStats)
+        thr_tmem_store_stats = tiled_tmem_store_stats.get_slice(thread_idx)
+        tTMEM_STOREcStats = thr_tmem_store_stats.partition_S(tOcStats)
+        tTMEM_STORErStats = cute.make_rmem_tensor(tTMEM_STOREcStats.shape, self.qk_acc_dtype)
+        tTMEM_STORErStats[0] = old_row_max
+        tTMEM_STORErStats[1] = row_max_safe
+        tTMEM_STOREtStats = thr_tmem_store_stats.partition_D(tOtStats)
+        cute.copy(tiled_tmem_store_stats, tTMEM_STORErStats, tTMEM_STOREtStats)
+        cute.arch.fence_view_async_tmem_store()
+        stats_handle.commit()
+
+        scale = scale_softmax_log2
+        minus_row_max_scale = (0.0 - row_max_safe) * scale
+        # Acquire P write slot early — overlaps any pipeline stall with exp2 compute
+        p_handle = p_mma_producer.acquire_and_advance()
+        # Fragment-based FMA + exp2 + bf16 conversion
+        # Trades SFU for FMA via polynomial emulation on a fraction of elements
+        ex2_frg_tile = 32
+        ex2_frg_cnt = cute.size(tTMEM_LOADrS) // ex2_frg_tile
+        tTMEM_LOADrS_ex2 = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(ex2_frg_tile))
+        tTMEM_STORErP = cute.make_rmem_tensor(tTMEM_LOADrS.shape, self.q_dtype)
+        tTMEM_STORErP_ex2 = cute.logical_divide(tTMEM_STORErP, cute.make_layout(ex2_frg_tile))
+        for j in cutlass.range_constexpr(ex2_frg_cnt):
+            for k in cutlass.range_constexpr(0, ex2_frg_tile, 2):
+                tTMEM_LOADrS_ex2[k, j], tTMEM_LOADrS_ex2[k + 1, j] = cute.arch.fma_packed_f32x2(
+                    (tTMEM_LOADrS_ex2[k, j], tTMEM_LOADrS_ex2[k + 1, j]),
+                    (scale, scale),
+                    (minus_row_max_scale, minus_row_max_scale),
+                )
+                if cutlass.const_expr(self.ex2_emu_freq == 0):
+                    tTMEM_LOADrS_ex2[k, j] = cute.math.exp2(tTMEM_LOADrS_ex2[k, j], fastmath=True)
+                    tTMEM_LOADrS_ex2[k + 1, j] = cute.math.exp2(
+                        tTMEM_LOADrS_ex2[k + 1, j], fastmath=True
+                    )
+                else:
+                    if cutlass.const_expr(
+                        k % self.ex2_emu_freq < self.ex2_emu_freq - self.ex2_emu_res
+                        or j >= ex2_frg_cnt - 1
+                        or j < self.ex2_emu_start_frg
+                    ):
+                        tTMEM_LOADrS_ex2[k, j] = cute.math.exp2(
+                            tTMEM_LOADrS_ex2[k, j], fastmath=True
+                        )
+                        tTMEM_LOADrS_ex2[k + 1, j] = cute.math.exp2(
+                            tTMEM_LOADrS_ex2[k + 1, j], fastmath=True
+                        )
+                    else:
+                        tTMEM_LOADrS_ex2[k, j], tTMEM_LOADrS_ex2[k + 1, j] = ex2_emulation_2(
+                            tTMEM_LOADrS_ex2[k, j], tTMEM_LOADrS_ex2[k + 1, j]
+                        )
+            tTMEM_STORErP_ex2[None, j].store(tTMEM_LOADrS_ex2[None, j].load().to(self.q_dtype))
+        tmem_store_atom = cute.make_copy_atom(
+            tcgen05.St32x32bOp(tcgen05.Repetition(32)), self.qk_acc_dtype
+        )
+        tilePlikeFP32 = tStS_slice.shape[1] // Float32.width * self.q_dtype.width
+        tStS_P_layout = cute.composition(
+            tStS_slice.layout, cute.make_layout((tStS_slice.shape[0], tilePlikeFP32))
+        )
+        tStS_P = cute.make_tensor(tStS_slice.iterator, tStS_P_layout)
+        tScS_P_layout = cute.composition(
+            tScS_slice.layout, cute.make_layout((tScS_slice.shape[0], tilePlikeFP32))
+        )
+        tScS_P = cute.make_tensor(tScS_slice.iterator, tScS_P_layout)
+        tmem_tiled_store = tcgen05.make_tmem_copy(tmem_store_atom, tStS_P)
+        thr_store = tmem_tiled_store.get_slice(thread_idx)
+        tTMEM_STOREtP = thr_store.partition_D(tStS_P)
+        tTMEM_STOREcS = thr_store.partition_S(tScS_P)
+        tTMEM_STORErP_ = cute.make_tensor(
+            cute.recast_ptr(tTMEM_STORErP.iterator, dtype=self.qk_acc_dtype),
+            tTMEM_STOREcS.shape,
+        )
+        cute.copy(tmem_tiled_store, tTMEM_STORErP_, tTMEM_STOREtP)
+        cute.arch.fence_view_async_tmem_store()
+
+        p_handle.commit()
+        acc_scale_ = scale * (old_row_max - row_max_safe)
+        acc_scale = cute.math.exp2(acc_scale_, fastmath=True) * 0.5
+        # TODO: calc row sum with TensorSSA
+        row_sum *= acc_scale
+        local_row_sum_0 = (row_sum, row_sum)
+        local_row_sum_1 = (0.0, 0.0)
+        local_row_sum_2 = (0.0, 0.0)
+        local_row_sum_3 = (0.0, 0.0)
+        reduction_unroll = 4
+        frg_tile = cute.size(tTMEM_LOADrS) // reduction_unroll
+        tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile))
+        for j in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS_frg, mode=[0]), 2):
+            local_row_sum_0 = cute.arch.add_packed_f32x2(
+                local_row_sum_0, (tTMEM_LOADrS_frg[j, 0], tTMEM_LOADrS_frg[j + 1, 0])
+            )
+            local_row_sum_1 = cute.arch.add_packed_f32x2(
+                local_row_sum_1, (tTMEM_LOADrS_frg[j, 1], tTMEM_LOADrS_frg[j + 1, 1])
+            )
+            local_row_sum_2 = cute.arch.add_packed_f32x2(
+                local_row_sum_2, (tTMEM_LOADrS_frg[j, 2], tTMEM_LOADrS_frg[j + 1, 2])
+            )
+            local_row_sum_3 = cute.arch.add_packed_f32x2(
+                local_row_sum_3, (tTMEM_LOADrS_frg[j, 3], tTMEM_LOADrS_frg[j + 1, 3])
+            )
+        local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_1)
+        local_row_sum_2 = cute.arch.add_packed_f32x2(local_row_sum_2, local_row_sum_3)
+        local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_2)
+        row_sum = local_row_sum_0[0] + local_row_sum_0[1]
+        return row_max, row_sum, mma_s_consumer, p_mma_producer, s_corr_producer
+
+    @cute.jit
+    def correction_rescale(
+        self,
+        scale_softmax_log2: Float32,
+        stats_args: tuple,
+        o_args: tuple,
+        epi_tile: cute.Tile,
+    ) -> pipeline.PipelineConsumer:
+        (s_corr_consumer, tStS, tScS) = stats_args
+        (mma_o_consumer, tOtO_staged, cO_staged) = o_args
+        tidx, _, _ = cute.arch.thread_idx()
+        thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids))
+
+        stats_handle = s_corr_consumer.wait_and_advance()
+        tStS_slice = tStS[(None, None), 0, 0, stats_handle.index]
+        tScS_slice = tScS[(None, None), 0, 0]
+        stats_layout = cute.composition(
+            tStS_slice.layout, cute.make_layout((tStS_slice.shape[0], 2))
+        )
+        stats_c_layout = cute.composition(
+            tScS_slice.layout, cute.make_layout((tScS_slice.shape[0], 2))
+        )
+        tOtStats = cute.make_tensor(tStS_slice.iterator + self.tilePlikeFP32, stats_layout)
+        tOcStats = cute.make_tensor(tScS_slice.iterator, stats_c_layout)
+        tmem_load_stats_atom = cute.make_copy_atom(
+            tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(2)),
+            self.qk_acc_dtype,
+        )
+        tiled_tmem_load_stats = tcgen05.make_tmem_copy(tmem_load_stats_atom, tOtStats)
+        thr_tmem_load_stats = tiled_tmem_load_stats.get_slice(thread_idx)
+        tTMEM_LOADtStats = thr_tmem_load_stats.partition_S(tOtStats)
+        tTMEM_LOADcStats = thr_tmem_load_stats.partition_D(tOcStats)
+        tTMEM_LOADrStats = cute.make_rmem_tensor(tTMEM_LOADcStats.shape, self.qk_acc_dtype)
+        cute.copy(tiled_tmem_load_stats, tTMEM_LOADtStats, tTMEM_LOADrStats)
+
+        scale = scale_softmax_log2 * (tTMEM_LOADrStats[0] - tTMEM_LOADrStats[1])
+        scale = cute.math.exp2(scale, fastmath=True)
+        stats_handle.release()
+        o_handle = mma_o_consumer.wait_and_advance()
+        for iter in cutlass.range(self.iterations_pv, unroll_full=True):
+            tOtO = tOtO_staged[(None, None), 0, 0, iter]
+            cO = cO_staged[None, None, iter]
+            tOtO_epi = cute.zipped_divide(tOtO, epi_tile)
+            cO_epi = cute.zipped_divide(cO, epi_tile)
+            tmem_load_atom = cute.make_copy_atom(
+                tcgen05.Ld32x32bOp(tcgen05.Repetition(16)),
+                self.pv_acc_dtype,
+            )
+            tmem_tiled_load = tcgen05.make_tmem_copy(tmem_load_atom, tOtO_epi)
+            thr_load = tmem_tiled_load.get_slice(thread_idx)
+            tmem_store_atom = cute.make_copy_atom(
+                tcgen05.St32x32bOp(tcgen05.Repetition(16)),
+                self.pv_acc_dtype,
+            )
+            tmem_store_atom = tcgen05.make_tmem_copy(tmem_store_atom, tOtO_epi)
+            thr_store = tmem_store_atom.get_slice(thread_idx)
+            tTMEM_LOADtO = thr_load.partition_S(tOtO_epi)
+            tTMEM_LOADcO = thr_load.partition_D(cO_epi)
+            tTMEM_STOREtO = thr_store.partition_D(tOtO_epi)
+            tTMrO = cute.make_rmem_tensor_like(
+                cute.append(
+                    cute.make_layout(tTMEM_LOADcO[None, 0, 0].shape),
+                    cute.make_layout(2, stride=cute.size(tTMEM_LOADcO[None, 0, 0].shape)),
+                ),
+                self.pv_acc_dtype,
+            )
+            tTMEM_LOADtO_0 = tTMEM_LOADtO[None, 0, 0]
+            cute.copy(tmem_tiled_load, tTMEM_LOADtO_0, tTMrO[None, 0])
+            iter_num = cute.size(tTMEM_LOADtO, mode=[1])
+            for i in cutlass.range(1, iter_num, unroll_full=True):
+                tTMEM_LOADtO_i = tTMEM_LOADtO[None, i, 0]
+                cute.copy(tmem_tiled_load, tTMEM_LOADtO_i, tTMrO[None, i % 2])
+                for j in cutlass.range(0, cute.size(tTMrO, mode=[0]), 2, unroll_full=True):
+                    tTMrO[j, (i - 1) % 2], tTMrO[j + 1, (i - 1) % 2] = cute.arch.mul_packed_f32x2(
+                        (tTMrO[j, (i - 1) % 2], tTMrO[j + 1, (i - 1) % 2]),
+                        (scale, scale),
+                    )
+                tTMEM_STOREtO_prev_i = tTMEM_STOREtO[None, i - 1, 0]
+                cute.copy(tmem_store_atom, tTMrO[None, (i - 1) % 2], tTMEM_STOREtO_prev_i)
+
+            for j in cutlass.range(0, cute.size(tTMrO, mode=[0]), 2, unroll_full=True):
+                tTMrO[j, (iter_num - 1) % 2], tTMrO[j + 1, (iter_num - 1) % 2] = (
+                    cute.arch.mul_packed_f32x2(
+                        (
+                            tTMrO[j, (iter_num - 1) % 2],
+                            tTMrO[j + 1, (iter_num - 1) % 2],
+                        ),
+                        (scale, scale),
+                    )
+                )
+            cute.copy(
+                tmem_store_atom,
+                tTMrO[None, (iter_num - 1) % 2],
+                tTMEM_STOREtO[None, iter_num - 1, 0],
+            )
+        cute.arch.fence_view_async_tmem_store()
+        o_handle.release()
+        return mma_o_consumer, s_corr_consumer
+
+    @cute.jit
+    def correction_epilog(
+        self,
+        value_args: Tuple,
+        sum_args: Tuple,
+        o_args: Tuple,
+        epi_tile: cute.Tile,
+    ) -> Tuple[pipeline.PipelineConsumer, pipeline.PipelineProducer]:
+        (seqlen_q, scale_output) = value_args
+        (sum_consumer, sSum) = sum_args
+        (mma_o_consumer, gO_staged, cO_staged, tOtO_staged) = o_args
+        tidx, _, _ = cute.arch.thread_idx()
+        thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids))
+        sum_handle = sum_consumer.wait_and_advance()
+        row_sum = sSum[thread_idx]
+        cute.arch.fence_view_async_shared()
+        sum_handle.release()
+        row_sum_is_zero_or_nan = row_sum == 0.0 or row_sum != row_sum
+        scale = scale_output / row_sum if not row_sum_is_zero_or_nan else 0.0
+        o_handle = mma_o_consumer.wait_and_advance()
+        for iter in cutlass.range(self.iterations_pv):
+            gO = gO_staged[None, None, iter]
+            cO = cO_staged[None, None, iter]
+            tOtO = tOtO_staged[(None, None), 0, 0, iter]
+            tOtO_epi = cute.zipped_divide(tOtO, epi_tile)
+            cO_epi = cute.zipped_divide(cO, epi_tile)
+            gO_epi = cute.zipped_divide(gO, epi_tile)
+            tidx, _, _ = cute.arch.thread_idx()
+            thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids))
+            tmem_copy_atom = cute.make_copy_atom(
+                tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.pv_acc_dtype
+            )
+            tiled_tmem_load = tcgen05.make_tmem_copy(tmem_copy_atom, tOtO_epi)
+            thr_tmem_load = tiled_tmem_load.get_slice(thread_idx)
+            tTMEM_LOADtO = thr_tmem_load.partition_S(tOtO_epi)
+            tTMEM_LOADgO = thr_tmem_load.partition_D(gO_epi)
+            tTMEM_LOADcO = thr_tmem_load.partition_D(cO_epi)
+            for i in cutlass.range(cute.size(tTMEM_LOADtO, mode=[1]), unroll_full=True):
+                tTMEM_LOADtO_i = tTMEM_LOADtO[None, i, 0]
+                tTMEM_LOADgO_i = tTMEM_LOADgO[None, i, 0]
+                tTMEM_LOADcO_i = tTMEM_LOADcO[None, i, 0]
+                tTMrO = cute.make_rmem_tensor(tTMEM_LOADcO[None, 0, i].shape, self.pv_acc_dtype)
+                cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO)
+                for j in cutlass.range(0, cute.size(tTMrO), 2, unroll_full=True):
+                    tTMrO[j], tTMrO[j + 1] = cute.arch.mul_packed_f32x2(
+                        (tTMrO[j], tTMrO[j + 1]),
+                        (scale, scale),
+                    )
+                tSMrO = cute.make_rmem_tensor(tTMrO.shape, self.o_dtype)
+                o_vec = tTMrO.load()
+                tSMrO.store(o_vec.to(self.o_dtype))
+                if cute.elem_less(tTMEM_LOADcO_i[0][0], seqlen_q):
+                    cute.autovec_copy(tSMrO, tTMEM_LOADgO_i)
+        o_handle.release()
+        return mma_o_consumer, sum_consumer
+
+    @cute.jit
+    def store_sum_max(
+        self,
+        row_max,
+        mLSE,
+        row_sum,
+        sSum,
+        sum_producer,
+        current_block_coord,
+        seqlen_q,
+        cum_seqlen_q,
+        cuseqlen_q,
+        scale_softmax,
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids))
+        sum_handle = sum_producer.acquire_and_advance()
+        sSum[thread_idx] = row_sum
+        cute.arch.fence_view_async_shared()
+        sum_handle.commit()
+        row_sum_is_zero_or_nan = row_sum == 0.0 or row_sum != row_sum
+
+        if cutlass.const_expr(mLSE is not None):
+            q_idx = current_block_coord[0] * self.cta_tiler[0] + tidx
+            hb_idx = (
+                (current_block_coord[2][0], Int32(0))
+                if cutlass.const_expr(cum_seqlen_q is not None)
+                else current_block_coord[2]
+            )
+            lse_value = (
+                scale_softmax * row_max + cute.math.log(row_sum, fastmath=True)
+                if not row_sum_is_zero_or_nan
+                else -Float32.inf
+            )
+            if cute.elem_less(q_idx, seqlen_q):
+                global_q_idx = (
+                    q_idx + cuseqlen_q if cutlass.const_expr(cum_seqlen_q is not None) else q_idx
+                )
+                mLSE[global_q_idx, hb_idx] = lse_value
+        return sum_producer
diff --git a/build/torch-cuda/softmax.py b/build/torch-cuda/softmax.py
index 41e6cba1688798dc75183a977946f6671d38c6e9..fd123a2c0de5d5a3a6f3e3f81c5b2fce225f3080 100644
--- a/build/torch-cuda/softmax.py
+++ b/build/torch-cuda/softmax.py
@@ -7,12 +7,86 @@ from dataclasses import dataclass
 
 import cutlass
 import cutlass.cute as cute
-from cutlass import Float32
+from cutlass import Float32, Boolean
 
 from .quack import layout_utils
-from . import utils
+from . import utils as utils
 from .quack.cute_dsl_utils import ParamsBase
 from .seqlen_info import SeqlenInfoQK
+from .utils import AuxData
+
+
+@cute.jit
+def call_score_mod(
+    score_mod: cutlass.Constexpr,
+    score,
+    batch_idx,
+    head_idx,
+    q_idx,
+    kv_idx,
+    seqlen_info,
+    aux_data: AuxData,
+):
+    aux_tensors = aux_data.tensors if aux_data.tensors is not None else ()
+    # Compatibility shim for pre-aux_scalars score_mod callables.
+    if cutlass.const_expr(aux_data.scalars is not None):
+        return score_mod(
+            score,
+            batch_idx,
+            head_idx,
+            q_idx=q_idx,
+            kv_idx=kv_idx,
+            seqlen_info=seqlen_info,
+            aux_tensors=aux_tensors,
+            aux_scalars=aux_data.scalars,
+        )
+    return score_mod(
+        score,
+        batch_idx,
+        head_idx,
+        q_idx=q_idx,
+        kv_idx=kv_idx,
+        seqlen_info=seqlen_info,
+        aux_tensors=aux_tensors,
+    )
+
+
+@cute.jit
+def call_score_mod_bwd(
+    score_mod_bwd: cutlass.Constexpr,
+    grad,
+    score,
+    batch_idx,
+    head_idx,
+    q_idx,
+    kv_idx,
+    seqlen_info,
+    aux_data: AuxData,
+):
+    aux_tensors = aux_data.tensors if aux_data.tensors is not None else ()
+    # Compatibility shim for pre-aux_scalars score_mod_bwd callables.
+    if cutlass.const_expr(aux_data.scalars is not None):
+        return score_mod_bwd(
+            grad,
+            score,
+            batch_idx,
+            head_idx,
+            q_idx=q_idx,
+            kv_idx=kv_idx,
+            seqlen_info=seqlen_info,
+            aux_tensors=aux_tensors,
+            aux_scalars=aux_data.scalars,
+        )
+    return score_mod_bwd(
+        grad,
+        score,
+        batch_idx,
+        head_idx,
+        q_idx=q_idx,
+        kv_idx=kv_idx,
+        seqlen_info=seqlen_info,
+        aux_tensors=aux_tensors,
+    )
 
 
 @dataclass
@@ -131,12 +205,13 @@ class Softmax(ParamsBase):
         row_scale = cute.make_fragment_like(row_max, Float32)
 
         for r in cutlass.range(cute.size(row_sum), unroll_full=True):
+            row_max_scaled = row_max[r] * scale_log2
             if cutlass.const_expr(sink_val is not None):
                 sink_val_cur = sink_val if not isinstance(sink_val, cute.Tensor) else sink_val[r]
                 LOG2_E = math.log2(math.e)
-                row_sum[r] += cute.math.exp2(
-                    sink_val_cur * LOG2_E - row_max[r] * scale_log2, fastmath=True
-                )
+                if row_max[r] == -Float32.inf:
+                    row_max_scaled = sink_val_cur * LOG2_E
+                row_sum[r] += cute.math.exp2(sink_val_cur * LOG2_E - row_max_scaled, fastmath=True)
 
             # if row_sum is zero or nan, set acc_O_mn_row to 1.0
             acc_O_mn_row_is_zero_or_nan = row_sum[r] == 0.0 or row_sum[r] != row_sum[r]
@@ -146,7 +221,7 @@ class Softmax(ParamsBase):
             row_sum_cur = row_sum[r]
             LN2 = math.log(2.0)
             row_sum[r] = (
-                (row_max[r] * scale_log2 + cute.math.log2(row_sum_cur, fastmath=True)) * LN2
+                (row_max_scaled + cute.math.log2(row_sum_cur, fastmath=True)) * LN2
                 if not acc_O_mn_row_is_zero_or_nan
                 else -Float32.inf
             )
@@ -169,12 +244,14 @@ class Softmax(ParamsBase):
 @dataclass
 class SoftmaxSm100(Softmax):
     rescale_threshold: cutlass.Constexpr[float] = 0.0
+    max_offset: cutlass.Constexpr[int] = 0
 
     @staticmethod
     def create(
         scale_log2: Float32,
         rescale_threshold: cutlass.Constexpr[float] = 0.0,
         softmax_scale: Float32 | None = None,
+        max_offset: cutlass.Constexpr[int] = 0,
     ):
         num_rows = 1
         arch = 100
@@ -188,8 +265,52 @@ class SoftmaxSm100(Softmax):
             arch,
             softmax_scale,
             rescale_threshold=rescale_threshold,
+            max_offset=max_offset,
         )
 
+    @cute.jit
+    def compute_row_max_local(self, acc_S_row: cute.TensorSSA, is_first: Boolean) -> Float32:
+        if cutlass.const_expr(is_first):
+            row_max_new = self._compute_row_max(acc_S_row)
+        else:
+            row_max_old = self.row_max[0]
+            row_max_new = self._compute_row_max(acc_S_row, init_val=row_max_old)
+        return row_max_new
+
+    @cute.jit
+    def update_row_max_from_local(
+        self,
+        row_max_new: Float32,
+        is_first: Boolean,
+    ) -> Tuple[Float32, Float32]:
+        if cutlass.const_expr(is_first):
+            row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0
+            acc_scale = 0.0
+        else:
+            row_max_old = self.row_max[0]
+            row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0
+            acc_scale_ = (row_max_old - row_max_safe) * self.scale_log2
+            acc_scale = cute.math.exp2(acc_scale_)
+            if cutlass.const_expr(self.rescale_threshold > 0.0):
+                if acc_scale_ >= -self.rescale_threshold:
+                    row_max_new = row_max_old
+                    row_max_safe = row_max_old
+                    acc_scale = 1.0
+        self.row_max[0] = row_max_new
+        return row_max_safe, acc_scale
+
+    @cute.jit
+    def update_row_max_precomputed(
+        self, hw_row_max: Float32, is_first: int
+    ) -> Tuple[Float32, Float32]:
+        """Row max already reduced in hardware (SM103 tcgen05.ld.red): skip the
+        software fmax tree — the TMEM controller computed the max during the S load."""
+        if cutlass.const_expr(is_first):
+            row_max_new = hw_row_max
+        else:
+            row_max_new = cute.arch.fmax(hw_row_max, self.row_max[0])
+        return self.update_row_max_from_local(row_max_new, is_first)
+
     @cute.jit
     def update_row_max(self, acc_S_row: cute.TensorSSA, is_first: int) -> Tuple[Float32, Float32]:
         if cutlass.const_expr(is_first):
@@ -227,11 +348,13 @@ class SoftmaxSm100(Softmax):
     ):
         assert cute.size(acc_S_row.shape) % 2 == 0, "acc_S_row must have an even number of elements"
         row_max_scaled = row_max * self.scale_log2
+        max_offset = Float32(self.max_offset)
+        bias = max_offset - row_max_scaled
         for i in cutlass.range(0, cute.size(acc_S_row.shape), 2, unroll_full=True):
             acc_S_row[i], acc_S_row[i + 1] = cute.arch.fma_packed_f32x2(
                 (acc_S_row[i], acc_S_row[i + 1]),
                 (self.scale_log2, self.scale_log2),
-                (-row_max_scaled, -row_max_scaled),
+                (bias, bias),
             )
 
     @cute.jit
@@ -350,7 +473,7 @@ def apply_score_mod_inner(
     softmax_scale,
     vec_size: cutlass.Constexpr,
     qk_acc_dtype: cutlass.Constexpr,
-    aux_tensors,
+    aux_data: AuxData,
     fastdiv_mods,
     seqlen_info: SeqlenInfoQK,
     constant_q_idx: cutlass.Constexpr,
@@ -369,6 +492,7 @@ def apply_score_mod_inner(
         vec_size: Vector size for processing elements
         qk_acc_dtype: Data type for accumulator
         aux_tensors: Optional aux_tensors for FlexAttention
+        aux_scalars: Optional runtime scalar captures for FlexAttention
         fastdiv_mods: Tuple of (seqlen_q_divmod, seqlen_k_divmod) for wrapping
         seqlen_info: Sequence length info
         constant_q_idx: If provided, use this constant for all q_idx values
@@ -398,7 +522,7 @@ def apply_score_mod_inner(
     q_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32)
 
     # For Pack-GQA with non-constant q_idx, we need per-element head indices
-    # since a thread my process multiple query head indices
+    # since a thread may process multiple query head indices
     if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None):
         head_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32)
 
@@ -415,7 +539,7 @@ def apply_score_mod_inner(
                 head_idx_vec[j] = head_idx * qhead_per_kvhead + head_offset
 
             # If we will do loads we mod, in order to not read OOB
-            if cutlass.const_expr(aux_tensors is not None and fastdiv_mods is not None):
+            if cutlass.const_expr(aux_data.tensors is not None and fastdiv_mods is not None):
                 if cutlass.const_expr(constant_q_idx is None):
                     seqlen_q_divmod, seqlen_k_divmod = fastdiv_mods
                     q_idx_floored = floor_if_packed(
@@ -450,18 +574,15 @@ def apply_score_mod_inner(
         else:
             head_idx_ssa = utils.scalar_to_ssa(head_idx, cutlass.Int32).broadcast_to((vec_size,))
 
-        aux_args = []
-        if cutlass.const_expr(aux_tensors is not None):
-            aux_args = aux_tensors
-
-        post_mod_scores = score_mod(
+        post_mod_scores = call_score_mod(
+            score_mod,
             score_ssa,
             batch_idx_ssa,
             head_idx_ssa,
-            q_idx=q_idx_ssa,
-            kv_idx=kv_idx_ssa,
-            seqlen_info=seqlen_info,
-            aux_tensors=aux_args,
+            q_idx_ssa,
+            kv_idx_ssa,
+            seqlen_info,
+            aux_data,
         )
 
         # Write back modified scores
@@ -481,7 +602,7 @@ def apply_score_mod_bwd_inner(
     softmax_scale,
     vec_size: cutlass.Constexpr,
     qk_acc_dtype: cutlass.Constexpr,
-    aux_tensors,
+    aux_data: AuxData,
     fastdiv_mods,
     seqlen_info,
     constant_q_idx: cutlass.Constexpr,
@@ -501,6 +622,7 @@ def apply_score_mod_bwd_inner(
         vec_size: Vector size for processing elements
         qk_acc_dtype: Data type for accumulator
         aux_tensors: Optional aux_tensors for FlexAttention
+        aux_scalars: Optional runtime scalar captures for FlexAttention
         fastdiv_mods: Tuple of (seqlen_q_divmod, seqlen_k_divmod) for wrapping
         seqlen_info: Sequence length info
         constant_q_idx: If provided, use this constant for all q_idx values
@@ -517,15 +639,15 @@ def apply_score_mod_bwd_inner(
         q_idx_pos = cutlass.const_expr(0)
         kv_idx_pos = cutlass.const_expr(1)
     n_vals = cutlass.const_expr(cute.size(grad_tensor.shape))
-    grad_vec = cute.make_fragment(vec_size, qk_acc_dtype)
-    score_vec = cute.make_fragment(vec_size, qk_acc_dtype)
-    kv_idx_vec = cute.make_fragment(vec_size, cutlass.Int32)
+    grad_vec = cute.make_rmem_tensor(vec_size, qk_acc_dtype)
+    score_vec = cute.make_rmem_tensor(vec_size, qk_acc_dtype)
+    kv_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32)
     batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32).broadcast_to((vec_size,))
-    q_idx_vec = cute.make_fragment(vec_size, cutlass.Int32)
+    q_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32)
 
     # For Pack-GQA with non-constant q_idx, we need per-element head indices
     if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None):
-        head_idx_vec = cute.make_fragment(vec_size, cutlass.Int32)
+        head_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32)
 
     for i in cutlass.range(0, n_vals, vec_size, unroll_full=True):
         for j in cutlass.range(vec_size, unroll_full=True):
@@ -539,7 +661,7 @@ def apply_score_mod_bwd_inner(
                 head_offset = q_idx_packed - q_idx_logical * qhead_per_kvhead
                 head_idx_vec[j] = head_idx * qhead_per_kvhead + head_offset
 
-            if cutlass.const_expr(aux_tensors is not None and fastdiv_mods is not None):
+            if cutlass.const_expr(aux_data.tensors is not None and fastdiv_mods is not None):
                 if cutlass.const_expr(constant_q_idx is None):
                     seqlen_q_divmod, seqlen_k_divmod = fastdiv_mods
                     q_idx_floored = floor_if_packed(
@@ -572,19 +694,16 @@ def apply_score_mod_bwd_inner(
         else:
             head_idx_ssa = utils.scalar_to_ssa(head_idx, cutlass.Int32).broadcast_to((vec_size,))
 
-        aux_args = []
-        if cutlass.const_expr(aux_tensors is not None):
-            aux_args = aux_tensors
-
-        grad_out_ssa = score_mod_bwd(
+        grad_out_ssa = call_score_mod_bwd(
+            score_mod_bwd,
             grad_ssa,
             score_ssa,
             batch_idx_ssa,
             head_idx_ssa,
-            q_idx=q_idx_ssa,
-            kv_idx=kv_idx_ssa,
-            seqlen_info=seqlen_info,
-            aux_tensors=aux_args,
+            q_idx_ssa,
+            kv_idx_ssa,
+            seqlen_info,
+            aux_data,
         )
 
         grad_vec.store(grad_out_ssa)
diff --git a/build/torch-cuda/testing.py b/build/torch-cuda/testing.py
index 6e3c40eb451239fe54dc84c901096df7daec9eb9..e28577be0508420d4bf8d371c9b971192686402d 100644
--- a/build/torch-cuda/testing.py
+++ b/build/torch-cuda/testing.py
@@ -91,20 +91,23 @@ def pad_input(hidden_states, indices, batch, seqlen):
     return rearrange(output, "(b s) ... -> b s ...", b=batch)
 
 
-def generate_random_padding_mask(max_seqlen, batch_size, device, mode="random", zero_lengths=False):
+def generate_random_padding_mask(
+    max_seqlen, batch_size, device, mode="random", zero_lengths=False, min_seqlen=None
+):
     assert mode in ["full", "random", "third"]
+    min_seqlen = min_seqlen if min_seqlen is not None else 0 if zero_lengths else 1
     if mode == "full":
         lengths = torch.full((batch_size, 1), max_seqlen, device=device, dtype=torch.int32)
     elif mode == "random":
         lengths = torch.randint(
-            max(0 if zero_lengths else 1, max_seqlen - 20),
+            max(min_seqlen, max_seqlen - 20),
             max_seqlen + 1,
             (batch_size, 1),
             device=device,
         )
     else:
         lengths = torch.randint(
-            max(0 if zero_lengths else 1, max_seqlen // 3),
+            max(min_seqlen, max_seqlen // 3),
             max_seqlen + 1,
             (batch_size, 1),
             device=device,
@@ -343,33 +346,48 @@ def attention_ref(
     upcast=True,
     reorder_ops=False,
     intermediate_dtype=None,
+    return_lse=False,
+    gather_kv_indices=None,
 ):
+    assert v is not None
+    has_qk = q is not None and k is not None
+    assert has_qk or qv is not None
     if causal:
         window_size = (window_size[0], 0)
-    dtype_og = q.dtype
+    dtype_og = v.dtype
+    q_shape = q.shape if q is not None else qv.shape
     if upcast:
-        q, k, v = q.float(), k.float(), v.float()
-        qv = qv.float() if qv is not None else None
+        q, k, v, qv, learnable_sink = [
+            t.float() if t is not None else None for t in (q, k, v, qv, learnable_sink)
+        ]
     if q_descale is not None:
-        q_descale = repeat(q_descale, "b h -> b 1 (h g) 1", g=q.shape[2] // k.shape[2])
-        q = (q.float() * q_descale).to(q.dtype)
-        qv = (qv.float() * q_descale).to(qv.dtype) if qv is not None else None
+        q_descale = repeat(q_descale, "b h -> b 1 (h g) 1", g=q_shape[2] // v.shape[2])
+        q, qv = [(t.float() * q_descale).to(t.dtype) if t is not None else None for t in (q, qv)]
     if k_descale is not None:
         k = (k.float() * rearrange(k_descale, "b h -> b 1 h 1")).to(dtype=k.dtype)
     if v_descale is not None:
         v = (v.float() * rearrange(v_descale, "b h -> b 1 h 1")).to(dtype=v.dtype)
-    seqlen_q, seqlen_k = q.shape[1], k.shape[1]
-    k = repeat(k, "b s h d -> b s (h g) d", g=q.shape[2] // k.shape[2])
-    v = repeat(v, "b s h d -> b s (h g) d", g=q.shape[2] // v.shape[2])
-    d = q.shape[-1]
+    seqlen_q, seqlen_k = q_shape[1], v.shape[1]
+    k, v = [
+        repeat(t, "b s h d -> b s (h g) d", g=q_shape[2] // t.shape[2]) if t is not None else None
+        for t in (k, v)
+    ]
+    d = q_shape[-1]  # == dv for qv
     dv = v.shape[-1]
-    softmax_scale = 1.0 / math.sqrt(d if qv is None else d + dv)
-    if not reorder_ops:
-        scores = torch.einsum("bthd,bshd->bhts", q * softmax_scale, k)
-    else:
-        scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
+    softmax_scale = 1.0 / math.sqrt(d if qv is None or q is None else d + dv)
+    if has_qk:
+        scores = torch.einsum(
+            "bthd,bshd->bhts",
+            q if reorder_ops else q * softmax_scale,
+            k * softmax_scale if reorder_ops else k,
+        )
     if qv is not None:
-        scores = scores + torch.einsum("bthd,bshd->bhts", qv * softmax_scale, v)
+        qv_scores = torch.einsum(
+            "bthd,bshd->bhts",
+            qv if reorder_ops else qv * softmax_scale,
+            v * softmax_scale if reorder_ops else v,
+        )
+        scores = qv_scores if not has_qk else scores + qv_scores
     if softcap > 0:
         scores = torch.tanh(scores / softcap) * softcap
     if key_padding_mask is not None:
@@ -384,7 +402,7 @@ def attention_ref(
             query_padding_mask,
             key_padding_mask,
             key_leftpad=key_leftpad,
-            device=q.device,
+            device=v.device,
         )
     if attention_chunk > 0:
         chunk_mask = construct_chunk_mask(
@@ -394,15 +412,29 @@ def attention_ref(
             query_padding_mask,
             key_padding_mask,
             key_leftpad=key_leftpad,
-            device=q.device,
+            device=v.device,
         )
         local_mask = (
             torch.logical_or(local_mask, chunk_mask) if local_mask is not None else chunk_mask
         )
+    if gather_kv_indices is not None:
+        batch = q_shape[0]
+        # -1 is a sentinel for invalid top-k slots; route sentinels (and any
+        # out-of-range index) to a dummy extra column so they never unmask a
+        # real key
+        gather_idx = gather_kv_indices.long()
+        gather_idx = gather_idx.masked_fill((gather_idx < 0) | (gather_idx >= seqlen_k), seqlen_k)
+        topk_index_mask = torch.full(
+            (batch, seqlen_q, seqlen_k + 1), False, device="cuda"
+        ).scatter_(-1, gather_idx, True)[..., :seqlen_k]
+        scores.masked_fill_(rearrange(~topk_index_mask, "b t s -> b 1 t s"), float("-inf"))
     if local_mask is not None:
         scores.masked_fill_(local_mask, float("-inf"))
     if attn_bias is not None:
         scores = scores + attn_bias
+    # After all masks are applied, before softmax:
+    # scores shape: [b, h, t, s]
+    lse = torch.logsumexp(scores, dim=-1)  # [b, h, t]
     if learnable_sink is None:
         attention = torch.softmax(scores, dim=-1).to(v.dtype)
     else:
@@ -414,6 +446,8 @@ def attention_ref(
         normalizer = unnormalized_scores.sum(dim=-1, keepdim=True) + torch.exp(
             learnable_sink - logits_or_sinks_max
         )
+        # LSE with sink: log(Z) = log(normalizer) + max
+        lse = (torch.log(normalizer.squeeze(-1)) + logits_or_sinks_max.squeeze(-1)).to(dtype_og)
         attention = (unnormalized_scores / normalizer).to(v.dtype)
     if query_padding_mask is not None:
         attention = attention.masked_fill(rearrange(~query_padding_mask, "b s -> b 1 s 1"), 0.0)
@@ -431,6 +465,8 @@ def attention_ref(
     output = torch.einsum("bhts,bshd->bthd", attention_drop, v * dropout_scaling)
     if query_padding_mask is not None:
         output.masked_fill_(rearrange(~query_padding_mask, "b s -> b s 1 1"), 0.0)
+    if return_lse:
+        return output.to(dtype_og), attention.to(dtype_og), lse.to(dtype_og)
     return output.to(dtype=dtype_og), attention.to(dtype=dtype_og)
 
 
diff --git a/build/torch-cuda/tile_scheduler.py b/build/torch-cuda/tile_scheduler.py
index f05709bdc9b0c2f9c61c2201e1b3b4d1ed5db767..3139147a0de0a546eace062b9cde08be86ea85b0 100644
--- a/build/torch-cuda/tile_scheduler.py
+++ b/build/torch-cuda/tile_scheduler.py
@@ -1,6 +1,7 @@
-# Copyright (c) 2025, Tri Dao.
+# Copyright (c) 2025, Tri Dao, Siyu Wang, Shengbin Di, Yuxi Chi, Johnsonms, Linfeng Zheng, Haoyan Huang, Lanbo Li, Yun Zhong, Man Yuan, Minmin Sun, Yong Li, Wei Lin.
 
-from typing import Optional, Tuple
+from enum import IntEnum, auto
+from typing import Optional, Tuple, Protocol, runtime_checkable
 from dataclasses import dataclass
 
 try:
@@ -9,17 +10,135 @@ except ImportError:  # Python < 3.12
     from typing_extensions import override
 
 import cutlass
+from cutlass.pipeline import PipelineClcFetchAsync, PipelineState
 from cutlass._mlir import ir
 import cutlass.cute as cute
 from cutlass import Int32, const_expr
 from cutlass.cute import FastDivmodDivisor
+from cutlass.utils import ClcDynamicPersistentTileScheduler, ClcDynamicPersistentTileSchedulerParams
+from cutlass.cute.typing import Boolean
+from cutlass.cutlass_dsl import (
+    min as dsl_min,
+    extract_mlir_values,
+    new_from_mlir_values,
+)
+from cutlass.utils.hardware_info import HardwareInfo
 
 from .quack.cute_dsl_utils import ParamsBase
 
-from . import utils
+from . import utils as utils
 from .fast_math import clz
 
 
+class SchedulingMode(IntEnum):
+    NONE = auto()
+    STATIC = auto()
+    DYNAMIC = auto()
+    CLC = auto()
+
+
+@dataclass
+class SchedulerState(ParamsBase):
+    """Runtime state shared by CLC and dynamic persistent tile schedulers:
+    the async pipeline and its producer/consumer states.
+
+    Main kernels construct this via `create_clc` / `create_dynamic_persistent`,
+    which return the appropriate concrete state (`ClcSchedulerState` or
+    `DynamicPersistentSchedulerState`). Schedulers consume it through the
+    `ctx: SchedulerState | None` parameter on their `__init__(...)`.
+    """
+
+    _pipeline: cutlass.pipeline.PipelineAsync
+    _consumer_state: PipelineState
+    _producer_state: PipelineState
+
+    @staticmethod
+    def create_clc(
+        *,
+        hw_scheduler: ClcDynamicPersistentTileScheduler,
+        pipeline: PipelineClcFetchAsync,
+        consumer_state: PipelineState,
+        producer_state: PipelineState,
+    ) -> "ClcSchedulerState":
+        return ClcSchedulerState(pipeline, consumer_state, producer_state, hw_scheduler)
+
+    @staticmethod
+    def create_dynamic_persistent(
+        *,
+        work_info: cute.Tensor,
+        pipeline: cutlass.pipeline.PipelineAsync,
+        consumer_state: PipelineState,
+        producer_state: PipelineState,
+    ) -> "DynamicPersistentSchedulerState":
+        return DynamicPersistentSchedulerState(pipeline, consumer_state, producer_state, work_info)
+
+    def consumer_wait(self, *, loc=None, ip=None):
+        self._pipeline.consumer_wait(self._consumer_state, loc=loc, ip=ip)
+
+    def consumer_release(self, *, loc=None, ip=None):
+        self._pipeline.consumer_release(self._consumer_state, loc=loc, ip=ip)
+        self._consumer_state.advance(loc=loc, ip=ip)
+
+    def advance_consumer_state(self, *, loc=None, ip=None):
+        self._consumer_state.advance(loc=loc, ip=ip)
+
+    def producer_tail(self, *, loc=None, ip=None):
+        self._pipeline.producer_tail(self._producer_state, loc=loc, ip=ip)
+
+
+@dataclass
+class ClcSchedulerState(SchedulerState):
+    """Owns the runtime state shared by CLC-capable tile schedulers.
+
+    `FlashAttentionForwardSm100` constructs this state because it owns the CLC
+    response buffer, mbarrier storage, and launch geometry needed to initialize
+    the hardware scheduler and async pipeline. Individual tile schedulers then
+    consume this state and map the returned hardware work tiles into their own
+    logical `WorkTileInfo` coordinates.
+
+    To add CLC support to a scheduler:
+    - implement `clc_problem_shape(params)` so the kernel can create the hardware scheduler
+    - map `ctx.initial_work_tile_info()` and `ctx.get_current_work()` into scheduler coordinates
+    """
+
+    _hw_scheduler: ClcDynamicPersistentTileScheduler
+
+    def initial_work_tile_info(self):
+        return self._hw_scheduler.initial_work_tile_info()
+
+    def get_current_work(self):
+        return self._hw_scheduler.get_current_work()
+
+    def prefetch_next_work(self, *, loc=None, ip=None):
+        self._pipeline.producer_acquire(self._producer_state, loc=loc, ip=ip)
+        mbarrier_addr = self._pipeline.producer_get_barrier(self._producer_state, loc=loc, ip=ip)
+        self._hw_scheduler.advance_to_next_work(mbarrier_addr, loc=loc, ip=ip)
+        self._producer_state.advance(loc=loc, ip=ip)
+
+
+@dataclass
+class DynamicPersistentSchedulerState(SchedulerState):
+    """Semaphore-backed: the scheduler class drives atomicAdd + warp-prefix-sum
+    and writes the resolved work tile via `write_work_info`."""
+
+    _work_info: cute.Tensor
+
+    def producer_acquire(self, *, loc=None, ip=None):
+        self._pipeline.producer_acquire(self._producer_state, loc=loc, ip=ip)
+
+    def producer_commit(self, *, loc=None, ip=None):
+        self._pipeline.producer_commit(self._producer_state, loc=loc, ip=ip)
+
+    def advance_producer_state(self, *, loc=None, ip=None):
+        self._producer_state.advance(loc=loc, ip=ip)
+
+    def write_work_info(self, block: Int32, head: Int32, batch: Int32, split: Int32):
+        self._work_info[0] = block
+        self._work_info[1] = head
+        self._work_info[2] = batch
+        self._work_info[3] = split
+
+
 class WorkTileInfo(cutlass.utils.WorkTileInfo):
     """Altered WorkTileInfo which includes four axes: (block, head, batch, split)"""
 
@@ -31,6 +150,43 @@ class WorkTileInfo(cutlass.utils.WorkTileInfo):
         return WorkTileInfo(new_tile_idx, new_is_valid_tile)
 
 
+@runtime_checkable
+class TileSchedulerProtocol(Protocol):
+    """Protocol defining the interface all tile schedulers must implement.
+
+    Schedulers are responsible for:
+    1. Coordinate mapping: linear tile index -> (m_block, head, batch, split)
+    2. Work distribution: how to get the next tile (static grid-stride vs dynamic)
+    """
+
+    def initial_work_tile_info(self) -> WorkTileInfo:
+        """Get the initial work tile for this CTA."""
+        ...
+
+    def advance_to_next_work(self, *, loc=None, ip=None):
+        """Consumer-side advance: move to next tile and return it.
+
+        For static schedulers: grid-stride increment + get_current_work.
+        For dynamic schedulers: consumer wait + get_current_work + consumer release + state advance.
+        """
+        ...
+
+    def prefetch_next_work(self, *, loc=None, ip=None) -> None:
+        """Producer-side prefetch of next work tile (no-op for static schedulers).
+
+        For dynamic schedulers: producer acquire (+ issue CLC query) + producer state advance.
+        Only called by the scheduler warp.
+        """
+        ...
+
+    def producer_tail(self, *, loc=None, ip=None) -> None:
+        """Producer-side cleanup after the last tile.
+
+        No-op for static schedulers. For dynamic schedulers: pipeline producer_tail.
+        """
+        ...
+
+
 @dataclass
 class TileSchedulerArguments(ParamsBase):
     num_block: Int32
@@ -51,6 +207,16 @@ class TileSchedulerArguments(ParamsBase):
     lpt: cutlass.Constexpr[bool] = False
     is_split_kv: cutlass.Constexpr[bool] = False
     head_swizzle: cutlass.Constexpr[bool] = False
+    use_cluster_idx: cutlass.Constexpr[bool] = False
+    num_splits_dynamic_ptr: Optional[cute.Tensor] = None
+    num_m_blocks_ptr: Optional[cute.Tensor] = None
+    virtual_batch_idx_ptr: Optional[cute.Tensor] = None
+    num_nheads_in_l2_ptr: Optional[cute.Tensor] = None
+    cu_total_m_blocks_ptr: Optional[cute.Tensor] = None
+    cu_total_splits_m_blocks_ptr: Optional[cute.Tensor] = None
+    blocks_to_batch_idx_ptr: Optional[cute.Tensor] = None
+    tile_count_semaphore: Optional[cute.Pointer] = None
+    persistent_cta_multiplier: cutlass.Constexpr[int] = 1
 
 
 class SingleTileScheduler:
@@ -63,6 +229,8 @@ class SingleTileScheduler:
         num_splits_divmod: FastDivmodDivisor
         is_split_kv: cutlass.Constexpr[bool] = False
         cluster_shape_mn: cutlass.Constexpr[Tuple[int, int]] = (1, 1)
+        use_cluster_idx: cutlass.Constexpr[bool] = False
+        num_splits_dynamic_ptr: Optional[cute.Tensor] = None
 
         @staticmethod
         def create(
@@ -76,6 +244,8 @@ class SingleTileScheduler:
                 FastDivmodDivisor(args.num_splits),
                 args.is_split_kv,
                 args.cluster_shape_mn,
+                args.use_cluster_idx,
+                args.num_splits_dynamic_ptr,
             )
 
     def __init__(self, params: Params, blk_coord: cute.Coord, *, loc=None, ip=None):
@@ -86,18 +256,26 @@ class SingleTileScheduler:
         self._ip = ip
 
     @staticmethod
-    def to_underlying_arguments(args: TileSchedulerArguments, *, loc=None, ip=None) -> Params:
+    def to_underlying_arguments(
+        args: TileSchedulerArguments,
+        *,
+        scheduling_mode: SchedulingMode = SchedulingMode.STATIC,
+        loc=None,
+        ip=None,
+    ) -> Params:
+        assert scheduling_mode == SchedulingMode.STATIC, (
+            f"SingleTileScheduler only supports STATIC, got {scheduling_mode!r}"
+        )
         return SingleTileScheduler.Params.create(args, loc=loc, ip=ip)
 
     @staticmethod
-    def create(params: Params, *, loc=None, ip=None) -> "SingleTileScheduler":
-        # if const_expr(cute.size(params.cluster_shape_mn) == 1):
-        #     blk_coord = cute.arch.block_idx()
-        # else:
-        #     # All CTAs in a cluster must get the same block coordinate
-        #     blk_coord = cute.arch.cluster_idx()
-        # Temporary set to block_idx until we sort out the best way to handle cluster
-        blk_coord = cute.arch.block_idx()
+    def create(
+        params: Params, ctx: SchedulerState | None = None, *, loc=None, ip=None
+    ) -> "SingleTileScheduler":
+        if const_expr(cute.size(params.cluster_shape_mn) == 1 or not params.use_cluster_idx):
+            blk_coord = cute.arch.block_idx()
+        else:
+            blk_coord = cute.arch.cluster_idx()
         return SingleTileScheduler(params, blk_coord, loc=loc, ip=ip)
 
     # called by host
@@ -110,21 +288,32 @@ class SingleTileScheduler:
     ) -> Tuple[Int32, Int32, Int32]:
         # TODO: this hard-codes the fact that we only use cluster = (1, 1) or (2, 1)
         assert params.cluster_shape_mn[1] == 1, "Only cluster_shape_mn[1] == 1 is supported"
+        if const_expr(params.use_cluster_idx):
+            # Grid must have num_block * cluster_m physical blocks so that there are num_block clusters
+            grid_x = params.num_block * params.cluster_shape_mn[0]
+        else:
+            grid_x = cute.round_up(params.num_block, params.cluster_shape_mn[0])
         return (
-            cute.round_up(params.num_block, params.cluster_shape_mn[0]),
+            grid_x,
             params.num_head * params.num_splits,
             params.num_batch,
         )
 
     def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo:
         block_idx, head_idx, batch_idx = self._blk_coord
+        is_valid = self._is_first_block
         if const_expr(self.params.is_split_kv):
             head_idx, split_idx = divmod(head_idx, self.params.num_splits_divmod)
         else:
             split_idx = Int32(0)
+        # Pack dynamic per-batch num_splits into high 16 bits of split_idx
+        if const_expr(self.params.is_split_kv and self.params.num_splits_dynamic_ptr is not None):
+            if is_valid:
+                num_splits = Int32(self.params.num_splits_dynamic_ptr[batch_idx])
+                split_idx = split_idx | (num_splits << 16)
         return WorkTileInfo(
             (block_idx, head_idx, batch_idx, split_idx),
-            self._is_first_block,
+            is_valid,
         )
 
     def initial_work_tile_info(self, *, loc=None, ip=None):
@@ -135,6 +324,10 @@ class SingleTileScheduler:
 
     def advance_to_next_work(self, *, loc=None, ip=None):
         self._is_first_block = False
+        return self.get_current_work()
+
+    def producer_tail(self, *, loc=None, ip=None):
+        pass
 
     def __extract_mlir_values__(self):
         values, self._values_pos = [], []
@@ -149,7 +342,11 @@ class SingleTileScheduler:
         for obj, n_items in zip([self.params, self._blk_coord], self._values_pos):
             obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items]))
             values = values[n_items:]
-        return SingleTileScheduler(*(tuple(obj_list)), loc=self._loc)
+        scheduler = SingleTileScheduler(*(tuple(obj_list)), loc=self._loc)
+        # Note: _is_first_block is a Python-only attribute omitted from MLIR values,
+        # so it must be restored explicitly after reconstruction.
+        scheduler._is_first_block = self._is_first_block
+        return scheduler
 
 
 class StaticPersistentTileScheduler:
@@ -180,18 +377,28 @@ class StaticPersistentTileScheduler:
         self._ip = ip
 
     @staticmethod
-    def to_underlying_arguments(args: TileSchedulerArguments, *, loc=None, ip=None) -> Params:
+    def to_underlying_arguments(
+        args: TileSchedulerArguments,
+        *,
+        scheduling_mode: SchedulingMode = SchedulingMode.STATIC,
+        loc=None,
+        ip=None,
+    ) -> Params:
+        assert scheduling_mode == SchedulingMode.STATIC, (
+            f"StaticPersistentTileScheduler only supports STATIC, got {scheduling_mode!r}"
+        )
         return StaticPersistentTileScheduler.Params.create(args, loc=loc, ip=ip)
 
     @staticmethod
-    def create(params: Params, *, loc=None, ip=None) -> "StaticPersistentTileScheduler":
+    def create(
+        params: Params, ctx: SchedulerState | None = None, *, loc=None, ip=None
+    ) -> "StaticPersistentTileScheduler":
         if const_expr(cute.size(params.cluster_shape_m) == 1):
             tile_idx = cute.arch.block_idx()[0]
         else:
             tile_idx = cute.arch.cluster_idx()[0]
         return StaticPersistentTileScheduler(params, tile_idx, loc=loc, ip=ip)
 
-    # called by host
     @staticmethod
     def get_grid_shape(
         params: Params,
@@ -201,18 +408,14 @@ class StaticPersistentTileScheduler:
     ) -> Tuple[Int32, Int32, Int32]:
         hardware_info = cutlass.utils.HardwareInfo()
         sm_count = hardware_info.get_device_multiprocessor_count()
-        # Grid must be a multiple of cluster_shape_m for CUDA cluster launch.
         max_ctas = (sm_count // params.cluster_shape_m) * params.cluster_shape_m
         grid_x = cutlass.min(max_ctas, params.total_blocks_cluster * params.cluster_shape_m)
         return (grid_x, Int32(1), Int32(1))
 
-    # @cute.jit
     def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo:
         hn_idx, block_idx = divmod(self._tile_idx, self.params.num_block_cluster_divmod)
         batch_idx, head_idx = divmod(hn_idx, self.params.num_head_divmod)
         is_valid = self._tile_idx < self.params.total_blocks_cluster
-        # if cute.arch.thread_idx()[0] == 0:
-        #     cute.printf("TileScheduler: tile_idx=%d, hn_idx=%d, block_idx=%d, batch_idx=%d, head_idx=%d, is_valid=%d", self._tile_idx, hn_idx, block_idx, batch_idx, head_idx, is_valid)
         return WorkTileInfo(
             (Int32(block_idx), Int32(head_idx), Int32(batch_idx), Int32(0)), is_valid
         )
@@ -228,6 +431,10 @@ class StaticPersistentTileScheduler:
             self._tile_idx += cute.arch.grid_dim()[0]
         else:
             self._tile_idx += cute.arch.cluster_dim()[0]
+        return self.get_current_work()
+
+    def producer_tail(self, *, loc=None, ip=None):
+        pass
 
     def __extract_mlir_values__(self):
         values, self._values_pos = [], []
@@ -254,32 +461,49 @@ class SingleTileLPTScheduler:
         total_blocks: Int32
         num_splits: Int32
         num_block: Int32
+        num_head: Int32
+        num_batch: Int32
         l2_minor: Int32
-        num_block_divmod: FastDivmodDivisor
         num_head_divmod: FastDivmodDivisor
         l2_minor_divmod: FastDivmodDivisor
         l2_major_divmod: FastDivmodDivisor
         l2_minor_residual_divmod: FastDivmodDivisor
         num_hb_quotient: Int32
+        num_splits_divmod: FastDivmodDivisor
         is_split_kv: cutlass.Constexpr[bool] = False
+        cluster_shape_m: cutlass.Constexpr[int] = 1
+        scheduling_mode: cutlass.Constexpr[SchedulingMode] = SchedulingMode.STATIC
+        lpt: cutlass.Constexpr[bool] = True
+        use_cluster_idx: cutlass.Constexpr[bool] = True
+        num_splits_dynamic_ptr: Optional[cute.Tensor] = None
 
         @staticmethod
         @cute.jit
         def create(
-            args: TileSchedulerArguments, *, loc=None, ip=None
+            args: TileSchedulerArguments,
+            *,
+            scheduling_mode: SchedulingMode = SchedulingMode.STATIC,
+            loc=None,
+            ip=None,
         ) -> "SingleTileLPTScheduler.Params":
-            # cute.printf(args.num_block, args.num_head, args.num_batch, args.seqlen_k, args.headdim, args.headdim_v, args.total_q, args.tile_shape_mn, args.qhead_per_kvhead_packgqa, args.element_size)
-            size_one_kv_head = args.seqlen_k * (args.headdim + args.headdim_v) * args.element_size
+            assert scheduling_mode in (SchedulingMode.STATIC, SchedulingMode.CLC), (
+                f"Only STATIC and CLC are supported, got {scheduling_mode!r}"
+            )
+            # int64: this product overflows int32 once seqlen_k * (headdim +
+            # headdim_v) * element_size > 2**31 (seqlen_k > ~4M for hdim-128 bf16).
+            size_one_kv_head = (
+                cutlass.Int64(args.seqlen_k) * (args.headdim + args.headdim_v) * args.element_size
+            )
             size_one_head = size_one_kv_head
             size_l2 = 50 * 1024 * 1024  # 40 MB for K & V
             # Swizzle is the size of each "section". Round swizzle to a power of 2
             # Need to be careful about the case where only one head will fit
             # swizzle is how many heads can fit in L2
-            # swizzle = 1 if size_l2 < size_one_head else (size_l2 // size_one_head)
-            # Seems faster if swizzle if a power of 2
+            # Seems faster if swizzle is a power of 2
             log2_floor = lambda n: 31 - clz(n)
-            swizzle = 1 if size_l2 < size_one_head else (1 << log2_floor(size_l2 // size_one_head))
-            # swizzle = 1 if size_l2 < size_one_head else (size_l2 // size_one_head)
+            swizzle = (
+                1 if size_l2 < size_one_head else (1 << log2_floor(Int32(size_l2 // size_one_head)))
+            )
             # If we're in the last section (called residual), we don't want to divide by
             # swizzle. Instead we want to divide by the remainder.
             num_hb_quotient = (args.num_head * args.num_batch) // swizzle
@@ -287,37 +511,91 @@ class SingleTileLPTScheduler:
             return SingleTileLPTScheduler.Params(
                 total_blocks=args.num_block * args.num_head * args.num_batch,
                 num_block=args.num_block,
+                num_head=args.num_head,
+                num_batch=args.num_batch,
                 l2_minor=Int32(swizzle),
-                num_block_divmod=FastDivmodDivisor(args.num_block),
                 num_head_divmod=FastDivmodDivisor(args.num_head),
                 l2_minor_divmod=FastDivmodDivisor(swizzle),
                 l2_major_divmod=FastDivmodDivisor(swizzle * args.num_block),
-                l2_minor_residual_divmod=FastDivmodDivisor(
-                    max(num_hb_remainder, 1)
-                ),  # don't divide by 0
+                l2_minor_residual_divmod=FastDivmodDivisor(max(num_hb_remainder, 1)),
                 num_hb_quotient=Int32(num_hb_quotient),
                 num_splits=args.num_splits,
+                num_splits_divmod=FastDivmodDivisor(args.num_splits),
                 is_split_kv=args.is_split_kv,
+                cluster_shape_m=args.cluster_shape_mn[0],
+                scheduling_mode=scheduling_mode,
+                lpt=args.lpt,
+                use_cluster_idx=args.use_cluster_idx,
+                num_splits_dynamic_ptr=args.num_splits_dynamic_ptr,
             )
 
-    def __init__(self, params: Params, tile_idx: Int32, split_idx: Int32, *, loc=None, ip=None):
+    def __init__(
+        self,
+        params: Params,
+        tile_idx: Int32,
+        split_idx: Int32,
+        ctx: SchedulerState | None = None,
+        *,
+        loc=None,
+        ip=None,
+    ):
         self.params = params
         self._tile_idx = tile_idx
         self._split_idx = split_idx
+        self._ctx = ctx
         self._loc = loc
         self._ip = ip
 
     @staticmethod
-    def to_underlying_arguments(args: TileSchedulerArguments, *, loc=None, ip=None) -> Params:
-        return SingleTileLPTScheduler.Params.create(args, loc=loc, ip=ip)
+    def to_underlying_arguments(
+        args: TileSchedulerArguments,
+        *,
+        scheduling_mode: SchedulingMode = SchedulingMode.STATIC,
+        loc=None,
+        ip=None,
+    ) -> Params:
+        return SingleTileLPTScheduler.Params.create(
+            args, scheduling_mode=scheduling_mode, loc=loc, ip=ip
+        )
+
+    @staticmethod
+    def _clc_grid_shape(params: Params):
+        num_batch_splits = (
+            params.num_batch * params.num_splits
+            if const_expr(params.is_split_kv)
+            else params.num_batch
+        )
+        if const_expr(params.use_cluster_idx):
+            # Grid must have num_block * cluster_m physical blocks so that there are num_block clusters
+            grid_x = params.num_block * params.cluster_shape_m
+        else:
+            grid_x = cute.round_up(params.num_block, params.cluster_shape_m)
+        return (
+            grid_x,
+            params.num_head,
+            num_batch_splits,
+        )
+
+    @staticmethod
+    @cute.jit
+    def clc_problem_shape(params: Params):
+        return ClcDynamicPersistentTileSchedulerParams(
+            problem_shape_ntile_mnl=SingleTileLPTScheduler._clc_grid_shape(params),
+            cluster_shape_mnk=(params.cluster_shape_m, 1, 1),
+        )
 
     @staticmethod
     @cute.jit
-    def create(params: Params, *, loc=None, ip=None) -> "SingleTileLPTScheduler":
+    def create(
+        params: Params, ctx: SchedulerState | None = None, *, loc=None, ip=None
+    ) -> "SingleTileLPTScheduler":
+        if const_expr(params.scheduling_mode == SchedulingMode.CLC):
+            return SingleTileLPTScheduler(
+                params, cute.arch.block_idx()[0], Int32(0), ctx, loc=loc, ip=ip
+            )
         tile_idx, split_idx, _ = cute.arch.block_idx()
         return SingleTileLPTScheduler(params, tile_idx, split_idx, loc=loc, ip=ip)
 
-    # called by host
     @staticmethod
     def get_grid_shape(
         params: Params,
@@ -325,10 +603,52 @@ class SingleTileLPTScheduler:
         loc=None,
         ip=None,
     ) -> Tuple[Int32, Int32, Int32]:
+        if const_expr(params.scheduling_mode == SchedulingMode.CLC):
+            return SingleTileLPTScheduler._clc_grid_shape(params)
         return (params.total_blocks, params.num_splits, Int32(1))
 
+    @cute.jit
+    def clc_work_to_coords(self, work) -> WorkTileInfo:
+        """Convert CLC response (block, head, batch_split) to WorkTileInfo.
+
+        CLC returns raw grid coordinates — no L2 swizzle (hardware decides order).
+        We only apply cluster division, optional LPT block reversal, and split_kv unpacking.
+        """
+        block_idx = work.tile_idx[0]
+        if const_expr(self.params.cluster_shape_m > 1):
+            block_idx = block_idx // self.params.cluster_shape_m
+        if const_expr(self.params.lpt):
+            # Longest-processing-time-first: reverse block order
+            if const_expr(self.params.cluster_shape_m > 1 and not self.params.use_cluster_idx):
+                num_block = self.params.num_block // self.params.cluster_shape_m
+            else:
+                num_block = self.params.num_block
+            block_idx = num_block - 1 - block_idx
+        split_idx = Int32(0)
+        if const_expr(self.params.is_split_kv):
+            batch_idx, split_idx = divmod(work.tile_idx[2], self.params.num_splits_divmod)
+        else:
+            batch_idx = work.tile_idx[2]
+        if const_expr(self.params.cluster_shape_m > 1 and not self.params.use_cluster_idx):
+            bidx_in_cluster = cute.arch.block_in_cluster_idx()
+            block_idx = block_idx * self.params.cluster_shape_m + bidx_in_cluster[0]
+        # Pack dynamic per-batch num_splits into high 16 bits of split_idx
+        if const_expr(self.params.is_split_kv and self.params.num_splits_dynamic_ptr is not None):
+            if work.is_valid_tile:
+                num_splits = Int32(self.params.num_splits_dynamic_ptr[batch_idx])
+                split_idx = split_idx | (num_splits << 16)
+        return WorkTileInfo(
+            (Int32(block_idx), Int32(work.tile_idx[1]), Int32(batch_idx), Int32(split_idx)),
+            work.is_valid_tile,
+        )
+
     @cute.jit
     def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo:
+        if const_expr(self.params.scheduling_mode == SchedulingMode.CLC):
+            work = self._ctx.get_current_work()
+            self._tile_idx = work.tile_idx[0]
+            return self.clc_work_to_coords(work)
+        # Static path: L2-swizzled coordinate mapping
         params = self.params
         # Implement LPT scheduling coordinate calculation
         bidhb, l2_mod = divmod(self._tile_idx, params.l2_major_divmod)
@@ -342,25 +662,51 @@ class SingleTileLPTScheduler:
         bidhb_actual = bidhb * params.l2_minor + bidhb_residual
         batch_idx, head_idx = divmod(bidhb_actual, params.num_head_divmod)
         # Longest-processing-time-first
-        block = params.num_block - 1 - block
+        if const_expr(params.lpt):
+            block = params.num_block - 1 - block
         is_valid = self._tile_idx < params.total_blocks
+        split_idx = self._split_idx
+        # Pack dynamic per-batch num_splits into high 16 bits of split_idx
+        if const_expr(params.is_split_kv and params.num_splits_dynamic_ptr is not None):
+            if is_valid:
+                num_splits = Int32(params.num_splits_dynamic_ptr[batch_idx])
+                split_idx = split_idx | (num_splits << 16)
         return WorkTileInfo(
-            (Int32(block), Int32(head_idx), Int32(batch_idx), Int32(self._split_idx)), is_valid
+            (Int32(block), Int32(head_idx), Int32(batch_idx), Int32(split_idx)), is_valid
         )
 
+    @cute.jit
     def initial_work_tile_info(self, *, loc=None, ip=None):
+        if const_expr(self.params.scheduling_mode == SchedulingMode.CLC):
+            work = self._ctx.initial_work_tile_info()
+            self._tile_idx = work.tile_idx[0]
+            return self.clc_work_to_coords(work)
         return self.get_current_work(loc=loc, ip=ip)
 
     def prefetch_next_work(self, *, loc=None, ip=None):
-        pass
+        if const_expr(self.params.scheduling_mode == SchedulingMode.CLC):
+            self._ctx.prefetch_next_work(loc=loc, ip=ip)
 
     def advance_to_next_work(self, *, loc=None, ip=None):
+        if const_expr(self.params.scheduling_mode == SchedulingMode.CLC):
+            self._ctx.consumer_wait(loc=loc, ip=ip)
+            work = self.get_current_work()
+            self._ctx.consumer_release(loc=loc, ip=ip)
+            return work
         # Single tile scheduler - set to invalid tile_idx to indicate no more work
         self._tile_idx = self.params.total_blocks
+        return self.get_current_work()
+
+    def producer_tail(self, *, loc=None, ip=None):
+        if const_expr(self.params.scheduling_mode == SchedulingMode.CLC):
+            self._ctx.producer_tail(loc=loc, ip=ip)
 
     def __extract_mlir_values__(self):
         values, self._values_pos = [], []
-        for obj in [self.params, self._tile_idx, self._split_idx]:
+        objs = [self.params, self._tile_idx, self._split_idx]
+        if const_expr(self.params.scheduling_mode == SchedulingMode.CLC):
+            objs += [self._ctx]
+        for obj in objs:
             obj_values = cutlass.extract_mlir_values(obj)
             values += obj_values
             self._values_pos.append(len(obj_values))
@@ -368,10 +714,13 @@ class SingleTileLPTScheduler:
 
     def __new_from_mlir_values__(self, values):
         obj_list = []
-        for obj, n_items in zip([self.params, self._tile_idx, self._split_idx], self._values_pos):
+        objs = [self.params, self._tile_idx, self._split_idx]
+        if const_expr(self.params.scheduling_mode == SchedulingMode.CLC):
+            objs += [self._ctx]
+        for obj, n_items in zip(objs, self._values_pos):
             obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items]))
             values = values[n_items:]
-        return self.__class__(*(tuple(obj_list)), loc=self._loc)
+        return self.__class__(*obj_list, loc=self._loc)
 
 
 class SingleTileLPTBwdScheduler:
@@ -394,12 +743,18 @@ class SingleTileLPTBwdScheduler:
             args: TileSchedulerArguments, *, loc=None, ip=None
         ) -> "SingleTileLPTBwdScheduler.Params":
             size_l2 = 50 * 1024 * 1024
-            size_one_qdo_head = args.seqlen_k * (args.headdim + args.headdim_v) * args.element_size
-            # size_one_dqaccum_head = args.seqlen_k * (args.headdim) * 4
-            size_one_dqaccum_head = 0
+            # int64: these products overflow int32 at large seqlen_k (> ~4M for
+            # hdim-128 bf16; the dqaccum *4 term wraps even sooner).
+            size_one_qdo_head = (
+                cutlass.Int64(args.seqlen_k) * (args.headdim + args.headdim_v) * args.element_size
+            )
+            size_one_dqaccum_head = cutlass.Int64(args.seqlen_k) * (args.headdim) * 4
+            # size_one_dqaccum_head = 0
             size_one_head = size_one_qdo_head + size_one_dqaccum_head
             log2_floor = lambda n: 31 - clz(n)
-            swizzle = 1 if size_l2 < size_one_head else (1 << log2_floor(size_l2 // size_one_head))
+            swizzle = (
+                1 if size_l2 < size_one_head else (1 << log2_floor(Int32(size_l2 // size_one_head)))
+            )
             # swizzle = 8
             # If we're in the last section (called residual), we don't want to divide by
             # swizzle. Instead we want to divide by the remainder.
@@ -430,7 +785,16 @@ class SingleTileLPTBwdScheduler:
         self._ip = ip
 
     @staticmethod
-    def to_underlying_arguments(args: TileSchedulerArguments, *, loc=None, ip=None) -> Params:
+    def to_underlying_arguments(
+        args: TileSchedulerArguments,
+        *,
+        scheduling_mode: SchedulingMode = SchedulingMode.STATIC,
+        loc=None,
+        ip=None,
+    ) -> Params:
+        assert scheduling_mode == SchedulingMode.STATIC, (
+            f"SingleTileLPTBwdScheduler only supports STATIC, got {scheduling_mode!r}"
+        )
         return SingleTileLPTBwdScheduler.Params.create(args, loc=loc, ip=ip)
 
     @staticmethod
@@ -481,6 +845,7 @@ class SingleTileLPTBwdScheduler:
     def advance_to_next_work(self, *, loc=None, ip=None):
         # Single tile scheduler - set to invalid tile_idx to indicate no more work
         self._tile_idx = self.params.total_blocks
+        return self.get_current_work()
 
     def __extract_mlir_values__(self):
         values, self._values_pos = [], []
@@ -498,66 +863,423 @@ class SingleTileLPTBwdScheduler:
         return self.__class__(*(tuple(obj_list)), loc=self._loc)
 
 
+@dataclass
+class VarlenDecoder(ParamsBase):
+    """Per-batch m-block lookup + warp-prefix-sum search-and-decode of the
+    varlen work tile. Composed into both `SingleTileVarlenScheduler.Params`
+    and `DynamicPersistentVarlenScheduler.Params`.
+
+    `fold_splits_into_scan` controls whether the prefix-sum scan folds per-batch
+    `num_splits` into the per-batch tile count (DynamicPersistent) or always
+    counts only m_blocks (SingleTileVarlen, where splits are dispatched at the
+    grid level and resolved post-scan).
+    """
+
+    num_head: Int32
+    num_batch: Int32
+    num_splits: Int32
+    max_kvblock_in_l2: Int32
+    tile_shape_mn: cutlass.Constexpr[Tuple[int, int]]
+    qhead_per_kvhead_packgqa: cutlass.Constexpr[int] = 1
+    is_split_kv: cutlass.Constexpr[bool] = False
+    lpt: cutlass.Constexpr[bool] = False
+    head_swizzle: cutlass.Constexpr[bool] = False
+    cluster_shape_m: cutlass.Constexpr[int] = 1
+    use_cluster_idx: cutlass.Constexpr[bool] = False
+    fold_splits_into_scan: cutlass.Constexpr[bool] = False
+    scheduling_mode: cutlass.Constexpr[SchedulingMode] = SchedulingMode.STATIC
+    mCuSeqlensQ: Optional[cute.Tensor] = None
+    mSeqUsedQ: Optional[cute.Tensor] = None
+    num_m_blocks_ptr: Optional[cute.Tensor] = None
+    num_splits_dynamic_ptr: Optional[cute.Tensor] = None
+    virtual_batch_idx_ptr: Optional[cute.Tensor] = None
+    num_nheads_in_l2_ptr: Optional[cute.Tensor] = None
+    cu_total_m_blocks_ptr: Optional[cute.Tensor] = None
+    cu_total_splits_m_blocks_ptr: Optional[cute.Tensor] = None
+    blocks_to_batch_idx_ptr: Optional[cute.Tensor] = None
+
+    @staticmethod
+    @cute.jit
+    def create(
+        args: TileSchedulerArguments,
+        *,
+        fold_splits_into_scan: bool,
+        head_swizzle: bool = False,
+        cluster_shape_m: int = 1,
+        scheduling_mode: SchedulingMode = SchedulingMode.STATIC,
+        loc=None,
+        ip=None,
+    ) -> "VarlenDecoder":
+        size_l2 = 50 * 1024 * 1024  # 50 MB for K & V
+        # if backward, this is qdo block size
+        kv_block_size = (args.headdim + args.headdim_v) * args.element_size * args.tile_shape_mn[1]
+        # if backward, add dqaccum block size to calculate swizzle
+        if head_swizzle:
+            kv_block_size += args.headdim * 4 * args.tile_shape_mn[1]
+        max_kvblock_in_l2 = size_l2 // kv_block_size
+        return VarlenDecoder(
+            num_head=args.num_head,
+            num_batch=args.num_batch,
+            num_splits=args.num_splits,
+            max_kvblock_in_l2=max_kvblock_in_l2,
+            tile_shape_mn=args.tile_shape_mn,
+            qhead_per_kvhead_packgqa=args.qhead_per_kvhead_packgqa,
+            is_split_kv=args.is_split_kv,
+            lpt=args.lpt,
+            head_swizzle=head_swizzle,
+            cluster_shape_m=cluster_shape_m,
+            use_cluster_idx=args.use_cluster_idx,
+            fold_splits_into_scan=fold_splits_into_scan,
+            scheduling_mode=scheduling_mode,
+            mCuSeqlensQ=args.mCuSeqlensQ,
+            mSeqUsedQ=args.mSeqUsedQ,
+            num_m_blocks_ptr=args.num_m_blocks_ptr,
+            num_splits_dynamic_ptr=args.num_splits_dynamic_ptr,
+            virtual_batch_idx_ptr=args.virtual_batch_idx_ptr,
+            num_nheads_in_l2_ptr=args.num_nheads_in_l2_ptr,
+            cu_total_m_blocks_ptr=args.cu_total_m_blocks_ptr,
+            cu_total_splits_m_blocks_ptr=args.cu_total_splits_m_blocks_ptr,
+            blocks_to_batch_idx_ptr=args.blocks_to_batch_idx_ptr,
+        )
+
+    @cute.jit
+    def _num_m_blocks(self, lane: Int32, bidb_start: Int32) -> Int32:
+        """Per-batch m-block count"""
+        batch_idx = lane + bidb_start
+        is_valid = batch_idx < self.num_batch and lane < cute.arch.WARP_SIZE - 1
+        if cutlass.const_expr(self.num_m_blocks_ptr is not None):
+            num_m_blocks_raw = Int32(0)
+            if is_valid:
+                if cutlass.const_expr(self.virtual_batch_idx_ptr is not None):
+                    real_batch_idx = self.virtual_batch_idx_ptr[batch_idx]
+                else:
+                    real_batch_idx = batch_idx
+                num_m_blocks_raw = Int32(self.num_m_blocks_ptr[real_batch_idx])
+            return cute.ceil_div(num_m_blocks_raw, self.cluster_shape_m) if is_valid else Int32(0)
+        if cutlass.const_expr(self.virtual_batch_idx_ptr is not None):
+            seqlen = Int32(0)
+            if is_valid:
+                real_batch_idx = self.virtual_batch_idx_ptr[batch_idx]
+                if cutlass.const_expr(self.mSeqUsedQ is not None):
+                    seqlen = self.mSeqUsedQ[real_batch_idx]
+                else:
+                    seqlen = self.mCuSeqlensQ[real_batch_idx + 1] - self.mCuSeqlensQ[real_batch_idx]
+            if cutlass.const_expr(self.qhead_per_kvhead_packgqa > 1):
+                seqlen *= self.qhead_per_kvhead_packgqa
+            return (
+                cute.ceil_div(cute.ceil_div(seqlen, self.tile_shape_mn[0]), self.cluster_shape_m)
+                if is_valid
+                else Int32(0)
+            )
+        if cutlass.const_expr(self.mSeqUsedQ is not None):
+            seqlen = Int32(0)
+            if batch_idx < self.num_batch:
+                seqlen = self.mSeqUsedQ[batch_idx]
+        else:
+            assert self.mCuSeqlensQ is not None
+            cur_cu_seqlen = Int32(0)
+            if batch_idx <= self.num_batch:
+                cur_cu_seqlen = self.mCuSeqlensQ[batch_idx]
+            next_cu_seqlen = cute.arch.shuffle_sync_down(cur_cu_seqlen, offset=1)
+            seqlen = next_cu_seqlen - cur_cu_seqlen
+        if cutlass.const_expr(self.qhead_per_kvhead_packgqa > 1):
+            seqlen *= self.qhead_per_kvhead_packgqa
+        return (
+            cute.ceil_div(cute.ceil_div(seqlen, self.tile_shape_mn[0]), self.cluster_shape_m)
+            if is_valid
+            else Int32(0)
+        )
+
+    @cute.jit
+    def _num_splits(self, lane: Int32, bidb_start: Int32) -> Int32:
+        if cutlass.const_expr(not self.fold_splits_into_scan):
+            return Int32(1)
+        batch_idx = lane + bidb_start
+        is_valid = batch_idx < self.num_batch and lane < cute.arch.WARP_SIZE - 1
+        if cutlass.const_expr(not self.is_split_kv):
+            return Int32(1)
+        elif cutlass.const_expr(self.num_splits_dynamic_ptr is not None):
+            num_splits = Int32(0)
+            if is_valid:
+                if cutlass.const_expr(self.virtual_batch_idx_ptr is not None):
+                    batch_idx = self.virtual_batch_idx_ptr[batch_idx]
+                num_splits = self.num_splits_dynamic_ptr[batch_idx]
+            return num_splits
+        else:
+            return Int32(0) if not is_valid else self.num_splits
+
+    @cute.jit
+    def decode(
+        self,
+        next_tile_idx: Int32,
+        bidb_start: Int32,
+        group_start_tile: Int32,
+    ) -> Tuple[Int32, Int32, Int32, Int32, Int32, Int32, Boolean]:
+        """Search varlen batches via warp-level prefix sums and decode the work tile.
+
+        Returns
+            - block
+            - head_idx
+            - batch_idx
+            - split_idx
+            - num_splits
+            - group_start_tile
+            - is_valid
+        """
+        # The scan counts m_blocks unless splits are folded into it, so the cumsum used as a
+        # hint must match: cu_total_splits_m_blocks only applies to the folded (persistent)
+        # layout, where SingleTileVarlen keeps splits in a separate grid dim.
+        if const_expr(self.fold_splits_into_scan):
+            cu_hint_ptr = self.cu_total_splits_m_blocks_ptr
+        else:
+            cu_hint_ptr = self.cu_total_m_blocks_ptr
+        # Both SingleTileVarlen STATIC and CLC; not DynamicPersistent (where
+        # warp-scan's _bidb_start resumption already amortizes per-call cost).
+        hint_mode_ok = const_expr(
+            cu_hint_ptr is not None
+            and (
+                self.scheduling_mode == SchedulingMode.STATIC
+                or self.scheduling_mode == SchedulingMode.CLC
+            )
+        )
+        # O(1) inverted index (flat block -> batch) replaces the warp scan outright. Needs the
+        # unfolded layout, where a batch owns a contiguous tile range given by the cumsum.
+        use_blocks_to_batch = const_expr(
+            hint_mode_ok
+            and self.blocks_to_batch_idx_ptr is not None
+            and not self.fold_splits_into_scan
+        )
+        use_cumsum_hint = const_expr(hint_mode_ok and not use_blocks_to_batch)
+
+        lane_idx = cute.arch.lane_idx()
+        if const_expr(use_blocks_to_batch):
+            batch_idx = Int32(self.blocks_to_batch_idx_ptr[next_tile_idx // self.num_head])
+            num_m_blocks, num_splits = Int32(0), Int32(1)
+            is_valid = batch_idx < self.num_batch
+            if is_valid:
+                cu_lo = cu_hint_ptr[batch_idx]
+                num_m_blocks = cu_hint_ptr[batch_idx + 1] - cu_lo
+                group_start_tile = cu_lo * self.num_head
+            else:
+                batch_idx = Int32(self.num_batch)
+        else:
+            if const_expr(use_cumsum_hint):
+                target = next_tile_idx // self.num_head
+                lo = utils.get_batch_from_cu_tensor(target, cu_hint_ptr)
+                group_size = Int32(cute.arch.WARP_SIZE - 1)
+                bidb_start = (lo // group_size) * group_size
+                group_start_tile = cu_hint_ptr[bidb_start] * self.num_head
+
+            num_m_blocks = self._num_m_blocks(lane_idx, bidb_start=bidb_start)
+            num_splits = self._num_splits(lane_idx, bidb_start=bidb_start)
+            per_batch = num_m_blocks * num_splits if const_expr(self.is_split_kv) else num_m_blocks
+            cumulative = utils.warp_prefix_sum(per_batch, lane_idx)
+            m_blocks_in_group = cute.arch.shuffle_sync(cumulative, cute.arch.WARP_SIZE - 1)
+            group_end_tile = m_blocks_in_group * self.num_head + group_start_tile
+
+            batch_idx = bidb_start
+            while group_end_tile <= next_tile_idx:
+                batch_idx += cute.arch.WARP_SIZE - 1
+                if batch_idx >= self.num_batch:
+                    batch_idx = Int32(self.num_batch)
+                    group_end_tile = next_tile_idx + 1
+                else:
+                    num_m_blocks = self._num_m_blocks(lane_idx, bidb_start=batch_idx)
+                    num_splits = self._num_splits(lane_idx, bidb_start=batch_idx)
+                    per_batch = (
+                        num_m_blocks * num_splits if const_expr(self.is_split_kv) else num_m_blocks
+                    )
+                    cumulative = utils.warp_prefix_sum(per_batch, lane_idx)
+                    m_blocks_in_group = cute.arch.shuffle_sync(cumulative, cute.arch.WARP_SIZE - 1)
+                    group_end_tile += m_blocks_in_group * self.num_head
+
+            is_valid = batch_idx < self.num_batch
+            if is_valid:
+                group_start_tile = group_end_tile - m_blocks_in_group * self.num_head
+                batch_idx_in_group = cute.arch.popc(
+                    cute.arch.vote_ballot_sync(
+                        group_start_tile + cumulative * self.num_head <= next_tile_idx
+                    )
+                )
+                batch_idx += batch_idx_in_group
+                num_m_blocks_prev_lane = (
+                    Int32(0)
+                    if batch_idx_in_group == 0
+                    else cute.arch.shuffle_sync(cumulative, batch_idx_in_group - 1)
+                )
+                group_start_tile += num_m_blocks_prev_lane * self.num_head
+                num_m_blocks = cute.arch.shuffle_sync(num_m_blocks, batch_idx_in_group)
+                if const_expr(self.is_split_kv):
+                    num_splits = cute.arch.shuffle_sync(num_splits, batch_idx_in_group)
+
+        block, head_idx, split_idx = Int32(0), Int32(0), Int32(0)
+        if is_valid:
+            mh_block = next_tile_idx - group_start_tile
+
+            if const_expr(self.lpt or self.head_swizzle):
+                # This is a version of the SingleTileLPTScheduler, complicated by the fact that
+                # the seqlen can vary per batch.
+                # TODO: is there any case where num_m_blocks is 0?
+                if const_expr(not self.is_split_kv) or num_splits == 1:
+                    if const_expr(self.num_nheads_in_l2_ptr is not None):
+                        if const_expr(self.virtual_batch_idx_ptr is not None):
+                            nheads_in_l2 = Int32(
+                                self.num_nheads_in_l2_ptr[self.virtual_batch_idx_ptr[batch_idx]]
+                            )
+                        else:
+                            nheads_in_l2 = Int32(self.num_nheads_in_l2_ptr[batch_idx])
+                    else:
+                        # TODO: by right we should read the seqlen_kv but we're assuming seqlen_q == seqlen_k here
+                        num_n_blocks = (
+                            num_m_blocks
+                            * self.tile_shape_mn[0]
+                            * self.cluster_shape_m
+                            // self.qhead_per_kvhead_packgqa
+                            // self.tile_shape_mn[1]
+                        )
+                        # Seems faster to have nheads_in_l2 be a power of 2
+                        nheads_in_l2 = (
+                            16
+                            if num_n_blocks * 16 <= self.max_kvblock_in_l2
+                            else (
+                                8
+                                if num_n_blocks * 8 <= self.max_kvblock_in_l2
+                                else (
+                                    4
+                                    if num_n_blocks * 4 <= self.max_kvblock_in_l2
+                                    else (2 if num_n_blocks * 2 <= self.max_kvblock_in_l2 else 1)
+                                )
+                            )
+                        )
+                        nheads_in_l2 = min(nheads_in_l2, self.num_head)
+                    mh_in_l2 = nheads_in_l2 * num_m_blocks
+                    section_idx = mh_block // mh_in_l2
+                    l2_mod = mh_block - section_idx * mh_in_l2
+                    nheads_in_this_section = (
+                        nheads_in_l2
+                        if nheads_in_l2 * (section_idx + 1) <= self.num_head
+                        else self.num_head - section_idx * nheads_in_l2
+                    )
+                    block = l2_mod // nheads_in_this_section
+                    head_idx_residual = l2_mod - block * nheads_in_this_section
+                    head_idx = section_idx * nheads_in_l2 + head_idx_residual
+                else:
+                    head_split_idx = mh_block // num_m_blocks
+                    block = mh_block - head_split_idx * num_m_blocks
+                    head_idx = head_split_idx // num_splits
+                    split_idx = head_split_idx - head_idx * num_splits
+                if const_expr(self.lpt):
+                    block = num_m_blocks - 1 - block
+            else:
+                head_split_idx = mh_block // num_m_blocks
+                block = mh_block - head_split_idx * num_m_blocks
+                if const_expr(self.is_split_kv):
+                    head_idx = head_split_idx // num_splits
+                    split_idx = head_split_idx - head_idx * num_splits
+                else:
+                    head_idx = head_split_idx
+
+            if const_expr(self.cluster_shape_m > 1 and not self.use_cluster_idx):
+                bidx_in_cluster = cute.arch.block_in_cluster_idx()
+                block = block * self.cluster_shape_m + bidx_in_cluster[0]
+
+        return block, head_idx, batch_idx, split_idx, num_splits, group_start_tile, is_valid
+
+
 class SingleTileVarlenScheduler:
     @dataclass
     class Params(ParamsBase):
-        num_head: Int32
-        num_batch: Int32
         total_q: Int32
-        num_splits: Int32
-        max_kvblock_in_l2: Int32
-        tile_shape_mn: cutlass.Constexpr[Tuple[int, int]]
-        mCuSeqlensQ: Optional[cute.Tensor] = None
-        mSeqUsedQ: Optional[cute.Tensor] = None
-        qhead_per_kvhead_packgqa: cutlass.Constexpr[int] = 1
-        lpt: cutlass.Constexpr[bool] = False
-        is_split_kv: cutlass.Constexpr[bool] = False
-        head_swizzle: cutlass.Constexpr[bool] = False
-        cluster_shape_m: cutlass.Constexpr[int] = 1
+        scheduling_mode: cutlass.Constexpr[SchedulingMode]
+        decoder: VarlenDecoder
 
         @staticmethod
         @cute.jit
         def create(
-            args: TileSchedulerArguments, *, loc=None, ip=None
+            args: TileSchedulerArguments,
+            *,
+            scheduling_mode: SchedulingMode = SchedulingMode.STATIC,
+            loc=None,
+            ip=None,
         ) -> "SingleTileVarlenScheduler.Params":
-            size_l2 = 50 * 1024 * 1024  # 50 MB for K & V
-            max_kvblock_in_l2 = size_l2 // (
-                (args.headdim + args.headdim_v) * args.element_size * args.tile_shape_mn[1]
+            assert scheduling_mode in (SchedulingMode.STATIC, SchedulingMode.CLC), (
+                f"Only STATIC and CLC are supported, got {scheduling_mode!r}"
             )
             assert args.mCuSeqlensQ is not None or args.mSeqUsedQ is not None, (
                 "At least one of mCuSeqlensQ or mSeqUsedQ must be provided"
             )
             assert args.cluster_shape_mn[1] == 1, "Only cluster_shape_mn[1] == 1 is supported"
-            return SingleTileVarlenScheduler.Params(
-                num_head=args.num_head,
-                num_batch=args.num_batch,
-                total_q=args.total_q,
-                num_splits=args.num_splits,
-                max_kvblock_in_l2=max_kvblock_in_l2,
-                tile_shape_mn=args.tile_shape_mn,
-                mCuSeqlensQ=args.mCuSeqlensQ,
-                mSeqUsedQ=args.mSeqUsedQ,
-                qhead_per_kvhead_packgqa=args.qhead_per_kvhead_packgqa,
-                lpt=args.lpt,
-                is_split_kv=args.is_split_kv,
+            decoder = VarlenDecoder.create(
+                args,
+                fold_splits_into_scan=False,
                 head_swizzle=args.head_swizzle,
                 cluster_shape_m=args.cluster_shape_mn[0],
+                scheduling_mode=scheduling_mode,
+                loc=loc,
+                ip=ip,
+            )
+            return SingleTileVarlenScheduler.Params(
+                total_q=args.total_q,
+                scheduling_mode=scheduling_mode,
+                decoder=decoder,
             )
 
-    def __init__(self, params: Params, tile_idx: Int32, split_idx: Int32, *, loc=None, ip=None):
+    def __init__(
+        self,
+        params: Params,
+        tile_idx: Int32,
+        split_idx: Int32,
+        ctx: SchedulerState | None = None,
+        *,
+        loc=None,
+        ip=None,
+    ):
         self.params = params
         self._tile_idx = tile_idx
         self._split_idx = split_idx
         self._is_first_block = True
+        self._ctx = ctx
         self._loc = loc
         self._ip = ip
 
     @staticmethod
-    def to_underlying_arguments(args: TileSchedulerArguments, *, loc=None, ip=None) -> Params:
-        return SingleTileVarlenScheduler.Params.create(args, loc=loc, ip=ip)
+    def to_underlying_arguments(
+        args: TileSchedulerArguments,
+        *,
+        scheduling_mode: SchedulingMode = SchedulingMode.STATIC,
+        loc=None,
+        ip=None,
+    ) -> Params:
+        return SingleTileVarlenScheduler.Params.create(
+            args, scheduling_mode=scheduling_mode, loc=loc, ip=ip
+        )
+
+    @staticmethod
+    @cute.jit
+    def clc_problem_shape(params: Params):
+        return ClcDynamicPersistentTileSchedulerParams(
+            problem_shape_ntile_mnl=SingleTileVarlenScheduler.get_grid_shape(params),
+            cluster_shape_mnk=(1, 1, 1),
+        )
 
     @staticmethod
-    def create(params: Params, *, loc=None, ip=None) -> "SingleTileVarlenScheduler":
+    @cute.jit
+    def create(
+        params: Params, ctx: SchedulerState | None = None, *, loc=None, ip=None
+    ) -> "SingleTileVarlenScheduler":
+        if const_expr(params.scheduling_mode == SchedulingMode.CLC):
+            block_idx = cute.arch.block_idx()
+            split_idx = Int32(0)
+            if const_expr(params.decoder.is_split_kv):
+                split_idx = block_idx[1]
+            return SingleTileVarlenScheduler(
+                params,
+                block_idx[0],
+                split_idx,
+                ctx,
+                loc=loc,
+                ip=ip,
+            )
         tile_idx, split_idx, _ = cute.arch.block_idx()
         return SingleTileVarlenScheduler(params, tile_idx, split_idx, loc=loc, ip=ip)
 
@@ -569,148 +1291,283 @@ class SingleTileVarlenScheduler:
         loc=None,
         ip=None,
     ) -> Tuple[Int32, Int32, Int32]:
+        d = params.decoder
         total_blocks_max = (
-            params.total_q
-            + params.num_batch * (params.cluster_shape_m * params.tile_shape_mn[0] - 1)
-        ) // params.tile_shape_mn[0]
-        # round down to nearest multiple of cluster since odd excess is always padding
-        total_blocks_max = total_blocks_max // params.cluster_shape_m * params.cluster_shape_m
-        return (total_blocks_max * params.num_head, params.num_splits, Int32(1))
+            params.total_q + d.num_batch * (d.cluster_shape_m * d.tile_shape_mn[0] - 1)
+        ) // d.tile_shape_mn[0]
+        # Round down to nearest multiple of cluster since odd excess is always padding.
+        total_blocks_max = total_blocks_max // d.cluster_shape_m * d.cluster_shape_m
+        return (total_blocks_max * d.num_head, d.num_splits, Int32(1))
 
     @cute.jit
-    def _get_num_m_blocks(self, lane: Int32, bidb_start: Int32) -> Int32:
-        params = self.params
-        batch_idx = lane + bidb_start
-        if cutlass.const_expr(params.mSeqUsedQ is not None):
-            seqlen = Int32(0)
-            if batch_idx < params.num_batch:
-                seqlen = params.mSeqUsedQ[batch_idx]
-        else:
-            assert params.mCuSeqlensQ is not None
-            cur_cu_seqlen = Int32(0)
-            if batch_idx <= params.num_batch:
-                cur_cu_seqlen = params.mCuSeqlensQ[batch_idx]
-            next_cu_seqlen = cute.arch.shuffle_sync_down(cur_cu_seqlen, offset=1)
-            seqlen = next_cu_seqlen - cur_cu_seqlen
-        if cutlass.const_expr(params.qhead_per_kvhead_packgqa > 1):
-            seqlen *= params.qhead_per_kvhead_packgqa
-        return (
-            cute.ceil_div(cute.ceil_div(seqlen, params.tile_shape_mn[0]), params.cluster_shape_m)
-            if batch_idx < params.num_batch and lane < cute.arch.WARP_SIZE - 1
-            else Int32(0)
+    def _decode_work_tile(self) -> WorkTileInfo:
+        """Map self._tile_idx to (block, head, batch, split) via warp-level prefix sums."""
+        d = self.params.decoder
+        next_tile_idx = self._tile_idx // d.cluster_shape_m
+        block, head_idx, batch_idx, _, _, _, is_valid = d.decode(next_tile_idx, Int32(0), Int32(0))
+        is_valid = is_valid and self._is_first_block
+        split_idx = self._split_idx if const_expr(d.is_split_kv) else Int32(0)
+        if const_expr(d.virtual_batch_idx_ptr is not None):
+            if is_valid:
+                batch_idx = d.virtual_batch_idx_ptr[batch_idx]
+        # Pack dynamic per-batch num_splits into high 16 bits of split_idx
+        if const_expr(d.is_split_kv and d.num_splits_dynamic_ptr is not None):
+            if is_valid:
+                num_splits = Int32(d.num_splits_dynamic_ptr[batch_idx])
+                split_idx = split_idx | (num_splits << 16)
+        return WorkTileInfo(
+            (Int32(block), Int32(head_idx), Int32(batch_idx), Int32(split_idx)),
+            is_valid,
         )
 
     @cute.jit
     def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo:
-        params = self.params
-        lane_idx = cute.arch.lane_idx()
-        num_m_blocks = self._get_num_m_blocks(lane_idx, bidb_start=0)
-        num_m_blocks_cumulative = utils.warp_prefix_sum(num_m_blocks, lane_idx)
-        # Total number of blocks for the next 31 batches
-        m_blocks_in_group = cute.arch.shuffle_sync(num_m_blocks_cumulative, cute.arch.WARP_SIZE - 1)
-        # Same for all lanes
-        group_end_tile = m_blocks_in_group * params.num_head
-        # if cute.arch.thread_idx()[0] == 128 + 31: cute.printf("SingleTileVarlenScheduler: tile_idx=%d, group_end_tile = %d, num_m_blocks=%d, num_m_blocks_cumulative = %d, m_blocks_in_group = %d", self._tile_idx, group_end_tile, num_m_blocks, num_m_blocks_cumulative, m_blocks_in_group)
-        block, head_idx, batch_idx = Int32(0), Int32(0), Int32(0)
-        next_tile_idx = self._tile_idx // params.cluster_shape_m
-        while group_end_tile <= next_tile_idx:
-            batch_idx += cute.arch.WARP_SIZE - 1
-            if batch_idx >= params.num_batch:
-                batch_idx = Int32(params.num_batch)
-                group_end_tile = next_tile_idx + 1
-            else:
-                num_m_blocks = self._get_num_m_blocks(lane_idx, bidb_start=batch_idx)
-                num_m_blocks_cumulative = utils.warp_prefix_sum(num_m_blocks, lane_idx)
-                m_blocks_in_group = cute.arch.shuffle_sync(
-                    num_m_blocks_cumulative, cute.arch.WARP_SIZE - 1
-                )
-                group_end_tile += m_blocks_in_group * params.num_head
-        is_valid = False
-        if batch_idx >= params.num_batch:
-            block, head_idx, batch_idx = Int32(0), Int32(0), Int32(params.num_batch)
-        else:
-            group_start_tile = group_end_tile - m_blocks_in_group * params.num_head
-            # if cute.arch.thread_idx()[0] == 128 + 31: cute.printf("SingleTileVarlenScheduler: tile_idx=%d, group_end_tile = %d, num_m_blocks=%d, batch_idx = %d", self._tile_idx, group_end_tile, num_m_blocks, batch_idx)
-            # The next problem to process is the first one that does not have ending tile position
-            # that is greater than or equal to tile index.
-            batch_idx_in_group = cute.arch.popc(
-                cute.arch.vote_ballot_sync(
-                    group_start_tile + num_m_blocks_cumulative * params.num_head <= next_tile_idx
-                )
-            )
-            batch_idx += batch_idx_in_group
-            num_m_blocks_prev_lane = (
-                0
-                if batch_idx_in_group == 0
-                else cute.arch.shuffle_sync(num_m_blocks_cumulative, batch_idx_in_group - 1)
-            )
-            num_m_blocks = cute.arch.shuffle_sync(num_m_blocks, batch_idx_in_group)
-            mh_block = next_tile_idx - group_start_tile - num_m_blocks_prev_lane * params.num_head
-            if cutlass.const_expr(params.lpt or params.head_swizzle):
-                # This is a version of the SingleTileLPTScheduler, complicated by the fact that
-                # the seqlen can vary per batch.
-                # TODO: is there any case where num_m_blocks is 0?
-                # TODO: by right we should read the seqlen_kv but we're assuming seqlen_q == seqlen_k here
-                num_n_blocks = (
-                    num_m_blocks
-                    * params.tile_shape_mn[0]
-                    // params.qhead_per_kvhead_packgqa
-                    // params.tile_shape_mn[1]
-                )
-                # nheads_in_l2 = min(max(self.max_kvblock_in_l2 // num_n_blocks, 1), self.num_head)
-                # Seems faster to have this be a power of 2
-                nheads_in_l2 = (
-                    16
-                    if num_n_blocks * 16 <= params.max_kvblock_in_l2
-                    else (
-                        8
-                        if num_n_blocks * 8 <= params.max_kvblock_in_l2
-                        else (
-                            4
-                            if num_n_blocks * 4 <= params.max_kvblock_in_l2
-                            else (2 if num_n_blocks * 2 <= params.max_kvblock_in_l2 else 1)
-                        )
-                    )
-                )
-                nheads_in_l2 = min(nheads_in_l2, params.num_head)
-                mh_in_l2 = nheads_in_l2 * num_m_blocks
-                section_idx = mh_block // mh_in_l2
-                l2_mod = mh_block - section_idx * mh_in_l2
-                # Deal with tail section
-                nheads_in_this_section = (
-                    nheads_in_l2
-                    if nheads_in_l2 * (section_idx + 1) <= params.num_head
-                    else params.num_head - section_idx * nheads_in_l2
-                )
-                block = l2_mod // nheads_in_this_section
-                head_idx_residual = l2_mod - block * nheads_in_this_section
-                head_idx = section_idx * nheads_in_l2 + head_idx_residual
-                if cutlass.const_expr(params.lpt):
-                    block = num_m_blocks - 1 - block
-            else:
-                head_idx = mh_block // num_m_blocks
-                block = mh_block - head_idx * num_m_blocks
-            is_valid = self._is_first_block and batch_idx < params.num_batch
-            if cutlass.const_expr(params.cluster_shape_m > 1):
-                bidx_in_cluster = cute.arch.block_in_cluster_idx()
-                block = block * params.cluster_shape_m + bidx_in_cluster[0]
-        # if cute.arch.thread_idx()[0] == 128: cute.printf("SingleTileVarlenScheduler: tile_idx=%d, batch_idx=%d, head_idx=%d, block=%d, is_valid = %d", self._tile_idx, batch_idx, head_idx, block, is_valid)
-        split_idx = self._split_idx if const_expr(params.is_split_kv) else Int32(0)
-        return WorkTileInfo((Int32(block), Int32(head_idx), Int32(batch_idx), split_idx), is_valid)
+        if const_expr(self.params.scheduling_mode == SchedulingMode.CLC):
+            clc_work = self._ctx.get_current_work()
+            # Default to grid_dim (one past last valid flat index) so _decode_work_tile
+            # returns is_valid=False when CLC is exhausted. CLC tile_idx is garbage when
+            # invalid, so we can't trust it. Local-then-assign avoids CuTe DSL structural
+            # mismatch on self inside the runtime if.
+            new_tile_idx = cute.arch.grid_dim()[0]
+            new_split_idx = Int32(0)
+            if clc_work.is_valid_tile:
+                new_tile_idx = clc_work.tile_idx[0]
+                if const_expr(self.params.decoder.is_split_kv):
+                    new_split_idx = clc_work.tile_idx[1]
+            self._tile_idx = new_tile_idx
+            self._split_idx = new_split_idx
+        return self._decode_work_tile()
 
+    @cute.jit
     def initial_work_tile_info(self, *, loc=None, ip=None):
-        return self.get_current_work(loc=loc, ip=ip)
+        if const_expr(self.params.scheduling_mode == SchedulingMode.CLC):
+            clc_work = self._ctx.initial_work_tile_info()
+            # See get_current_work for why grid_dim and local-then-assign.
+            new_tile_idx = cute.arch.grid_dim()[0]
+            new_split_idx = Int32(0)
+            if clc_work.is_valid_tile:
+                new_tile_idx = clc_work.tile_idx[0]
+                if const_expr(self.params.decoder.is_split_kv):
+                    new_split_idx = clc_work.tile_idx[1]
+            self._tile_idx = new_tile_idx
+            self._split_idx = new_split_idx
+        return self._decode_work_tile()
 
     def prefetch_next_work(self, *, loc=None, ip=None):
-        pass
+        if const_expr(self.params.scheduling_mode == SchedulingMode.CLC):
+            self._ctx.prefetch_next_work(loc=loc, ip=ip)
 
     def advance_to_next_work(self, *, loc=None, ip=None):
-        # Single tile scheduler - set to invalid tile_idx to indicate no more work
+        if const_expr(self.params.scheduling_mode == SchedulingMode.CLC):
+            self._ctx.consumer_wait(loc=loc, ip=ip)
+            work = self.get_current_work()
+            self._ctx.consumer_release(loc=loc, ip=ip)
+            return work
         self._is_first_block = False
+        return self.get_current_work()
+
+    def producer_tail(self, *, loc=None, ip=None):
+        if const_expr(self.params.scheduling_mode == SchedulingMode.CLC):
+            self._ctx.producer_tail(loc=loc, ip=ip)
 
     def __extract_mlir_values__(self):
         values, self._values_pos = [], []
-        for obj in [self.params, self._tile_idx, self._split_idx]:
+        objs = [self.params, self._tile_idx, self._split_idx]
+        if const_expr(self.params.scheduling_mode == SchedulingMode.CLC):
+            objs += [self._ctx]
+        for obj in objs:
+            obj_values = cutlass.extract_mlir_values(obj)
+            values += obj_values
+            self._values_pos.append(len(obj_values))
+        return values
+
+    def __new_from_mlir_values__(self, values):
+        obj_list = []
+        objs = [self.params, self._tile_idx, self._split_idx]
+        if const_expr(self.params.scheduling_mode == SchedulingMode.CLC):
+            objs += [self._ctx]
+        for obj, n_items in zip(objs, self._values_pos):
+            obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items]))
+            values = values[n_items:]
+        scheduler = self.__class__(*obj_list, loc=self._loc)
+        # See the note on Python-only attributes in SingleTileScheduler.
+        scheduler._is_first_block = self._is_first_block
+        return scheduler
+
+
+class DynamicPersistentVarlenScheduler:
+    @dataclass
+    class Params(ParamsBase):
+        total_q: Int32
+        decoder: VarlenDecoder
+        tile_count_semaphore: Optional[cute.Pointer] = None
+        persistent_cta_multiplier: cutlass.Constexpr[int] = 1
+
+        @staticmethod
+        @cute.jit
+        def create(
+            args: TileSchedulerArguments, *, loc=None, ip=None
+        ) -> "DynamicPersistentVarlenScheduler.Params":
+            assert args.mCuSeqlensQ is not None or args.mSeqUsedQ is not None, (
+                "At least one of mCuSeqlensQ or mSeqUsedQ must be provided"
+            )
+            # TODO: support non-trivial cluster shapes in a follow-on PR
+            assert args.cluster_shape_mn[0] == 1 and args.cluster_shape_mn[1] == 1, (
+                "DynamicPersistentVarlenScheduler currently requires cluster_shape_mn == (1, 1)"
+            )
+            decoder = VarlenDecoder.create(
+                args,
+                fold_splits_into_scan=True,
+                scheduling_mode=SchedulingMode.DYNAMIC,
+                loc=loc,
+                ip=ip,
+            )
+            return DynamicPersistentVarlenScheduler.Params(
+                total_q=args.total_q,
+                decoder=decoder,
+                tile_count_semaphore=args.tile_count_semaphore,
+                persistent_cta_multiplier=args.persistent_cta_multiplier,
+            )
+
+    def __init__(
+        self,
+        params: Params,
+        ctx: SchedulerState,
+        bidb_start: Int32,
+        group_start_tile: Int32,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        self.params = params
+        self._ctx = ctx
+        self._bidb_start = bidb_start
+        self._group_start_tile = group_start_tile
+        self._loc = loc
+        self._ip = ip
+
+    @staticmethod
+    def to_underlying_arguments(
+        args: TileSchedulerArguments,
+        *,
+        scheduling_mode: SchedulingMode = SchedulingMode.DYNAMIC,
+        loc=None,
+        ip=None,
+    ) -> Params:
+        assert scheduling_mode == SchedulingMode.DYNAMIC, (
+            f"DynamicPersistentVarlenScheduler only supports DYNAMIC, got {scheduling_mode!r}"
+        )
+        return DynamicPersistentVarlenScheduler.Params.create(args, loc=loc, ip=ip)
+
+    @staticmethod
+    @cute.jit
+    def create(
+        params: Params,
+        ctx: SchedulerState,
+        *,
+        loc=None,
+        ip=None,
+    ) -> "DynamicPersistentVarlenScheduler":
+        return DynamicPersistentVarlenScheduler(params, ctx, Int32(0), Int32(0), loc=loc, ip=ip)
+
+    # called by host
+    @staticmethod
+    def get_grid_shape(
+        params: Params,
+        *,
+        loc=None,
+        ip=None,
+    ) -> Tuple[Int32, Int32, Int32]:
+        d = params.decoder
+        total_blocks_max = (
+            params.total_q + d.num_batch * (d.tile_shape_mn[0] - 1)
+        ) // d.tile_shape_mn[0]
+        total_blocks = total_blocks_max * d.num_head * d.num_splits
+        hardware_info = HardwareInfo()
+        sm_count = (
+            hardware_info.get_device_multiprocessor_count() * params.persistent_cta_multiplier
+        )
+        return (cutlass.min(sm_count, total_blocks), Int32(1), Int32(1))
+
+    @cute.jit
+    def get_current_work(
+        self,
+        next_tile_idx: Int32,
+        bidb_start: Int32,
+        group_start_tile: Int32,
+        *,
+        loc=None,
+        ip=None,
+    ) -> Tuple[WorkTileInfo, Int32]:
+        d = self.params.decoder
+        block, head_idx, batch_idx, split_idx, num_splits, group_start_tile, is_valid = d.decode(
+            next_tile_idx, bidb_start, group_start_tile
+        )
+        if const_expr(d.is_split_kv and d.num_splits_dynamic_ptr is not None):
+            if is_valid:
+                split_idx = split_idx | (num_splits << 16)
+        if const_expr(d.virtual_batch_idx_ptr is not None):
+            if is_valid:
+                batch_idx = d.virtual_batch_idx_ptr[batch_idx]
+        return (
+            WorkTileInfo(
+                (Int32(block), Int32(head_idx), Int32(batch_idx), Int32(split_idx)),
+                is_valid,
+            ),
+            group_start_tile,
+        )
+
+    @cute.jit
+    def prefetch_next_work(self, *, loc=None, ip=None):
+        ctx = self._ctx
+        next_tile_idx = Int32(0)
+        if cute.arch.lane_idx() == 0:
+            next_tile_idx = cute.arch.grid_dim()[0] + utils.atomic_add_i32(
+                1,
+                self.params.tile_count_semaphore,
+            )
+        next_tile_idx = cute.arch.shuffle_sync(next_tile_idx, 0)
+        work_info, new_group_start_tile = self.get_current_work(
+            next_tile_idx, self._bidb_start, self._group_start_tile
+        )
+        # Advance scan state so the next prefetch resumes from this tile's batch
+        # group instead of restarting at batch 0.
+        self._bidb_start = Int32(work_info.tile_idx[2])
+        self._group_start_tile = new_group_start_tile
+        ctx.producer_acquire()
+        with cute.arch.elect_one():
+            block, head_idx, batch_idx, split_idx = work_info.tile_idx
+            ctx.write_work_info(block, head_idx, batch_idx, split_idx)
+            ctx.producer_commit()
+        ctx.advance_producer_state()
+
+    @cute.jit
+    def advance_to_next_work(self, *, loc=None, ip=None) -> WorkTileInfo:
+        ctx = self._ctx
+        ctx.consumer_wait()
+        block = ctx._work_info[0]
+        head_idx = ctx._work_info[1]
+        batch_idx = ctx._work_info[2]
+        split_idx = ctx._work_info[3]
+        is_valid = batch_idx < self.params.decoder.num_batch
+        work_info = WorkTileInfo((block, head_idx, batch_idx, split_idx), is_valid)
+        ctx.consumer_release()
+        return work_info
+
+    @cute.jit
+    def initial_work_tile_info(self, *, loc=None, ip=None) -> WorkTileInfo:
+        cta_tile_idx, _, _ = cute.arch.block_idx()
+        work_info, new_group_start_tile = self.get_current_work(cta_tile_idx, Int32(0), Int32(0))
+        self._bidb_start = Int32(work_info.tile_idx[2])
+        self._group_start_tile = new_group_start_tile
+        return work_info
+
+    def producer_tail(self, *, loc=None, ip=None):
+        self._ctx.producer_tail(loc=loc, ip=ip)
+
+    def __extract_mlir_values__(self):
+        values, self._values_pos = [], []
+        for obj in [self.params, self._ctx, self._bidb_start, self._group_start_tile]:
             obj_values = cutlass.extract_mlir_values(obj)
             values += obj_values
             self._values_pos.append(len(obj_values))
@@ -719,9 +1576,545 @@ class SingleTileVarlenScheduler:
     def __new_from_mlir_values__(self, values):
         obj_list = []
         for obj, n_items in zip(
-            [self.params, self._tile_idx, self._split_idx],
+            [self.params, self._ctx, self._bidb_start, self._group_start_tile],
             self._values_pos,
         ):
             obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items]))
             values = values[n_items:]
-        return SingleTileVarlenScheduler(*(tuple(obj_list)), loc=self._loc)
+        return self.__class__(*obj_list, loc=self._loc)
+
+
+# -----------------------------------------------------------------------------
+# SM100 FMHA-specific schedulers (kept separate from generic schedulers).
+# -----------------------------------------------------------------------------
+
+
+class Sm100FmhaStaticTileSchedulerParams:
+    """A class to represent parameters for the FMHA (Fused Multi-Head Attention) static tile scheduler.
+
+    This class holds the configuration parameters needed to initialize and configure
+    the tile scheduler for FMHA operations.
+
+    :ivar is_persistent: Whether to use persistent kernel mode.
+    :type is_persistent: bool
+    :ivar problem_shape_mbh: Problem shape in (M, B, H) format.
+    :type problem_shape_mbh: cute.Shape
+    """
+
+    def __init__(
+        self,
+        is_persistent: bool,
+        problem_shape_mbh: cute.Shape,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        """
+        Initializes the Sm100FmhaStaticTileSchedulerParams with the given parameters.
+
+        :param is_persistent: Whether to use persistent kernel mode.
+        :type is_persistent: bool
+        :param problem_shape_mbh: Problem shape in (M, B, H) format.
+        :type problem_shape_mbh: cute.Shape
+        """
+        self.is_persistent = is_persistent
+        self.problem_shape_mbh = problem_shape_mbh
+        self._loc = loc
+        self._ip = ip
+
+    def __extract_mlir_values__(self):
+        values, self._values_pos = [], []
+        for obj in [self.problem_shape_mbh]:
+            obj_values = extract_mlir_values(obj)
+            values += obj_values
+            self._values_pos.append(len(obj_values))
+        return values
+
+    def __new_from_mlir_values__(self, values):
+        obj_list = []
+        for obj, n_items in zip([self.problem_shape_mbh], self._values_pos):
+            obj_list.append(new_from_mlir_values(obj, values[:n_items]))
+            values = values[n_items:]
+        return Sm100FmhaStaticTileSchedulerParams(
+            self.is_persistent, *(tuple(obj_list)), loc=self._loc
+        )
+
+
+class Sm100FmhaStaticTileScheduler:
+    """A static tile scheduler for FMHA (Fused Multi-Head Attention) operations.
+
+    This class manages the scheduling of work tiles for FMHA kernels, supporting
+    both persistent and non-persistent kernel modes. It tracks the current work
+    position and advances through the problem space efficiently.
+
+    :ivar _params: Scheduler parameters.
+    :type _params: Sm100FmhaStaticTileSchedulerParams
+    :ivar _blk_coord: Block coordinates.
+    :type _blk_coord: cute.Coord
+    :ivar _grid_shape: Grid shape for the kernel.
+    :type _grid_shape: cute.Shape
+    :ivar _is_persistent: Whether to use persistent kernel mode.
+    :type _is_persistent: bool
+    :ivar _current_work_linear_idx: Current linear work index.
+    :type _current_work_linear_idx: Int32
+    :ivar _problem_shape_mbh: Problem shape in (M, B, H) format.
+    :type _problem_shape_mbh: cute.Layout
+    :ivar _num_blocks: Number of blocks in the problem.
+    :type _num_blocks: Int32
+    :ivar _is_first_block: Whether this is the first block.
+    :type _is_first_block: bool
+    :ivar num_persistent_sm: Number of persistent SMs.
+    :type num_persistent_sm: Int32
+    """
+
+    def __init__(
+        self,
+        params: Sm100FmhaStaticTileSchedulerParams,
+        current_work_linear_idx: Int32,
+        blk_coord: cute.Coord,
+        grid_shape: cute.Shape,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        """
+        Initializes the Sm100FmhaStaticTileScheduler with the given parameters.
+
+        :param params: Scheduler parameters.
+        :type params: Sm100FmhaStaticTileSchedulerParams
+        :param current_work_linear_idx: Current linear work index.
+        :type current_work_linear_idx: Int32
+        :param blk_coord: Block coordinates.
+        :type blk_coord: cute.Coord
+        :param grid_shape: Grid shape for the kernel.
+        :type grid_shape: cute.Shape
+        """
+        self._params = params
+        self._blk_coord = blk_coord
+        self._grid_shape = grid_shape
+        self._is_persistent = params.is_persistent
+        self._current_work_linear_idx = current_work_linear_idx
+        self._problem_shape_mbh = cute.make_layout(params.problem_shape_mbh, loc=loc, ip=ip)
+        self._num_blocks = cute.size(self._problem_shape_mbh, loc=loc, ip=ip)
+        self._is_first_block = True
+        self.num_persistent_sm = cute.size(grid_shape, loc=loc, ip=ip)
+        self._loc = loc
+        self._ip = ip
+
+    # called by host
+    @staticmethod
+    def get_grid_shape(
+        params: Sm100FmhaStaticTileSchedulerParams,
+        *,
+        loc=None,
+        ip=None,
+    ) -> cute.Shape:
+        """
+        Determine the grid shape for the FMHA kernel.
+
+        For persistent kernels, the grid shape is limited by the number of SMs
+        (Streaming Multiprocessors) available on the device. For non-persistent
+        kernels, the grid shape matches the problem shape.
+
+        :param params: Scheduler parameters.
+        :type params: Sm100FmhaStaticTileSchedulerParams
+
+        :return: Grid shape as (M, B, H) tuple.
+        :rtype: cute.Shape
+        """
+        if params.is_persistent:
+            hardware_info = HardwareInfo()
+            sm_count = hardware_info.get_device_multiprocessor_count()
+            return (
+                dsl_min(sm_count, cute.size(params.problem_shape_mbh, loc=loc, ip=ip)),
+                1,
+                1,
+            )
+        else:
+            return params.problem_shape_mbh
+
+    @staticmethod
+    def check_valid_work_for_seqlen_q(
+        q_tiler: int,
+        current_idx: Int32,
+        seqlen_q: Int32,
+    ) -> Boolean:
+        """
+        Check if the current work index is valid for the given query sequence length.
+
+        This method verifies that the current work tile index multiplied by the
+        query tiler size is within the bounds of the query sequence length.
+
+        :param q_tiler: Query tiler size.
+        :type q_tiler: int
+        :param current_idx: Current work index.
+        :type current_idx: Int32
+        :param seqlen_q: Query sequence length.
+        :type seqlen_q: Int32
+
+        :return: True if the work is valid, False otherwise.
+        :rtype: Boolean
+        """
+        return current_idx * q_tiler < seqlen_q
+
+    def get_current_work(self, *, loc=None, ip=None) -> cutlass.utils.WorkTileInfo:
+        """
+        Get information about the current work tile.
+
+        Determines if the current work is valid and computes the tile coordinates
+        based on whether the kernel is persistent or non-persistent.
+
+        :return: WorkTileInfo containing tile coordinates and validity flag.
+        :rtype: WorkTileInfo
+        """
+        is_valid = (
+            self._current_work_linear_idx < self._num_blocks
+            if self._is_persistent
+            else self._is_first_block
+        )
+
+        blk_coord = (0, 0, 0)
+        if self._is_persistent:
+            blk_coord = self._problem_shape_mbh.get_hier_coord(
+                self._current_work_linear_idx, loc=loc, ip=ip
+            )
+        else:
+            blk_coord = self._blk_coord
+
+        # cur_tile_coord is (mid, 0, (bid, hid))
+        cur_tile_coord = (
+            blk_coord[0],
+            0,
+            (blk_coord[1], blk_coord[2]),
+        )
+
+        return cutlass.utils.WorkTileInfo(cur_tile_coord, is_valid)
+
+    def initial_work_tile_info(self, *, loc=None, ip=None):
+        """
+        Get the initial work tile information.
+
+        :return: Initial WorkTileInfo.
+        :rtype: WorkTileInfo
+        """
+        return self.get_current_work(loc=loc, ip=ip)
+
+    def advance_to_next_work(self, *, advance_count=1, loc=None, ip=None):
+        """
+        Advance to the next work tile and return it.
+
+        For persistent kernels, advances by the number of persistent SMs.
+        For non-persistent kernels, marks that the first block has been processed.
+        """
+        if self._is_persistent:
+            self._current_work_linear_idx += advance_count * self.num_persistent_sm
+        self._is_first_block = False
+        return self.get_current_work()
+
+    def prefetch_next_work(self, *, loc=None, ip=None):
+        """No-op for static scheduler."""
+        pass
+
+    def producer_tail(self, *, loc=None, ip=None):
+        """No-op for static scheduler."""
+        pass
+
+    def __extract_mlir_values__(self):
+        values = extract_mlir_values(self._params)
+        values.extend(extract_mlir_values(self._current_work_linear_idx))
+        values.extend(extract_mlir_values(self._blk_coord))
+        values.extend(extract_mlir_values(self._grid_shape))
+        return values
+
+    def __new_from_mlir_values__(self, values):
+        assert len(values) == 10
+        new_params = new_from_mlir_values(self._params, values[0:3])
+        new_current_work_linear_idx = new_from_mlir_values(
+            self._current_work_linear_idx, [values[3]]
+        )
+        new_blk_coord = new_from_mlir_values(self._blk_coord, values[4:7])
+        new_grid_shape = new_from_mlir_values(self._grid_shape, values[7:])
+        scheduler = Sm100FmhaStaticTileScheduler(
+            new_params, new_current_work_linear_idx, new_blk_coord, new_grid_shape
+        )
+        # See the note on Python-only attributes in SingleTileScheduler.
+        scheduler._is_first_block = self._is_first_block
+        return scheduler
+
+
+def compute_sm100_fmha_grid(
+    o_shape: cute.Shape,
+    cta_tiler: Tuple[int, int, int],
+    is_persistent: bool,
+) -> Tuple[Sm100FmhaStaticTileSchedulerParams, Tuple[int, int, int]]:
+    """Compute grid parameters for FMHA (static scheduler).
+
+    The output tensor o has shape (s, d, ((h_r, h_k), b)).
+    """
+    tile_sched_params = Sm100FmhaStaticTileSchedulerParams(
+        is_persistent,
+        (
+            cute.ceil_div(cute.size(o_shape[0]), cta_tiler[0]),
+            cute.size(o_shape[2][0]),
+            cute.size(o_shape[2][1]),
+        ),
+    )
+    grid = Sm100FmhaStaticTileScheduler.get_grid_shape(tile_sched_params)
+    return tile_sched_params, grid
+
+
+##############################################################################
+# Fmha CLC dynamic tile scheduler
+##############################################################################
+
+
+class Sm100FmhaClcDynamicTileSchedulerParams:
+    """Parameters for FMHA CLC dynamic persistent tile scheduler.
+
+    This class manages the layout of tiles for CLC (Cluster Launch Control)
+    based dynamic scheduling, adapted for FMHA's (M, B, H) problem shape.
+
+    :ivar problem_shape_mbh: Problem shape in (M, B, H) format.
+    :type problem_shape_mbh: cute.Shape
+    :ivar cluster_shape_mnk: Cluster shape in (M, N, K) format.
+    :type cluster_shape_mnk: cute.Shape
+    """
+
+    def __init__(
+        self,
+        problem_shape_mbh: cute.Shape,
+        cluster_shape_mnk: cute.Shape,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        self.problem_shape_mbh = problem_shape_mbh
+        self._cluster_shape_mnk = cluster_shape_mnk
+        self.cluster_shape_mn = cluster_shape_mnk[:2]
+        self._loc = loc
+        self._ip = ip
+
+        # FMHA uses linear indexing over (M, B, H), convert to (M, N, L) style
+        # For FMHA: M dim is tile count along sequence, N=1, L=(B*H)
+        self.problem_shape_ntile_mnl = (
+            problem_shape_mbh[0],  # M tiles
+            1,  # N tiles (always 1 for FMHA)
+            problem_shape_mbh[1] * problem_shape_mbh[2],  # L = B * H
+        )
+
+        # Create layout for cluster-to-tile mapping
+        self.problem_layout_ncluster_mnl = cute.make_layout(
+            cute.ceil_div(self.problem_shape_ntile_mnl, cluster_shape_mnk[:2]),
+            loc=loc,
+            ip=ip,
+        )
+
+    def __extract_mlir_values__(self):
+        values, self._values_pos = [], []
+        for obj in [
+            self.problem_shape_mbh,
+            self._cluster_shape_mnk,
+        ]:
+            obj_values = extract_mlir_values(obj)
+            values += obj_values
+            self._values_pos.append(len(obj_values))
+        return values
+
+    def __new_from_mlir_values__(self, values):
+        obj_list = []
+        values_copy = list(values)
+        for obj, n_items in zip(
+            [self.problem_shape_mbh, self._cluster_shape_mnk],
+            self._values_pos,
+        ):
+            obj_list.append(new_from_mlir_values(obj, values_copy[:n_items]))
+            values_copy = values_copy[n_items:]
+        return Sm100FmhaClcDynamicTileSchedulerParams(*(tuple(obj_list)), loc=self._loc)
+
+    def get_grid_shape(self, *, loc=None, ip=None) -> Tuple[int, int, int]:
+        """Compute grid shape aligned with cluster shape."""
+        return cute.round_up(self.problem_shape_ntile_mnl, self._cluster_shape_mnk)
+
+    def clc_hw_params(self) -> ClcDynamicPersistentTileSchedulerParams:
+        """Return params for the upstream CLC hardware scheduler."""
+        return ClcDynamicPersistentTileSchedulerParams(
+            problem_shape_ntile_mnl=self.problem_shape_ntile_mnl,
+            cluster_shape_mnk=self._cluster_shape_mnk,
+        )
+
+
+class Sm100FmhaClcDynamicTileScheduler:
+    """CLC dynamic persistent tile scheduler for FMHA.
+
+    This scheduler uses Blackwell's Cluster Launch Control hardware mechanism
+    for dynamic tile distribution, providing automatic load balancing.
+    Adapted for FMHA's (M, B, H) problem shape.
+    """
+
+    def __init__(
+        self,
+        params: Sm100FmhaClcDynamicTileSchedulerParams,
+        cta_id_in_cluster: cute.Coord,
+        num_tiles_executed: Int32,
+        clc_response_ptr: cute.Pointer,
+        block_idx: Tuple,
+        clc: ClcSchedulerState = None,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        self.params = params
+        self.cta_id_in_cluster = cta_id_in_cluster
+        self._num_tiles_executed = num_tiles_executed
+        self._clc_response_ptr = clc_response_ptr
+        self._block_idx = block_idx
+        self.clc = clc
+        self._loc = loc
+        self._ip = ip
+
+    def __extract_mlir_values__(self):
+        values = extract_mlir_values(self.cta_id_in_cluster)
+        values.extend(extract_mlir_values(self._num_tiles_executed))
+        values.extend(extract_mlir_values(self._clc_response_ptr))
+        values.extend(extract_mlir_values(self._block_idx))
+        if self.clc is not None:
+            values.extend(extract_mlir_values(self.clc))
+        return values
+
+    def __new_from_mlir_values__(self, values):
+        new_cta_id_in_cluster = new_from_mlir_values(self.cta_id_in_cluster, values[0:3])
+        new_num_tiles_executed = new_from_mlir_values(self._num_tiles_executed, [values[3]])
+        new_clc_response_ptr = new_from_mlir_values(self._clc_response_ptr, [values[4]])
+        new_block_idx = new_from_mlir_values(self._block_idx, values[5:8])
+        new_clc = None
+        if self.clc is not None:
+            new_clc = new_from_mlir_values(self.clc, values[8:])
+        return Sm100FmhaClcDynamicTileScheduler(
+            self.params,
+            new_cta_id_in_cluster,
+            new_num_tiles_executed,
+            new_clc_response_ptr,
+            new_block_idx,
+            new_clc,
+        )
+
+    @staticmethod
+    def create(
+        params: Sm100FmhaClcDynamicTileSchedulerParams,
+        block_idx: Tuple,
+        grid_dim: Tuple,
+        clc_response_ptr: cute.Pointer,
+        clc: ClcSchedulerState = None,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        """Create a CLC dynamic tile scheduler instance."""
+        bidx, bidy, bidz = block_idx
+
+        # CTA id in cluster
+        cta_id_in_cluster = (
+            Int32(bidx % params.cluster_shape_mn[0]),
+            Int32(bidy % params.cluster_shape_mn[1]),
+            Int32(0),
+        )
+
+        num_tiles_executed = Int32(0)
+
+        return Sm100FmhaClcDynamicTileScheduler(
+            params,
+            cta_id_in_cluster,
+            num_tiles_executed,
+            clc_response_ptr,
+            block_idx,
+            clc,
+        )
+
+    @staticmethod
+    def get_grid_shape(
+        params: Sm100FmhaClcDynamicTileSchedulerParams,
+        *,
+        loc=None,
+        ip=None,
+    ) -> Tuple[int, int, int]:
+        """Get grid shape for kernel launch."""
+        return params.get_grid_shape(loc=loc, ip=ip)
+
+    def work_tile_info_from_clc_response(self, result_addr: cute.Pointer, *, loc=None, ip=None):
+        """Parse CLC response and convert to FMHA tile coordinates."""
+        m_idx, n_idx, l_idx, vld = cute.arch.clc_response(result_addr, loc=loc, ip=ip)
+        cute.arch.fence_proxy("async.shared", space="cta")
+
+        # CLC returns first CTA coordinates: m_idx=x, l_idx=z
+        # l_idx is the L (batch) dimension; decode to (bid, hid)
+        hid = l_idx % self.params.problem_shape_mbh[2]
+        bid = l_idx // self.params.problem_shape_mbh[2]
+
+        cta_idx_in_cluster, cta_idy_in_cluster, _ = self.cta_id_in_cluster
+        cur_tile_coord = (
+            m_idx + cta_idx_in_cluster,  # M dimension
+            0,  # N always 0 for FMHA
+            (bid, hid),  # (B, H) packed
+        )
+
+        return cutlass.utils.WorkTileInfo(cur_tile_coord, vld)
+
+    def get_current_work(self, *, loc=None, ip=None):
+        """Get current work tile from CLC response."""
+        return self.work_tile_info_from_clc_response(self._clc_response_ptr, loc=loc, ip=ip)
+
+    def initial_work_tile_info(self, *, loc=None, ip=None):
+        """Get initial work tile based on block index."""
+        bidx, bidy, bidz = self._block_idx
+        # bidz is the L (batch) dimension; decode to (bid, hid)
+        hid = bidz % self.params.problem_shape_mbh[2]
+        bid = bidz // self.params.problem_shape_mbh[2]
+        return cutlass.utils.WorkTileInfo((bidx, 0, (bid, hid)), True)
+
+    def advance_to_next_work(self, *, loc=None, ip=None):
+        """Consumer-side advance: wait for next tile, read coordinates, release."""
+        self.clc.consumer_wait(loc=loc, ip=ip)
+        work = self.get_current_work(loc=loc, ip=ip)
+        self.clc.consumer_release(loc=loc, ip=ip)
+        self._num_tiles_executed += Int32(1)
+        return work
+
+    def prefetch_next_work(self, *, loc=None, ip=None):
+        """Producer-side: issue CLC query for next tile."""
+        self.clc.prefetch_next_work(loc=loc, ip=ip)
+
+    def producer_tail(self, *, loc=None, ip=None):
+        """Producer-side cleanup after last tile."""
+        self.clc.producer_tail(loc=loc, ip=ip)
+
+    @property
+    def num_tiles_executed(self) -> Int32:
+        return self._num_tiles_executed
+
+
+def compute_sm100_fmha_grid_clc(
+    o_shape: cute.Shape,
+    cta_tiler: Tuple[int, int, int],
+    cluster_shape_mnk: Tuple[int, int, int],
+) -> Tuple[Sm100FmhaClcDynamicTileSchedulerParams, Tuple[int, int, int]]:
+    """Compute grid parameters for FMHA with CLC dynamic scheduling."""
+    problem_shape_mbh = (
+        cute.ceil_div(cute.size(o_shape[0]), cta_tiler[0]),
+        cute.size(o_shape[2][0]),
+        cute.size(o_shape[2][1]),
+    )
+    tile_sched_params = Sm100FmhaClcDynamicTileSchedulerParams(problem_shape_mbh, cluster_shape_mnk)
+    grid = Sm100FmhaClcDynamicTileScheduler.get_grid_shape(tile_sched_params)
+    return tile_sched_params, grid
+
+
+##############################################################################
+# Fused Mask
+##############################################################################
+
+
+def make_sm100_thread_cooperative_group(size: int):
+    return cutlass.pipeline.CooperativeGroup(cutlass.pipeline.Agent.Thread, size)
+
+
+SM100_TMEM_CAPACITY_COLUMNS = 512
diff --git a/build/torch-cuda/topk_gather_kv.py b/build/torch-cuda/topk_gather_kv.py
new file mode 100644
index 0000000000000000000000000000000000000000..b6d1572c389f9b39b6c0d5bffca73e18639c21d3
--- /dev/null
+++ b/build/torch-cuda/topk_gather_kv.py
@@ -0,0 +1,278 @@
+from typing import Type, Optional
+from dataclasses import dataclass
+import operator
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+from cutlass.cute.nvgpu import cpasync
+from cutlass import Int32, Uint32, const_expr, Boolean
+
+from . import utils
+from .utils import warp_reduce
+from .quack.cute_dsl_utils import ParamsBase
+
+import math
+
+
+@dataclass
+class CpasyncGatherKVManager(ParamsBase):
+    mIndexTopk: cute.Tensor
+    sBitmask: Optional[cute.Tensor]
+
+    cta_rank_in_cluster: Int32
+    thread_idx: Int32
+    warp_idx: Int32
+
+    topk_length: Int32
+    seqlen_k_limit: Int32
+    tile_n: Int32
+    num_threads: cutlass.Constexpr[Int32]
+    hdim: cutlass.Constexpr[Int32]
+    hdim_v: cutlass.Constexpr[Int32]
+    num_hdimv_splits: cutlass.Constexpr[Int32]
+    cta_group_size: cutlass.Constexpr[Int32]
+
+    gmem_threads_per_row: cutlass.Constexpr[Int32]
+    topk_indices_per_thread: Int32
+    async_copy_elems: Int32
+
+    gmem_tiled_copy_KV: cute.TiledCopy
+    gmem_thr_copy_KV: cute.TiledCopy
+
+    rTopk: cute.Tensor
+    rTopkHalf: cute.Tensor
+    # for bitmask
+    rTopk_NonInterleaved: cute.Tensor
+
+    pipeline_bitmask: Optional[pipeline.PipelineAsync]
+    cpasync_barrier: Optional[pipeline.NamedBarrier]
+
+    disable_bitmask: cutlass.Constexpr[Boolean]
+
+    @staticmethod
+    def create(
+        mIndexTopk: cute.Tensor,
+        cta_rank_in_cluster: Int32,
+        thread_idx: Int32,
+        warp_idx: Int32,
+        topk_length: Int32,
+        seqlen_k_limit: Int32,
+        tile_n: cutlass.Constexpr[Int32],
+        hdim: cutlass.Constexpr[Int32],
+        hdim_v: cutlass.Constexpr[Int32],
+        num_hdimv_splits: cutlass.Constexpr[Int32],
+        num_threads: cutlass.Constexpr[Int32],
+        dtype: Type[cutlass.Numeric],
+        cta_group_size: cutlass.Constexpr[Int32],
+        cpasync_barrier: Optional[pipeline.NamedBarrier] = None,
+        disable_bitmask: cutlass.Constexpr[Boolean] = False,
+        sBitmask: Optional[cute.Tensor] = None,
+        pipeline_bitmask: Optional[pipeline.PipelineAsync] = None,
+    ):
+        assert tile_n % num_threads == 0
+        assert num_threads == 128
+        assert hdim % 64 == 0
+        assert (hdim_v // num_hdimv_splits // cta_group_size) % 64 == 0
+        assert num_threads % cute.arch.WARP_SIZE == 0
+        universal_copy_bits = 128
+        async_copy_elems = universal_copy_bits // dtype.width
+        dtype_bytes = dtype.width // 8
+        # assumes hdim is never part of transposed operand
+        gmem_k_block_size = math.gcd(
+            hdim,
+            hdim_v // num_hdimv_splits // cta_group_size,
+            128 // dtype_bytes,
+        )
+        assert gmem_k_block_size % async_copy_elems == 0
+        gmem_threads_per_row = gmem_k_block_size // async_copy_elems
+        assert cute.arch.WARP_SIZE % gmem_threads_per_row == 0
+        atom_async_copy = cute.make_copy_atom(
+            cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
+            dtype,
+            num_bits_per_copy=universal_copy_bits,
+        )
+        thr_layout = cute.make_ordered_layout(
+            (num_threads // gmem_threads_per_row, gmem_threads_per_row),
+            order=(1, 0),
+        )
+        val_layout = cute.make_layout((1, 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)
+        topk_indices_per_thread = tile_n // num_threads
+
+        rTopk = cute.make_rmem_tensor((topk_indices_per_thread,), Int32)
+        rTopkHalf = cute.make_rmem_tensor((topk_indices_per_thread,), Int32)
+        rTopk_NonInterleaved = cute.make_rmem_tensor((topk_indices_per_thread,), Int32)
+
+        return CpasyncGatherKVManager(
+            mIndexTopk,
+            sBitmask,
+            cta_rank_in_cluster,
+            thread_idx,
+            warp_idx,
+            topk_length,
+            seqlen_k_limit,
+            tile_n,
+            num_threads,
+            hdim,
+            hdim_v,
+            num_hdimv_splits,
+            cta_group_size,
+            gmem_threads_per_row,
+            topk_indices_per_thread,
+            async_copy_elems,
+            gmem_tiled_copy_KV,
+            gmem_thr_copy_KV,
+            rTopk,
+            rTopkHalf,
+            rTopk_NonInterleaved,
+            pipeline_bitmask,
+            cpasync_barrier,
+            disable_bitmask,
+        )
+
+    @cute.jit
+    def load_index_topk(
+        self,
+        n_block: Int32,
+        transpose: bool,
+    ):
+        entries_per_thread = self.topk_indices_per_thread
+        rTopk = self.rTopk if const_expr(transpose) else self.rTopkHalf
+
+        for i in cutlass.range_constexpr(entries_per_thread):
+            row = (
+                i * self.num_threads
+                + (self.thread_idx % self.gmem_threads_per_row)
+                * (self.num_threads // self.gmem_threads_per_row)
+                + (self.thread_idx // self.gmem_threads_per_row)
+            )
+            # need this if not offset in load_X
+            # if const_expr(not transpose):
+            #     row += self.cta_rank_in_cluster * (self.tile_n//self.cta_group_size)
+            #     row = row % self.tile_n
+            row_idx = n_block * self.tile_n + row
+            rTopk[i] = self.mIndexTopk[row_idx]
+
+            if const_expr(not transpose and not self.disable_bitmask):
+                row_non_interleaved = i * self.num_threads + self.thread_idx
+                row_idx_non_interleaved = n_block * self.tile_n + row_non_interleaved
+                self.rTopk_NonInterleaved[0] = self.mIndexTopk[row_idx_non_interleaved]
+
+    @cute.jit
+    def compute_bitmask(
+        self,
+        producer_state_bitmask,
+    ):
+        assert self.pipeline_bitmask is not None, "pipeline_bitmask not provided"
+        assert self.cpasync_barrier is not None, "cpasync barrier not provided"
+
+        lane_idx = cute.arch.lane_idx()
+        assert cute.size(self.rTopk_NonInterleaved) == 1
+        bitmask = Uint32(0)
+
+        # Step 1. Construct per-thread bitmask
+        topk_idx = self.rTopk_NonInterleaved[0]
+        is_valid = topk_idx >= 0 and topk_idx < self.seqlen_k_limit
+        if is_valid:
+            bitmask = Uint32(1 << lane_idx)
+
+        # Step 2. Warp shuffle bitwise OR = add since indices are exclusive.
+        bitmask = warp_reduce(bitmask, operator.add)
+
+        self.pipeline_bitmask.producer_acquire(producer_state_bitmask)
+        # store to smem and sync threads
+        if lane_idx == 0:
+            self.sBitmask[self.warp_idx, producer_state_bitmask.index] = bitmask
+        self.cpasync_barrier.arrive_and_wait()
+
+        self.pipeline_bitmask.producer_commit(producer_state_bitmask)
+        producer_state_bitmask.advance()
+        return producer_state_bitmask
+
+    @cute.jit
+    def compute_X_ptr(
+        self,
+        mX: cute.Tensor,
+        transpose: bool,
+        d_offset: int = 0,
+    ):
+        entries_per_thread = self.topk_indices_per_thread
+        tPrXPtr = cute.make_rmem_tensor((entries_per_thread,), cutlass.Int64)
+        tPrRowValid = cute.make_rmem_tensor((entries_per_thread,), cutlass.Int32)
+        rTopk = self.rTopk if const_expr(transpose) else self.rTopkHalf
+
+        for i in cutlass.range_constexpr(entries_per_thread):
+            topk_idx = rTopk[i]
+            if const_expr(not self.disable_bitmask):
+                row_valid = topk_idx >= 0 and topk_idx < self.seqlen_k_limit
+                tPrRowValid[i] = row_valid
+            if const_expr(not transpose):
+                tPrXPtr[i] = utils.elem_pointer(mX, (topk_idx, d_offset)).toint()
+            else:
+                tPrXPtr[i] = utils.elem_pointer(mX, (d_offset, topk_idx)).toint()
+
+        return tPrXPtr, tPrRowValid
+
+    @cute.jit
+    def load_X(
+        self,
+        mX: cute.Tensor,
+        sX: cute.Tensor,
+        transpose: bool,
+        K_or_V: str,
+        d_offset: int = 0,
+    ):
+        assert K_or_V in ("K", "V")
+        cta_tile_n = self.tile_n if const_expr(transpose) else self.tile_n // self.cta_group_size
+        head_dim = self.hdim if const_expr(K_or_V == "K") else self.hdim_v // self.num_hdimv_splits
+        if const_expr(transpose):
+            head_dim = head_dim // self.cta_group_size
+        order = (1, 0) if const_expr(transpose) else (0, 1)
+
+        sX_nd_layout = cute.make_ordered_layout((cta_tile_n, head_dim), order=order)
+        sX_nd = cute.composition(sX, sX_nd_layout)
+
+        cX = cute.make_identity_tensor((cta_tile_n, head_dim))
+        tXsX = self.gmem_thr_copy_KV.partition_D(sX_nd)
+        tXcX = self.gmem_thr_copy_KV.partition_S(cX)
+
+        tPrXPtr, tPrRowValid = self.compute_X_ptr(mX, transpose, d_offset)
+
+        if const_expr(not transpose):
+            offset = self.cta_rank_in_cluster * (self.gmem_threads_per_row // self.cta_group_size)
+        else:
+            offset = 0
+
+        for m in cutlass.range_constexpr(cute.size(tXsX, mode=[1])):
+            if const_expr(not self.disable_bitmask):
+                row_valid = utils.shuffle_sync(
+                    tPrRowValid[m // self.gmem_threads_per_row],
+                    (m + offset) % self.gmem_threads_per_row,
+                    width=self.gmem_threads_per_row,
+                )
+                should_load = cute.make_fragment_like(tXsX[(0, None), m, 0], Boolean)
+                should_load.fill(Boolean(row_valid))
+            x_ptr_i64 = utils.shuffle_sync(
+                tPrXPtr[m // self.gmem_threads_per_row],
+                (m + offset) % self.gmem_threads_per_row,
+                width=self.gmem_threads_per_row,
+            )
+            x_gmem_ptr = cute.make_ptr(
+                mX.element_type, x_ptr_i64, cute.AddressSpace.gmem, assumed_align=16
+            )
+            mX_cur = cute.make_tensor(x_gmem_ptr, cute.make_layout((head_dim,)))
+            mX_cur_copy = cute.tiled_divide(mX_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_cur_copy_ki = mX_cur_copy[None, ki]
+                tXsX_k = tXsX[None, m, k]
+                mX_cur_copy_ki = cute.make_tensor(mX_cur_copy_ki.iterator, tXsX_k.layout)
+                cute.copy(
+                    self.gmem_tiled_copy_KV,
+                    mX_cur_copy_ki,
+                    tXsX_k,
+                    pred=should_load if const_expr(not self.disable_bitmask) else None,
+                )
diff --git a/build/torch-cuda/utils.py b/build/torch-cuda/utils.py
index a05305b287b48aded27d4cd11562d1951cfed8b6..be01b9447ed5356351283a45829eb55cb5e0196a 100644
--- a/build/torch-cuda/utils.py
+++ b/build/torch-cuda/utils.py
@@ -3,21 +3,28 @@
 import math
 import hashlib
 import inspect
-from typing import Type, Callable, Optional, Tuple, overload
+import os
+from functools import partial
+from typing import Type, Callable, Optional, Tuple, overload, NamedTuple
 
 import cutlass
 import cutlass.cute as cute
 
-from cutlass import Float32, const_expr
+from cutlass import Float32, Int32, const_expr
+from cutlass.cute import FastDivmodDivisor
 from cutlass.cutlass_dsl import T, dsl_user_op
 from cutlass._mlir.dialects import nvvm, llvm
 from cutlass.cute.runtime import from_dlpack
 
 
-from .quack import activation
-
 _MIXER_ATTRS = ("__vec_size__",)
 
+
+class AuxData(NamedTuple):
+    tensors: tuple | list | None = None
+    scalars: tuple | None = None
+
+
 # Obtained from sollya:
 # fpminimax(exp(x * log(2.0)), 1, [|1,24...|],[0;1],relative);
 POLY_EX2 = {
@@ -54,9 +61,44 @@ POLY_EX2 = {
     ),
 }
 
+_fa_clc_enabled: bool = os.environ.get("FA_CLC", "0") == "1"
+_fa_disable_2cta_enabled: bool = os.environ.get("FA_DISABLE_2CTA", "0") == "1"
+
+
+def _is_cuda_12() -> bool:
+    """Check if the CUDA toolkit version is 12.x.
+
+    2CTA forward non-causal has a codegen regression on CUDA 12 that causes
+    ~18% slowdown compared to 1CTA. This is fixed in CUDA 13.x.
+    """
+    try:
+        import torch
+
+        cuda_version = torch.version.cuda
+        if cuda_version is not None:
+            major = cuda_version.split(".")[0]
+            return int(major) == 12
+    except Exception:
+        pass
+    return False
+
+
+_fa_disable_2cta_cuda12: bool = _is_cuda_12()
+
+
+def _get_use_clc_scheduler_default() -> bool:
+    return _fa_clc_enabled
+
+
+def _get_disable_2cta_default(is_fwd: bool = False) -> bool:
+    if is_fwd:
+        return _fa_disable_2cta_enabled or _fa_disable_2cta_cuda12
+    else:
+        return _fa_disable_2cta_enabled
+
 
 def _compute_base_hash(func: Callable) -> str:
-    """Compute hash from source code or bytecode and closure values."""
+    """Compute a hash from callable code and captured compile-time values."""
     try:
         data = inspect.getsource(func).encode()
     except (OSError, TypeError):
@@ -71,13 +113,22 @@ def _compute_base_hash(func: Callable) -> str:
         for cell in func.__closure__:
             hasher.update(repr(cell.cell_contents).encode())
 
+    # Factory-generated score/mask modifiers often capture specialization values
+    # as defaults instead of closure cells. Those values affect generated code and
+    # therefore must distinguish entries in the compile cache as well.
+    for attr in ("__defaults__", "__kwdefaults__"):
+        values = getattr(func, attr, None)
+        if values is not None:
+            hasher.update(attr.encode())
+            hasher.update(repr(values).encode())
+
     return hasher.hexdigest()
 
 
 def hash_callable(
     func: Callable, mixer_attrs: Tuple[str] = _MIXER_ATTRS, set_cute_hash: bool = True
 ) -> str:
-    """Hash a callable based on the source code or bytecode and closure values.
+    """Hash a callable based on its code and captured compile-time values.
     Fast-path: if the callable (or its __wrapped__ base) has a ``__cute_hash__``
     attribute, that value is returned immediately as the base hash, then
     metadata dunders are mixed in to produce the final dict-key hash.
@@ -106,23 +157,69 @@ def hash_callable(
 
     hasher = hashlib.sha256(base_hash.encode())
 
-    for attr, val in zip(_MIXER_ATTRS, mixer_values):
+    for attr, val in zip(mixer_attrs, mixer_values):
         hasher.update(f"{attr}={val!r}".encode())
 
     return hasher.hexdigest()
 
 
 def create_softcap_scoremod(softcap_val):
-    inv_softcap = 1.0 / softcap_val
-
     @cute.jit
-    def scoremod_premask_fn(acc_S_SSA, batch_idx, head_idx, q_idx, kv_idx, aux_tensors):
-        scores = acc_S_SSA * inv_softcap
-        return scores * cute.math.tanh(scores, fastmath=True)
+    def scoremod_premask_fn(
+        acc_S_SSA, batch_idx, head_idx, q_idx, kv_idx, seqlen_info, aux_tensors
+    ):
+        scores = acc_S_SSA / softcap_val
+        return softcap_val * cute.math.tanh(scores, fastmath=True)
 
     return scoremod_premask_fn
 
 
+def create_softcap_scoremod_bwd(softcap_val):
+    @cute.jit
+    def scoremod_bwd_fn(
+        grad_out_SSA, score_SSA, batch_idx, head_idx, q_idx, kv_idx, seqlen_info, aux_tensors
+    ):
+        scores = score_SSA / softcap_val
+        tanh_scores = cute.math.tanh(scores, fastmath=True)
+        return grad_out_SSA * (1.0 - tanh_scores * tanh_scores)
+
+    return scoremod_bwd_fn
+
+
+LOG2_E = math.log2(math.e)
+
+
+def compute_softmax_scale_log2(softmax_scale, score_mod):
+    """Compute softmax_scale_log2 and adjusted softmax_scale based on whether score_mod is used.
+
+    When score_mod is None, fold the log2(e) factor into softmax_scale_log2 and set softmax_scale
+    to None. When score_mod is present, keep softmax_scale separate so it can be applied before
+    the score_mod, and set softmax_scale_log2 to just the change-of-base constant.
+
+    Returns (softmax_scale_log2, softmax_scale).
+    """
+    if const_expr(score_mod is None):
+        return softmax_scale * LOG2_E, None
+    else:
+        return LOG2_E, softmax_scale
+
+
+def compute_fastdiv_mods(mQ, mK, qhead_per_kvhead, pack_gqa, aux_tensors, mPageTable=None):
+    """Compute FastDivmodDivisor pairs for aux_tensors index computation.
+
+    Returns a (seqlen_q_divmod, seqlen_k_divmod) tuple, or None if aux_tensors is None.
+    """
+    if const_expr(aux_tensors is None):
+        return None
+    seqlen_q = cute.size(mQ.shape[0]) // (qhead_per_kvhead if const_expr(pack_gqa) else 1)
+    seqlen_k = (
+        cute.size(mK.shape[0])
+        if const_expr(mPageTable is None)
+        else mK.shape[0] * mPageTable.shape[1]
+    )
+    return (FastDivmodDivisor(seqlen_q), FastDivmodDivisor(seqlen_k))
+
+
 def convert_from_dlpack(x, leading_dim, alignment=16, divisibility=1) -> cute.Tensor:
     return (
         from_dlpack(x, assumed_align=alignment)
@@ -133,6 +230,34 @@ def convert_from_dlpack(x, leading_dim, alignment=16, divisibility=1) -> cute.Te
     )
 
 
+def convert_from_dlpack_compact_dynamic(
+    x,
+    *,
+    dynamic_modes: tuple[int, ...],
+    alignment: int = 16,
+    stride_order=None,
+    divisibility: int = 1,
+    enable_tvm_ffi: bool = False,
+) -> cute.Tensor:
+    """Convert via DLPack and mark selected compact dimensions as dynamic."""
+    if isinstance(dynamic_modes, int):
+        dynamic_modes = (dynamic_modes,)
+    if stride_order is None:
+        stride_order = x.dim_order()
+    t = (
+        from_dlpack(x, assumed_align=alignment, enable_tvm_ffi=True)
+        if enable_tvm_ffi
+        else from_dlpack(x, assumed_align=alignment)
+    )
+    for mode in dynamic_modes:
+        t = t.mark_compact_shape_dynamic(
+            mode=mode,
+            stride_order=stride_order,
+            divisibility=divisibility,
+        )
+    return t
+
+
 def convert_from_dlpack_leading_static(
     x, leading_dim, alignment=16, static_modes=None, stride_order=None
 ) -> cute.Tensor:
@@ -164,7 +289,7 @@ def make_tiled_copy_B(
 
 
 def mma_make_fragment_A(
-    smem: cute.Tensor, thr_mma: cute.core.ThrMma, swapAB: cutlass.Constexpr[bool] = False
+    smem: cute.Tensor, thr_mma: cute.ThrMma, swapAB: cutlass.Constexpr[bool] = False
 ) -> cute.Tensor:
     if const_expr(swapAB):
         return mma_make_fragment_B(smem, thr_mma)
@@ -173,7 +298,7 @@ def mma_make_fragment_A(
 
 
 def mma_make_fragment_B(
-    smem: cute.Tensor, thr_mma: cute.core.ThrMma, swapAB: cutlass.Constexpr[bool] = False
+    smem: cute.Tensor, thr_mma: cute.ThrMma, swapAB: cutlass.Constexpr[bool] = False
 ) -> cute.Tensor:
     if const_expr(swapAB):
         return mma_make_fragment_A(smem, thr_mma)
@@ -204,7 +329,7 @@ def warp_reduce(
     width: cutlass.Constexpr[int] = cute.arch.WARP_SIZE,
 ) -> cute.TensorSSA | cute.Numeric:
     if const_expr(isinstance(val, cute.TensorSSA)):
-        res = cute.make_fragment(val.shape, val.dtype)
+        res = cute.make_rmem_tensor(val.shape, val.dtype)
         res.store(val)
         for i in cutlass.range_constexpr(cute.size(val.shape)):
             res[i] = warp_reduce(res[i], op, width)
@@ -215,36 +340,34 @@ def warp_reduce(
     return val
 
 
+@dsl_user_op
+def smid(*, loc=None, ip=None) -> Int32:
+    return Int32(
+        llvm.inline_asm(
+            T.i32(),
+            [],
+            "mov.u32 $0, %smid;",
+            "=r",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
 @dsl_user_op
 def fmax(
     a: float | Float32, b: float | Float32, c: float | Float32 | None = None, *, loc=None, ip=None
 ) -> Float32:
-    from cutlass import CUDA_VERSION
-
-    # * NVVM call based on nvvm version
-    if CUDA_VERSION.major == 12 and CUDA_VERSION.minor == 9:
-        # Old API: requires explicit result type as first positional argument
-        return Float32(
-            nvvm.fmax(
-                T.f32(),
-                Float32(a).ir_value(loc=loc, ip=ip),
-                Float32(b).ir_value(loc=loc, ip=ip),
-                c=Float32(c).ir_value(loc=loc, ip=ip) if c is not None else None,
-                loc=loc,
-                ip=ip,
-            )
-        )
-    else:
-        # New API: infers result type automatically
-        return Float32(
-            nvvm.fmax(
-                Float32(a).ir_value(loc=loc, ip=ip),
-                Float32(b).ir_value(loc=loc, ip=ip),
-                c=Float32(c).ir_value(loc=loc, ip=ip) if c is not None else None,
-                loc=loc,
-                ip=ip,
-            )
+    return Float32(
+        nvvm.fmax(
+            Float32(a).ir_value(loc=loc, ip=ip),
+            Float32(b).ir_value(loc=loc, ip=ip),
+            c=Float32(c).ir_value(loc=loc, ip=ip) if c is not None else None,
+            loc=loc,
+            ip=ip,
         )
+    )
 
 
 @cute.jit
@@ -255,7 +378,7 @@ def fmax_reduce(
         # if const_expr(init_val is None):
         #     init_val = -cutlass.Float32.if
         # return x.reduce(cute.ReductionOp.MAX, init_val, 0)
-        res = cute.make_fragment(x.shape, Float32)
+        res = cute.make_rmem_tensor(x.shape, Float32)
         res.store(x)
         # local_max = [res[0], res[1]]
         # for i in cutlass.range_constexpr(2, cute.size(x.shape), 2):
@@ -276,7 +399,7 @@ def fmax_reduce(
     else:
         # [2025-06-15] x.reduce only seems to use 50% 3-input max and 50% 2-input max
         # We instead force the 3-input max.
-        res = cute.make_fragment(x.shape, Float32)
+        res = cute.make_rmem_tensor(x.shape, Float32)
         res.store(x)
         local_max_0 = (
             fmax(init_val, res[0], res[1])
@@ -306,7 +429,7 @@ def fadd_reduce(
         if const_expr(init_val is None):
             init_val = Float32.zero
         return x.reduce(cute.ReductionOp.ADD, init_val, 0)
-        # res = cute.make_fragment(x.shape, Float32)
+        # res = cute.make_rmem_tensor(x.shape, Float32)
         # res.store(x)
         # local_sum = [res[0], res[1], res[2], res[3]]
         # for i in cutlass.range_constexpr(4, cute.size(x.shape), 4):
@@ -319,7 +442,7 @@ def fadd_reduce(
         # local_sum[0] += local_sum[2]
         # return local_sum[0] if const_expr(init_val is None) else local_sum[0] + init_val
     else:
-        res = cute.make_fragment(x.shape, Float32)
+        res = cute.make_rmem_tensor(x.shape, Float32)
         res.store(x)
         local_sum_0 = (
             cute.arch.add_packed_f32x2((init_val, 0.0), (res[0], res[1]))
@@ -339,6 +462,19 @@ def fadd_reduce(
         return local_sum[0][0] + local_sum[0][1]
 
 
+@dsl_user_op
+def atomic_add_i32(a: int | Int32, ptr: cute.Pointer, *, loc=None, ip=None) -> Int32:
+    return Int32(
+        nvvm.atomicrmw(
+            op=nvvm.AtomicOpKind.ADD,
+            ptr=ptr.llvm_ptr,
+            a=Int32(a).ir_value(loc=loc, ip=ip),
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
 @dsl_user_op
 def atomic_add_fp32(a: float | Float32, gmem_ptr: cute.Pointer, *, loc=None, ip=None) -> None:
     # gmem_ptr_i64 = gmem_ptr.toint(loc=loc, ip=ip).ir_value()
@@ -356,9 +492,7 @@ def atomic_add_fp32(a: float | Float32, gmem_ptr: cute.Pointer, *, loc=None, ip=
     #     is_align_stack=False,
     #     asm_dialect=llvm.AsmDialect.AD_ATT,
     # )
-    nvvm.atomicrmw(
-        res=T.f32(), op=nvvm.AtomicOpKind.FADD, ptr=gmem_ptr.llvm_ptr, a=Float32(a).ir_value()
-    )
+    nvvm.atomicrmw(op=nvvm.AtomicOpKind.FADD, ptr=gmem_ptr.llvm_ptr, a=Float32(a).ir_value())
 
 
 @dsl_user_op
@@ -369,7 +503,7 @@ def elem_pointer(x: cute.Tensor, coord: cute.Coord, *, loc=None, ip=None) -> cut
 @cute.jit
 def predicate_k(tAcA: cute.Tensor, limit: cutlass.Int32) -> cute.Tensor:
     # Only compute predicates for the "k" dimension. For the mn dimension, we will use "if"
-    tApA = cute.make_fragment(
+    tApA = cute.make_rmem_tensor(
         cute.make_layout(
             (cute.size(tAcA, mode=[0, 1]), cute.size(tAcA, mode=[1]), cute.size(tAcA, mode=[2])),
             stride=(cute.size(tAcA, mode=[2]), 0, 1),
@@ -429,8 +563,48 @@ def shuffle_sync(
     return val[0]
 
 
+@dsl_user_op
+def shl_u32(val: cutlass.Uint32, shift: cutlass.Uint32, *, loc=None, ip=None) -> cutlass.Uint32:
+    """
+    Left-shift val by shift bits using PTX shl.b32 (sign-agnostic).
+
+    Named ``shl_u32`` (not ``shl_b32``) because python type annotations
+    distinguish signed/unsigned.
+
+    PTX semantics (§9.7.8.8): "Shift amounts greater than the register width N
+    are clamped to N."  So ``shl.b32 d, a, 32`` is well-defined and yields 0.
+
+    This differs from C/C++ and LLVM IR, where shifting by >= the type width is
+    undefined behavior.  CuTeDSL compiles through MLIR -> LLVM IR, so a plain
+    Python-level ``Uint32(x) << Uint32(n)`` inherits LLVM's UB: the optimizer
+    may treat the result as poison and eliminate dependent code.  Inline PTX
+    bypasses the LLVM IR shift entirely — the instruction is emitted verbatim
+    into PTX where clamping makes it safe for all shift amounts.
+    """
+    return cutlass.Uint32(
+        llvm.inline_asm(
+            T.i32(),
+            [
+                cutlass.Uint32(val).ir_value(loc=loc, ip=ip),
+                cutlass.Uint32(shift).ir_value(loc=loc, ip=ip),
+            ],
+            "shl.b32 $0, $1, $2;",
+            "=r,r,r",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
 @dsl_user_op
 def shr_u32(val: cutlass.Uint32, shift: cutlass.Uint32, *, loc=None, ip=None) -> cutlass.Uint32:
+    """
+    Unsigned right-shift val by shift bits using PTX shr.u32 (zero-fills).
+
+    See ``shl_u32`` docstring for why inline PTX is used instead of plain
+    CuTeDSL shift operators (LLVM shift-by-type-width UB).
+    """
     return cutlass.Uint32(
         llvm.inline_asm(
             T.i32(),
@@ -438,7 +612,7 @@ def shr_u32(val: cutlass.Uint32, shift: cutlass.Uint32, *, loc=None, ip=None) ->
                 cutlass.Uint32(val).ir_value(loc=loc, ip=ip),
                 cutlass.Uint32(shift).ir_value(loc=loc, ip=ip),
             ],
-            "shr.s32 $0, $1, $2;",
+            "shr.u32 $0, $1, $2;",
             "=r,r,r",
             has_side_effects=False,
             is_align_stack=False,
@@ -502,7 +676,7 @@ def cvt_f16(src: cute.Tensor, dst_or_dtype):
     if const_expr(isinstance(dst_or_dtype, type)):
         # dtype variant: create new tensor and call the tensor variant
         dtype = dst_or_dtype
-        dst = cute.make_fragment(src.shape, dtype)
+        dst = cute.make_rmem_tensor(src.shape, dtype)
         cvt_f16(src, dst)
         return dst
     else:
@@ -613,10 +787,8 @@ def ex2_emulation_2(
     xy_rounded = cute.arch.add_packed_f32x2(xy_clamped, (fp32_round_int, fp32_round_int), rnd="rm")
     # The integer floor of x & y are now in the last 8 bits of xy_rounded
     # We want the next 2 ops to round to nearest even. The rounding mode is important.
-    xy_rounded_back = activation.sub_packed_f32x2(
-        xy_rounded, (fp32_round_int, fp32_round_int)
-    )
-    xy_frac = activation.sub_packed_f32x2(xy_clamped, xy_rounded_back)
+    xy_rounded_back = cute.arch.sub_packed_f32x2(xy_rounded, (fp32_round_int, fp32_round_int))
+    xy_frac = cute.arch.sub_packed_f32x2(xy_clamped, xy_rounded_back)
     xy_frac_ex2 = evaluate_polynomial_2(*xy_frac, POLY_EX2[poly_degree], loc=loc, ip=ip)
     x_out = combine_int_frac_ex2(xy_rounded[0], xy_frac_ex2[0], loc=loc, ip=ip)
     y_out = combine_int_frac_ex2(xy_rounded[1], xy_frac_ex2[1], loc=loc, ip=ip)
@@ -685,10 +857,96 @@ def domain_offset_aligned(
     return cute.make_tensor(new_ptr, tensor.layout)
 
 
+@dsl_user_op
+def warp_reduction(
+    val: cute.Numeric, op: Callable, *, threads_in_group: int = 32, loc=None, ip=None
+) -> cute.Numeric:
+    """Warp-wide reduction helper for a custom binary op."""
+    offset = threads_in_group // 2
+    while offset > 0:
+        val = op(
+            val,
+            cute.arch.shuffle_sync_bfly(
+                val, offset=offset, mask=-1, mask_and_clamp=31, loc=loc, ip=ip
+            ),
+        )
+        offset //= 2
+    return val
+
+
+warp_reduction_max = partial(
+    warp_reduction, op=lambda x, y: fmax(x, y) if isinstance(x, Float32) else max(x, y)
+)
+warp_reduction_sum = partial(warp_reduction, op=lambda x, y: x + y)  # noqa: FURB118
+
+
+@dsl_user_op
+def make_cotiled_copy(
+    atom: cute.CopyAtom, atom_layout_tv: cute.Layout, data_layout: cute.Layout, *, loc=None, ip=None
+) -> cute.TiledCopy:
+    """Compatibility wrapper for deprecated CuTeDSL `make_cotiled_copy`."""
+    assert cute.is_static(atom_layout_tv.type), "atom_layout_tv must be static"
+    assert cute.is_static(data_layout.type), "data_layout must be static"
+
+    inv_layout_ = cute.left_inverse(data_layout, loc=loc, ip=ip)
+    inv_data_layout = cute.make_layout(
+        (inv_layout_.shape, (1)), stride=(inv_layout_.stride, (0)), loc=loc, ip=ip
+    )
+    layout_tv_data = cute.composition(inv_data_layout, atom_layout_tv, loc=loc, ip=ip)
+
+    atom_layout_v_to_check = cute.coalesce(
+        cute.make_layout(atom_layout_tv.shape[1], stride=atom_layout_tv.stride[1], loc=loc, ip=ip),
+        loc=loc,
+        ip=ip,
+    )
+    data_layout_v_to_check = cute.coalesce(
+        cute.composition(
+            data_layout,
+            cute.make_layout(
+                layout_tv_data.shape[1], stride=layout_tv_data.stride[1], loc=loc, ip=ip
+            ),
+            loc=loc,
+            ip=ip,
+        ),
+        loc=loc,
+        ip=ip,
+    )
+    assert data_layout_v_to_check == atom_layout_v_to_check, (
+        "the memory pointed to by atom_layout_tv does not exist in the data_layout."
+    )
+
+    flat_data_shape = cute.product_each(data_layout.shape, loc=loc, ip=ip)
+    tiler = tuple(
+        cute.filter(
+            cute.composition(
+                cute.make_layout(
+                    flat_data_shape,
+                    stride=tuple(0 if j != i else 1 for j in range(cute.rank(flat_data_shape))),
+                    loc=loc,
+                    ip=ip,
+                ),
+                layout_tv_data,
+                loc=loc,
+                ip=ip,
+            ),
+            loc=loc,
+            ip=ip,
+        )
+        for i in range(cute.rank(flat_data_shape))
+    )
+    tile2data = cute.composition(
+        cute.make_layout(flat_data_shape, loc=loc, ip=ip), tiler, loc=loc, ip=ip
+    )
+    layout_tv = cute.composition(
+        cute.left_inverse(tile2data, loc=loc, ip=ip), layout_tv_data, loc=loc, ip=ip
+    )
+    return cute.make_tiled_copy(atom, layout_tv, tiler, loc=loc, ip=ip)
+
+
 @cute.jit
 def scalar_to_ssa(a: cute.Numeric, dtype) -> cute.TensorSSA:
     """Convert a scalar to a cute TensorSSA of shape (1,) and given dtype"""
-    vec = cute.make_fragment(1, dtype)
+    vec = cute.make_rmem_tensor(1, dtype)
     vec[0] = a
     return vec.load()
 
@@ -696,3 +954,93 @@ def scalar_to_ssa(a: cute.Numeric, dtype) -> cute.TensorSSA:
 def ssa_to_scalar(val):
     """Could inline but nice for reflecting the above api"""
     return val[0]
+
+
+@cute.jit
+def get_batch_from_cu_tensor(idx: Int32, cu_tensor: cute.Tensor) -> Int32:
+    """Binary search to determine batch from packed index in a cumulative tensor"""
+    batch_size = cute.size(cu_tensor) - 1
+    lo = Int32(0)
+    hi = batch_size
+
+    while lo < hi:
+        mid = (lo + hi) // 2
+        if cu_tensor[mid + 1] <= idx:
+            lo = mid + 1
+        else:
+            hi = mid
+
+    return lo
+
+
+def as_bshkrd_tensor(
+    tensor: cute.Tensor,
+    h_k: Int32,
+    h_r: Int32,
+    varlen: bool,
+) -> cute.Tensor:
+    """Normalize (B,S,H,D)/(S,H,D) tensors to (B,S,H_k,H_r,D) view."""
+    if cutlass.const_expr(cute.rank(tensor.layout) == 5):
+        return tensor
+    if cutlass.const_expr(cute.rank(tensor.layout) == 4):
+        return cute.make_tensor(
+            tensor.iterator,
+            cute.make_layout(
+                (tensor.shape[0], tensor.shape[1], h_k, h_r, tensor.shape[3]),
+                stride=(
+                    tensor.stride[0],
+                    tensor.stride[1],
+                    tensor.stride[2] * h_r,
+                    tensor.stride[2],
+                    tensor.stride[3],
+                ),
+            ),
+        )
+    assert cutlass.const_expr(cute.rank(tensor.layout) == 3), "Expected rank-3 varlen tensor"
+    assert cutlass.const_expr(varlen), "Rank-3 input is only valid for varlen"
+    return cute.make_tensor(
+        tensor.iterator,
+        cute.make_layout(
+            (1, tensor.shape[0], h_k, h_r, tensor.shape[2]),
+            stride=(
+                0,
+                tensor.stride[0],
+                tensor.stride[1] * h_r,
+                tensor.stride[1],
+                tensor.stride[2],
+            ),
+        ),
+    )
+
+
+def as_shhb_tensor(
+    tensor: cute.Tensor,
+    h_k: Int32,
+    h_r: Int32,
+    b: Int32,
+    varlen: bool,
+) -> cute.Tensor:
+    """Normalize (B,H,S)/(H,S) tensors to (S, ((H_r, H_k), B)) view."""
+    if cutlass.const_expr(cute.rank(tensor.layout) == 3):
+        return cute.make_tensor(
+            tensor.iterator,
+            cute.make_layout(
+                (tensor.shape[2], ((h_r, h_k), tensor.shape[0])),
+                stride=(
+                    tensor.stride[2],
+                    ((tensor.stride[1], tensor.stride[1] * h_r), tensor.stride[0]),
+                ),
+            ),
+        )
+    assert cutlass.const_expr(cute.rank(tensor.layout) == 2), "Expected rank-2 varlen tensor"
+    assert cutlass.const_expr(varlen), "Rank-2 input is only valid for varlen"
+    return cute.make_tensor(
+        tensor.iterator,
+        cute.make_layout(
+            (tensor.shape[1], ((h_r, h_k), b)),
+            stride=(
+                tensor.stride[1],
+                ((tensor.stride[0], tensor.stride[0] * h_r), 0),
+            ),
+        ),
+    )