# Copyright (c) 2026, Colfax International.

import math
import time
from functools import partial
from typing import Callable, Optional

import cuda.bindings.driver as cuda
import cutlass
import cutlass.cute as cute
import cutlass.pipeline as pipeline
import cutlass.utils.blackwell_helpers as sm100_utils
import torch
import torch.utils.benchmark as benchmark
from cutlass import Boolean, Float32, Int32, Int64, Uint32, const_expr
from cutlass.cute import FastDivmodDivisor
from cutlass.cute.nvgpu import cpasync, tcgen05
from cutlass.cute.runtime import from_dlpack
from cutlass.utils import ClcDynamicPersistentTileScheduler
from quack import copy_utils

import sglang.kernels.ops.attention.flash_attn.cute.blackwell_helpers as fa_sm100_utils
from sglang.kernels.ops.attention.flash_attn.cute import utils as fa_utils
from sglang.kernels.ops.attention.flash_attn.cute.block_info import BlockInfo
from sglang.kernels.ops.attention.flash_attn.cute.cute_dsl_utils import (
    dump_kernel_attributes,
)
from sglang.kernels.ops.attention.flash_attn.cute.fa_logging import fa_log, fa_printf
from sglang.kernels.ops.attention.flash_attn.cute.mask import AttentionMask
from sglang.kernels.ops.attention.flash_attn.cute.named_barrier import (
    NamedBarrierFwdSm100_MLA2CTA,
)
from sglang.kernels.ops.attention.flash_attn.cute.pack_gqa import (
    make_packgqa_tiled_tma_atom,
    pack_gqa_layout,
)
from sglang.kernels.ops.attention.flash_attn.cute.paged_kv import PagedKVManager
from sglang.kernels.ops.attention.flash_attn.cute.seqlen_info import SeqlenInfoQK
from sglang.kernels.ops.attention.flash_attn.cute.softmax import SoftmaxSm100
from sglang.kernels.ops.attention.flash_attn.cute.testing import attention_ref
from sglang.kernels.ops.attention.flash_attn.cute.tile_scheduler import (
    ClcState,
    ParamsBase,
    SchedulingMode,
    SingleTileLPTScheduler,
    SingleTileScheduler,
    SingleTileVarlenScheduler,
    TileSchedulerArguments,
    TileSchedulerProtocol,
)
from sglang.kernels.ops.attention.flash_attn.cute.topk_gather_kv import (
    CpasyncGatherKVManager,
)
from sglang.kernels.ops.attention.flash_attn.cute.utils import smid


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,
        is_varlen_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
        self.nheads_kv = nheads_kv
        self.is_varlen_q = is_varlen_q
        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 and not is_varlen_q
        self.sched_stages = 1
        self.scheduling_mode = (
            SchedulingMode.CLC if self.use_clc_scheduler else SchedulingMode.STATIC
        )

        if const_expr(is_varlen_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 = 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 = 80
            self.num_regs_mma = 80
            self.num_regs_softmax = 208
            self.num_regs_epilogue = 128
            self.num_regs_cpasync = 96 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 = 4
        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,
            ]
        )
        mbar_ptr_tmem_dealloc_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
            mbar_ptr_tmem_dealloc: mbar_ptr_tmem_dealloc_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

        tile_sched_args = TileSchedulerArguments(
            num_block=cute.ceil_div(cute.size(mQv.shape[0]), self.cta_tile_m),
            num_head=cute.size(mQv.shape[2]),
            num_batch=cute.size(mQv.shape[3])
            if const_expr(mCuSeqlensQ is None)
            else cute.size(mCuSeqlensQ.shape[0] - 1),
            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=False,
        )
        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_m_block, head_idx, batch_idx = cute.arch.block_idx()
        cluster_m_block = cta_m_block // self.cta_group_size
        mma_tile_coord_v = cta_m_block % 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,
            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.mbar_ptr_tmem_dealloc,
        )

        # ==== 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
            )
            clc = ClcState.create(
                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, clc=clc)
        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_K,
                    pipeline_V,
                    pipeline_K_cpasync,
                    pipeline_V_cpasync,
                    sO_empty_mbar_ptr,
                    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,
                    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,
                )

        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,
            )

        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,
            )
            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)
            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,
            )
            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,
            )
            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_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,
    ):
        # ==== 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:
            cta_m_block, head_idx, batch_idx, _ = work_tile.tile_idx
            cluster_m_block = cta_m_block // self.cta_group_size

            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,
    ):
        # ==== 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:
            cta_m_block, head_idx, batch_idx, _ = work_tile.tile_idx
            cluster_m_block = cta_m_block // self.cta_group_size
            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
                    mIndexTopk_cur = mIndexTopk[None, m_idx + offset_q]

                if const_expr(self.is_causal):
                    seqlen_k_limit = m_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,
    ):
        # ==== 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:
            cta_m_block, head_idx, batch_idx, _ = work_tile.tile_idx
            cluster_m_block = cta_m_block // self.cta_group_size
            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,
    ):
        # ==== 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 = not self.is_topk_gather
        use_ptx_gemm_QvV = not self.is_topk_gather
        use_ptx_gemm_PVt = not self.is_topk_gather

        # 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:
            cta_m_block, head_idx, batch_idx, _ = work_tile.tile_idx
            cluster_m_block = cta_m_block // self.cta_group_size

            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,
    ):
        # ==== 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
        )

        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:
            cta_m_block, head_idx, batch_idx, _ = work_tile.tile_idx
            cluster_m_block = cta_m_block // self.cta_group_size
            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,
    ):
        ### ==== 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:
            cta_m_block, head_idx, batch_idx, _ = work_tile.tile_idx
            cluster_m_block = cta_m_block // self.cta_group_size

            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)

            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:
                    row_max = sRowMax[tidx, 0]
                    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()


