# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:

# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.

# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.

# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.

# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


import math
from typing import Type, Tuple, Optional
from types import SimpleNamespace

import cuda.bindings.driver as cuda

import cutlass
import cutlass.cute as cute
import cutlass.cute.nvgpu.tcgen05 as tcgen05

# TODO: Remove this hook helper function after nvidia-cutlass-dsl 4.3.x is no longer supported.
# Compat shim: setmaxregister_{decrease,increase} added in cutlass-dsl 4.4;
# older versions only have the deprecated warpgroup_reg_{dealloc,alloc}.
_setmaxregister_decrease = getattr(
    cute.arch,
    "setmaxregister_decrease",
    getattr(cute.arch, "warpgroup_reg_dealloc", None),
)
_setmaxregister_increase = getattr(
    cute.arch,
    "setmaxregister_increase",
    getattr(cute.arch, "warpgroup_reg_alloc", None),
)

# Compat shim: get_max_tmem_alloc_cols added in cutlass-dsl 4.4;
# older versions don't have it, so we provide a fallback implementation.
_TMEM_MAX_ALLOC_COLUMNS_MAP = {"sm_100": 512, "sm_103": 512, "sm_120": 512}


# TODO: Remove this hook helper function after nvidia-cutlass-dsl 4.3.x is no longer supported.
def _get_max_tmem_alloc_cols(compute_capability: str) -> int:
    if hasattr(cute.arch, "get_max_tmem_alloc_cols"):
        return cute.arch.get_max_tmem_alloc_cols(compute_capability)
    if compute_capability not in _TMEM_MAX_ALLOC_COLUMNS_MAP:
        raise ValueError(f"Unsupported compute capability: {compute_capability}")
    return _TMEM_MAX_ALLOC_COLUMNS_MAP[compute_capability]


from cutlass.cute.nvgpu.tcgen05 import OperandMajorMode
import cutlass.cute.nvgpu.cpasync as cpasync
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
import cutlass.utils.blackwell_helpers as sm100_utils
from cutlass.base_dsl.arch import Arch
from cutlass.cutlass_dsl import BaseDSL


from .mla_helpers import (
    ceil_div,
    compute_q_tile_layout,
    MAX_SPLITS,
    LOG2_E,
    MLAStaticTileScheduler,
    MLAStaticTileSchedulerParams,
    create_mla_static_tile_scheduler,
    create_mla_static_tile_scheduler_params,
    get_variable_query_tile_info,
)

"""
A Multi-Head Latent Attention (MLA) example with FP16 data type for the NVIDIA Blackwell SM100 architecture using CUTE DSL

This example demonstrates an implementation of inference of multi-head latent attention using a TMA + Blackwell
SM100 TensorCore warp-specialized persistent kernel. The implementation integrates the (Qc + Qr)*(Kc + Kr)^T
matrix multiplication, softmax normalization, and softmax((Qc + Qr)*(Kc + Kr)^T)*Vc into a single kernel.
The kernel provides support for page table storage and variable-length KV cache sequences. It implements KV splitting
functionality to minimize latency when processing long KV sequences.

The kernel implements key optimizations including:
- Warp specialization for different computation phases (load, MMA, softmax, correction, epilogue)
- Pipeline stages between different warps for overlapping computation and memory access
- Support for different precision data types
- Two sub-kernels (split KV kernel and reduction kernel) that enable split KV processing

This file defines the BF16/FP16 monolithic kernel class only; production access
goes through ``flashinfer.cute_dsl.attention.cute_dsl_mla_decode`` (the dispatcher
in front of the modular and monolithic implementations).  See
``flashinfer/cute_dsl/attention/monolithic/mla_decode.py`` for the standalone
launcher and ``flashinfer/cute_dsl/attention/mla_dispatch.py`` for impl selection.

Constraints:
* Data type requirements:
  - Input: Float16 or BFloat16
  - Output: Float16 or BFloat16
  - Accumulation and LSE: Float32
* Fixed architecture parameters:
  - Number of attention heads: 1-128
  - Latent dimension: 512
  - RoPE dimension: 64
* Query storage is fixed [BatchSize, SeqLenQ, NumHeads, Dim] or compact
  [TotalQueryTokens, NumHeads, Dim]; both map token/head rows into M128 tiles
* Input kv latent/rope modes should be (SeqLenK, LatentDim/RopeDim, BatchSize)
* Query sequence length must be positive; query (token, head) rows are packed
  continuously into 128-row M tiles with a safely padded final tile
* Only supports 2-CTA instructions
* Variable sequence length requires page table storage enabled
"""