def test_mla_kernel(
    seqlen_q=2048,
    seqlen_k=2048,
    topk_length=2048,
    nheads=1,
    batch=1,
    iter=0,
    compile_cache=dict(),
    validate=True,
    seed=0,
    gather_kv=True,
    pack_gqa=False,
    is_causal=False,
    varlen_q=False,
    varlen_k=False,
    disable_bitmask=False,
    has_qk=True,
    store_P=False,
):
    torch.manual_seed(seed)
    hdim = 64
    hdimv = 512
    softmax_scale = 1.0 / math.sqrt(hdim + hdimv) if has_qk else 1.0 / math.sqrt(hdimv)

    nheads_kv = 1
    qhead_per_kvhead = nheads
    seqlen_k_rounded = (seqlen_k + 128 - 1) // 128 * 128
    P_k_length = seqlen_k_rounded if not gather_kv else topk_length

    torch_stream = torch.cuda.current_stream()
    stream = cuda.CUstream(torch_stream.cuda_stream)

    compile_key = (
        is_causal,
        gather_kv,
        topk_length if gather_kv else None,
        pack_gqa,
        qhead_per_kvhead,
        nheads_kv,
        varlen_q,
        varlen_k,
        disable_bitmask,
        has_qk,
    )
    if compile_key not in compile_cache:
        total_q_dummy = batch * seqlen_q
        total_k_dummy = batch * seqlen_k

        if varlen_q:
            Q = torch.randn(
                total_q_dummy, nheads, hdim, dtype=torch.bfloat16, device="cuda"
            )
            Qv = torch.randn(
                total_q_dummy, nheads, hdimv, dtype=torch.bfloat16, device="cuda"
            )
            O = torch.empty(
                total_q_dummy, nheads, hdimv, dtype=torch.bfloat16, device="cuda"
            )
            P = torch.empty(
                total_q_dummy, nheads, P_k_length, dtype=torch.bfloat16, device="cuda"
            )
            lse = torch.empty(total_q_dummy, nheads, dtype=torch.float32, device="cuda")
            row_max = torch.empty(
                total_q_dummy,
                nheads,
                P_k_length // 128,
                dtype=torch.float32,
                device="cuda",
            )
            index_topk = (
                torch.rand(total_q_dummy, topk_length, device="cuda")
                .argsort(dim=-1)
                .to(torch.int32)
            )
            cu_seqlens_q_dummy = torch.arange(
                0, (batch + 1) * seqlen_q, seqlen_q, dtype=torch.int32, device="cuda"
            )
        else:
            Q = torch.randn(
                batch, seqlen_q, nheads, hdim, dtype=torch.bfloat16, device="cuda"
            )
            Qv = torch.randn(
                batch, seqlen_q, nheads, hdimv, dtype=torch.bfloat16, device="cuda"
            )
            O = torch.empty(
                batch, seqlen_q, nheads, hdimv, dtype=torch.bfloat16, device="cuda"
            )
            P = torch.empty(
                batch, seqlen_q, nheads, P_k_length, dtype=torch.bfloat16, device="cuda"
            )
            lse = torch.empty(
                batch, seqlen_q, nheads, dtype=torch.float32, device="cuda"
            )
            row_max = torch.empty(
                batch,
                seqlen_q,
                nheads,
                P_k_length // 128,
                dtype=torch.float32,
                device="cuda",
            )
            index_topk = (
                torch.rand(batch, seqlen_q, topk_length, device="cuda")
                .argsort(dim=-1)
                .to(torch.int32)
            )

        if varlen_k:
            K = torch.randn(
                total_k_dummy, nheads_kv, hdim, dtype=torch.bfloat16, device="cuda"
            )
            V = torch.randn(
                total_k_dummy, nheads_kv, hdimv, dtype=torch.bfloat16, device="cuda"
            )
            cu_seqlens_k_dummy = torch.arange(
                0, (batch + 1) * seqlen_k, seqlen_k, dtype=torch.int32, device="cuda"
            )
        else:
            K = torch.randn(
                batch, seqlen_k, nheads_kv, hdim, dtype=torch.bfloat16, device="cuda"
            )
            V = torch.randn(
                batch, seqlen_k, nheads_kv, hdimv, dtype=torch.bfloat16, device="cuda"
            )

        mQ = from_dlpack(Q, assumed_align=16).mark_layout_dynamic(
            leading_dim=Q.ndim - 1
        )
        mQv = from_dlpack(Qv, assumed_align=16).mark_layout_dynamic(
            leading_dim=Qv.ndim - 1
        )
        mK = from_dlpack(K, assumed_align=16).mark_layout_dynamic(
            leading_dim=K.ndim - 1
        )
        mV = from_dlpack(V, assumed_align=16).mark_layout_dynamic(
            leading_dim=V.ndim - 1
        )
        mO = from_dlpack(O, assumed_align=16).mark_layout_dynamic(
            leading_dim=O.ndim - 1
        )
        mP = from_dlpack(P, assumed_align=16).mark_layout_dynamic(
            leading_dim=P.ndim - 1
        )
        mLSE = from_dlpack(lse, assumed_align=4).mark_layout_dynamic(
            leading_dim=lse.ndim - 1
        )
        mRowMax = from_dlpack(row_max, assumed_align=4).mark_layout_dynamic(
            leading_dim=row_max.ndim - 1
        )
        if gather_kv:
            mIndexTopk = from_dlpack(index_topk, assumed_align=16).mark_layout_dynamic(
                leading_dim=index_topk.ndim - 1
            )
        else:
            mIndexTopk = None

        compile_kwargs = dict(mIndexTopk=mIndexTopk)
        if varlen_q:
            compile_kwargs["mCuSeqlensQ"] = from_dlpack(
                cu_seqlens_q_dummy, assumed_align=4
            )
        if varlen_k:
            compile_kwargs["mCuSeqlensK"] = from_dlpack(
                cu_seqlens_k_dummy, assumed_align=4
            )

        if not has_qk:
            mQ = mK = None

        if store_P is False:
            mP = mRowMax = None

        kernel = cute.compile(
            FlashAttentionMLAForwardSm100(
                is_causal=is_causal,
                use_cpasync_load_KV=gather_kv,
                topk_length=topk_length if gather_kv else 2048,
                is_topk_gather=gather_kv,
                pack_gqa=pack_gqa,
                qhead_per_kvhead=qhead_per_kvhead,
                nheads_kv=nheads_kv,
                is_varlen_q=varlen_q,
                disable_bitmask=disable_bitmask,
                has_qk=has_qk,
            ),
            mQ,
            mQv,
            mK,
            mV,
            mO,
            mLSE,
            softmax_scale,
            mP,
            mRowMax,
            **compile_kwargs,
            stream=stream,
            options="--keep-ptx --keep-cubin --generate-line-info",
        )
        dump_kernel_attributes(kernel)
        compile_cache[compile_key] = kernel

    # ================================================================
    # ---- Generate variable seqlens for this run ----
    if varlen_q:
        torch.manual_seed(seed + 1000)
        # When causal without varlen_k, every per-batch seqlen_q must not exceed seqlen_k.
        max_seqlen_q = seqlen_k if (is_causal and not varlen_k) else seqlen_q
        seqlens_q = torch.randint(1, max_seqlen_q + 1, (batch,), dtype=torch.int32)
        cu_seqlens_q = torch.zeros(batch + 1, dtype=torch.int32, device="cuda")
        cu_seqlens_q[1:] = seqlens_q.cumsum(0).to(torch.int32).cuda()
        total_q = cu_seqlens_q[-1].item()
    else:
        seqlens_q = torch.full((batch,), seqlen_q, dtype=torch.int32)
        total_q = None  # unused

    if varlen_k:
        torch.manual_seed(seed + 2000)
        # Each batch item must have at least topk_length keys so topk gather is valid.
        min_seqlen_k = topk_length if gather_kv else 1
        seqlens_k = torch.randint(
            min_seqlen_k, seqlen_k + 1, (batch,), dtype=torch.int32
        )
        # When causal, every batch item needs seqlens_k[b] >= seqlens_q[b].
        if is_causal:
            seqlens_k = torch.maximum(seqlens_k, seqlens_q)
        cu_seqlens_k = torch.zeros(batch + 1, dtype=torch.int32, device="cuda")
        cu_seqlens_k[1:] = seqlens_k.cumsum(0).to(torch.int32).cuda()
        total_k = cu_seqlens_k[-1].item()
    else:
        seqlens_k = torch.full((batch,), seqlen_k, dtype=torch.int32)
        total_k = None  # unused

    torch.manual_seed(seed)  # restore main seed before drawing actual tensors

    # ---- Allocate Q / Qv / O / lse ----
    if varlen_q:
        Q = torch.randn(total_q, nheads, hdim, dtype=torch.bfloat16, device="cuda")
        Qv = torch.randn(total_q, nheads, hdimv, dtype=torch.bfloat16, device="cuda")
        O = torch.empty(total_q, nheads, hdimv, dtype=torch.bfloat16, device="cuda")
        P = torch.empty(
            total_q, nheads, P_k_length, dtype=torch.bfloat16, device="cuda"
        )
        lse = torch.empty(total_q, nheads, dtype=torch.float32, device="cuda")
        row_max = torch.empty(
            total_q_dummy, P_k_length // 128, nheads, dtype=torch.float32, device="cuda"
        )
    else:
        Q = torch.randn(
            batch, seqlen_q, nheads, hdim, dtype=torch.bfloat16, device="cuda"
        )
        Qv = torch.randn(
            batch, seqlen_q, nheads, hdimv, dtype=torch.bfloat16, device="cuda"
        )
        O = torch.empty(
            batch, seqlen_q, nheads, hdimv, dtype=torch.bfloat16, device="cuda"
        )
        P = torch.empty(
            batch, seqlen_q, nheads, P_k_length, dtype=torch.bfloat16, device="cuda"
        )
        lse = torch.empty(batch, seqlen_q, nheads, dtype=torch.float32, device="cuda")
        row_max = torch.empty(
            batch,
            seqlen_q,
            P_k_length // 128,
            nheads,
            dtype=torch.float32,
            device="cuda",
        )

    # ---- Allocate K / V ----
    if varlen_k:
        K = torch.randn(total_k, nheads_kv, hdim, dtype=torch.bfloat16, device="cuda")
        V = torch.randn(total_k, nheads_kv, hdimv, dtype=torch.bfloat16, device="cuda")
    else:
        K = torch.randn(
            batch, seqlen_k, nheads_kv, hdim, dtype=torch.bfloat16, device="cuda"
        )
        V = torch.randn(
            batch, seqlen_k, nheads_kv, hdimv, dtype=torch.bfloat16, device="cuda"
        )

    # ---- Generate index_topk with per-batch valid ranges when varlen_k ----
    # index_topk shape: (total_q, topk_length) if varlen_q else (batch, seqlen_q, topk_length)
    if gather_kv:
        topk_parts = []
        for b in range(batch):
            sl_q_b = seqlens_q[b].item()
            sl_k_b = seqlens_k[b].item()
            # Draw topk_length unique indices from [0, sl_k_b) for each query in this batch item.
            topk_b = (
                torch.rand(sl_q_b, sl_k_b, device="cuda")
                .argsort(dim=-1)[..., :topk_length]
                .to(torch.int32)
            )  # (sl_q_b, topk_length), all < sl_k_b
            topk_parts.append(topk_b)

        if varlen_q:
            index_topk = torch.cat(topk_parts, dim=0)  # (total_q, topk_length)
        else:
            index_topk = torch.stack(
                topk_parts, dim=0
            )  # (batch, seqlen_q, topk_length)
    else:
        index_topk = None

    # ---- Reference computation (per-batch loop covers all four varlen combos) ----
    O_ref_list, O_pt_list, lse_ref_list, lse_pt_list = [], [], [], []
    for b in range(batch):
        qs = cu_seqlens_q[b].item() if varlen_q else b * seqlen_q
        qe = cu_seqlens_q[b + 1].item() if varlen_q else (b + 1) * seqlen_q
        ks = cu_seqlens_k[b].item() if varlen_k else b * seqlen_k
        ke = cu_seqlens_k[b + 1].item() if varlen_k else (b + 1) * seqlen_k

        Q_b = (
            Q[qs:qe].unsqueeze(0) if varlen_q else Q[b : b + 1]
        )  # (1, sl_q, nheads, hdim)
        Qv_b = (
            Qv[qs:qe].unsqueeze(0) if varlen_q else Qv[b : b + 1]
        )  # (1, sl_q, nheads, hdimv)
        K_b = (
            K[ks:ke].unsqueeze(0) if varlen_k else K[b : b + 1]
        )  # (1, sl_k, nheads_kv, hdim)
        V_b = (
            V[ks:ke].unsqueeze(0) if varlen_k else V[b : b + 1]
        )  # (1, sl_k, nheads_kv, hdimv)
        if gather_kv:
            topk_b = (
                index_topk[qs:qe].unsqueeze(0) if varlen_q else index_topk[b : b + 1]
            )
        else:
            topk_b = None

        O_b, _, lse_b = attention_ref(
            Q_b if has_qk else None,
            K_b if has_qk else None,
            V_b,
            qv=Qv_b,
            causal=is_causal,
            return_lse=True,
            gather_kv_indices=topk_b,
        )
        O_pt_b, _, lse_pt_b = attention_ref(
            Q_b if has_qk else None,
            K_b if has_qk else None,
            V_b,
            qv=Qv_b,
            causal=is_causal,
            upcast=False,
            reorder_ops=True,
            return_lse=True,
            gather_kv_indices=topk_b,
        )
        O_ref_list.append(O_b.squeeze(0))
        O_pt_list.append(O_pt_b.squeeze(0))
        lse_ref_list.append(lse_b.squeeze(0))
        lse_pt_list.append(lse_pt_b.squeeze(0))

    cat_dim_o = 0 if (varlen_q) else 0  # always 0: leading token/batch dim
    cat_dim_lse = -1 if (varlen_q) else -1  # always last: token dim

    if varlen_q:
        O_ref = torch.cat(O_ref_list, dim=0)  # (total_q, nheads, hdimv)
        O_pt = torch.cat(O_pt_list, dim=0)
        lse_ref = torch.cat(lse_ref_list, dim=-1)  # (nheads, total_q)
        lse_pt = torch.cat(lse_pt_list, dim=-1)
    else:
        O_ref = torch.stack(O_ref_list, dim=0)  # (batch, seqlen_q, nheads, hdimv)
        O_pt = torch.stack(O_pt_list, dim=0)
        lse_ref = torch.stack(lse_ref_list, dim=0)  # (batch, nheads, seqlen_q)
        lse_pt = torch.stack(lse_pt_list, dim=0)

    rtol = 2
    atol = 2 * (O_ref + 0.3 - 0.3 - O_ref).abs().max().item()

    # ---- CuTe tensor wrappers ----
    mQ = from_dlpack(Q, assumed_align=16).mark_layout_dynamic(leading_dim=Q.ndim - 1)
    mQv = from_dlpack(Qv, assumed_align=16).mark_layout_dynamic(leading_dim=Qv.ndim - 1)
    mK = from_dlpack(K, assumed_align=16).mark_layout_dynamic(leading_dim=K.ndim - 1)
    mV = from_dlpack(V, assumed_align=16).mark_layout_dynamic(leading_dim=V.ndim - 1)
    mO = from_dlpack(O, assumed_align=16).mark_layout_dynamic(leading_dim=O.ndim - 1)
    mP = from_dlpack(P, assumed_align=16).mark_layout_dynamic(leading_dim=P.ndim - 1)
    mLSE = from_dlpack(lse, assumed_align=4).mark_layout_dynamic(
        leading_dim=lse.ndim - 1
    )
    mRowMax = from_dlpack(row_max, assumed_align=4).mark_layout_dynamic(
        leading_dim=row_max.ndim - 1
    )
    if index_topk is not None:
        mIndexTopk = from_dlpack(index_topk, assumed_align=16).mark_layout_dynamic(
            leading_dim=index_topk.ndim - 1
        )
    else:
        mIndexTopk = None

    run_kwargs = dict(mIndexTopk=mIndexTopk)
    if varlen_q:
        run_kwargs["mCuSeqlensQ"] = from_dlpack(cu_seqlens_q, assumed_align=4)
    if varlen_k:
        run_kwargs["mCuSeqlensK"] = from_dlpack(cu_seqlens_k, assumed_align=4)

    if not has_qk:
        mQ = mK = None

    if store_P is False:
        mP = mRowMax = None

    # ---- Run kernel ----
    compile_cache[compile_key](
        mQ,
        mQv,
        mK,
        mV,
        mO,
        mLSE,
        softmax_scale,
        mP,
        mRowMax,
        **run_kwargs,
        stream=stream,
    )

    O_ref_max = O_ref.abs().max().item()
    O_max = O.abs().max().item()
    print(f"Pytorch O max = {O_ref_max} and our O max = {O_max}")
    print(f"Pytorch max O diff: {(O_pt - O_ref).abs().max().item()}")
    print(f"Pytorch mean O diff: {(O_pt - O_ref).abs().mean().item()}")
    print(f"Max abs diff O, O_ref: {(O - O_ref).abs().max().item()}")
    print(f"Mean abs diff O, O_ref: {(O - O_ref).abs().mean().item()}")

    lse = lse.transpose(-1, -2)
    lse_ref_max = lse_ref.abs().max().item()
    lse_max = lse.abs().max().item()
    print(f"Pytorch LSE max = {lse_ref_max} and our LSE max = {lse_max}")
    print(f"Pytorch LSE max diff: {(lse_pt - lse_ref).abs().max().item()}")
    print(f"Pytorch LSE mean diff: {(lse_pt - lse_ref).abs().mean().item()}")
    print(f"Max abs diff LSE: {(lse - lse_ref).abs().max().item()}")
    print(f"Mean abs diff LSE: {(lse - lse_ref).abs().mean().item()}")

    if validate:
        assert (O - O_ref).abs().max().item() <= rtol * (
            O_pt - O_ref
        ).abs().max().item() + atol
        varlen_tag = ""
        if varlen_q:
            varlen_tag += f", total_q:{total_q}"
        if varlen_k:
            varlen_tag += f", total_k:{total_k}"
        print(
            f"batch:{batch:3d}, nheads:{nheads:3d}, seqlen_q:{seqlen_q:5d}, seqlen_k:{seqlen_k:5d}"
            f"{varlen_tag}, iter:{iter:2d} PASSED"
        )
    else:
        print(mO)
        print(
            f"batch:{batch:3d}, nheads:{nheads:3d}, seqlen_q:{seqlen_q:5d}, seqlen_k:{seqlen_k:5d}"
            f", iter:{iter:2d} RUN (NOT TESTING CORRECTNESS)"
        )

    return None


def timeit(fn, *args, **kwargs):
    # Synchronize before timing
    torch.cuda.synchronize()

    # Warmup
    for _ in range(10):
        fn(*args, **kwargs)

    # Benchmark using PyTorch's Timer
    t = benchmark.Timer(
        stmt="fn(*args, **kwargs)", globals={"fn": fn, "args": args, "kwargs": kwargs}
    )

    # Time it multiple runs
    measurement = t.timeit(20)  # 20 repeats
    avg_time = measurement.mean  # Average time in seconds

    time.sleep(1)

    return avg_time


def benchmark_mla_kernel(
    batch=1,
    seqlen_q=2048,
    seqlen_k=2048,
    topk_length=2048,
    nheads=128,
    hdim=64,
    hdimv=512,
    compile_cache=dict(),
    gather_kv=True,
    is_causal=False,
    disable_bitmask=False,
    store_P=False,
):
    assert hdim == 64, "hdim must be 64"
    assert hdimv == 512, "hdimv must be 512"

    qhead_per_kvhead = nheads
    nheads_kv = 1
    pack_gqa = True
    softmax_scale = 1.0 / math.sqrt(hdim + hdimv)
    seqlen_k_rounded = (seqlen_k + 128 - 1) // 128 * 128
    P_k_length = seqlen_k_rounded if not gather_kv else topk_length

    torch_stream = torch.cuda.current_stream()
    stream = cuda.CUstream(torch_stream.cuda_stream)

    compile_key = (
        is_causal,
        gather_kv,
        topk_length if gather_kv else None,
        pack_gqa,
        qhead_per_kvhead,
        nheads_kv,
        disable_bitmask,
    )
    if compile_key not in compile_cache:
        Q = torch.randn(
            batch, seqlen_q, nheads, hdim, dtype=torch.bfloat16, device="cuda"
        )
        Qv = torch.randn(
            batch, seqlen_q, nheads, hdimv, dtype=torch.bfloat16, device="cuda"
        )
        K = torch.randn(
            batch, seqlen_k, nheads_kv, hdim, dtype=torch.bfloat16, device="cuda"
        )
        V = torch.randn(
            batch, seqlen_k, nheads_kv, hdimv, dtype=torch.bfloat16, device="cuda"
        )
        O = torch.empty(
            batch, seqlen_q, nheads, hdimv, dtype=torch.bfloat16, device="cuda"
        )
        P = torch.empty(
            batch, seqlen_q, nheads, P_k_length, dtype=torch.bfloat16, device="cuda"
        )
        index_topk = (
            torch.rand(batch, seqlen_q, topk_length, device="cuda")
            .argsort(dim=-1)
            .to(torch.int32)
        )

        mQ = from_dlpack(Q, assumed_align=16).mark_layout_dynamic(
            leading_dim=Q.ndim - 1
        )
        mQv = from_dlpack(Qv, assumed_align=16).mark_layout_dynamic(
            leading_dim=Qv.ndim - 1
        )
        mK = from_dlpack(K, assumed_align=16).mark_layout_dynamic(
            leading_dim=K.ndim - 1
        )
        mV = from_dlpack(V, assumed_align=16).mark_layout_dynamic(
            leading_dim=V.ndim - 1
        )
        mO = from_dlpack(O, assumed_align=16).mark_layout_dynamic(
            leading_dim=O.ndim - 1
        )
        mP = from_dlpack(P, assumed_align=16).mark_layout_dynamic(
            leading_dim=P.ndim - 1
        )
        if gather_kv:
            mIndexTopk = from_dlpack(index_topk, assumed_align=16).mark_layout_dynamic(
                leading_dim=index_topk.ndim - 1
            )
        else:
            mIndexTopk = None

        mLSE = None

        if store_P is False:
            mP = None

        kernel = cute.compile(
            FlashAttentionMLAForwardSm100(
                is_causal=is_causal,
                use_cpasync_load_KV=gather_kv,
                topk_length=topk_length if gather_kv else 2048,
                is_topk_gather=gather_kv,
                pack_gqa=pack_gqa,
                qhead_per_kvhead=qhead_per_kvhead,
                nheads_kv=nheads_kv,
                disable_bitmask=disable_bitmask,
            ),
            mQ,
            mQv,
            mK,
            mV,
            mO,
            mLSE,
            softmax_scale,
            mP=mP,
            mIndexTopk=mIndexTopk,
            stream=stream,
        )
        compile_cache[compile_key] = kernel

    Q = torch.randn(batch, seqlen_q, nheads, hdim, dtype=torch.bfloat16, device="cuda")
    Qv = torch.randn(
        batch, seqlen_q, nheads, hdimv, dtype=torch.bfloat16, device="cuda"
    )
    K = torch.randn(
        batch, seqlen_k, nheads_kv, hdim, dtype=torch.bfloat16, device="cuda"
    )
    V = torch.randn(
        batch, seqlen_k, nheads_kv, hdimv, dtype=torch.bfloat16, device="cuda"
    )
    O = torch.empty(batch, seqlen_q, nheads, hdimv, dtype=torch.bfloat16, device="cuda")
    P = torch.empty(
        batch, seqlen_q, nheads, P_k_length, dtype=torch.bfloat16, device="cuda"
    )

    index_topk = (
        torch.rand(batch, seqlen_q, topk_length, device="cuda")
        .argsort(dim=-1)
        .to(torch.int32)
    )

    mQ = from_dlpack(Q, assumed_align=16).mark_layout_dynamic(leading_dim=Q.ndim - 1)
    mQv = from_dlpack(Qv, assumed_align=16).mark_layout_dynamic(leading_dim=Qv.ndim - 1)
    mK = from_dlpack(K, assumed_align=16).mark_layout_dynamic(leading_dim=K.ndim - 1)
    mV = from_dlpack(V, assumed_align=16).mark_layout_dynamic(leading_dim=V.ndim - 1)
    mO = from_dlpack(O, assumed_align=16).mark_layout_dynamic(leading_dim=O.ndim - 1)
    mP = from_dlpack(P, assumed_align=16).mark_layout_dynamic(leading_dim=P.ndim - 1)
    if gather_kv:
        mIndexTopk = from_dlpack(index_topk, assumed_align=16).mark_layout_dynamic(
            leading_dim=index_topk.ndim - 1
        )
    else:
        mIndexTopk = None
    mLSE = None

    if store_P is False:
        mP = None

    exec_time_in_s = timeit(
        compile_cache[compile_key],
        mQ,
        mQv,
        mK,
        mV,
        mO,
        mLSE,
        softmax_scale,
        mP=mP,
        mIndexTopk=mIndexTopk,
        stream=stream,
    )

    seqlen_k_eff = topk_length if gather_kv else seqlen_k

    FLOPs = 2 * batch * nheads * seqlen_q * seqlen_k_eff * (hdim + 2 * hdimv)
    if is_causal and not gather_kv:
        FLOPs /= 2

    TFLOPS = FLOPs / exec_time_in_s / 1e12

    q_bytes = 2 * batch * nheads * seqlen_q * hdim
    qv_bytes = 2 * batch * nheads * seqlen_q * hdimv
    k_bytes = 2 * batch * nheads_kv * seqlen_k_eff * hdim
    v_bytes = 2 * batch * nheads_kv * seqlen_k_eff * hdimv
    o_bytes = 2 * batch * nheads * seqlen_q * hdimv
    total_bytes = q_bytes + qv_bytes + k_bytes + v_bytes + o_bytes
    TBs = total_bytes / exec_time_in_s / 1e12

    print(
        f"batch: {batch}, seqlen_q: {seqlen_q}, seqlen_k: {seqlen_k}, nheads: {nheads}, -> {exec_time_in_s * 1e3:.2f} ms, {TFLOPS:.2f} TFLOPS, {TBs:.2f} TBs"
    )