class BlackwellMultiHeadLatentAttentionForwardFP16:
    def __init__(
        self,
        acc_dtype: Type[cutlass.Numeric],
        lse_dtype: Type[cutlass.Numeric],
        mma_qk_tiler_mn: Tuple[int, int],
        mma_pv_tiler_mn: Tuple[int, int],
        max_active_clusters: int,
        page_size: int,
        skip_correction_threshold: float,
        is_persistent: bool,
        is_var_seq: bool,
        is_var_split_kv: bool,
        enable_pdl: bool,
        is_var_q: bool = False,
        num_heads: int = 128,
        seq_len_q: int = 1,
        reducer_d_tiles: int = 1,
        reducer_max_splits: int = MAX_SPLITS,
        enable_dcp: bool = False,
        cp_world: int = 1,
    ):
        """Initializes the configuration for a Blackwell Multi-Head Latent Attention (MLA) kernel.

        :param acc_dtype: Data type for accumulation S and O
        :type acc_dtype: Type[cutlass.Numeric]
        :param lse_dtype: Data type for output LSE
        :type lse_dtype: Type[cutlass.Numeric]
        :param mma_s_tiler: The (H, K) tile shape of the MMA instruction for S
        :type mma_s_tiler: Tuple[int, int]
        :param mma_p_tiler: The (H, D) tile shape of the MMA instruction for P
        :type mma_p_tiler: Tuple[int, int]
        :param max_active_clusters: Maximum number of active clusters
        :type max_active_clusters: int
        :param page_size: The page size of the page table
        :type page_size: int
        :param skip_correction_threshold: Threshold to skip correction
        :type skip_correction_threshold: float
        :param is_persistent: Whether to use persistent kernel mode
        :type is_persistent: bool
        :param is_var_seq: Whether to use variable sequence length
        :type is_var_seq: bool
        :param is_var_split_kv: Whether to use variable split KV
        :type is_var_split_kv: bool
        :param enable_pdl: Whether to use PDL
        :type enable_pdl: bool
        :param is_var_q: Whether query tensors use compact variable-length
            storage described by ``cum_seq_lens_q``.
        :type is_var_q: bool
        :param num_heads: Number of attention heads. Defines the flattened
            ``(query_token, head)`` row geometry and division-free causal mask.
        :type num_heads: int
        :param seq_len_q: Fixed query length, or the per-request query-length
            capacity when ``is_var_q`` is enabled. Combined with ``num_heads``
            to define the scheduler's flattened query-row extent.
        :type seq_len_q: int
        :param reducer_d_tiles: Number of independent D bands reduced per row.
        :type reducer_d_tiles: int
        :param reducer_max_splits: Compile-time reducer capacity. Direct users
            retain the generic 256-split capacity by default; callers choosing
            a smaller specialization must cap runtime split-KV accordingly.
        :type reducer_max_splits: int
        :param enable_dcp: Whether to compile the decode-context-parallel
            global-coordinate causal mask.
        :type enable_dcp: bool
        :param cp_world: Number of cyclic context-parallel shards. This is a
            compile-time parameter when DCP is enabled.
        :type cp_world: int
        """

        self.latent_dim = 512
        self.rope_dim = 64
        self.acc_dtype = acc_dtype
        self.lse_dtype = lse_dtype
        self.mma_qk_tiler_mn = mma_qk_tiler_mn
        self.mma_pv_tiler_mn = mma_pv_tiler_mn
        self.max_active_clusters = max_active_clusters
        self.skip_correction_threshold = skip_correction_threshold
        self.is_persistent = is_persistent
        self.page_size = page_size
        self.is_var_seq = is_var_seq
        self.is_var_split_kv = is_var_split_kv
        self.is_var_q = is_var_q
        if not 1 <= reducer_max_splits <= MAX_SPLITS:
            raise ValueError(
                f"reducer_max_splits must be in [1, {MAX_SPLITS}], "
                f"got {reducer_max_splits}"
            )
        if is_var_split_kv and reducer_max_splits != MAX_SPLITS:
            raise ValueError(
                "variable split-KV requires the generic reducer capacity "
                f"of {MAX_SPLITS}, got {reducer_max_splits}"
            )
        self.reducer_max_splits = reducer_max_splits
        self.enable_pdl = enable_pdl
        # ``seq_len_q`` is the compile-time query capacity.  Fixed-query
        # specializations use it as the request length; variable-query
        # specializations obtain each request's actual length from its indptr.
        # Both flatten (query token, head) into the same physical M128 mode.
        if cp_world < 1:
            raise ValueError(f"cp_world must be positive, got {cp_world}")
        if not enable_dcp and cp_world != 1:
            raise ValueError(
                f"cp_world must be 1 when DCP is disabled, got cp_world={cp_world}"
            )
        self.enable_dcp = enable_dcp
        self.cp_world = cp_world
        self.num_heads = num_heads
        self.seq_len_q = seq_len_q
        (
            self.total_q_rows,
            self.num_q_tiles,
            self.tail_q_rows,
        ) = compute_q_tile_layout(num_heads, seq_len_q, mma_qk_tiler_mn[0])
        self.cluster_shape_mnk = (2, 1, 1)
        self.use_2cta_instrs = True
        # When using 2 CTAs with m=128: warps 0-1 handle accumulation for first half [0, n/2),
        # while warps 2-3 handle accumulation for second half [n/2, n)
        self.warps_in_n = 2
        self.num_compute_warps = 4
        self.threads_per_warp = 32
        if reducer_d_tiles not in (1, 2, 4):
            raise ValueError(f"unsupported reducer_d_tiles={reducer_d_tiles}")
        self.reducer_d_tiles = reducer_d_tiles
        self.reducer_d_tile = self.latent_dim // reducer_d_tiles
        mma_qk_tiler_k = self.rope_dim if self.seq_len_q == 1 else self.rope_dim * 2
        self.mma_qk_tiler = (
            self.mma_qk_tiler_mn[0],
            self.mma_qk_tiler_mn[1],
            mma_qk_tiler_k,
        )
        self.mma_qk_rope_tiler = (
            self.mma_qk_tiler_mn[0],
            self.mma_qk_tiler_mn[1],
            self.rope_dim,
        )
        self.mma_pv_tiler = (
            self.mma_pv_tiler_mn[0],
            self.mma_pv_tiler_mn[1],
            self.mma_qk_tiler[1] * self.mma_qk_tiler[2] // self.mma_pv_tiler_mn[1],
        )
        self.iterations_qk_latent = self.latent_dim // self.mma_qk_tiler[2]
        self.iterations_qk_rope = mma_qk_tiler_k // self.mma_qk_tiler[2]
        self.iterations_qk = self.iterations_qk_latent + self.iterations_qk_rope
        self.iterations_pv_k = self.mma_qk_tiler[1] // self.mma_pv_tiler[2]
        self.iterations_pv_n = self.latent_dim // self.mma_pv_tiler[1]

        # Set specialized warp ids
        self.compute_warp_ids = (0, 1, 2, 3)
        self.correction_warp_ids = (4, 5, 6, 7)
        self.mma_warp_id = 8

        self.load_tma_warp_id = 9
        self.load_pt_warp_id = 10
        self.empty_warp_ids = (11,)
        self.threads_per_cta = self.threads_per_warp * len(
            (
                self.mma_warp_id,
                self.load_tma_warp_id,
                self.load_pt_warp_id,
                *self.compute_warp_ids,
                *self.correction_warp_ids,
                *self.empty_warp_ids,
            )
        )

        # register settings
        self.softmax_reg_num = 192
        self.correction_reg_num = 208
        self.other_reg_num = 96
        # Named barriers
        self.tmem_ptr_sync_bar = pipeline.NamedBarrier(
            barrier_id=1,
            num_threads=(
                self.threads_per_warp
                + self.threads_per_warp * self.num_compute_warps * 2
            ),
        )
        self.softmax_exchange_sync_bar = pipeline.NamedBarrier(
            barrier_id=2, num_threads=(self.threads_per_warp * self.num_compute_warps)
        )
        self.epilogue_exchange_sync_bar = pipeline.NamedBarrier(
            barrier_id=3, num_threads=(self.threads_per_warp * self.num_compute_warps)
        )

    def _setup_attributes(self):
        """Set up configurations and parameters for the MLA kernel operation.

        This method initializes and configures various attributes required for the
        execution of the multi-head latent attention kernel, mainly about the pipeline stages:

        - Sets up staging parameters for Q, K, V inputs and accumulator data
        - Configures pipeline stages for softmax, correction, and epilogue operations
        """

        self.load_q_stage = 1
        self.load_kv_stage = 15 if self.seq_len_q == 1 else 7
        self.mma_s_stage = 2
        self.p_mma_stage = 2
        self.p_cor_stage = 2
        self.mma_o_stage = 1
        self.load_pt_stage = 4

        self.tmem_o_offset = self.mma_s_stage * self.mma_qk_tiler[1] // self.warps_in_n
        self.correction_factor_offset = (
            self.tmem_o_offset + self.latent_dim // self.warps_in_n
        )

    @cute.jit
    def __call__(
        self,
        q_latent: cute.Tensor,
        q_rope: cute.Tensor,
        c_latent: cute.Tensor,
        c_rope: cute.Tensor,
        page_table: cute.Tensor,
        o: cute.Tensor,
        lse: cute.Tensor,
        workspace: cute.Tensor,
        split_kv: cutlass.Int32,
        cache_seqs: Optional[cute.Tensor],
        cum_seq_lens_q: Optional[cute.Tensor],
        causal_seqlens_kv_global: Optional[cute.Tensor],
        cp_rank: cutlass.Int32,
        block_split_kvs: Optional[cute.Tensor],
        softmax_scale: cutlass.Float32,
        output_scale: cutlass.Float32,
        stream: cuda.CUstream,
    ):
        """Execute the Multi-Head Latent Attention operation on the provided tensors.

        The method handles:
        1. Initialization of workspace for temporary split KV buffers
        2. Validation of tensor data types
        3. Initialization of hardware-specific parameters and memory layouts
        4. Configuration of TMA (Tensor Memory Access) operations
        5. Grid and work scheduling computation
        6. Kernel launch(split KV kernel and reduction kernel) with appropriate parameters

        :param q_latent: Fixed-Q query with shape
            [batch_size, seq_len_q, num_head, latent_dim], or compact variable-Q
            query with shape [total_q, num_head, latent_dim] (contiguous).
        :type q_latent: cute.Tensor
        :param q_rope: Fixed-Q query RoPE with shape
            [batch_size, seq_len_q, num_head, rope_dim], or compact variable-Q
            query RoPE with shape [total_q, num_head, rope_dim] (contiguous).
        :type q_rope: cute.Tensor
        :param c_latent: The key tensor with shape [num_pages, page_size, latent_dim] (contiguous)
        :type c_latent: cute.Tensor
        :param c_rope: The key RoPE tensor with shape [num_pages, page_size, rope_dim] (contiguous)
        :type c_rope: cute.Tensor
        :param page_table: The page table tensor with shape [batch_size, page_count] (contiguous)
        :type page_table: cute.Tensor
        :param o: Fixed-Q output with shape
            [batch_size, seq_len_q, num_head, latent_dim], or compact variable-Q
            output with shape [total_q, num_head, latent_dim] (contiguous).
        :type o: cute.Tensor
        :param lse: Fixed-Q LSE with shape
            [batch_size, seq_len_q, num_head], or compact variable-Q LSE with
            shape [total_q, num_head] (contiguous).
        :type lse: cute.Tensor
        :param workspace: The workspace tensor with 1-d shape prepared for acc_o and acc_lse
        :type workspace: cute.Tensor
        :param split_kv: The scalar factor for split KV
        :type split_kv: cutlass.Int32
        :param cache_seqs: The cache sequences tensor with shape [batch_size]
        :type cache_seqs: cute.Tensor
        :param cum_seq_lens_q: Query indptr with shape [batch_size + 1] for
            compact variable-Q tensors; unused for fixed Q.
        :type cum_seq_lens_q: cute.Tensor
        :param causal_seqlens_kv_global: Global exclusive causal KV bound for
            the newest query on each batch item. Present only for DCP kernels.
        :type causal_seqlens_kv_global: cute.Tensor
        :param cp_rank: Runtime context-parallel rank.
        :type cp_rank: cutlass.Int32
        :param block_split_kvs: The block split KV tensor with shape [batch_size]
        :type block_split_kvs: cute.Tensor
        :param softmax_scale: The scale factor for softmax
        :type softmax_scale: cutlass.Float32
        :param output_scale: The scale factor for the output
        :type output_scale: cutlass.Float32
        :param stream: The CUDA stream to execute the kernel on
        :type stream: cuda.CUstream

        :raises TypeError: If tensor data types don't match or aren't supported
        """

        # setup static attributes before smem/grid/tma computation
        self.q_dtype = q_latent.element_type
        self.k_dtype = c_latent.element_type
        self.v_dtype = c_latent.element_type
        self.o_dtype = o.element_type

        # check type consistency
        if cutlass.const_expr(
            self.q_dtype != self.k_dtype or self.q_dtype != self.v_dtype
        ):
            raise TypeError(
                f"Type mismatch: {self.q_dtype} != {self.k_dtype} or {self.q_dtype} != {self.v_dtype}"
            )

        def _flatten_compact_q_rows(t):
            total_tokens = t.shape[0]
            num_heads = t.shape[1]
            head_dim = t.shape[2]
            token_stride = t.stride[0]
            head_stride = t.stride[1]
            dim_stride = t.stride[2]
            return cute.make_tensor(
                t.iterator,
                cute.make_layout(
                    (total_tokens * num_heads, head_dim, 1),
                    stride=(
                        head_stride,
                        dim_stride,
                        token_stride * total_tokens,
                    ),
                ),
            )

        def _make_compact_unpacked_o(t):
            total_tokens = t.shape[0]
            num_heads = t.shape[1]
            head_dim = t.shape[2]
            token_stride = t.stride[0]
            head_stride = t.stride[1]
            dim_stride = t.stride[2]
            return cute.make_tensor(
                t.iterator,
                cute.make_layout(
                    (num_heads, head_dim, total_tokens),
                    stride=(head_stride, dim_stride, token_stride),
                ),
            )

        def _reinterpret_4d(t):
            return cute.make_tensor(
                t.iterator,
                cute.make_layout(
                    (t.shape[2], t.shape[3], t.shape[1], t.shape[0]),
                    stride=(
                        t.stride[2],
                        t.stride[3],
                        t.stride[1],
                        t.stride[0],
                    ),
                ),
            )

        def _flatten_fixed_q_rows(t):
            return cute.make_tensor(
                t.iterator,
                cute.make_layout(
                    (self.total_q_rows, t.shape[1], t.shape[3]),
                    stride=(t.stride[0], t.stride[1], t.stride[3]),
                ),
            )

        if cutlass.const_expr(self.is_var_q):
            # Compact variable-Q tensors use one global affine row space:
            # [total_q, H, D] -> [global_token * H + head, D, 1].
            q_latent = _flatten_compact_q_rows(q_latent)
            q_rope = _flatten_compact_q_rows(q_rope)
            o_unpacked = _make_compact_unpacked_o(o)
            o = _flatten_compact_q_rows(o)
        else:
            # Fixed Q retains an explicit batch mode in its descriptors.
            q_latent = _reinterpret_4d(q_latent)
            q_rope = _reinterpret_4d(q_rope)
            o = _reinterpret_4d(o)
            o_unpacked = o

        # Reinterpret contiguous [num_pages, page_size, D] as [page_size, D, num_pages]
        # Input stride: (PS*D, D, 1) → Target: (D, 1, PS*D)
        def _reinterpret_3d_kv(t):
            return cute.make_tensor(
                t.iterator,
                cute.make_layout(
                    (t.shape[1], t.shape[2], t.shape[0]),
                    stride=(t.stride[1], t.stride[2], t.stride[0]),
                ),
            )

        c_latent = _reinterpret_3d_kv(c_latent)
        c_rope = _reinterpret_3d_kv(c_rope)

        # Reinterpret contiguous [B, page_count] as [page_count, B]
        page_table = cute.make_tensor(
            page_table.iterator,
            cute.make_layout(
                (page_table.shape[1], page_table.shape[0]),
                stride=(page_table.stride[1], page_table.stride[0]),
            ),
        )

        m_tile = self.mma_qk_tiler_mn[0]
        runtime_batch_size = page_table.shape[1]
        if cutlass.const_expr(self.is_var_q):
            lse_total_tokens = lse.shape[0]
            lse_num_heads = lse.shape[1]
            lse_token_stride = lse.stride[0]
            lse_head_stride = lse.stride[1]
            lse_unpacked = cute.make_tensor(
                lse.iterator,
                cute.make_layout(
                    (lse_num_heads, lse_total_tokens),
                    stride=(lse_head_stride, lse_token_stride),
                ),
            )
            lse = cute.make_tensor(
                lse.iterator,
                cute.make_layout(
                    (lse_total_tokens * lse_num_heads, 1, 1),
                    stride=(
                        lse_head_stride,
                        lse_token_stride * lse_total_tokens,
                        lse_token_stride * lse_total_tokens,
                    ),
                ),
            )

            runtime_num_q_tiles = cute.ceil_div(
                o_unpacked.shape[0] * self.seq_len_q, m_tile
            )
            acc_o, acc_lse = self.initialize_workspace(
                cutlass.Int32(m_tile),
                o.shape[1],
                runtime_num_q_tiles,
                runtime_batch_size,
                split_kv,
                self.acc_dtype,
                workspace,
            )
        else:
            # Reinterpret contiguous [B, S_q, H] as [H, S_q, B].
            lse = cute.make_tensor(
                lse.iterator,
                cute.make_layout(
                    (lse.shape[2], lse.shape[1], lse.shape[0]),
                    stride=(lse.stride[2], lse.stride[1], lse.stride[0]),
                ),
            )
            lse_unpacked = lse

            # Flatten token-major [H, D, S_q, B] query storage while retaining
            # the fixed-Q batch mode.
            q_latent = _flatten_fixed_q_rows(q_latent)
            q_rope = _flatten_fixed_q_rows(q_rope)

            runtime_num_q_tiles = cute.ceil_div(o.shape[0] * o.shape[2], m_tile)
            o = cute.make_tensor(
                o.iterator,
                cute.make_layout(
                    (m_tile, o.shape[1], runtime_num_q_tiles, o.shape[3]),
                    stride=(
                        o.stride[0],
                        o.stride[1],
                        o.stride[0] * m_tile,
                        o.stride[3],
                    ),
                ),
            )
            lse = cute.make_tensor(
                lse.iterator,
                cute.make_layout(
                    (m_tile, runtime_num_q_tiles, lse.shape[2]),
                    stride=(
                        lse.stride[0],
                        lse.stride[0] * m_tile,
                        lse.stride[2],
                    ),
                ),
            )
            acc_o, acc_lse = self.initialize_workspace(
                o.shape[0],
                o.shape[1],
                o.shape[2],
                o.shape[3],
                split_kv,
                self.acc_dtype,
                workspace,
            )

        c_latent_tranpose_layout = cute.select(c_latent.layout, mode=[1, 0, 2])
        c_latent_transpose = cute.make_tensor(
            c_latent.iterator, c_latent_tranpose_layout
        )

        self.q_major_mode = tcgen05.OperandMajorMode.K
        self.k_major_mode = tcgen05.OperandMajorMode.K
        self.v_major_mode = tcgen05.OperandMajorMode.MN

        self._setup_attributes()

        cta_group = tcgen05.CtaGroup.TWO
        # the intermediate tensor p is from smem & k-major
        p_major_mode = tcgen05.OperandMajorMode.K
        qk_tiled_mma = sm100_utils.make_trivial_tiled_mma(
            self.q_dtype,
            self.q_major_mode,
            self.k_major_mode,
            self.acc_dtype,
            cta_group,
            self.mma_qk_tiler[:2],
        )
        pv_tiled_mma = sm100_utils.make_trivial_tiled_mma(
            self.v_dtype,
            p_major_mode,
            self.v_major_mode,
            self.acc_dtype,
            cta_group,
            self.mma_pv_tiler[:2],
        )

        cta_layout_vmnk = cute.tiled_divide(
            cute.make_layout(self.cluster_shape_mnk),
            (qk_tiled_mma.thr_id.shape,),
        )

        self.epi_tile = self.mma_pv_tiler[:2]

        q_latent_smem_layout_staged = sm100_utils.make_smem_layout_a(
            qk_tiled_mma,
            self.mma_qk_tiler,
            self.q_dtype,
            (self.iterations_qk_latent * self.load_q_stage),
        )
        q_latent_smem_layout_staged = cute.logical_divide(
            q_latent_smem_layout_staged, (None, None, None, self.iterations_qk_latent)
        )
        q_rope_smem_layout_staged = sm100_utils.make_smem_layout_a(
            qk_tiled_mma,
            self.mma_qk_rope_tiler,
            self.q_dtype,
            self.load_q_stage,
        )

        # rope reuse the same smem layout as latent
        kc_smem_layout_staged = sm100_utils.make_smem_layout_b(
            qk_tiled_mma,
            self.mma_qk_tiler,
            self.k_dtype,
            self.load_kv_stage,
        )
        kc_page_tile_size = min(
            self.page_size, qk_tiled_mma.op.shape_mnk[0] // qk_tiled_mma.thr_id.shape
        )

        kc_smem_layout_for_tma = sm100_utils.make_smem_layout(
            OperandMajorMode.K,
            (self.mma_qk_tiler[0] // qk_tiled_mma.thr_id.shape, self.mma_qk_tiler[2]),
            self.k_dtype,
            self.load_kv_stage,
        )
        kc_smem_layout_for_tma = cute.tiled_divide(
            kc_smem_layout_for_tma, (kc_page_tile_size, self.mma_qk_tiler[2])
        )

        p_smem_layout_staged = sm100_utils.make_smem_layout_a(
            pv_tiled_mma,
            self.mma_pv_tiler,
            self.q_dtype,
            (self.iterations_pv_k * self.p_mma_stage),
        )
        p_smem_layout_staged = cute.logical_divide(
            p_smem_layout_staged, (None, None, None, self.iterations_pv_k)
        )

        vc_smem_layout_staged = sm100_utils.make_smem_layout_b(
            pv_tiled_mma,
            self.mma_pv_tiler,
            self.v_dtype,
            self.load_kv_stage,
        )
        vc_page_tile_size = min(self.page_size, self.mma_pv_tiler[2])
        vc_smem_layout_for_tma = sm100_utils.make_smem_layout(
            OperandMajorMode.MN,
            (self.mma_pv_tiler[1] // pv_tiled_mma.thr_id.shape, self.mma_pv_tiler[2]),
            self.v_dtype,
            self.load_kv_stage,
        )
        vc_smem_layout_for_tma = cute.tiled_divide(
            vc_smem_layout_for_tma,
            (
                pv_tiled_mma.op.shape_mnk[1] // pv_tiled_mma.thr_id.shape,
                vc_page_tile_size,
            ),
        )
        # TMA load for Q latent and rope
        tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group)

        q_latent_smem_layout = cute.select(q_latent_smem_layout_staged, mode=[0, 1, 2])
        tma_atom_q_latent, tma_tensor_q_latent = cute.nvgpu.make_tiled_tma_atom_A(
            tma_load_op,
            q_latent,
            q_latent_smem_layout,
            self.mma_qk_tiler,
            qk_tiled_mma,
            cta_layout_vmnk.shape,
        )
        q_rope_smem_layout = cute.select(q_rope_smem_layout_staged, mode=[0, 1, 2])
        tma_atom_q_rope, tma_tensor_q_rope = cute.nvgpu.make_tiled_tma_atom_A(
            tma_load_op,
            q_rope,
            q_rope_smem_layout,
            self.mma_qk_rope_tiler,
            qk_tiled_mma,
            cta_layout_vmnk.shape,
        )
        # TMA load for c latent and k rope
        kc_smem_layout = cute.select(kc_smem_layout_for_tma, mode=[0])
        tma_atom_c_latent, tma_tensor_c_latent = self.make_paged_tiled_tma_atom(
            tma_load_op,
            c_latent,
            kc_smem_layout,
            (self.mma_qk_tiler[1], self.mma_qk_tiler[2]),
            qk_tiled_mma,
            is_k_load=True,
        )
        tma_atom_c_rope, tma_tensor_c_rope = self.make_paged_tiled_tma_atom(
            tma_load_op,
            c_rope,
            kc_smem_layout,
            (self.mma_qk_tiler[1], self.mma_qk_tiler[2]),
            qk_tiled_mma,
            is_k_load=True,
        )
        # TMA load for c latent transpose
        vc_smem_layout = cute.select(vc_smem_layout_for_tma, mode=[0])
        tma_atom_c_latent_transpose, tma_tensor_c_latent_transpose = (
            self.make_paged_tiled_tma_atom(
                tma_load_op,
                c_latent_transpose,
                vc_smem_layout,
                (self.mma_pv_tiler[1], self.mma_pv_tiler[2]),
                pv_tiled_mma,
                is_k_load=False,
            )
        )

        q_latent_copy_size = (
            cute.size_in_bytes(self.q_dtype, q_latent_smem_layout)
            * cute.size(qk_tiled_mma.thr_id.shape)
            * self.iterations_qk_latent
        )
        q_rope_copy_size = (
            cute.size_in_bytes(self.q_dtype, q_rope_smem_layout)
            * cute.size(qk_tiled_mma.thr_id.shape)
            * self.iterations_qk_rope
        )
        q_copy_size = q_latent_copy_size + q_rope_copy_size
        kc_copy_size = cute.size_in_bytes(
            self.k_dtype, cute.select(kc_smem_layout_staged, mode=[0, 1, 2])
        ) * cute.size(qk_tiled_mma.thr_id.shape)
        vc_copy_size = cute.size_in_bytes(
            self.v_dtype, cute.select(vc_smem_layout_staged, mode=[0, 1, 2])
        ) * cute.size(pv_tiled_mma.thr_id.shape)
        assert kc_copy_size == vc_copy_size, (
            "kc_copy_size and vc_copy_size must be the same"
        )

        self.tma_copy_q_bytes = q_copy_size
        self.tma_copy_kc_bytes = kc_copy_size

        tile_sched_params, grid = self._compute_grid(
            runtime_batch_size,
            runtime_num_q_tiles,
            split_kv,
            self.cluster_shape_mnk,
            self.max_active_clusters,
            self.is_persistent,
        )

        @cute.struct
        class SplitKVKernelSharedStorage:
            # Pipeline barriers
            load_q_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_q_stage * 2]
            load_kv_mbar_ptr: cute.struct.MemRange[
                cutlass.Int64, self.load_kv_stage * 2
            ]
            mma_s_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.mma_s_stage * 2]
            p_mma_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.p_mma_stage * 2]
            p_cor_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.p_cor_stage * 2]
            mma_o_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.mma_o_stage * 2]
            load_pt_mbar_ptr: cute.struct.MemRange[
                cutlass.Int64, self.load_pt_stage * 2
            ]
            # Tmem dealloc cluster barrier
            tmem_dealloc_mbar_ptr: cutlass.Int64

            # Tmem holding buffer
            tmem_holding_buf: cutlass.Int32
            # Smem tensors
            softmax_smem_exchange: cute.struct.MemRange[
                self.acc_dtype, self.num_compute_warps * self.threads_per_warp
            ]
            epilogue_smem_exchange: cute.struct.MemRange[
                self.acc_dtype, self.num_compute_warps * self.threads_per_warp
            ]
            smem_q_latent: cute.struct.Align[
                cute.struct.MemRange[
                    self.q_dtype, cute.cosize(q_latent_smem_layout_staged)
                ],
                1024,
            ]
            smem_q_rope: cute.struct.Align[
                cute.struct.MemRange[
                    self.q_dtype, cute.cosize(q_rope_smem_layout_staged)
                ],
                1024,
            ]
            smem_kc: cute.struct.Align[
                cute.struct.MemRange[self.k_dtype, cute.cosize(kc_smem_layout_staged)],
                1024,
            ]
            smem_p: cute.struct.Align[
                cute.struct.MemRange[self.q_dtype, cute.cosize(p_smem_layout_staged)],
                1024,
            ]
            smem_page_table: cute.struct.MemRange[
                cutlass.Int32, self.load_pt_stage * self.mma_qk_tiler[1] // 2
            ]

        softmax_scale_log2 = softmax_scale * LOG2_E
        self.split_kv_kernel(
            qk_tiled_mma,
            pv_tiled_mma,
            tma_atom_q_latent,
            tma_tensor_q_latent,
            tma_atom_q_rope,
            tma_tensor_q_rope,
            tma_atom_c_latent,
            tma_tensor_c_latent,
            tma_atom_c_rope,
            tma_tensor_c_rope,
            tma_atom_c_latent_transpose,
            tma_tensor_c_latent_transpose,
            page_table,
            o,
            lse,
            acc_o,
            acc_lse,
            split_kv,
            cache_seqs,
            cum_seq_lens_q,
            causal_seqlens_kv_global,
            cp_rank,
            block_split_kvs,
            softmax_scale_log2,
            output_scale,
            q_latent_smem_layout_staged,
            q_rope_smem_layout_staged,
            kc_smem_layout_staged,
            p_smem_layout_staged,
            vc_smem_layout_staged,
            kc_smem_layout_for_tma,
            vc_smem_layout_for_tma,
            cta_layout_vmnk,
            tile_sched_params,
            SplitKVKernelSharedStorage,
        ).launch(
            grid=grid,
            block=[self.threads_per_cta, 1, 1],
            cluster=self.cluster_shape_mnk,
            smem=SplitKVKernelSharedStorage.size_in_bytes(),  # type: ignore[attr-defined]
            stream=stream,
            min_blocks_per_mp=1,
            use_pdl=self.enable_pdl,
        )
        if cutlass.const_expr(acc_o is not None):
            self.reduction_kernel(
                o_unpacked,
                lse_unpacked,
                acc_o,
                acc_lse,
                split_kv,
                cache_seqs,
                cum_seq_lens_q,
                block_split_kvs,
            ).launch(
                grid=(
                    o_unpacked.shape[0] * self.reducer_d_tiles,
                    self.seq_len_q,
                    runtime_batch_size,
                ),
                block=[self.threads_per_warp * self.num_compute_warps, 1, 1],
                smem=self.reducer_max_splits * self.acc_dtype.width // 8,
                stream=stream,
                min_blocks_per_mp=1,
                use_pdl=self.enable_pdl,
            )

    @cute.jit
    def make_paged_tiled_tma_atom(
        self,
        tma_load_op: cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp,
        gmem: cute.Tensor,
        smem_layout: cute.Layout,
        mma_tiler,
        tiled_mma: cute.TiledMma,
        is_k_load: bool,
    ):
        ident = cute.make_identity_layout(gmem.shape)
        g_tile = cute.composition(ident, mma_tiler)
        cta_mn = mma_tiler[0] // tiled_mma.thr_id.shape
        cta_v_map = cute.flat_divide(g_tile, (cta_mn,))
        cta_v_map = cute.select(cta_v_map, mode=[0, 2])
        page_tile_size = (
            min(self.page_size, cta_mn)
            if is_k_load
            else min(self.page_size, mma_tiler[1])
        )
        cta_v_map = cute.zipped_divide(
            cta_v_map,
            (page_tile_size, mma_tiler[1]) if is_k_load else (cta_mn, page_tile_size),
        )
        cta_v_map = cute.select(cta_v_map, mode=[0])
        from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir

        res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_load(
            gmem.value,
            smem_layout.value,
            cta_v_map,
            tma_load_op._to_ir(),
            num_multicast=1,
        )
        return cute.CopyAtom(
            tma_load_op, cpasync.CopyBulkTensorTileG2SNonExecTrait(res[0])
        ), res[1]

    @cute.kernel
    def split_kv_kernel(
        self,
        tiled_mma_qk: cute.TiledMma,
        tiled_mma_pv: cute.TiledMma,
        tma_atom_q_latent: Optional[cute.CopyAtom],
        mQL: cute.Tensor,
        tma_atom_q_rope: Optional[cute.CopyAtom],
        mQR: cute.Tensor,
        tma_atom_c_latent: Optional[cute.CopyAtom],
        mCL: cute.Tensor,
        tma_atom_c_rope: Optional[cute.CopyAtom],
        mKR: cute.Tensor,
        tma_atom_c_latent_transpose: Optional[cute.CopyAtom],
        mCLT: cute.Tensor,
        mPT: cute.Tensor,
        mO: Optional[cute.Tensor],
        mLSE: Optional[cute.Tensor],
        mAccO: Optional[cute.Tensor],
        mAccLSE: Optional[cute.Tensor],
        split_kv: cutlass.Int32,
        cache_seqs: cute.Tensor,
        cum_seq_lens_q: Optional[cute.Tensor],
        causal_seqlens_kv_global: Optional[cute.Tensor],
        cp_rank: cutlass.Int32,
        block_split_kvs: cute.Tensor,
        softmax_scale_log2: cutlass.Float32,
        output_scale: cutlass.Float32,
        q_latent_smem_layout_staged: cute.ComposedLayout,
        q_rope_smem_layout_staged: cute.ComposedLayout,
        kc_smem_layout_staged: cute.ComposedLayout,
        p_smem_layout_staged: cute.ComposedLayout,
        vc_smem_layout_staged: cute.ComposedLayout,
        kc_smem_layout_for_tma: cute.ComposedLayout,
        vc_smem_layout_for_tma: cute.ComposedLayout,
        cta_layout_vmnk: cute.Layout,
        tile_sched_params: MLAStaticTileSchedulerParams,
        SharedStorage: cutlass.Constexpr,
    ):
        """The device split_kv kernel implementation of the Multi-Head Latent Attention.

        This kernel coordinates multiple specialized warps to perform different phases of the MLA computation:
        1. Load warp: Loads Q/C latent/rope data from global memory to shared memory using TMA
        2. MMA warp: Performs matrix multiplications (Q*K^T and P*V)
        3. Compute warps: Compute softmax and do rescaling on accumulators, and store the intermediate/final results
        to global memory

        The kernel produces either intermediate or final results of the MLA computation based on the split_kv parameter.
        When split_kv is 1, the kernel generates the final results directly. Otherwise, it produces intermediate results
        that will later be combined by a reduction kernel.

        The kernel implements a complex pipeline with overlapping computation and memory operations,
        using tensor memory access (TMA) for efficient data loading, warp specialization for different
        computation phases.

        :param tiled_mma_qk: Tiled MMA for Q*K^T
        :type tiled_mma_qk: cute.TiledMma
        :param tiled_mma_pv: Tiled MMA for P*V
        :type tiled_mma_pv: cute.TiledMma
        :param tma_atom_q_latent: TMA copy atom for query latent tensor
        :type tma_atom_q_latent: cute.CopyAtom
        :param mQL: query latent tensor
        :type mQL: cute.Tensor
        :param tma_atom_q_rope: TMA copy atom for query rope tensor
        :type tma_atom_q_rope: cute.CopyAtom
        :param mQR: query rope tensor
        :type mQR: cute.Tensor
        :param tma_atom_c_latent: TMA copy atom for c latent tensor
        :type tma_atom_c_latent: cute.CopyAtom
        :param mCL: Compressed latent tensor
        :type mCL: cute.Tensor
        :param tma_atom_c_rope: TMA copy atom for c rope tensor
        :type tma_atom_c_rope: cute.CopyAtom
        :param mCLT: Compressed latent transpose tensor
        :type mCLT: cute.Tensor
        :param mPT: Page table tensor
        :type mPT: cute.Tensor
        :param mO: Output tensor
        :type mO: cute.Tensor
        :param mLSE: Log-sum-exp tensor
        :type mLSE: cute.Tensor
        :param mAccO: Intermediate accumulator output tensor
        :type mAccO: cute.Tensor
        :param mAccLSE: Intermediate accumulator log-sum-exp tensor
        :type mAccLSE: cute.Tensor
        :param split_kv: The split_kv parameter
        :type split_kv: cutlass.Int32
        :param cache_seqs: The variable sequence length tensor
        :type cache_seqs: cute.Tensor
        :param causal_seqlens_kv_global: Per-batch global exclusive causal KV
            bound for the newest query, present only when DCP is enabled.
        :type causal_seqlens_kv_global: cute.Tensor
        :param cp_rank: Runtime context-parallel rank.
        :type cp_rank: cutlass.Int32
        :param block_split_kvs: The per-block split_kv values tensor
        :type block_split_kvs: cute.Tensor
        :param softmax_scale_log2: The log2 scale factor for softmax
        :type softmax_scale_log2: cutlass.Float32
        :param output_scale: The scale factor for the output
        :type output_scale: cutlass.Float32
        :param q_latent_smem_layout_staged: Shared memory layout for query latent tensor
        :type q_latent_smem_layout_staged: cute.ComposedLayout
        :param q_rope_smem_layout_staged: Shared memory layout for query rope tensor
        :type q_rope_smem_layout_staged: cute.ComposedLayout
        :param kc_smem_layout_staged: Shared memory layout for key/value latent/rope tensor
        :type kc_smem_layout_staged: cute.ComposedLayout
        :param p_smem_layout_staged: Shared memory layout for probability matrix
        :type p_smem_layout_staged: cute.ComposedLayout
        :param vc_smem_layout_staged: Shared memory layout for value tensor
        :type vc_smem_layout_staged: cute.ComposedLayout
        :param kc_smem_layout_for_tma: Shared memory layout for key/value latent tensor for TMA
        :type kc_smem_layout_for_tma: cute.ComposedLayout
        :param vc_smem_layout_for_tma: Shared memory layout for value tensor for TMA
        :type vc_smem_layout_for_tma: cute.ComposedLayout
        :param cta_layout_vmnk: Layout for compute threads
        :type cta_layout_vmnk: cute.Layout
        :param tile_sched_params: Scheduling parameters for work distribution
        :type tile_sched_params: MLAStaticTileSchedulerParams
        :param SharedStorage: Shared storage for the kernel
        :type SharedStorage: cutlass.Constexpr
        """

        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())

        tidx, _, _ = cute.arch.thread_idx()
        bidx, _, _ = cute.arch.block_idx()
        mma_tile_coord_v = bidx % cute.size(tiled_mma_qk.thr_id.shape)
        is_leader_cta = mma_tile_coord_v == 0

        # Prefetch tma descriptor
        if warp_idx == self.mma_warp_id:
            cpasync.prefetch_descriptor(tma_atom_q_latent)
            cpasync.prefetch_descriptor(tma_atom_q_rope)
            cpasync.prefetch_descriptor(tma_atom_c_latent)
            cpasync.prefetch_descriptor(tma_atom_c_rope)
            cpasync.prefetch_descriptor(tma_atom_c_latent_transpose)

        # Alloc
        smem = utils.SmemAllocator()
        storage = smem.allocate(SharedStorage)

        # Tensor memory dealloc barrier init
        tmem = utils.TmemAllocator(
            storage.tmem_holding_buf.ptr,
            barrier_for_retrieve=self.tmem_ptr_sync_bar,
            allocator_warp_id=self.mma_warp_id,
            is_two_cta=self.use_2cta_instrs,
            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr.ptr,
        )

        load_q_pipeline = self.make_and_init_load_qkv_pipeline(
            storage.load_q_mbar_ptr.data_ptr(),
            cta_layout_vmnk,
            self.load_q_stage,
            self.tma_copy_q_bytes,
        )
        load_kv_pipeline = self.make_and_init_load_qkv_pipeline(
            storage.load_kv_mbar_ptr.data_ptr(),
            cta_layout_vmnk,
            self.load_kv_stage,
            self.tma_copy_kc_bytes,
        )
        mma_s_pipeline = self.make_and_init_mma_s_pipeline(
            storage.mma_s_mbar_ptr.data_ptr(), cta_layout_vmnk
        )
        p_mma_pipeline = self.make_and_init_p_mma_pipeline(
            storage.p_mma_mbar_ptr.data_ptr(), cta_layout_vmnk
        )
        p_cor_pipeline = self.make_and_init_p_cor_pipeline(
            storage.p_cor_mbar_ptr.data_ptr()
        )
        mma_o_pipeline = self.make_and_init_mma_o_pipeline(
            storage.mma_o_mbar_ptr.data_ptr(), cta_layout_vmnk
        )
        load_pt_pipeline = self.make_and_init_load_pt_pipeline(
            storage.load_pt_mbar_ptr.data_ptr()
        )

        # Cluster arrive after barrier init
        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, is_relaxed=True)

        # Generate smem tensor Q/KC/VC/exchange
        # (MMA, MMA_H, MMA_R, PIPE)
        sQ = storage.smem_q_latent.get_tensor(
            q_latent_smem_layout_staged.outer, swizzle=q_latent_smem_layout_staged.inner
        )
        sQ_rope = storage.smem_q_rope.get_tensor(
            q_rope_smem_layout_staged.outer, swizzle=q_rope_smem_layout_staged.inner
        )
        # (MMA, MMA_K, MMA_R, PIPE)
        sKC = storage.smem_kc.get_tensor(
            kc_smem_layout_staged.outer, swizzle=kc_smem_layout_staged.inner
        )
        sKC_for_tma = storage.smem_kc.get_tensor(
            kc_smem_layout_for_tma.outer,
            swizzle=kc_smem_layout_for_tma.inner,
        )
        # (MMA, MMA_D, MMA_K, PIPE)
        # reuse smem
        sVC_ptr = cute.recast_ptr(sKC.iterator, vc_smem_layout_staged.inner)
        sVC = cute.make_tensor(sVC_ptr, vc_smem_layout_staged.outer)
        sVC_for_tma = cute.make_tensor(sVC_ptr, vc_smem_layout_for_tma.outer)
        # (MMA, MMA_H, MMA_K)
        sP = storage.smem_p.get_tensor(
            p_smem_layout_staged.outer, swizzle=p_smem_layout_staged.inner
        )
        sPT = storage.smem_page_table.get_tensor(
            cute.make_layout((self.mma_qk_tiler[1] // 2, self.load_pt_stage))
        )
        # (compute_threads,)
        softmax_smem_exchange = storage.softmax_smem_exchange.get_tensor(
            cute.make_layout(self.num_compute_warps * self.threads_per_warp)
        )
        epilogue_smem_exchange = storage.epilogue_smem_exchange.get_tensor(
            cute.make_layout(self.num_compute_warps * self.threads_per_warp)
        )

        #
        # Cluster wait before tensor memory alloc
        #
        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk)

        if cutlass.const_expr(self.enable_pdl):
            cute.arch.griddepcontrol_wait()

        # ///////////////////////////////////////////////////////////////////////////////
        #  Load warps, including page table and data tensors
        # ///////////////////////////////////////////////////////////////////////////////

        if warp_idx >= self.empty_warp_ids[0] and warp_idx <= self.empty_warp_ids[-1]:
            _setmaxregister_decrease(self.other_reg_num)
        if warp_idx == self.load_pt_warp_id:
            _setmaxregister_decrease(self.other_reg_num)
            load_pt_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.load_pt_stage
            )
            tile_sched = create_mla_static_tile_scheduler(
                tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
            )
            work_tile = tile_sched.initial_work_tile_info()
            while work_tile.is_valid_tile:
                blk_coord = work_tile.tile_idx
                (
                    k_index,
                    k_tile_count,
                    local_split_kv,
                    _,
                    _,
                    _,
                ) = self.get_k_tile_count(
                    split_kv,
                    cache_seqs,
                    cum_seq_lens_q,
                    block_split_kvs,
                    blk_coord,
                )
                if k_tile_count > 0:
                    load_pt_common_params = SimpleNamespace(
                        blk_coord=blk_coord,
                        load_pt_pipeline=load_pt_pipeline,
                        mPT=mPT,
                        sPT=sPT,
                        tidx=tidx,
                        page_size=mCL.shape[0],
                    )
                    load_pt_producer_state = self.load_page_table(
                        load_pt_common_params,
                        k_index,
                        k_tile_count,
                        load_pt_producer_state,
                    )
                tile_sched.advance_to_next_work()
                work_tile = tile_sched.get_current_work()
            load_pt_pipeline.producer_tail(load_pt_producer_state)
        if warp_idx == self.load_tma_warp_id:
            _setmaxregister_decrease(self.other_reg_num)
            load_q_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.load_q_stage
            )
            load_kv_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.load_kv_stage
            )
            load_pt_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.load_pt_stage
            )
            load_pt_release_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.load_pt_stage
            )
            tile_sched = create_mla_static_tile_scheduler(
                tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
            )
            work_tile = tile_sched.initial_work_tile_info()
            while work_tile.is_valid_tile:
                blk_coord = work_tile.tile_idx
                (
                    k_index,
                    k_tile_count,
                    local_split_kv,
                    q_begin,
                    _,
                    _,
                ) = self.get_k_tile_count(
                    split_kv,
                    cache_seqs,
                    cum_seq_lens_q,
                    block_split_kvs,
                    blk_coord,
                )
                if k_tile_count > 0:
                    # Construct fixed common/tma_qk/tma_pv params for load_tma
                    tma_common_params = SimpleNamespace(
                        blk_coord=blk_coord,
                        local_split_kv=local_split_kv,
                        q_begin=q_begin,
                        load_q_pipeline=load_q_pipeline,
                        load_kv_pipeline=load_kv_pipeline,
                        mPT=mPT,
                        sPT=sPT,
                        load_pt_pipeline=load_pt_pipeline,
                    )
                    tma_qk_params = SimpleNamespace(
                        tiled_mma_qk=tiled_mma_qk,
                        tma_atom_q_latent=tma_atom_q_latent,
                        tma_atom_q_rope=tma_atom_q_rope,
                        tma_atom_c_latent=tma_atom_c_latent,
                        tma_atom_c_rope=tma_atom_c_rope,
                        mQL=mQL,
                        mQR=mQR,
                        mCL=mCL,
                        mKR=mKR,
                        sQ=sQ,
                        sQ_rope=sQ_rope,
                        sKC=sKC_for_tma,
                    )
                    tma_pv_params = SimpleNamespace(
                        tiled_mma_pv=tiled_mma_pv,
                        tma_atom_c_latent_transpose=tma_atom_c_latent_transpose,
                        mCL=mCL,
                        mKR=mKR,
                        mCLT=mCLT,
                        sVC=sVC_for_tma,
                    )
                    # Load tma
                    (
                        load_q_producer_state,
                        load_kv_producer_state,
                        load_pt_consumer_state,
                        load_pt_release_state,
                    ) = self.load_tma(
                        tma_common_params,
                        tma_qk_params,
                        tma_pv_params,
                        k_index,
                        k_tile_count,
                        load_q_producer_state,
                        load_kv_producer_state,
                        load_pt_consumer_state,
                        load_pt_release_state,
                    )
                tile_sched.advance_to_next_work()
                work_tile = tile_sched.get_current_work()

            load_q_pipeline.producer_tail(load_q_producer_state)
            load_kv_pipeline.producer_tail(load_kv_producer_state)

        # ///////////////////////////////////////////////////////////////////////////////
        #  MMA warp
        # ///////////////////////////////////////////////////////////////////////////////
        if warp_idx == self.mma_warp_id:
            _setmaxregister_decrease(self.other_reg_num)
            # Alloc tensor memory buffer
            tmem.allocate(_get_max_tmem_alloc_cols("sm_100"))
            tmem.wait_for_alloc()
            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)

            load_q_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.load_q_stage
            )
            load_kv_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.load_kv_stage
            )
            mma_s_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.mma_s_stage
            )
            p_mma_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.p_mma_stage
            )
            mma_o_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.mma_o_stage
            )
            tile_sched = create_mla_static_tile_scheduler(
                tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
            )
            work_tile = tile_sched.initial_work_tile_info()
            while work_tile.is_valid_tile:
                blk_coord = work_tile.tile_idx
                (
                    k_index,
                    k_tile_count,
                    local_split_kv,
                    _,
                    _,
                    _,
                ) = self.get_k_tile_count(
                    split_kv,
                    cache_seqs,
                    cum_seq_lens_q,
                    block_split_kvs,
                    blk_coord,
                )
                if k_tile_count > 0:
                    mma_common_params = SimpleNamespace(
                        blk_coord=blk_coord,
                        local_split_kv=local_split_kv,
                        load_q_pipeline=load_q_pipeline,
                        load_kv_pipeline=load_kv_pipeline,
                        tmem_ptr=tmem_ptr,
                        is_leader_cta=is_leader_cta,
                        L=mCL.shape[1],
                    )
                    mma_qk_params = SimpleNamespace(
                        mma_s_pipeline=mma_s_pipeline,
                        sQ=sQ,
                        sQ_rope=sQ_rope,
                        sKC=sKC,
                    )
                    mma_pv_params = SimpleNamespace(
                        p_mma_pipeline=p_mma_pipeline,
                        mma_o_pipeline=mma_o_pipeline,
                        sP=sP,
                        sVC=sVC,
                    )
                    (
                        tiled_mma_qk,
                        tiled_mma_pv,
                        load_q_consumer_state,
                        load_kv_consumer_state,
                        mma_s_producer_state,
                        p_mma_consumer_state,
                        mma_o_producer_state,
                    ) = self.mma(
                        mma_common_params,
                        mma_qk_params,
                        mma_pv_params,
                        k_tile_count,
                        tiled_mma_qk,
                        tiled_mma_pv,
                        load_q_consumer_state,
                        load_kv_consumer_state,
                        mma_s_producer_state,
                        p_mma_consumer_state,
                        mma_o_producer_state,
                    )
                tile_sched.advance_to_next_work()
                work_tile = tile_sched.get_current_work()

            mma_s_pipeline.producer_tail(mma_s_producer_state)
            mma_o_pipeline.producer_tail(mma_o_producer_state)

            tmem.relinquish_alloc_permit()
            tmem.free(tmem_ptr)
            if cutlass.const_expr(self.enable_pdl):
                cute.arch.griddepcontrol_launch_dependents()

        # ///////////////////////////////////////////////////////////////////////////////
        #  Compute warp
        # ///////////////////////////////////////////////////////////////////////////////
        if (
            warp_idx >= self.compute_warp_ids[0]
            and warp_idx <= self.compute_warp_ids[-1]
        ):
            _setmaxregister_increase(self.softmax_reg_num)
            mma_s_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.mma_s_stage
            )
            p_mma_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.p_mma_stage
            )
            p_cor_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.p_cor_stage
            )
            mma_o_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.mma_o_stage
            )
            # sync with mma warp before retrieving tmem ptr
            tmem.wait_for_alloc()

            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)

            tile_sched = create_mla_static_tile_scheduler(
                tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
            )
            work_tile = tile_sched.initial_work_tile_info()
            while work_tile.is_valid_tile:
                blk_coord = work_tile.tile_idx
                (
                    k_index,
                    k_tile_count,
                    local_split_kv,
                    _,
                    q_len,
                    _,
                ) = self.get_k_tile_count(
                    split_kv,
                    cache_seqs,
                    cum_seq_lens_q,
                    block_split_kvs,
                    blk_coord,
                )
                if k_tile_count > 0:
                    causal_global = cutlass.Int32(0)
                    if cutlass.const_expr(self.enable_dcp):
                        causal_global = causal_seqlens_kv_global[blk_coord[2]]
                    compute_common_params = SimpleNamespace(
                        blk_coord=blk_coord,
                        split_kv=split_kv,
                        local_split_kv=local_split_kv,
                        smem_exchange=softmax_smem_exchange,
                        mAccO=mAccO,
                        mO=mO,
                        K=cache_seqs[blk_coord[2]],
                        q_len=q_len,
                        causal_global=causal_global,
                        cp_rank=cp_rank,
                        L=mCL.shape[1],
                        tmem_ptr=tmem_ptr,
                        tidx=tidx,
                        p_cor_pipeline=p_cor_pipeline,
                    )
                    compute_softmax_params = SimpleNamespace(
                        tiled_mma_qk=tiled_mma_qk,
                        sP=sP,
                        mma_s_pipeline=mma_s_pipeline,
                        p_mma_pipeline=p_mma_pipeline,
                        softmax_scale_log2=softmax_scale_log2,
                    )
                    mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state = (
                        self.compute(
                            compute_common_params,
                            compute_softmax_params,
                            k_index=k_index,
                            k_tile_count=k_tile_count,
                            mma_s_consumer_state=mma_s_consumer_state,
                            p_mma_producer_state=p_mma_producer_state,
                            p_cor_producer_state=p_cor_producer_state,
                        )
                    )
                tile_sched.advance_to_next_work()
                work_tile = tile_sched.get_current_work()
            p_cor_pipeline.producer_tail(p_cor_producer_state)

        # ///////////////////////////////////////////////////////////////////////////////
        #  Correction warp
        # ///////////////////////////////////////////////////////////////////////////////
        if (
            warp_idx >= self.correction_warp_ids[0]
            and warp_idx <= self.correction_warp_ids[-1]
        ):
            _setmaxregister_increase(self.correction_reg_num)
            p_cor_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.p_cor_stage
            )
            mma_o_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.mma_o_stage
            )
            # sync with mma warp before retrieving tmem ptr
            tmem.wait_for_alloc()

            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)

            tile_sched = create_mla_static_tile_scheduler(
                tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
            )
            work_tile = tile_sched.initial_work_tile_info()
            while work_tile.is_valid_tile:
                blk_coord = work_tile.tile_idx
                (
                    k_index,
                    k_tile_count,
                    local_split_kv,
                    q_begin,
                    q_len,
                    valid_q_rows,
                ) = self.get_k_tile_count(
                    split_kv,
                    cache_seqs,
                    cum_seq_lens_q,
                    block_split_kvs,
                    blk_coord,
                )
                if k_tile_count > 0:
                    causal_global = cutlass.Int32(0)
                    if cutlass.const_expr(self.enable_dcp):
                        causal_global = causal_seqlens_kv_global[blk_coord[2]]
                    compute_common_params = SimpleNamespace(
                        blk_coord=blk_coord,
                        split_kv=split_kv,
                        local_split_kv=local_split_kv,
                        smem_exchange=epilogue_smem_exchange,
                        mAccO=mAccO,
                        mO=mO,
                        K=cache_seqs[blk_coord[2]],
                        q_begin=q_begin,
                        q_len=q_len,
                        causal_global=causal_global,
                        cp_rank=cp_rank,
                        split_start_key=k_index * self.mma_qk_tiler[1],
                        L=mCL.shape[1],
                        H=valid_q_rows,
                        tmem_ptr=tmem_ptr,
                        tidx=tidx,
                        tiled_mma_pv=tiled_mma_pv,
                        p_cor_pipeline=p_cor_pipeline,
                        mma_o_pipeline=mma_o_pipeline,
                    )
                    compute_epilogue_params = SimpleNamespace(
                        output_scale=output_scale,
                        softmax_scale_log2=softmax_scale_log2,
                        mAccLSE=mAccLSE,
                        mLSE=mLSE,
                    )
                    p_cor_consumer_state, mma_o_consumer_state = self.correction(
                        compute_common_params,
                        compute_epilogue_params,
                        k_tile_count=k_tile_count,
                        p_cor_consumer_state=p_cor_consumer_state,
                        mma_o_consumer_state=mma_o_consumer_state,
                    )
                elif cutlass.const_expr(self.enable_dcp):
                    self.store_empty_dcp_work(
                        mO,
                        mLSE,
                        mAccO,
                        mAccLSE,
                        blk_coord,
                        self.get_valid_q_rows(blk_coord[1]),
                        tidx,
                    )
                tile_sched.advance_to_next_work()
                work_tile = tile_sched.get_current_work()

        return

    @cute.jit
    def store_empty_dcp_work(
        self,
        mO: Optional[cute.Tensor],
        mLSE: Optional[cute.Tensor],
        mAccO: Optional[cute.Tensor],
        mAccLSE: Optional[cute.Tensor],
        blk_coord: cute.Coord,
        valid_q_rows: cutlass.Int32,
        tidx: cutlass.Int32,
    ):
        """Materialize the neutral attention state for a split with no K tile."""
        cta_m_rows = self.mma_qk_tiler[0] // self.cluster_shape_mnk[0]
        compute_threads = self.num_compute_warps * self.threads_per_warp
        local_tidx = tidx % compute_threads
        cta_row_base = blk_coord[0] * cta_m_rows

        for linear_idx in cutlass.range(
            local_tidx, cta_m_rows * self.latent_dim, compute_threads
        ):
            cta_row = linear_idx // self.latent_dim
            d_idx = linear_idx % self.latent_dim
            q_row = cta_row_base + cta_row
            if cute.elem_less(q_row, valid_q_rows):
                if cutlass.const_expr(mAccO is None):
                    mO[q_row, d_idx, blk_coord[1], blk_coord[2]] = self.o_dtype(0.0)
                else:
                    mAccO[
                        q_row,
                        blk_coord[3],
                        d_idx,
                        blk_coord[1],
                        blk_coord[2],
                    ] = self.acc_dtype(0.0)

        if cute.elem_less(local_tidx, cta_m_rows):
            q_row = cta_row_base + local_tidx
            if cute.elem_less(q_row, valid_q_rows):
                if cutlass.const_expr(mAccLSE is None):
                    mLSE[q_row, blk_coord[1], blk_coord[2]] = -self.lse_dtype.inf
                else:
                    mAccLSE[
                        q_row, blk_coord[3], blk_coord[1], blk_coord[2]
                    ] = -self.lse_dtype.inf

    @cute.jit
    def get_valid_q_rows(self, q_tile_idx: cutlass.Int32) -> cutlass.Int32:
        """Number of valid flattened query rows in one fixed-Q M tile."""
        valid_rows = self.mma_qk_tiler_mn[0]
        if q_tile_idx == self.num_q_tiles - 1:
            valid_rows = self.tail_q_rows
        return valid_rows

    @cute.kernel
    def reduction_kernel(
        self,
        mO: cute.Tensor,
        mLSE: cute.Tensor,
        mAccO: cute.Tensor,
        mAccLSE: cute.Tensor,
        split_kv: cutlass.Int32,
        cache_seqs: cute.Tensor,
        cum_seq_lens_q: Optional[cute.Tensor],
        block_split_kvs: cute.Tensor,
    ):
        """The reduction kernel for Multi-Head Latent Attention (MLA) that combines intermediate results
        from multiple split_kv blocks into final outputs.

        :param mO: Output tensor for storing final results
        :type mO: cute.Tensor
        :param mLSE: Log-sum-exp tensor for storing final LSE values
        :type mLSE: cute.Tensor
        :param mAccO: Accumulated output tensor from split_kv blocks
        :type mAccO: cute.Tensor
        :param mAccLSE: Accumulated LSE tensor from split_kv blocks
        :type mAccLSE: cute.Tensor
        :param split_kv: Number of split_kv blocks
        :type split_kv: cutlass.Int32
        :param cache_seqs: Cache sequence lengths tensor
        :type cache_seqs: cute.Tensor
        :param block_split_kvs: Per-block split_kv values tensor (for variable split_kv)
        :type block_split_kvs: cute.Tensor
        """
        bidx, bidy, bidz = cute.arch.block_idx()
        tidx, _, _ = cute.arch.thread_idx()
        # Reducer blocks cover one D band of one logical output row.
        d_tile_idx = bidx % self.reducer_d_tiles
        head_idx = bidx // self.reducer_d_tiles
        q_begin = cutlass.Int32(0)
        q_len = cutlass.Int32(self.seq_len_q)
        is_active = True
        if cutlass.const_expr(self.is_var_q):
            q_begin, q_len, _ = get_variable_query_tile_info(
                self.num_heads,
                cutlass.Int32(0),
                bidz,
                cum_seq_lens_q,
                self.mma_qk_tiler[0],
            )
            is_active = bidy < q_len

        # In the variable-Q specialization, inactive CTAs still execute the
        # balanced PDL protocol but do not touch workspace or outputs.
        smem = utils.SmemAllocator()
        storage = smem.allocate(self.reducer_max_splits * self.acc_dtype.width // 8, 16)
        lse_scale_ptr = cute.recast_ptr(storage, dtype=self.acc_dtype)
        smem_lse_scale = cute.make_tensor(
            lse_scale_ptr, cute.make_layout(self.reducer_max_splits)
        )

        if cutlass.const_expr(self.enable_pdl):
            cute.arch.griddepcontrol_wait()
        if is_active:
            flat_q_row = bidy * self.num_heads + head_idx
            # The physical M tile is fixed at 128 rows, so map with shift/mask.
            q_tile = flat_q_row >> 7
            q_tile_row = flat_q_row & 127
            local_split_kv = split_kv
            if cutlass.const_expr(self.is_var_split_kv):
                local_split_kv = block_split_kvs[bidz]
            k_tile_total = cute.ceil_div(cache_seqs[bidz], self.mma_qk_tiler[1])
            if cutlass.const_expr(self.enable_dcp):
                k_tile_per_cta = cutlass.max(
                    cute.ceil_div(k_tile_total, local_split_kv), cutlass.Int32(1)
                )
                local_split_kv = cutlass.max(
                    cute.ceil_div(k_tile_total, k_tile_per_cta), cutlass.Int32(1)
                )
            else:
                k_tile_per_cta = cute.ceil_div(k_tile_total, local_split_kv)
                local_split_kv = cute.ceil_div(k_tile_total, k_tile_per_cta)

            gLSE = mAccLSE[q_tile_row, None, q_tile, bidz]
            warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
            if warp_idx == 0:
                # Calculate the global LSE and exp2(local_lse - global_lse).
                lse_per_thread = cute.ceil_div(
                    self.reducer_max_splits, self.threads_per_warp
                )

                local_lse = cute.make_rmem_tensor(
                    cute.make_layout(lse_per_thread), self.lse_dtype
                )
                lse_max = -self.lse_dtype.inf
                for i in cutlass.range_constexpr(lse_per_thread):
                    split_kv_idx = tidx + i * self.threads_per_warp
                    local_lse[i] = (
                        gLSE[split_kv_idx]
                        if cute.elem_less(split_kv_idx, local_split_kv)
                        else -self.lse_dtype.inf
                    )
                    lse_max = cute.arch.fmax(lse_max, local_lse[i])
                lse_max = cute.arch.warp_reduction_max(lse_max)
                if cutlass.const_expr(self.enable_dcp):
                    has_valid_lse = lse_max != -self.lse_dtype.inf
                    lse_max = lse_max if has_valid_lse else 0.0
                else:
                    lse_max = lse_max if lse_max != -self.lse_dtype.inf else 0.0
                sum_lse = 0.0
                for i in cutlass.range_constexpr(lse_per_thread):
                    sum_lse += cute.math.exp2(local_lse[i] - lse_max, fastmath=True)
                sum_lse = cute.arch.warp_reduction_sum(sum_lse)
                if cutlass.const_expr(self.enable_dcp):
                    global_lse = (
                        lse_max + cute.math.log2(sum_lse, fastmath=True)
                        if has_valid_lse
                        else -self.lse_dtype.inf
                    )
                else:
                    global_lse = (
                        lse_max + cute.math.log2(sum_lse, fastmath=True)
                        if sum_lse != self.lse_dtype(0.0) or sum_lse != sum_lse
                        else self.lse_dtype.inf
                    )
                if d_tile_idx == 0:
                    if tidx == 0:
                        # Convert the internal log2 value to natural log.
                        if cutlass.const_expr(self.is_var_q):
                            mLSE[head_idx, q_begin + bidy] = global_lse * (1.0 / LOG2_E)
                        else:
                            mLSE[head_idx, bidy, bidz] = global_lse * (1.0 / LOG2_E)
                for i in cutlass.range_constexpr(lse_per_thread):
                    split_kv_idx = tidx + i * self.threads_per_warp
                    if cute.elem_less(split_kv_idx, local_split_kv):
                        if cutlass.const_expr(self.enable_dcp):
                            smem_lse_scale[split_kv_idx] = (
                                cute.math.exp2(local_lse[i] - global_lse, fastmath=True)
                                if has_valid_lse
                                else self.acc_dtype(0.0)
                            )
                        else:
                            smem_lse_scale[split_kv_idx] = cute.math.exp2(
                                local_lse[i] - global_lse, fastmath=True
                            )

            pipeline.sync(barrier_id=4)

            elements_per_thread = cute.ceil_div(
                self.reducer_d_tile,
                self.threads_per_warp * self.num_compute_warps,
            )
            gAccO = mAccO[q_tile_row, None, None, q_tile, bidz]
            rAccO = cute.make_rmem_tensor(
                cute.make_layout(elements_per_thread), self.acc_dtype
            )
            rO = cute.make_rmem_tensor(
                cute.make_layout(elements_per_thread), self.o_dtype
            )
            rAccO.fill(0.0)
            for i in range(local_split_kv):
                for j in cutlass.range_constexpr(elements_per_thread):
                    element_idx = (
                        d_tile_idx * self.reducer_d_tile
                        + tidx
                        + j * self.threads_per_warp * self.num_compute_warps
                    )
                    rAccO[j] += gAccO[i, element_idx] * smem_lse_scale[i]
            rO.store(rAccO.load().to(self.o_dtype))
            for j in cutlass.range_constexpr(elements_per_thread):
                element_idx = (
                    d_tile_idx * self.reducer_d_tile
                    + tidx
                    + j * self.threads_per_warp * self.num_compute_warps
                )
                if cutlass.const_expr(self.is_var_q):
                    mO[head_idx, element_idx, q_begin + bidy] = rO[j]
                else:
                    mO[head_idx, element_idx, bidy, bidz] = rO[j]
        if cutlass.const_expr(self.enable_pdl):
            cute.arch.griddepcontrol_launch_dependents()
        return

    @staticmethod
    def get_split_kv(
        B: int, S: int, K: int, mma_qk_tiler_mn: tuple, max_active_blocks: int
    ) -> int:
        """Get the proper split_kv value for the MLA kernel based on parameters.

        :param B: Batch size
        :type B: int
        :param S: Sequence length
        :type S: int
        :param K: Sequence length
        :type K: int
        :param mma_qk_tiler_mn: MLA tiling parameters
        :type mma_qk_tiler_mn: tuple
        :param max_active_blocks: Maximum number of active blocks
        :type max_active_blocks: int
        :return: Split_kv value
        :rtype: int
        """
        max_splits = ceil_div(K, mma_qk_tiler_mn[1])
        blocks_per_batch = max(1, max_active_blocks // B // (S * 2))
        split_heur = min(max_splits, blocks_per_batch)
        # {$nv-internal-release begin}
        # TODO: figure out the error of make_tile with dynamic int_tuple
        # {$nv-internal-release end}
        k_waves = ceil_div(max_splits, split_heur)
        split_wave_aware = ceil_div(max_splits, k_waves)
        max_split_kv = 32
        return min(split_wave_aware, max_split_kv)

    @staticmethod
    def get_split_kv_simplified(B: int, S: int, max_active_blocks: int) -> int:
        blocks_per_batch = max(1, max_active_blocks // B // (S * 2))
        max_split_kv = 32
        return min(blocks_per_batch, max_split_kv)

    @cute.jit
    def get_k_tile_count(
        self,
        split_kv: cutlass.Int32,
        cache_seqs: cute.Tensor,
        cum_seq_lens_q: Optional[cute.Tensor],
        block_split_kvs: cute.Tensor,
        blk_coord: cute.Coord,
    ) -> tuple[
        cutlass.Int32,
        cutlass.Int32,
        cutlass.Int32,
        cutlass.Int32,
        cutlass.Int32,
        cutlass.Int32,
    ]:
        """Get K split and request-local query metadata for one work item.

        :param split_kv: Split_kv value
        :type split_kv: cutlass.Int32
        :param cache_seqs: Cache sequence lengths tensor
        :type cache_seqs: cute.Tensor
        :param cum_seq_lens_q: Compact-query indptr, or None for fixed Q
        :type cum_seq_lens_q: cute.Tensor
        :param block_split_kvs: Per-block split_kv values tensor
        :type block_split_kvs: cute.Tensor
        :param blk_coord: Block coordinate
        :type blk_coord: cute.Coord
        :return: K range, effective split count, and request-local Q metadata
        """
        K = cache_seqs[blk_coord[2]]
        if cutlass.const_expr(self.is_var_split_kv):
            split_kv = block_split_kvs[blk_coord[2]]

        k_tile_total = cute.ceil_div(K, self.mma_qk_tiler[1])
        # {$nv-internal-release begin}
        # TODO: figure out the error of make_tile with dynamic int_tuple
        # {$nv-internal-release end}
        k_tile_per_cta = cute.ceil_div(k_tile_total, split_kv)
        k_index = blk_coord[3] * k_tile_per_cta
        k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index)
        q_begin = cutlass.Int32(0)
        q_len = cutlass.Int32(self.seq_len_q)
        if cutlass.const_expr(self.is_var_q):
            q_begin, q_len, valid_q_rows = get_variable_query_tile_info(
                self.num_heads,
                blk_coord[1],
                blk_coord[2],
                cum_seq_lens_q,
                self.mma_qk_tiler[0],
            )
            if valid_q_rows == 0:
                # Keep the rectangular scheduler slot valid.  A persistent
                # cluster may encounter a later active request after this
                # no-op slot.
                k_tile_count = cutlass.Int32(0)
        else:
            valid_q_rows = self.get_valid_q_rows(blk_coord[1])
        return (
            k_index,
            k_tile_count,
            split_kv,
            q_begin,
            q_len,
            valid_q_rows,
        )

    @cute.jit
    def load_page_table(
        self,
        common_params: SimpleNamespace,
        k_index: cutlass.Int32,
        k_tile_count: cutlass.Int32,
        load_pt_producer_state: pipeline.PipelineState,
    ) -> pipeline.PipelineState:
        """Load warp to load page table. Updates the load pt producer state.

        :param common_params: The common parameters
        :type common_params: SimpleNamespace
        :param k_index: The k index
        :type k_index: cutlass.Int32
        :param k_tile_count: The k tile count
        :type k_tile_count: cutlass.Int32
        :param load_pt_producer_state: The load pt producer state
        :type load_pt_producer_state: pipeline.PipelineState

        :return: The load pt producer state
        :rtype: pipeline.PipelineState
        """
        mPT = common_params.mPT[None, common_params.blk_coord[2]]
        page_per_tile = self.mma_qk_tiler[1] // self.page_size
        tidx = common_params.tidx % self.threads_per_warp

        load_pt_pipeline = common_params.load_pt_pipeline
        while k_tile_count > 0:
            load_pt_pipeline.producer_acquire(load_pt_producer_state)

            elem_per_thread = cute.ceil_div(page_per_tile, self.threads_per_warp)

            # atom_async_copy: async copy atom for page table load
            atom_async_copy = cute.make_copy_atom(
                cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.ALWAYS),
                cutlass.Int32,
                num_bits_per_copy=cutlass.Int32.width,
            )
            mPT_for_copy = cute.flat_divide(mPT, (1,))
            sPT_for_copy = cute.flat_divide(common_params.sPT, (1,))
            # elem_per_thread is a dynamic value depends on the page_size setting.
            for i in range(elem_per_thread):
                idx = i * self.threads_per_warp + tidx
                if cute.elem_less(
                    k_index * page_per_tile + idx, mPT.shape[0]
                ) and cute.elem_less(idx, page_per_tile):
                    cute.copy(
                        atom_async_copy,
                        mPT_for_copy[None, k_index * page_per_tile + idx],
                        sPT_for_copy[None, idx, load_pt_producer_state.index],
                    )
                else:
                    sPT_for_copy[None, idx, load_pt_producer_state.index].fill(0)
            mbar_ptr = load_pt_pipeline.producer_get_barrier(load_pt_producer_state)  # noqa: F841
            load_pt_pipeline.producer_commit(load_pt_producer_state)
            load_pt_producer_state.advance()
            k_index += 1
            k_tile_count -= 1

        return load_pt_producer_state

    @cute.jit
    def load_tma(
        self,
        common_params: SimpleNamespace,
        qk_params: SimpleNamespace,
        v_params: SimpleNamespace,
        k_index: cutlass.Int32,
        k_tile_count: cutlass.Int32,
        load_q_producer_state: pipeline.PipelineState,
        load_kv_producer_state: pipeline.PipelineState,
        load_pt_consumer_state: pipeline.PipelineState,
        load_pt_release_state: pipeline.PipelineState,
    ) -> tuple[
        pipeline.PipelineState,
        pipeline.PipelineState,
        pipeline.PipelineState,
        pipeline.PipelineState,
    ]:
        """Load wrap to load Q/C latent/rope tensors. Updates the load qkv producer state.

        :param common_params: The common parameters
        :type common_params: SimpleNamespace
        :param qk_params: The qk parameters
        :type qk_params: SimpleNamespace
        :param v_params: The v parameters
        :type v_params: SimpleNamespace
        :param k_index: The k index
        :type k_index: cutlass.Int32
        :param k_tile_count: The k tile count
        :type k_tile_count: cutlass.Int32
        :param load_q_producer_state: The load q producer state
        :type load_q_producer_state: pipeline.PipelineState
        :param load_kv_producer_state: The load kv producer state
        :type load_kv_producer_state: pipeline.PipelineState
        :param load_pt_consumer_state: The load pt consumer state
        :type load_pt_consumer_state: pipeline.PipelineState
        :param load_pt_release_state: The load pt release state
        :type load_pt_release_state: pipeline.PipelineState

        :return: The load q producer state, load kv producer state, load pt consumer state, and load pt release state
        :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState]
        """
        # page table
        mPT = common_params.mPT[None, common_params.blk_coord[2]]

        if cutlass.const_expr(self.is_var_q):
            # Offset the compact global descriptor to this request.  A middle
            # request's tail can read the next request's rows; output stores
            # remain predicated by this request's valid-row count.
            mQL = cute.domain_offset(
                (common_params.q_begin * self.num_heads, 0, 0),
                qk_params.mQL,
            )
            mQR = cute.domain_offset(
                (common_params.q_begin * self.num_heads, 0, 0),
                qk_params.mQR,
            )
        else:
            mQL = qk_params.mQL
            mQR = qk_params.mQR
        # (bM, bK, rM, rK, rL)
        mma_qk_tiler_mk = cute.select(self.mma_qk_tiler, mode=[0, 2])
        gQL = cute.flat_divide(mQL, mma_qk_tiler_mk)
        mma_qk_tiler_mk_rope = cute.select(self.mma_qk_rope_tiler, mode=[0, 2])
        gQR = cute.flat_divide(mQR, mma_qk_tiler_mk_rope)

        thr_mma_qk = qk_params.tiled_mma_qk.get_slice(
            common_params.blk_coord[0] % cute.size(qk_params.tiled_mma_qk.thr_id)
        )
        tSgQL = thr_mma_qk.partition_A(gQL)
        tSgQR = thr_mma_qk.partition_A(gQR)
        cta_m = min(
            qk_params.tiled_mma_qk.op.shape_mnk[0]
            // qk_params.tiled_mma_qk.thr_id.shape,
            self.page_size,
        )
        page_tile_size = min(self.page_size, cta_m)
        gCL = cute.tiled_divide(qk_params.mCL, (page_tile_size, self.mma_qk_tiler[2]))
        tSgCL = (
            gCL[
                None,
                common_params.blk_coord[0] % qk_params.tiled_mma_qk.thr_id.shape,
                None,
                None,
            ]
            if cta_m < self.page_size
            else gCL[None, 0, None, None]
        )
        gKR = cute.tiled_divide(qk_params.mKR, (page_tile_size, self.mma_qk_tiler[2]))
        tSgKR = (
            gKR[
                None,
                common_params.blk_coord[0] % qk_params.tiled_mma_qk.thr_id.shape,
                None,
                None,
            ]
            if cta_m < self.page_size
            else gKR[None, 0, None, None]
        )

        # tma partition for q, k latent/rope
        # smem: ((atom_v, rest_v), STAGE)
        # gmem: ((atom_v, rest_v), RestM, RestK, RestL)
        tQsQ, tQLgQL_mkl = cpasync.tma_partition(
            qk_params.tma_atom_q_latent,
            0,
            cute.make_layout(1),
            cute.group_modes(qk_params.sQ, 0, 3),
            cute.group_modes(tSgQL, 0, 3),
        )

        tQsQ_rope, tQRgQR_mkl = cpasync.tma_partition(
            qk_params.tma_atom_q_rope,
            0,
            cute.make_layout(1),
            cute.group_modes(qk_params.sQ_rope, 0, 3),
            cute.group_modes(tSgQR, 0, 3),
        )
        tKCsKC, tCLgCL = cpasync.tma_partition(
            qk_params.tma_atom_c_latent,
            0,
            cute.make_layout(1),
            qk_params.sKC,
            tSgCL,
        )

        _, tKRgKR = cpasync.tma_partition(
            qk_params.tma_atom_c_rope,
            0,
            cute.make_layout(1),
            qk_params.sKC,
            tSgKR,
        )

        if cutlass.const_expr(self.is_var_q):
            tQLgQL = tQLgQL_mkl[None, common_params.blk_coord[1], None, 0]
            tQRgQR = tQRgQR_mkl[None, common_params.blk_coord[1], None, 0]
        else:
            # Fixed Q retains its descriptor batch mode.
            tQLgQL = tQLgQL_mkl[
                None,
                common_params.blk_coord[1],
                None,
                common_params.blk_coord[2],
            ]
            tQRgQR = tQRgQR_mkl[
                None,
                common_params.blk_coord[1],
                None,
                common_params.blk_coord[2],
            ]

        # Flatten divide and partition global tensors for V TMA load
        page_tile_size = min(self.page_size, self.mma_pv_tiler[2])
        gCLT = cute.flat_divide(v_params.mCLT, (self.mma_pv_tiler[1], page_tile_size))
        cta_n = self.mma_pv_tiler[1] // v_params.tiled_mma_pv.thr_id.shape
        gCLT = cute.logical_divide(gCLT, (cta_n,))[
            (None, common_params.blk_coord[0]), None, None, None, None
        ]
        tOgCLT = cute.tiled_divide(gCLT, (cta_n, page_tile_size))
        tOgCLT = tOgCLT[None, 0, 0, None, None, None]

        # tma partition for vc
        # smem: ((atom_v, rest_v), STAGE)
        # gmem: ((atom_v, rest_v), RestM, RestK, RestL)
        tVCsVC, tCLTgCLT = cpasync.tma_partition(
            v_params.tma_atom_c_latent_transpose,
            0,
            cute.make_layout(1),
            v_params.sVC,
            tOgCLT,
        )

        # set extra params
        common_params.mPT = mPT
        qk_params.tQLgQL = tQLgQL
        qk_params.tQRgQR = tQRgQR
        qk_params.tCLgCL = tCLgCL
        qk_params.tKRgKR = tKRgKR
        qk_params.tQsQ = tQsQ
        qk_params.tQsQ_rope = tQsQ_rope
        qk_params.tKCsKC = tKCsKC
        v_params.tCLTgCLT = tCLTgCLT
        v_params.tVCsVC = tVCsVC

        load_q_producer_state, load_kv_producer_state, load_pt_consumer_state = (
            self.load_tma_qk_one_k_tile(
                common_params,
                qk_params,
                k_index,
                k_tile_count,
                load_q_producer_state,
                load_kv_producer_state,
                load_pt_consumer_state,
                load_q=True,
            )
        )
        k_index += 1
        k_tile_count -= 1
        while k_tile_count > 0:
            # {$nv-internal-release begin}
            # TODO: figure out how to support SingleNamespace/struct in ast
            # {$nv-internal-release end}
            load_q_producer_state, load_kv_producer_state, load_pt_consumer_state = (
                self.load_tma_qk_one_k_tile(
                    common_params,
                    qk_params,
                    k_index,
                    k_tile_count,
                    load_q_producer_state,
                    load_kv_producer_state,
                    load_pt_consumer_state,
                    load_q=False,
                )
            )
            load_kv_producer_state, load_pt_release_state = self.load_tma_v_one_k_tile(
                common_params,
                v_params,
                k_index - 1,
                load_kv_producer_state,
                load_pt_release_state,
            )
            k_index += 1
            k_tile_count -= 1

        # load last v tile
        load_kv_producer_state, load_pt_release_state = self.load_tma_v_one_k_tile(
            common_params,
            v_params,
            k_index - 1,
            load_kv_producer_state,
            load_pt_release_state,
        )
        return (
            load_q_producer_state,
            load_kv_producer_state,
            load_pt_consumer_state,
            load_pt_release_state,
        )

    @cute.jit
    def issue_q_tma_load(
        self,
        tma_atom_q_latent: cute.CopyAtom,
        tQLgQL: cute.Tensor,
        tQsQ: cute.Tensor,
        tma_atom_q_rope: cute.CopyAtom,
        tQRgQR: cute.Tensor,
        tQsQ_rope: cute.Tensor,
        tma_bar_ptr,
    ):
        """Issue one bounded flattened query-tile load into shared Q stages."""
        for i in cutlass.range(self.iterations_qk_latent):
            cute.copy(
                tma_atom_q_latent,
                tQLgQL[None, i],
                tQsQ[None, (i, 0)],
                tma_bar_ptr=tma_bar_ptr,
            )
        for i in cutlass.range(self.iterations_qk_rope):
            cute.copy(
                tma_atom_q_rope,
                tQRgQR[None, i],
                tQsQ_rope[None, i],
                tma_bar_ptr=tma_bar_ptr,
            )

    @cute.jit
    def load_tma_qk_one_k_tile(
        self,
        common_params: SimpleNamespace,
        qk_params: SimpleNamespace,
        k_index: cutlass.Int32,
        k_tile_count: cutlass.Int32,
        load_q_producer_state: pipeline.PipelineState,
        load_kv_producer_state: pipeline.PipelineState,
        load_pt_consumer_state: pipeline.PipelineState,
        load_q: bool,
    ) -> tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState]:
        """Load one k-tile of Q/C latent/rope tensors. Updates the load qkv producer state.

        :param common_params: The common parameters
        :type common_params: SimpleNamespace
        :param qk_params: The qk parameters
        :type qk_params: SimpleNamespace
        :param k_index: The k index
        :type k_index: cutlass.Int32
        :param k_tile_count: The k tile count
        :type k_tile_count: cutlass.Int32
        :param load_q_producer_state: The load q producer state
        :type load_q_producer_state: pipeline.PipelineState
        :param load_kv_producer_state: The load kv producer state
        :type load_kv_producer_state: pipeline.PipelineState
        :param load_pt_consumer_state: The load pt consumer state
        :type load_pt_consumer_state: pipeline.PipelineState
        :param load_q: Whether to load q
        :type load_q: bool

        :return: The load q producer state, load kv producer state, and load pt consumer state
        :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState]
        """
        page_per_tile = ceil_div(
            self.mma_qk_tiler[1] // self.page_size, qk_params.tiled_mma_qk.thr_id.shape
        )
        common_params.load_pt_pipeline.consumer_wait(load_pt_consumer_state)
        page_table_stage = load_pt_consumer_state.index
        load_pt_consumer_state.advance()
        k_idx = cute.make_rmem_tensor(cute.make_layout(page_per_tile), cutlass.Int32)
        for i in cutlass.range_constexpr(page_per_tile):
            k_idx[i] = (
                common_params.sPT[0, page_table_stage]
                if self.mma_qk_tiler[1] // self.page_size == 1
                else common_params.sPT[
                    i + common_params.blk_coord[0] * page_per_tile, page_table_stage
                ]
            )
        # load q once at first iteration
        if cutlass.const_expr(load_q):
            common_params.load_q_pipeline.producer_acquire(load_q_producer_state)
            # get the mbar ptr from pipeline.
            tma_bar_ptr = common_params.load_q_pipeline.producer_get_barrier(
                load_q_producer_state
            )
            self.issue_q_tma_load(
                qk_params.tma_atom_q_latent,
                qk_params.tQLgQL,
                qk_params.tQsQ,
                qk_params.tma_atom_q_rope,
                qk_params.tQRgQR,
                qk_params.tQsQ_rope,
                tma_bar_ptr,
            )
            load_q_producer_state.advance()
        load_kv_pipeline = common_params.load_kv_pipeline
        tma_bar_ptr = load_kv_pipeline.producer_get_barrier(load_kv_producer_state)
        for i in cutlass.range(self.iterations_qk_latent):
            # get the mbar ptr from pipeline.
            tma_bar_ptr = load_kv_pipeline.producer_get_barrier(load_kv_producer_state)
            load_kv_pipeline.producer_acquire(load_kv_producer_state)
            for k in cutlass.range(page_per_tile):
                # load k latent
                cute.copy(
                    qk_params.tma_atom_c_latent,
                    qk_params.tCLgCL[None, i, k_idx[k]],
                    qk_params.tKCsKC[None, k, 0, load_kv_producer_state.index],
                    tma_bar_ptr=tma_bar_ptr,
                )
            load_kv_producer_state.advance()

        for i in cutlass.range(self.iterations_qk_rope):
            # get the mbar ptr from pipeline.
            tma_bar_ptr = load_kv_pipeline.producer_get_barrier(load_kv_producer_state)
            load_kv_pipeline.producer_acquire(load_kv_producer_state)
            for k in cutlass.range(page_per_tile):
                # load k rope
                cute.copy(
                    qk_params.tma_atom_c_rope,
                    qk_params.tKRgKR[None, i, k_idx[k]],
                    qk_params.tKCsKC[None, k, 0, load_kv_producer_state.index],
                    tma_bar_ptr=tma_bar_ptr,
                )
            load_kv_producer_state.advance()

        return load_q_producer_state, load_kv_producer_state, load_pt_consumer_state

    @cute.jit
    def load_tma_v_one_k_tile(
        self,
        common_params: SimpleNamespace,
        v_params: SimpleNamespace,
        k_index: cutlass.Int32,
        load_kv_producer_state: pipeline.PipelineState,
        load_pt_release_state: pipeline.PipelineState,
    ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]:
        """Load one k-tile of compressed latent transpose tensor(v). Updates the load qkv producer state.

        :param common_params: The common parameters
        :type common_params: SimpleNamespace
        :param v_params: The load tma v parameters
        :type v_params: SimpleNamespace
        :param k_index: The k index
        :type k_index: cutlass.Int32
        :param load_kv_producer_state: The load qkv producer state
        :type load_kv_producer_state: pipeline.PipelineState
        :param load_pt_release_state: The load pt release state
        :type load_pt_release_state: pipeline.PipelineState

        :return: The load kv producer state and load pt release state
        :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState]
        """
        page_per_tile = self.mma_pv_tiler[2] * self.iterations_pv_k // self.page_size
        page_per_subtile = ceil_div(page_per_tile, self.iterations_pv_k)
        k_idx = cute.make_rmem_tensor(cute.make_layout(page_per_tile), cutlass.Int32)
        page_table_stage = load_pt_release_state.index
        for i in cutlass.range(page_per_tile):
            k_idx[i] = (
                common_params.sPT[0, page_table_stage]
                if page_per_tile == 1
                else common_params.sPT[i, page_table_stage]
            )
        common_params.load_pt_pipeline.consumer_release(load_pt_release_state)
        load_pt_release_state.advance()
        load_kv_pipeline = common_params.load_kv_pipeline
        tma_bar_ptr = load_kv_pipeline.producer_get_barrier(load_kv_producer_state)
        for i in cutlass.range(self.iterations_pv_k):
            for j in cutlass.range(self.iterations_pv_n):
                # get the mbar ptr from pipeline.
                tma_bar_ptr = load_kv_pipeline.producer_get_barrier(
                    load_kv_producer_state
                )
                load_kv_pipeline.producer_acquire(load_kv_producer_state)
                for k in cutlass.range(page_per_subtile):
                    k_idx_i = k_idx[
                        k
                        + i
                        // ceil_div(self.iterations_pv_k, page_per_tile)
                        * page_per_subtile
                    ]
                    cute.copy(
                        v_params.tma_atom_c_latent_transpose,
                        v_params.tCLTgCLT[
                            None,
                            j,
                            i % ceil_div(self.iterations_pv_k, page_per_tile),
                            k_idx_i,
                        ],
                        v_params.tVCsVC[None, 0, k, load_kv_producer_state.index],
                        tma_bar_ptr=tma_bar_ptr,
                    )

                load_kv_producer_state.advance()
        return load_kv_producer_state, load_pt_release_state

    @cute.jit
    def mma(
        self,
        common_params: SimpleNamespace,
        qk_params: SimpleNamespace,
        pv_params: SimpleNamespace,
        k_tile_count: cutlass.Int32,
        tiled_mma_qk: cute.TiledMma,
        tiled_mma_pv: cute.TiledMma,
        load_q_consumer_state: pipeline.PipelineState,
        load_kv_consumer_state: pipeline.PipelineState,
        mma_s_producer_state: pipeline.PipelineState,
        p_mma_consumer_state: pipeline.PipelineState,
        mma_o_producer_state: pipeline.PipelineState,
    ) -> tuple[
        cute.TiledMma,
        cute.TiledMma,
        pipeline.PipelineState,
        pipeline.PipelineState,
        pipeline.PipelineState,
        pipeline.PipelineState,
    ]:
        """MMA warp to compute the result of Q*K^T and P*V. Updates the tiled mma and pipeline states.

        :param common_params: The common parameters for mma qk and pv
        :type common_params: SimpleNamespace
        :param qk_params: The mma qk parameters
        :type qk_params: SimpleNamespace
        :param pv_params: The mma pv parameters
        :type pv_params: SimpleNamespace
        :param k_tile_count: The k tile count
        :type k_tile_count: cutlass.Int32
        :param tiled_mma_qk: The tiled mma qk
        :type tiled_mma_qk: cute.TiledMma
        :param tiled_mma_pv: The tiled mma pv
        :type tiled_mma_pv: cute.TiledMma
        :param load_q_consumer_state: The load q consumer state
        :type load_q_consumer_state: pipeline.PipelineState
        :param load_kv_consumer_state: The load kv consumer state
        :type load_kv_consumer_state: pipeline.PipelineState
        :param mma_s_producer_state: The mma s producer state
        :type mma_s_producer_state: pipeline.PipelineState
        :param p_mma_consumer_state: The p mma consumer state
        :type p_mma_consumer_state: pipeline.PipelineState
        :param mma_o_producer_state: The mma o producer state
        :type mma_o_producer_state: pipeline.PipelineState

        :return: The tiled mma qk, the tiled mma pv, the load q consumer state, the load kv consumer state, the mma s producer state, the p mma consumer state, and the mma o producer state
        :rtype: tuple[cute.TiledMma, cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState]
        """

        tSrQ = tiled_mma_qk.make_fragment_A(qk_params.sQ)
        tSrQ_rope = tiled_mma_qk.make_fragment_A(qk_params.sQ_rope)
        tSrKC = tiled_mma_qk.make_fragment_B(qk_params.sKC)
        tOrP = tiled_mma_pv.make_fragment_A(pv_params.sP)
        tOrVC = tiled_mma_pv.make_fragment_B(pv_params.sVC)

        tStS_shape = tiled_mma_qk.partition_shape_C(
            cute.select(self.mma_qk_tiler, mode=[0, 1])
        )
        tStS_staged_fake = tiled_mma_qk.make_fragment_C(
            cute.append(tStS_shape, self.mma_s_stage)
        )
        # use real tmem ptr for tStS
        tStS_staged = cute.make_tensor(common_params.tmem_ptr, tStS_staged_fake.layout)
        tOtO_shape = tiled_mma_pv.partition_shape_C(
            cute.select(self.mma_pv_tiler, mode=[0, 1])
        )
        # mma O has 1 stage.
        tOtO = tiled_mma_pv.make_fragment_C(tOtO_shape)
        tOtO_layout = cute.append(
            tOtO.layout,
            cute.make_layout(
                common_params.L // self.mma_pv_tiler[1],
                stride=self.mma_pv_tiler[1] // self.warps_in_n,
            ),
        )
        tOtO_staged = cute.make_tensor(
            tStS_staged.iterator + self.tmem_o_offset, tOtO_layout
        )

        # set more parameters
        qk_params.tSrQ = tSrQ
        qk_params.tSrQ_rope = tSrQ_rope
        qk_params.tSrKC = tSrKC
        qk_params.tStS_staged = tStS_staged
        pv_params.tOrP = tOrP
        pv_params.tOrVC = tOrVC
        pv_params.tOtO_staged = tOtO_staged

        # mma O accumulates on K, so the accumlate flag is set to False once before all K blocks.
        tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, False)
        load_q_pipeline = common_params.load_q_pipeline
        if common_params.is_leader_cta:
            load_q_release_state = load_q_consumer_state.clone()

            (
                tiled_mma_qk,
                load_q_consumer_state,
                load_kv_consumer_state,
                mma_s_producer_state,
            ) = self.mma_qk(
                common_params,
                qk_params,
                tiled_mma_qk,
                load_q_consumer_state,
                load_kv_consumer_state,
                mma_s_producer_state,
                wait_q=True,
            )
            k_tile_count -= 1
            while k_tile_count > 0:
                (
                    tiled_mma_qk,
                    load_q_consumer_state,
                    load_kv_consumer_state,
                    mma_s_producer_state,
                ) = self.mma_qk(
                    common_params,
                    qk_params,
                    tiled_mma_qk,
                    load_q_consumer_state,
                    load_kv_consumer_state,
                    mma_s_producer_state,
                    wait_q=False,
                )
                (
                    tiled_mma_pv,
                    load_kv_consumer_state,
                    p_mma_consumer_state,
                    mma_o_producer_state,
                ) = self.mma_pv(
                    common_params,
                    pv_params,
                    tiled_mma_pv,
                    load_kv_consumer_state,
                    p_mma_consumer_state,
                    mma_o_producer_state,
                )
                k_tile_count -= 1

            # release q consumer states
            load_q_pipeline.consumer_release(load_q_release_state)
            load_q_release_state.advance()
            (
                tiled_mma_pv,
                load_kv_consumer_state,
                p_mma_consumer_state,
                mma_o_producer_state,
            ) = self.mma_pv(
                common_params,
                pv_params,
                tiled_mma_pv,
                load_kv_consumer_state,
                p_mma_consumer_state,
                mma_o_producer_state,
            )

        return (  # type: ignore[return-value]
            tiled_mma_qk,
            tiled_mma_pv,
            load_q_consumer_state,
            load_kv_consumer_state,
            mma_s_producer_state,
            p_mma_consumer_state,
            mma_o_producer_state,
        )

    @cute.jit
    def mma_qk(
        self,
        common_params: SimpleNamespace,
        qk_params: SimpleNamespace,
        tiled_mma_qk: cute.TiledMma,
        load_q_consumer_state: pipeline.PipelineState,
        load_kv_consumer_state: pipeline.PipelineState,
        mma_s_producer_state: pipeline.PipelineState,
        wait_q: bool,
    ) -> tuple[
        cute.TiledMma,
        pipeline.PipelineState,
        pipeline.PipelineState,
        pipeline.PipelineState,
    ]:
        """Compute one k-tile of mma for Q*K^T. Updates the tiled MMA QK and pipeline states.

        :param qk_params: The qk parameters
        :type qk_params: SimpleNamespace
        :param tiled_mma_qk: The tiled mma qk
        :type tiled_mma_qk: cute.TiledMma
        :param load_q_consumer_state: The load q consumer state
        :type load_q_consumer_state: pipeline.PipelineState
        :param load_kv_consumer_state: The load kv consumer state
        :type load_kv_consumer_state: pipeline.PipelineState
        :param mma_s_producer_state: The mma s producer state
        :type mma_s_producer_state: pipeline.PipelineState

        :return: The tiled mma qk, the load q consumer state, the load kv consumer state, and the mma s producer state
        :rtype: tuple[cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState]
        """
        tStS = qk_params.tStS_staged[None, None, None, mma_s_producer_state.index]

        qk_params.mma_s_pipeline.producer_acquire(mma_s_producer_state)
        tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, False)
        load_q_pipeline = common_params.load_q_pipeline
        load_kv_pipeline = common_params.load_kv_pipeline
        if cutlass.const_expr(wait_q):
            load_q_pipeline.consumer_wait(load_q_consumer_state)
            load_q_consumer_state.advance()
        for q_stage in range(self.iterations_qk_latent):
            load_kv_pipeline.consumer_wait(load_kv_consumer_state)
            kc_stage = load_kv_consumer_state.index
            for k_block in cutlass.range(cute.size(qk_params.tSrQ.shape[2])):
                cute.gemm(
                    tiled_mma_qk,
                    tStS,
                    qk_params.tSrQ[None, None, k_block, q_stage],
                    qk_params.tSrKC[None, None, k_block, kc_stage],
                    tStS,
                )
                tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, True)
            load_kv_pipeline.consumer_release(load_kv_consumer_state)
            load_kv_consumer_state.advance()
        for q_stage in range(self.iterations_qk_rope):
            load_kv_pipeline.consumer_wait(load_kv_consumer_state)
            kc_stage = load_kv_consumer_state.index
            for k_block in cutlass.range(self.rope_dim // tiled_mma_qk.shape_mnk[2]):
                cute.gemm(
                    tiled_mma_qk,
                    tStS,
                    qk_params.tSrQ_rope[None, None, k_block, q_stage],
                    qk_params.tSrKC[None, None, k_block, kc_stage],
                    tStS,
                )
                tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, True)
            load_kv_pipeline.consumer_release(load_kv_consumer_state)
            load_kv_consumer_state.advance()

        qk_params.mma_s_pipeline.producer_commit(mma_s_producer_state)
        mma_s_producer_state.advance()
        return (
            tiled_mma_qk,
            load_q_consumer_state,
            load_kv_consumer_state,
            mma_s_producer_state,
        )

    @cute.jit
    def mma_pv(
        self,
        common_params: SimpleNamespace,
        pv_params: SimpleNamespace,
        tiled_mma_pv: cute.TiledMma,
        load_kv_consumer_state: pipeline.PipelineState,
        p_mma_consumer_state: pipeline.PipelineState,
        mma_o_producer_state: pipeline.PipelineState,
    ) -> tuple[
        cute.TiledMma,
        pipeline.PipelineState,
        pipeline.PipelineState,
        pipeline.PipelineState,
    ]:
        """Compute one k-tile of mma for P*V. Updates the tiled mma pv and pipeline states.

        :param common_params: The common parameters
        :type common_params: SimpleNamespace
        :param pv_params: The pv parameters
        :type pv_params: SimpleNamespace
        :param tiled_mma_pv: The tiled mma pv
        :type tiled_mma_pv: cute.TiledMma
        :param load_kv_consumer_state: The load kv consumer state
        :type load_kv_consumer_state: pipeline.PipelineState
        :param p_mma_consumer_state: The P MMA consumer state
        :type p_mma_consumer_state: pipeline.PipelineState
        :param mma_o_producer_state: The MMA o producer state
        :type mma_o_producer_state: pipeline.PipelineState

        :return: The tiled mma pv, the load qkv consumer state, the P MMA consumer state, and the MMA o producer state
        :rtype: tuple[cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState]
        """

        pv_params.mma_o_pipeline.producer_acquire(mma_o_producer_state)
        pv_params.p_mma_pipeline.consumer_wait(p_mma_consumer_state)
        load_kv_pipeline = common_params.load_kv_pipeline
        for p_stage in range(self.iterations_pv_k):
            accumulate_flag = tiled_mma_pv.get(tcgen05.Field.ACCUMULATE)
            for acc_stage in range(self.iterations_pv_n):
                load_kv_pipeline.consumer_wait(load_kv_consumer_state)
                tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, accumulate_flag)
                vc_stage = load_kv_consumer_state.index
                tOtO = pv_params.tOtO_staged[None, None, None, acc_stage]
                for k_block in cutlass.range(pv_params.tOrP.shape[2]):
                    cute.gemm(
                        tiled_mma_pv,
                        tOtO,
                        pv_params.tOrP[
                            None,
                            None,
                            k_block,
                            (p_stage, p_mma_consumer_state.index),
                        ],
                        pv_params.tOrVC[None, None, k_block, vc_stage],
                        tOtO,
                    )
                    tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, True)
                load_kv_pipeline.consumer_release(load_kv_consumer_state)
                load_kv_consumer_state.advance()
        pv_params.p_mma_pipeline.consumer_release(p_mma_consumer_state)
        p_mma_consumer_state.advance()
        pv_params.mma_o_pipeline.producer_commit(mma_o_producer_state)
        mma_o_producer_state.advance()

        return (
            tiled_mma_pv,
            load_kv_consumer_state,
            p_mma_consumer_state,
            mma_o_producer_state,
        )

    @cute.jit
    def compute(
        self,
        common_params: SimpleNamespace,
        softmax_params: SimpleNamespace,
        k_index: cutlass.Int32,
        k_tile_count: cutlass.Int32,
        mma_s_consumer_state: pipeline.PipelineState,
        p_mma_producer_state: pipeline.PipelineState,
        p_cor_producer_state: pipeline.PipelineState,
    ) -> tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState]:
        """Compute warp to compute the result of softmax, rescale, and epilogue. Updates the related pipeline states.

        :param common_params: The common parameters
        :type common_params: SimpleNamespace
        :param softmax_params: The softmax parameters
        :type softmax_params: SimpleNamespace
        :param k_index: The index of the k-tile
        :type k_index: cutlass.Int32
        :param k_tile_count: The number of k-tiles
        :type k_tile_count: cutlass.Int32
        :param mma_s_consumer_state: The MMA s consumer state
        :type mma_s_consumer_state: pipeline.PipelineState
        :param p_mma_producer_state: The P MMA producer state
        :type p_mma_producer_state: pipeline.PipelineState
        :param p_cor_producer_state: The P correction producer state
        :type p_cor_producer_state: pipeline.PipelineState

        :return: The MMA s consumer state, the P MMA producer state, and the P correction producer state
        :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState]
        """

        row_max = -self.acc_dtype.inf
        row_sum = self.acc_dtype(0)
        correction_factor = self.acc_dtype(1)
        common_params.p_cor_pipeline.producer_acquire(p_cor_producer_state)

        # The first tile that can contain a key at or beyond this query tile's
        # earliest causal bound.  A flat M tile starts at row q_tile * M, so
        # floor(row / H) is its earliest query token.  Later query tiles can
        # therefore keep more fully dense K tiles on the compile-time unmasked
        # path.  The token division is outside the per-score loop (and is
        # replicated by each participating CTA/compute group); tile_n is a
        # static power of two, so the remaining division is a shift.
        tile_n = self.mma_qk_tiler[1]
        first_q_token = cutlass.Int32(0)
        if cutlass.const_expr(self.num_q_tiles > 1):
            first_q_token = (
                common_params.blk_coord[1] * self.mma_qk_tiler[0]
            ) // self.num_heads
        if cutlass.const_expr(self.enable_dcp):
            # The earliest token in this full M128 query tile determines the
            # dense-prefix boundary shared by both CTAs.  Keep the physical
            # local K bound separate so a partial final local tile is never
            # misclassified as dense.
            earliest_bound_numer = (
                common_params.causal_global
                - common_params.cp_rank
                - (common_params.q_len - 1)
                + first_q_token
            )
            earliest_local_bound = cute.ceil_div(
                cutlass.max(earliest_bound_numer, cutlass.Int32(0)),
                self.cp_world,
            )
            effective_local_bound = cutlass.min(common_params.K, earliest_local_bound)
            first_mask_tile_idx = effective_local_bound // tile_n
        else:
            first_mask_tile_idx = cutlass.max(
                (common_params.K - common_params.q_len + 1 + first_q_token) // tile_n,
                cutlass.Int32(0),
            )

        # Phase 1: pure unmasked bulk tiles (all columns strictly < min k_bound).
        while k_tile_count > 1 and k_index < first_mask_tile_idx:
            (
                mma_s_consumer_state,
                p_mma_producer_state,
                p_cor_producer_state,
                row_max,
                row_sum,
                correction_factor,
            ) = self.softmax(
                common_params,
                softmax_params,
                k_index,
                mma_s_consumer_state,
                p_mma_producer_state,
                p_cor_producer_state,
                row_max,
                row_sum,
                correction_factor,
                False,
                False,
            )
            k_index = k_index + 1
            k_tile_count = k_tile_count - 1

        # Phase 2: intermediate tiles that overlap the causal/K-bound region
        # but are not this work-split's final tile.
        while k_tile_count > 1:
            (
                mma_s_consumer_state,
                p_mma_producer_state,
                p_cor_producer_state,
                row_max,
                row_sum,
                correction_factor,
            ) = self.softmax(
                common_params,
                softmax_params,
                k_index,
                mma_s_consumer_state,
                p_mma_producer_state,
                p_cor_producer_state,
                row_max,
                row_sum,
                correction_factor,
                True,
                False,
            )
            k_index = k_index + 1
            k_tile_count = k_tile_count - 1

        # Phase 3: this work-split's final tile.  Mask it only when its global
        # index reaches the causal/K-bound region.
        (
            mma_s_consumer_state,
            p_mma_producer_state,
            p_cor_producer_state,
            row_max,
            row_sum,
            correction_factor,
        ) = self.softmax(
            common_params,
            softmax_params,
            k_index,
            mma_s_consumer_state,
            p_mma_producer_state,
            p_cor_producer_state,
            row_max,
            row_sum,
            correction_factor,
            k_index >= first_mask_tile_idx,
            True,
        )

        return mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state

    @cute.jit
    def correction(
        self,
        common_params: SimpleNamespace,
        epilogue_params: SimpleNamespace,
        k_tile_count: cutlass.Int32,
        p_cor_consumer_state: pipeline.PipelineState,
        mma_o_consumer_state: pipeline.PipelineState,
    ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]:
        """Compute warp to compute the result of softmax, rescale, and epilogue. Updates the related pipeline states.

        :param common_params: The common parameters
        :type common_params: SimpleNamespace
        :param epilogue_params: The epilogue parameters
        :type epilogue_params: SimpleNamespace
        :param k_index: The index of the k-tile
        :type k_index: cutlass.Int32
        :param k_tile_count: The number of k-tiles
        :type k_tile_count: cutlass.Int32
        :param p_cor_consumer_state: The P correction consumer state
        :type p_cor_consumer_state: pipeline.PipelineState
        :param mma_o_consumer_state: The MMA o consumer state
        :type mma_o_consumer_state: pipeline.PipelineState

        :return: The P correction consumer state, and the MMA o consumer state
        :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState]
        """

        k_tile_count_init = k_tile_count
        while k_tile_count > 0:
            p_cor_consumer_state, row_sum, row_max, correction_factor, no_correction = (
                self.get_correction_factor(common_params, p_cor_consumer_state)
            )
            if k_tile_count_init != k_tile_count:
                mma_o_consumer_state = self.rescale(
                    common_params,
                    mma_o_consumer_state,
                    correction_factor,
                    no_correction,
                )
            k_tile_count = k_tile_count - 1
            if k_tile_count == 0:
                mma_o_consumer_state = self.epilogue(
                    common_params,
                    epilogue_params,
                    mma_o_consumer_state,
                    row_sum,
                    row_max,
                )

        return p_cor_consumer_state, mma_o_consumer_state

    @cute.jit
    def exchange_p_cor_metadata(
        self,
        common_params: SimpleNamespace,
        softmax_params: SimpleNamespace,
        correction_factor: cutlass.Float32,
        row_sum: cutlass.Float32,
        row_max: cutlass.Float32,
        row_max_new: cutlass.Float32,
        tAcc: cute.Tensor,
        tidx: cutlass.Int32,
        p_cor_producer_state: pipeline.PipelineState,
    ) -> pipeline.PipelineState:
        """Compute the correction factor for the last k tile."""
        no_correction = 0
        if (
            row_max_new - row_max
        ) * softmax_params.softmax_scale_log2 <= self.skip_correction_threshold:
            no_correction = 1
            row_max_new = row_max

        # pad for 4x32b
        corr_layout = cute.make_layout(
            (tAcc.shape[0], (4, tAcc.shape[1][1]), self.mma_s_stage),
            stride=(tAcc.stride[0], (1, tAcc.stride[1][1]), 4),
        )
        tCor = cute.make_tensor(
            common_params.tmem_ptr + self.correction_factor_offset,
            corr_layout,
        )
        cCor = cute.make_identity_tensor(tCor.shape)
        corr_tmem_store_atom = cute.make_copy_atom(
            tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(4)), self.acc_dtype
        )
        corr_tmem_store_tiled_copy = tcgen05.make_tmem_copy(corr_tmem_store_atom, tCor)
        corr_tmem_store_thr_copy = corr_tmem_store_tiled_copy.get_slice(tidx)
        cCor_for_copy = corr_tmem_store_thr_copy.partition_S(cCor)
        tCor_for_copy = corr_tmem_store_thr_copy.partition_D(tCor)
        rCor = cute.make_fragment_like(
            cCor_for_copy[None, None, None, 0], self.acc_dtype
        )
        rCor_int = cute.make_tensor(
            cute.recast_ptr(rCor.iterator, dtype=cutlass.Int32), rCor.layout
        )
        rCor[0] = row_sum
        rCor[1] = row_max_new
        rCor[2] = correction_factor
        rCor_int[3] = no_correction

        cute.copy(
            corr_tmem_store_tiled_copy,
            rCor,
            tCor_for_copy[None, None, None, p_cor_producer_state.index],
        )
        # fence between tmem store and correction warp
        cute.arch.fence_view_async_tmem_store()
        common_params.p_cor_pipeline.producer_commit(p_cor_producer_state)
        p_cor_producer_state.advance()
        return p_cor_producer_state, row_max_new

    @cute.jit
    def softmax(
        self,
        common_params: SimpleNamespace,
        softmax_params: SimpleNamespace,
        k_index: cutlass.Int32,
        mma_s_consumer_state: pipeline.PipelineState,
        p_mma_producer_state: pipeline.PipelineState,
        p_cor_producer_state: pipeline.PipelineState,
        row_max: cutlass.Float32,
        row_sum: cutlass.Float32,
        correction_factor: cutlass.Float32,
        apply_mask: bool,
        is_local_last_tile: cutlass.Boolean,
    ) -> tuple[
        pipeline.PipelineState,
        pipeline.PipelineState,
        pipeline.PipelineState,
        cutlass.Float32,
        cutlass.Float32,
        cutlass.Float32,
    ]:
        """Softmax for one k-tile. Updates the related pipeline states and returns the computed results.

        :param common_params: The common parameters
        :type common_params: SimpleNamespace
        :param softmax_params: The softmax parameters
        :type softmax_params: SimpleNamespace
        :param k_index: The index of the k-tile
        :type k_index: cutlass.Int32
        :param mma_s_consumer_state: The MMA s consumer state
        :type mma_s_consumer_state: pipeline.PipelineState
        :param p_mma_producer_state: The P MMA producer state
        :type p_mma_producer_state: pipeline.PipelineState
        :param p_cor_producer_state: The P correction producer state
        :type p_cor_producer_state: pipeline.PipelineState
        :param row_max: The row max
        :type row_max: cutlass.Float32
        :param row_sum: The row sum
        :type row_sum: cutlass.Float32
        :param correction_factor: The correction factor
        :type correction_factor: cutlass.Float32
        :param apply_mask: Whether the tile needs K-bound / causal masking (Python bool
            for the unmasked/masked bulk loops; runtime cutlass.Boolean for a
            work-split final tile at the exact per-query-tile mask boundary).
        :type apply_mask: bool | cutlass.Boolean
        :param is_local_last_tile: Whether the last tile is local
        :type is_local_last_tile: cutlass.Boolean

        :return: The MMA s consumer state, the P MMA producer state, the P correction producer state, the row max, the row sum, and the correction factor
        :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, cutlass.Float32, cutlass.Float32, cutlass.Float32]
        """

        softmax_params.p_mma_pipeline.producer_acquire(p_mma_producer_state)
        softmax_params.mma_s_pipeline.consumer_wait(mma_s_consumer_state)

        # load S from tmem
        tStS_shape = softmax_params.tiled_mma_qk.partition_shape_C(
            cute.select(self.mma_qk_tiler, mode=[0, 1])
        )
        tStS_staged_fake = softmax_params.tiled_mma_qk.make_fragment_C(
            cute.append(tStS_shape, self.mma_s_stage)
        )
        tStS_staged = cute.make_tensor(common_params.tmem_ptr, tStS_staged_fake.layout)
        tStS = tStS_staged[None, None, None, mma_s_consumer_state.index]

        tAcc = tStS[(None, None), 0, 0]
        cta_qk_tiler = (
            self.mma_qk_tiler[0] // self.cluster_shape_mnk[0],
            self.mma_qk_tiler[1],
            self.mma_qk_tiler[2],
        )
        cS = cute.make_identity_tensor(cute.select(cta_qk_tiler, mode=[0, 1]))

        tmem_load_atom = cute.make_copy_atom(
            tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype
        )
        tmem_tiled_copy = tcgen05.make_tmem_copy(tmem_load_atom, tAcc)

        tidx = common_params.tidx % (self.num_compute_warps * self.threads_per_warp)

        tmem_thr_copy = tmem_tiled_copy.get_slice(tidx)
        tTR_tAcc = tmem_thr_copy.partition_S(tAcc)
        tTR_tS = tmem_thr_copy.partition_D(cS)

        tTR_rAcc = cute.make_fragment_like(tTR_tS, self.acc_dtype)

        row_max_new = row_max
        # Spec-decoding (MTP) causal mask.  A flattened row r corresponds to
        # q_token=floor(r/H), whose last valid key is K-S_q+q_token.  Avoid the
        # per-element integer division using the equivalent integer predicate:
        #
        #   H * (key_pos - K + S_q) <= r
        #
        # This arithmetic remains inside `if apply_mask`, so the compile-time
        # dense path emits no row-mask work.  Masked positions use a large
        # negative sentinel (not -inf) to avoid all-masked-row NaNs.
        cta_m_rows = self.mma_qk_tiler[0] // self.cluster_shape_mnk[0]
        arch = BaseDSL._get_dsl().get_arch_enum()
        if cutlass.const_expr(arch >= Arch.sm_100 and arch <= Arch.sm_100f):
            cute.copy(tmem_tiled_copy, tTR_tAcc, tTR_rAcc)
            for i in cutlass.range_constexpr(cute.size(tTR_rAcc)):
                if apply_mask:
                    flat_q_row = (
                        common_params.blk_coord[1] * self.mma_qk_tiler[0]
                        + common_params.blk_coord[0] * cta_m_rows
                        + tTR_tS[i][0]
                    )
                    key_pos = tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index
                    if cutlass.const_expr(self.enable_dcp):
                        mask_threshold = self.num_heads * (
                            self.cp_world * key_pos
                            + common_params.cp_rank
                            - common_params.causal_global
                            + common_params.q_len
                        )
                        tTR_rAcc[i] = (
                            tTR_rAcc[i]
                            if cute.elem_less(key_pos, common_params.K)
                            and not cute.elem_less(flat_q_row, mask_threshold)
                            else self.acc_dtype(-1.0e6)
                        )
                    else:
                        mask_threshold = self.num_heads * (
                            key_pos - common_params.K + common_params.q_len
                        )
                        tTR_rAcc[i] = (
                            tTR_rAcc[i]
                            if not cute.elem_less(flat_q_row, mask_threshold)
                            else self.acc_dtype(-1.0e6)
                        )
            # reduction for row_max
            row_max_new = tTR_rAcc.load().reduce(cute.ReductionOp.MAX, row_max_new, 0)

        elif cutlass.const_expr(arch >= Arch.sm_103 and arch <= Arch.sm_103f):
            tmem_load_red_atom = cute.make_copy_atom(
                tcgen05.copy.LdRed32x32bOp(
                    tcgen05.copy.Repetition(64), redOp=tcgen05.TmemLoadRedOp.MAX
                ),
                self.acc_dtype,
            )
            tmem_red_tiled_copy = tcgen05.make_tmem_copy(tmem_load_red_atom, tAcc)
            tmem_red_thr_copy = tmem_red_tiled_copy.get_slice(tidx)
            tTR_tAcc_red = tmem_red_thr_copy.partition_S(tAcc)
            tTR_tS_red = tmem_red_thr_copy.partition_D(cS)
            tTR_rAcc_red = cute.make_fragment_like(tTR_tS_red, self.acc_dtype)
            tTR_rMax = cute.make_rmem_tensor(
                cute.make_layout((1, tTR_tS_red.shape[1], tTR_tS_red.shape[2])),
                self.acc_dtype,
            )
            cute.copy(
                tmem_red_tiled_copy,
                tTR_tAcc_red,
                (tTR_rAcc_red, tTR_rMax),
            )
            tTR_rAcc = cute.make_tensor(tTR_rAcc_red.iterator, tTR_rAcc.layout)
            if apply_mask:
                for i in cutlass.range_constexpr(cute.size(tTR_rAcc)):
                    flat_q_row = (
                        common_params.blk_coord[1] * self.mma_qk_tiler[0]
                        + common_params.blk_coord[0] * cta_m_rows
                        + tTR_tS[i][0]
                    )
                    key_pos = tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index
                    if cutlass.const_expr(self.enable_dcp):
                        mask_threshold = self.num_heads * (
                            self.cp_world * key_pos
                            + common_params.cp_rank
                            - common_params.causal_global
                            + common_params.q_len
                        )
                        tTR_rAcc[i] = (
                            tTR_rAcc[i]
                            if cute.elem_less(key_pos, common_params.K)
                            and not cute.elem_less(flat_q_row, mask_threshold)
                            else self.acc_dtype(-1.0e6)
                        )
                    else:
                        mask_threshold = self.num_heads * (
                            key_pos - common_params.K + common_params.q_len
                        )
                        tTR_rAcc[i] = (
                            tTR_rAcc[i]
                            if not cute.elem_less(flat_q_row, mask_threshold)
                            else self.acc_dtype(-1.0e6)
                        )
                # reduction for row_max after manual masking
                row_max_new = tTR_rAcc.load().reduce(
                    cute.ReductionOp.MAX, row_max_new, 0
                )
            else:
                row_max_new = cute.arch.fmax(row_max_new, tTR_rMax[0])

        # if warps in N is 2, reduce row_max across warps (0, 1) and (2, 3)
        if cutlass.const_expr(self.warps_in_n == 2):
            common_params.smem_exchange[tidx] = row_max_new
            self.softmax_exchange_sync_bar.wait()
            row_max_new = cute.arch.fmax(
                row_max_new,
                common_params.smem_exchange[
                    (tidx + 64) % (self.num_compute_warps * self.threads_per_warp)
                ],
            )

        # find correction factor
        correction_factor = cute.math.exp2(
            (row_max - row_max_new) * softmax_params.softmax_scale_log2, fastmath=True
        )
        # split kv case
        if cutlass.const_expr(not is_local_last_tile):
            p_cor_producer_state, row_max_new = self.exchange_p_cor_metadata(
                common_params,
                softmax_params,
                correction_factor,
                row_sum,
                row_max,
                row_max_new,
                tAcc,
                tidx,
                p_cor_producer_state,
            )

        # softmax
        fma_b = softmax_params.softmax_scale_log2
        fma_c = (0.0 - row_max_new) * softmax_params.softmax_scale_log2

        for i in cutlass.range(cute.size(tTR_rAcc), vectorize=True, unroll_full=True):
            tTR_rAcc[i] = tTR_rAcc[i] * fma_b + fma_c
            tTR_rAcc[i] = cute.math.exp2(tTR_rAcc[i], fastmath=True)

        tTR_rS = cute.make_fragment_like(tTR_tS, self.q_dtype)

        # quantize
        tTR_rS.store(tTR_rAcc.load().to(self.q_dtype))

        # create sP
        sP = softmax_params.sP[None, None, None, (None, p_mma_producer_state.index)]
        sP_mk_view = cute.make_tensor(
            sP.iterator,
            cute.make_layout(
                (
                    (sP.shape[0][0], sP.shape[1]),
                    (sP.shape[0][1], sP.shape[2], sP.shape[3]),
                ),
                stride=(
                    (sP.stride[0][0], sP.stride[1]),
                    (sP.stride[0][1], sP.stride[2], sP.stride[3]),
                ),
            ),
        )
        # {$nv-internal-release begin}
        # TODO: figure out if we could use A tmem for pv.
        # {$nv-internal-release end}
        # change to PISL
        sP_wo_swizzle_iter = cute.recast_ptr(sP.iterator, swizzle_=None)
        swizzle_bits = (
            int(math.log2(self.mma_pv_tiler[2] * self.q_dtype.width // 8 // 32)) + 1
        )
        swizzle_base = 3 if self.q_dtype.width == 16 else 4
        sP_swizzle = cute.make_swizzle(swizzle_bits, swizzle_base, 3)
        sP_mk_view = cute.make_tensor(
            sP_wo_swizzle_iter,
            cute.make_composed_layout(sP_swizzle, 0, sP_mk_view.layout),
        )
        universal_copy_bits = 128
        smem_copy_atom = cute.make_copy_atom(
            cute.nvgpu.CopyUniversalOp(),
            self.q_dtype,
            num_bits_per_copy=universal_copy_bits,
        )
        smem_tiled_copy = cute.make_tiled_copy_D(smem_copy_atom, tmem_tiled_copy)
        smem_thr_copy = smem_tiled_copy.get_slice(tidx)
        rP_copy_view = smem_thr_copy.retile(tTR_rS)
        sP_copy_view = smem_thr_copy.partition_D(sP_mk_view)
        cute.copy(smem_tiled_copy, rP_copy_view, sP_copy_view)

        # fence between smem store and mma o
        cute.arch.fence_view_async_shared()
        softmax_params.p_mma_pipeline.producer_commit(p_mma_producer_state)
        p_mma_producer_state.advance()

        # row_sum, using `add_packed_f32x2` to reduce the number of instructions
        row_sum = row_sum * correction_factor
        row_sum_vec = (0.0, 0.0)
        for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc), 2):
            row_sum_vec = cute.arch.add_packed_f32x2(
                row_sum_vec, (tTR_rAcc[i], tTR_rAcc[i + 1])
            )
        row_sum = row_sum_vec[0] + row_sum_vec[1] + row_sum

        # split kv case
        if cutlass.const_expr(is_local_last_tile):
            p_cor_producer_state, row_max_new = self.exchange_p_cor_metadata(
                common_params,
                softmax_params,
                correction_factor,
                row_sum,
                row_max,
                row_max_new,
                tAcc,
                tidx,
                p_cor_producer_state,
            )

        # store correction factor/row_sum/row_max to tmem for correction warp
        common_params.p_cor_pipeline.producer_acquire(p_cor_producer_state)

        # fence between tmem load and mma s
        cute.arch.fence_view_async_tmem_load()

        softmax_params.mma_s_pipeline.consumer_release(mma_s_consumer_state)
        mma_s_consumer_state.advance()

        return (
            mma_s_consumer_state,
            p_mma_producer_state,
            p_cor_producer_state,
            row_max_new,
            row_sum,
            correction_factor,
        )

    @cute.jit
    def _tmem_load_partition(
        self, common_params: SimpleNamespace, tiled_mma_pv: cute.TiledMma, iter_n: int
    ) -> tuple[
        cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma
    ]:
        """Tensor memory load partition for rescale and epilogue.

        :param common_params: The common parameters
        :type common_params: SimpleNamespace
        :param tiled_mma_pv: The tiled mma pv
        :type tiled_mma_pv: cute.TiledMma
        :param iter_n: The iteration number
        :type iter_n: int

        :return: The tiled mma pv, the tiled mma pv, the tiled mma pv, the tiled mma pv, the tiled mma pv
        :rtype: tuple[cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma]
        """

        tOtO_shape = tiled_mma_pv.partition_shape_C(
            cute.select(self.mma_pv_tiler, mode=[0, 1])
        )
        tOtO = tiled_mma_pv.make_fragment_C(tOtO_shape)
        tOtO_layout = cute.append(
            tOtO.layout,
            cute.make_layout(
                common_params.L // self.mma_pv_tiler[1],
                stride=self.mma_pv_tiler[1] // self.warps_in_n,
            ),
        )
        tOtO = cute.make_tensor(
            common_params.tmem_ptr + self.tmem_o_offset, tOtO_layout
        )
        tOtO = tOtO[None, None, None, iter_n]

        tAcc = tOtO[(None, None), 0, 0]

        tmem_load_atom = cute.make_copy_atom(
            tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype
        )
        tmem_load_tiled_copy = tcgen05.make_tmem_copy(tmem_load_atom, tAcc)
        # {$nv-internal-release begin}
        # TODO: supports size() on tiled copy.
        # {$nv-internal-release end}
        tmem_load_thr_copy = tmem_load_tiled_copy.get_slice(
            common_params.tidx % (self.num_compute_warps * self.threads_per_warp)
        )

        cta_pv_tiler = (
            self.mma_pv_tiler[0] // self.cluster_shape_mnk[0],
            self.mma_pv_tiler[1],
            self.mma_pv_tiler[2],
        )
        # Flatten divide and partition global tensors for O
        cta_pv_tiler_mn = cute.select(cta_pv_tiler, mode=[0, 1])

        gO = None
        if cutlass.const_expr(common_params.mAccO is not None):
            gO = cute.local_tile(
                common_params.mAccO[None, common_params.blk_coord[3], None, None, None],
                cta_pv_tiler_mn,
                (
                    common_params.blk_coord[0],
                    iter_n,
                    common_params.blk_coord[1],
                    common_params.blk_coord[2],
                ),
            )
            cO = cute.local_tile(
                cute.make_identity_tensor(
                    common_params.mAccO[
                        None, common_params.blk_coord[3], None, None, None
                    ].shape
                ),
                cta_pv_tiler_mn,
                (
                    common_params.blk_coord[0],
                    iter_n,
                    common_params.blk_coord[1],
                    common_params.blk_coord[2],
                ),
            )
        else:
            if cutlass.const_expr(self.is_var_q):
                request_o = cute.domain_offset(
                    (common_params.q_begin * self.num_heads, 0, 0),
                    common_params.mO,
                )
                request_cta_m = (
                    common_params.blk_coord[1] * self.cluster_shape_mnk[0]
                    + common_params.blk_coord[0]
                )
                gO = cute.local_tile(
                    request_o,
                    cta_pv_tiler_mn,
                    (
                        request_cta_m,
                        iter_n,
                        0,
                    ),
                )
                cO = cute.local_tile(
                    # Keep coordinates request-local so ``coord < valid_rows``
                    # remains correct for every q_tile and q_begin.
                    cute.make_identity_tensor(
                        (self.mma_qk_tiler[0], request_o.shape[1], 1)
                    ),
                    cta_pv_tiler_mn,
                    (
                        common_params.blk_coord[0],
                        iter_n,
                        0,
                    ),
                )
            else:
                gO = cute.local_tile(
                    common_params.mO,
                    cta_pv_tiler_mn,
                    (
                        common_params.blk_coord[0],
                        iter_n,
                        common_params.blk_coord[1],
                        common_params.blk_coord[2],
                    ),
                )
                cO = cute.local_tile(
                    cute.make_identity_tensor(common_params.mO.shape),
                    cta_pv_tiler_mn,
                    (
                        common_params.blk_coord[0],
                        iter_n,
                        common_params.blk_coord[1],
                        common_params.blk_coord[2],
                    ),
                )
        tTR_tAcc = tmem_load_thr_copy.partition_S(tAcc)
        tTR_gO = tmem_load_thr_copy.partition_D(gO)
        tTR_cO = tmem_load_thr_copy.partition_D(cO)
        tTR_rAcc = cute.make_fragment_like(tTR_gO, self.acc_dtype)
        return tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc  # type: ignore[return-value]

    def get_correction_factor(
        self,
        common_params: SimpleNamespace,
        p_cor_consumer_state: pipeline.PipelineState,
    ) -> tuple[
        pipeline.PipelineState,
        cutlass.Float32,
        cutlass.Float32,
        cutlass.Float32,
        cutlass.Int32,
    ]:
        """Get the correction factor from the P correction consumer state.

        :param common_params: The common parameters
        :type common_params: SimpleNamespace
        :param p_cor_consumer_state: The P correction consumer state
        :type p_cor_consumer_state: pipeline.PipelineState

        :return: The P correction consumer state, the row_sum, the row_max, and the correction factor
        :rtype: tuple[pipeline.PipelineState, cutlass.Float32, cutlass.Float32, cutlass.Float32, cutlass.Int32]
        """
        common_params.p_cor_pipeline.consumer_wait(p_cor_consumer_state)
        tidx = common_params.tidx % (self.num_compute_warps * self.threads_per_warp)
        # load correction factor
        _, tAcc, _, _, _, _ = self._tmem_load_partition(
            common_params, common_params.tiled_mma_pv, 0
        )
        corr_layout = cute.make_layout(
            (tAcc.shape[0], (4, tAcc.shape[1][1]), self.p_cor_stage),
            stride=(tAcc.stride[0], (1, tAcc.stride[1][1]), 4),
        )
        tCor = cute.make_tensor(
            common_params.tmem_ptr + self.correction_factor_offset, corr_layout
        )
        cCor = cute.make_identity_tensor(tCor.shape)
        corr_tmem_load_atom = cute.make_copy_atom(
            tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(4)), self.acc_dtype
        )
        corr_tmem_load_tiled_copy = tcgen05.make_tmem_copy(corr_tmem_load_atom, tCor)
        corr_tmem_load_thr_copy = corr_tmem_load_tiled_copy.get_slice(tidx)
        tCor_for_copy = corr_tmem_load_thr_copy.partition_S(tCor)
        cCor_for_copy = corr_tmem_load_thr_copy.partition_D(cCor)
        rCor = cute.make_fragment_like(
            cCor_for_copy[None, None, None, 0], self.acc_dtype
        )
        rCor_int = cute.make_tensor(
            cute.recast_ptr(rCor.iterator, dtype=cutlass.Int32), rCor.layout
        )
        cute.copy(
            corr_tmem_load_tiled_copy,
            tCor_for_copy[None, None, None, p_cor_consumer_state.index],
            rCor,
        )
        row_sum = rCor[0]
        row_max = rCor[1]
        correction_factor = rCor[2]
        no_correction = rCor_int[3]

        common_params.p_cor_pipeline.consumer_release(p_cor_consumer_state)
        p_cor_consumer_state.advance()
        return p_cor_consumer_state, row_sum, row_max, correction_factor, no_correction

    @cute.jit
    def rescale(
        self,
        common_params: SimpleNamespace,
        mma_o_consumer_state: pipeline.PipelineState,
        correction_factor: cutlass.Float32,
        no_correction: cutlass.Int32,
    ) -> pipeline.PipelineState:
        """Rescale for one k-tile. Updates the related pipeline state.

        :param common_params: The common parameters
        :type common_params: SimpleNamespace
        :param mma_o_consumer_state: The mma o consumer state
        :type mma_o_consumer_state: pipeline.PipelineState
        :param correction_factor: The correction factor
        :type correction_factor: cutlass.Float32
        :param no_correction: Whether to apply correction factor
        :type no_correction: cutlass.Int32

        :return: The MMA o consumer state
        :rtype: pipeline.PipelineState
        """
        skip_correction = cute.arch.vote_all_sync(no_correction == 1)
        common_params.mma_o_pipeline.consumer_wait(mma_o_consumer_state)
        if not skip_correction:
            for iter_n in cutlass.range_constexpr(self.iterations_pv_n):
                # tmem load tiled copy and partition results.
                tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc = (
                    self._tmem_load_partition(
                        common_params, common_params.tiled_mma_pv, iter_n
                    )
                )

                # tmem store tiled copy
                tmem_store_atom = cute.make_copy_atom(
                    tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype
                )
                tmem_store_tiled_copy = tcgen05.make_tmem_copy(tmem_store_atom, tAcc)

                # load o
                cute.copy(tmem_load_tiled_copy, tTR_tAcc, tTR_rAcc)
                # rescale, using `mul_packed_f32x2` to reduce the number of instructions
                for i in cutlass.range(
                    cute.size(tTR_rAcc), vectorize=True, unroll_full=True
                ):
                    tTR_rAcc[i] = tTR_rAcc[i] * correction_factor

                # store o to tensor memory for next k tile
                cute.copy(tmem_store_tiled_copy, tTR_rAcc, tTR_tAcc)

        cute.arch.fence_view_async_tmem_store()
        common_params.mma_o_pipeline.consumer_release(mma_o_consumer_state)
        mma_o_consumer_state.advance()

        return mma_o_consumer_state

    @cute.jit
    def epilogue(
        self,
        common_params: SimpleNamespace,
        epilogue_params: SimpleNamespace,
        mma_o_consumer_state: pipeline.PipelineState,
        row_sum: cutlass.Float32,
        row_max: cutlass.Float32,
    ) -> pipeline.PipelineState:
        """Epilogue for one k-tile. Updates the related pipeline state.

        :param common_params: The common parameters
        :type common_params: SimpleNamespace
        :param epilogue_params: The epilogue parameters
        :type epilogue_params: SimpleNamespace
        :param mma_o_consumer_state: The mma o consumer state
        :type mma_o_consumer_state: pipeline.PipelineState
        :param row_sum: The row sum
        :type row_sum: cutlass.Float32
        :param row_max: The row max
        :type row_max: cutlass.Float32

        :return: The MMA o consumer state
        :rtype: pipeline.PipelineState
        """

        tidx = common_params.tidx % (self.num_compute_warps * self.threads_per_warp)

        # exchange row_sum between warps (0, 1) and (2, 3)
        if cutlass.const_expr(self.warps_in_n == 2):
            common_params.smem_exchange[tidx] = row_sum
            self.epilogue_exchange_sync_bar.wait()
            # (64, 2)
            row_sum = (
                row_sum
                + common_params.smem_exchange[
                    (tidx + 64) % (self.num_compute_warps * self.threads_per_warp)
                ]
            )
        # mma_o pipeline consumer wait
        common_params.mma_o_pipeline.consumer_wait(mma_o_consumer_state)
        for iter_n in cutlass.range_constexpr(self.iterations_pv_n):
            # tmem load tiled copy and partition results.
            tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc = (
                self._tmem_load_partition(
                    common_params, common_params.tiled_mma_pv, iter_n
                )
            )

            # load o
            cute.copy(tmem_load_tiled_copy, tTR_tAcc, tTR_rAcc)

            row_has_key = True
            if cutlass.const_expr(self.enable_dcp):
                # The coordinate tensor returned by local_tile already carries
                # the CTA's offset within the full M128 query tile.
                flat_q_row = (
                    common_params.blk_coord[1] * self.mma_qk_tiler[0] + tTR_cO[0][0]
                )
                first_key_threshold = self.num_heads * (
                    self.cp_world * common_params.split_start_key
                    + common_params.cp_rank
                    - common_params.causal_global
                    + common_params.q_len
                )
                row_has_key = cute.elem_less(
                    common_params.split_start_key, common_params.K
                ) and not cute.elem_less(flat_q_row, first_key_threshold)

            # apply output scale and normalize by row_sum
            if cutlass.const_expr(self.enable_dcp):
                normalization_scale = (
                    epilogue_params.output_scale * cute.arch.rcp_approx(row_sum)
                )
                normalization_scale = (
                    normalization_scale if row_has_key else self.acc_dtype(0.0)
                )
                for i in cutlass.range(
                    cute.size(tTR_rAcc), vectorize=True, unroll_full=True
                ):
                    tTR_rAcc[i] = tTR_rAcc[i] * normalization_scale
            else:
                for i in cutlass.range(
                    cute.size(tTR_rAcc), vectorize=True, unroll_full=True
                ):
                    tTR_rAcc[i] = (
                        tTR_rAcc[i]
                        * epilogue_params.output_scale
                        * cute.arch.rcp_approx(row_sum)
                    )

            # store o to global memory
            tR2G_rO_src = None
            tR2G_rO_dst = tTR_gO
            if cutlass.const_expr(common_params.mAccO is None):
                tR2G_rO_src = cute.make_fragment_like(tTR_gO, self.o_dtype)
                # using final output dtype for o
                tR2G_rO_src.store(tTR_rAcc.load().to(self.o_dtype))
            else:
                # using accumulate dtype for o
                tR2G_rO_src = tTR_rAcc

            if cute.elem_less(tTR_cO[0][0], common_params.H):
                cute.autovec_copy(
                    tR2G_rO_src,
                    tR2G_rO_dst,
                    l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.NO_ALLOCATE,
                )

            # store the lse to global memory
            cta_pv_tiler = (
                self.mma_pv_tiler[0] // self.cluster_shape_mnk[0],
                self.mma_pv_tiler[1],
                self.mma_pv_tiler[2],
            )
            gLSE = None
            cLSE = None
            if cutlass.const_expr(epilogue_params.mAccLSE is None):
                if cutlass.const_expr(self.is_var_q):
                    request_lse = cute.domain_offset(
                        (common_params.q_begin * self.num_heads, 0, 0),
                        epilogue_params.mLSE,
                    )
                    request_cta_m = (
                        common_params.blk_coord[1] * self.cluster_shape_mnk[0]
                        + common_params.blk_coord[0]
                    )
                    gLSE = cute.local_tile(
                        request_lse,
                        (cta_pv_tiler[0], 1, 1),
                        (
                            request_cta_m,
                            0,
                            0,
                        ),
                        (1, 1, 1),
                    )
                    cLSE = cute.local_tile(
                        cute.make_identity_tensor((self.mma_qk_tiler[0], 1, 1)),
                        (cta_pv_tiler[0], 1, 1),
                        (
                            common_params.blk_coord[0],
                            0,
                            0,
                        ),
                        (1, 1, 1),
                    )
                else:
                    gLSE = cute.local_tile(
                        epilogue_params.mLSE,
                        (cta_pv_tiler[0], 1, 1),
                        (
                            common_params.blk_coord[0],
                            common_params.blk_coord[1],
                            common_params.blk_coord[2],
                        ),
                        (1, 1, 1),
                    )
                    cLSE = cute.local_tile(
                        cute.make_identity_tensor(epilogue_params.mLSE.shape),
                        (cta_pv_tiler[0], 1, 1),
                        (
                            common_params.blk_coord[0],
                            common_params.blk_coord[1],
                            common_params.blk_coord[2],
                        ),
                        (1, 1, 1),
                    )

            else:
                gLSE = cute.local_tile(
                    epilogue_params.mAccLSE[
                        None, common_params.blk_coord[3], None, None
                    ],
                    (cta_pv_tiler[0], 1, 1),
                    (
                        common_params.blk_coord[0],
                        common_params.blk_coord[1],
                        common_params.blk_coord[2],
                    ),
                    (1, 1, 1),
                )
                cLSE = cute.local_tile(
                    cute.make_identity_tensor(
                        epilogue_params.mAccLSE[
                            None, common_params.blk_coord[3], None, None
                        ].shape
                    ),
                    (cta_pv_tiler[0], 1, 1),
                    (
                        common_params.blk_coord[0],
                        common_params.blk_coord[1],
                        common_params.blk_coord[2],
                    ),
                    (1, 1, 1),
                )
            lse = (
                cute.math.log2(row_sum, fastmath=True)
                + epilogue_params.softmax_scale_log2 * row_max
            )
            if cutlass.const_expr(self.enable_dcp):
                lse = lse if row_has_key else -self.lse_dtype.inf
            # When writing directly to the user-facing mLSE (single-tile,
            # no split-KV merge), convert from log2 base to natural log.
            # When writing the per-split intermediate (mAccLSE branch), keep
            # log2 base so the merge code above can use exp2 / log2 ops.
            if cutlass.const_expr(epilogue_params.mAccLSE is None):
                lse = lse * (1.0 / LOG2_E)
            if cutlass.const_expr(self.warps_in_n == 2):
                if cute.elem_less(cLSE[tidx][0], common_params.H):
                    gLSE[tidx] = lse

        cute.arch.fence_view_async_tmem_load()
        common_params.mma_o_pipeline.consumer_release(mma_o_consumer_state)
        mma_o_consumer_state.advance()

        return mma_o_consumer_state

    def make_and_init_load_pt_pipeline(self, load_pt_mbar_ptr):
        """Create and initialize the load page table pipeline.

        :param load_pt_mbar_ptr: The load page table mbar pointer
        :type load_pt_mbar_ptr: cute.Tensor

        :return: The load page table pipeline
        :rtype: pipeline.PipelineAsync
        """
        load_pt_producer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread,
            self.threads_per_warp * len([self.load_pt_warp_id]),
        )
        load_pt_consumer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread,
            self.threads_per_warp * len([self.load_tma_warp_id]),
        )
        return pipeline.PipelineCpAsync.create(
            barrier_storage=load_pt_mbar_ptr,
            num_stages=self.load_pt_stage,
            producer_group=load_pt_producer_group,
            consumer_group=load_pt_consumer_group,
            defer_sync=True,
        )

    def make_and_init_load_qkv_pipeline(
        self, load_qkv_mbar_ptr, cta_layout_vmnk, load_stages, tx_count
    ) -> pipeline.PipelineTmaUmma:
        """Create and initialize the tma load qkv pipeline.

        :param load_qkv_mbar_ptr: The load qkv mbar pointer
        :type load_qkv_mbar_ptr: cute.Tensor
        :param cta_layout_vmnk: The cta layout vmnk
        :type cta_layout_vmnk: tuple[int, int, int]
        :param load_stages: The load stages
        :type load_stages: list[int]
        :param tx_count: The tx count
        :type tx_count: int

        :return: The tma load qkv pipeline
        :rtype: pipeline.PipelineTmaUmma
        """
        load_qkv_producer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread, len([self.load_tma_warp_id])
        )
        load_qkv_consumer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread, len([self.mma_warp_id])
        )
        return pipeline.PipelineTmaUmma.create(
            barrier_storage=load_qkv_mbar_ptr,
            num_stages=load_stages,
            producer_group=load_qkv_producer_group,
            consumer_group=load_qkv_consumer_group,
            tx_count=tx_count,
            cta_layout_vmnk=cta_layout_vmnk,
            defer_sync=True,
        )

    def make_and_init_mma_s_pipeline(
        self, mma_s_mbar_ptr, cta_layout_vmnk
    ) -> pipeline.PipelineUmmaAsync:
        """Create and initialize the mma s pipeline.

        :param mma_s_mbar_ptr: The mma s mbar pointer
        :type mma_s_mbar_ptr: cute.Tensor
        :param cta_layout_vmnk: The cta layout vmnk
        :type cta_layout_vmnk: tuple[int, int, int]

        :return: The mma s pipeline
        :rtype: pipeline.PipelineUmmaAsync
        """

        mma_s_producer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread, len([self.mma_warp_id])
        )
        consumer_thread_size = (
            self.threads_per_warp
            * len(self.compute_warp_ids)
            * self.cluster_shape_mnk[0]
        )
        mma_s_consumer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread,
            consumer_thread_size,
        )
        return pipeline.PipelineUmmaAsync.create(
            barrier_storage=mma_s_mbar_ptr,
            num_stages=self.mma_s_stage,
            producer_group=mma_s_producer_group,
            consumer_group=mma_s_consumer_group,
            cta_layout_vmnk=cta_layout_vmnk,
            defer_sync=True,
        )

    def make_and_init_p_mma_pipeline(
        self, p_mma_mbar_ptr, cta_layout_vmnk
    ) -> pipeline.PipelineAsyncUmma:
        """Create and initialize the p mma pipeline.

        :param p_mma_mbar_ptr: The p mma mbar pointer
        :type p_mma_mbar_ptr: cute.Tensor
        :param cta_layout_vmnk: The cta layout vmnk
        :type cta_layout_vmnk: tuple[int, int, int]

        :return: The p mma pipeline
        :rtype: pipeline.PipelineAsyncUmma
        """

        producer_thread_size = (
            self.threads_per_warp
            * len(self.compute_warp_ids)
            * self.cluster_shape_mnk[0]
        )
        p_mma_producer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread,
            producer_thread_size,
        )
        p_mma_consumer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread, len([self.mma_warp_id])
        )
        return pipeline.PipelineAsyncUmma.create(
            barrier_storage=p_mma_mbar_ptr,
            num_stages=self.p_mma_stage,
            producer_group=p_mma_producer_group,
            consumer_group=p_mma_consumer_group,
            cta_layout_vmnk=cta_layout_vmnk,
            defer_sync=True,
        )

    def make_and_init_p_cor_pipeline(
        self, p_cor_mbar_ptr
    ) -> pipeline.PipelineAsyncUmma:
        """Create and initialize the p correction pipeline.

        :param p_cor_mbar_ptr: The p correction mbar pointer
        :type p_cor_mbar_ptr: cute.Tensor

        :return: The p correction pipeline
        :rtype: pipeline.PipelineAsyncUmma
        """

        producer_thread_size = self.threads_per_warp * len(self.compute_warp_ids)
        p_cor_producer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread,
            producer_thread_size,
        )
        p_cor_consumer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread,
            producer_thread_size,
        )
        return pipeline.PipelineAsync.create(
            barrier_storage=p_cor_mbar_ptr,
            num_stages=self.p_cor_stage,
            producer_group=p_cor_producer_group,
            consumer_group=p_cor_consumer_group,
            defer_sync=True,
        )

    def make_and_init_mma_o_pipeline(
        self, mma_o_mbar_ptr, cta_layout_vmnk
    ) -> pipeline.PipelineUmmaAsync:
        """Create and initialize the mma o pipeline.

        :param mma_o_mbar_ptr: The mma o mbar pointer
        :type mma_o_mbar_ptr: cute.Tensor
        :param cta_layout_vmnk: The cta layout vmnk
        :type cta_layout_vmnk: tuple[int, int, int]

        :return: The mma o pipeline
        :rtype: pipeline.PipelineUmmaAsync
        """

        mma_o_producer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread, len([self.mma_warp_id])
        )
        consumer_thread_size = (
            self.threads_per_warp
            * len(self.compute_warp_ids)
            * self.cluster_shape_mnk[0]
        )
        mma_o_consumer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread,
            consumer_thread_size,
        )
        return pipeline.PipelineUmmaAsync.create(
            barrier_storage=mma_o_mbar_ptr,
            num_stages=self.mma_o_stage,
            producer_group=mma_o_producer_group,
            consumer_group=mma_o_consumer_group,
            cta_layout_vmnk=cta_layout_vmnk,
            defer_sync=True,
        )

    @staticmethod
    def _compute_grid(
        batch_size: cutlass.Int32,
        num_q_tiles: cutlass.Int32,
        split_kv: cutlass.Int32,
        cluster_shape_mnk: Tuple[int, int, int],
        max_active_clusters: int,
        is_persistent: bool,
    ) -> Tuple[MLAStaticTileSchedulerParams, Tuple[int, int, int]]:
        """Compute grid shape for the output tensor C.

        :param c: The output tensor C
        :type c: cute.Tensor
        :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile.
        :type cta_tile_shape_mnk: tuple[int, int, int]
        :param cluster_shape_mn: Shape of each cluster in M, N dimensions.
        :type cluster_shape_mn: tuple[int, int]

        :return: Tile scheduler parameters and grid shape.
        :rtype: tuple[MLAStaticTileSchedulerParams, tuple[int, int, int]]
        """
        tile_sched_params = create_mla_static_tile_scheduler_params(
            is_persistent,
            batch_size,
            num_q_tiles,
            cluster_shape_mnk,
            split_kv,
        )
        grid = MLAStaticTileScheduler.get_grid_shape(
            tile_sched_params, max_active_clusters
        )

        return tile_sched_params, grid

    @staticmethod
    def get_workspace_size(
        H: int,
        S: int,
        D: int,
        B: int,
        split_kv: int,
        acc_dtype: Type[cutlass.Numeric],
    ) -> int:
        """Get the extra workspace(device memory) size for the MLA kernel when split_kv is not 1.

        :param H: The height of the output tensor C
        :type H: int
        :param S: The sequence length of the output tensor C
        :type S: int
        :param D: The depth of the output tensor C
        :type D: int
        :param B: The batch size of the output tensor C
        :type B: int
        :param split_kv: The split key-value of the output tensor C
        :type split_kv: int
        :param acc_dtype: The data type of the output tensor C
        :type acc_dtype: Type[cutlass.Numeric]

        :return: The workspace size for the MLA kernel
        :rtype: int
        """
        if split_kv == 1:
            return 0
        # The first workspace mode is the physical MMA-M row extent. Direct
        # callers may still pass a logical H < 128, so preserve the M128 floor
        # used by initialize_workspace(). The flat internal path passes 128.
        workspace_rows = max(H, 128)
        return B * workspace_rows * S * split_kv * (D + 1) * acc_dtype.width // 8

    @cute.jit
    def initialize_workspace(
        self,
        H: cutlass.Int32,
        D: cutlass.Int32,
        S: cutlass.Int32,
        B: cutlass.Int32,
        split_kv: cutlass.Int32,
        acc_dtype: Type[cutlass.Numeric],
        workspace: cute.Tensor,
    ) -> tuple[cute.Tensor, cute.Tensor]:
        """Initialize the workspace for the MLA kernel. Construct the intermediate tensors
        acc_o and acc_lse.

        :param H: The height of the output tensor C
        :type H: cutlass.Int32
        :param D: The depth of the output tensor C
        :type D: cutlass.Int32
        :param S: The sequence length of the output tensor C
        :type S: cutlass.Int32
        :param B: The batch size of the output tensor C
        :type B: cutlass.Int32
        :param split_kv: The split key-value of the output tensor C
        :type split_kv: cutlass.Int32
        :param acc_dtype: The data type of the output tensor C
        :type acc_dtype: Type[cutlass.Numeric]
        :param workspace: The workspace tensor
        :type workspace: cute.Tensor

        :return: The output tensor C and the workspace tensor
        :rtype: tuple[cute.Tensor, cute.Tensor]
        """
        acc_o, acc_lse = None, None
        if cutlass.const_expr(workspace is not None):
            # The first workspace mode is the physical MMA-M row extent.
            # Preserve the M128 floor for direct callers that pass H < 128.
            workspace_rows = cutlass.max(H, cutlass.Int32(128))
            align = 256 // self.q_dtype.width
            acc_o_layout = cute.make_layout(
                (workspace_rows, split_kv, D, S, B),
                stride=(
                    cute.assume(split_kv * D, align),
                    cute.assume(D, align),
                    1,
                    cute.assume(split_kv * workspace_rows * D, align),
                    cute.assume(workspace_rows * split_kv * S * D, align),
                ),
            )
            acc_o_iter = cute.recast_ptr(workspace.iterator, dtype=acc_dtype)
            acc_o = cute.make_tensor(acc_o_iter, acc_o_layout)
            acc_lse_layout = cute.make_layout(
                (workspace_rows, split_kv, S, B),
                stride=(
                    split_kv,
                    1,
                    workspace_rows * split_kv,
                    workspace_rows * split_kv * S,
                ),
            )
            acc_lse_iter = cute.recast_ptr(
                workspace.iterator + cute.cosize(acc_o_layout) * acc_dtype.width // 8,
                dtype=acc_dtype,
            )
            acc_lse = cute.make_tensor(acc_lse_iter, acc_lse_layout)
        return acc_o, acc_lse

    @staticmethod
    def can_implement(
        B: int,
        S: int,
        K: int,
        H: int,
        L: int,
        R: int,
        in_dtype: Type[cutlass.Numeric],
        out_dtype: Type[cutlass.Numeric],
        acc_dtype: Type[cutlass.Numeric],
        lse_dtype: Type[cutlass.Numeric],
        mma_qk_tiler_mn: Tuple[int, int],
        mma_pv_tiler_mn: Tuple[int, int],
        is_persistent: bool,
        is_var_seq: bool,
        is_var_split_kv: bool,
        page_size: int,
    ) -> bool:
        """Check if the MLA kernel can be implemented.

        :param B: The batch size of the output tensor C
        :type B: int
        :param S: The sequence length of the output tensor C
        :type S: int
        :param K: The width of the output tensor KV
        :type K: int
        :param H: The number of heads of the output tensor C
        :type H: int
        :param L: The number of latent dimensions of the tensor KV
        :type L: int
        :param R: The number of rope dimensions of the tensor C_rope
        :type R: int
        :param in_dtype: The data type of the input tensor
        :type in_dtype: Type[cutlass.Numeric]
        :param out_dtype: The data type of the output tensor
        :type out_dtype: Type[cutlass.Numeric]
        :param acc_dtype: The data type of the accumulator
        :type acc_dtype: Type[cutlass.Numeric]
        :param lse_dtype: The data type of the log-sum-exp
        :type lse_dtype: Type[cutlass.Numeric]
        :param mma_qk_tiler_mn: The tile shape of the query-key matrix multiplication
        :type mma_qk_tiler_mn: Tuple[int, int]
        :param mma_pv_tiler_mn: The tile shape of the probability-value matrix multiplication
        :type mma_pv_tiler_mn: Tuple[int, int]
        :param is_persistent: Whether to use persistent kernel optimization
        :type is_persistent: bool
        :param is_var_seq: Whether to use variable sequence length
        :type is_var_seq: bool
        :param is_var_split_kv: Whether to use variable split_kv
        :type is_var_split_kv: bool
        :param page_size: The page size of the page table
        :type page_size: int

        :return: Whether the MLA kernel can be implemented
        :rtype: bool
        """
        if L != 512 or R != 64:
            return False
        if in_dtype not in [cutlass.Float16, cutlass.BFloat16]:
            return False
        if out_dtype not in [cutlass.Float16, cutlass.BFloat16]:
            return False
        if acc_dtype != cutlass.Float32 or lse_dtype != cutlass.Float32:
            return False
        # page size equals 1 is prohibited by tma specification, not 128B aligned.
        if mma_qk_tiler_mn[1] % page_size != 0 or page_size == 1:
            return False
        if mma_qk_tiler_mn[0] != mma_pv_tiler_mn[0] or mma_qk_tiler_mn[0] != 128:
            return False
        if is_var_split_kv and not is_var_seq:
            return False
        if H <= 0 or mma_qk_tiler_mn[0] < H:
            return False
        # Query (token, head) rows are packed continuously into 128-row M
        # tiles.  A partial final tile is supported, so any H <= M_tile and
        # S >= 1 is valid.
        if S < 1:
            return False
        if K <= 0:
            return False
        return True