if __name__ == "__main__":
    run_test = True
    run_benchmark = True
    gather_kv = True
    is_causal = False
    pack_gqa = True
    topk_length = 2048
    varlen_q = False
    varlen_k = False
    disable_bitmask = False
    validate = True
    has_qk = True

    if run_test:
        if not gather_kv:
            seqlen_q_test_values = range(1, 4002, 400)
            seqlen_k_test_values = range(1, 4002, 400)
        else:
            seqlen_q_test_values = range(1, 1001, 200)
            seqlen_k_test_values = range(topk_length, 9001, 2000)
        seqlen_q_test_values = [4096]
        seqlen_k_test_values = [4096]
        nheads_test_values = [128]
        batch_test_values = [1]
        test_configs = [
            (
                batch,
                nheads,
                seqlen_q,
                seqlen_k,
            )
            for batch in batch_test_values
            for nheads in nheads_test_values
            for seqlen_q in seqlen_q_test_values
            for seqlen_k in seqlen_k_test_values
        ]
        iters_per_config = 1
        compile_cache = dict()
        print("=" * 40)
        print("Testing MLA Kernel")
        print("=" * 40)
        for config in test_configs:
            batch, nheads, seqlen_q, seqlen_k = config
            # if is_causal and seqlen_k < seqlen_q:
            #     continue
            for iter in range(iters_per_config):
                test_mla_kernel(
                    seqlen_q=seqlen_q,
                    seqlen_k=seqlen_k,
                    topk_length=topk_length,
                    nheads=nheads,
                    batch=batch,
                    iter=iter,
                    compile_cache=compile_cache,
                    validate=validate,
                    seed=0,
                    gather_kv=gather_kv,
                    pack_gqa=pack_gqa,
                    is_causal=is_causal,
                    varlen_q=varlen_q,
                    varlen_k=varlen_k,
                    disable_bitmask=disable_bitmask,
                    has_qk=has_qk,
                )
    if run_benchmark:
        if gather_kv:
            seqlen_q_benchmark_values = [1]
            seqlen_k_benchmark_values = [8192 * 2]
            nheads_benchmark_values = [128]
            batch_benchmark_values = [128]
        else:
            seqlen_q_benchmark_values = [1]
            seqlen_k_benchmark_values = [8192]
            nheads_benchmark_values = [128]
            batch_benchmark_values = [128]
        # seqlen_q_benchmark_values = [4096]
        # seqlen_k_benchmark_values = [4096]
        # nheads_benchmark_values = [16]
        # batch_benchmark_values = [8]
        benchmark_configs = [
            (
                batch,
                nheads,
                seqlen_q,
                seqlen_k,
            )
            for batch in batch_benchmark_values
            for nheads in nheads_benchmark_values
            for seqlen_q in seqlen_q_benchmark_values
            for seqlen_k in seqlen_k_benchmark_values
        ]
        compile_cache = dict()
        print("=" * 40)
        print("Benchmarking MLA Kernel")
        print("=" * 40)
        for config in benchmark_configs:
            batch, nheads, seqlen_q, seqlen_k = config
            benchmark_mla_kernel(
                batch=batch,
                seqlen_q=seqlen_q,
                seqlen_k=seqlen_k,
                topk_length=topk_length,
                nheads=nheads,
                gather_kv=gather_kv,
                is_causal=is_causal,
                disable_bitmask=disable_bitmask,
                compile_cache=compile_cache,
            )
