# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
from typing import Tuple, Type, Optional

import cuda.bindings.driver as cuda
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass
import cutlass.cute as cute
from cutlass.cute.typing import Int32, Float32, Int64

from cutlass._mlir.dialects import cute_nvgpu

"""
A NSA(Native Sparse Attention) Top-K Reduction Forward Pass for NVIDIA Blackwell SM100 architecture using Cute DSL.
"""


class FineGrainedReductionQK:
    def __init__(
        self,
        element_dtype: Type[cutlass.Numeric],
        acc_dtype: Type[cutlass.Numeric],
        k_value: int,
        selection_block_size: int,
        compress_block_sliding_stride: int,
        mma_tiler: Tuple[int, int, int],
        is_causal: bool = False,
    ):
        self.element_dtype = element_dtype
        self.acc_dtype = acc_dtype
        self.indices_dtype = Int32
        self.k_value = k_value
        self.selection_block_size = selection_block_size
        self.compress_block_sliding_stride = compress_block_sliding_stride
        self.num_elem_for_reduction = selection_block_size // compress_block_sliding_stride

        self.cluster_shape_mn = (1, 1)
        self.mma_tiler = mma_tiler
        self.is_causal = is_causal

        self.compute_warp_id = (0, 1, 2, 3)
        # self.reduce_warp_id = (4, 5, 6, 7)
        self.load_warp_id = 4
        self.mma_warp_id = 5
        self.epi_warp_id = 6
        self.threads_per_warp = 32
        self.num_compute_warps = len(self.compute_warp_id)
        # self.num_reduce_warps = len(self.reduce_warp_id)

        self.threads_per_cta = self.threads_per_warp * (self.num_compute_warps + 4)

        self.cta_sync_bar_id = 0
        self.tmem_alloc_sync_bar_id = 1
        self.compute_sync_bar_id = 2
        # self.reduce_sync_bar_id = 3
        # self.num_regs_compute = 240
        # self.num_regs_reduce = 96
        # self.num_regs_other = 72
        self.num_regs_compute = 256
        # self.num_regs_reduce = 96
        self.num_regs_other = 96

    @cute.jit
    def __call__(
        self,
        problem_size: Tuple[Int32, Int32, Int32, Int32, Int32, Int32],
        Q: cute.Tensor,
        K: cute.Tensor,
        LSE: cute.Tensor,
        Topk_scores: cute.Tensor,
        Topk_indices: cute.Tensor,
        softmax_scale_log2_e: Float32,
        cumulative_s_q: Optional[cute.Tensor],
        cumulative_s_k: Optional[cute.Tensor],
        stream: cuda.CUstream,
    ):
        b, s_q_max, s_k_max, _, _, _ = problem_size
        s_q, s_k = Q.shape[2], K.shape[2]
        h_q, h_k = Q.shape[1], K.shape[1]
        head_dim = Q.shape[3]
        h_r = h_q // h_k

        stride_b_q = s_q * head_dim * h_k * h_r if cumulative_s_q is None else 0
        stride_b_k = s_k * head_dim * h_k if cumulative_s_k is None else 0
        stride_b_lse = s_q * h_r * h_k if cumulative_s_q is None else 0
        stride_b_out = s_q * (s_k_max // self.num_elem_for_reduction) * h_k if cumulative_s_q is None else 0
        stride_b_topk_scores = s_q * self.k_value * h_k if cumulative_s_q is None else 0
        stride_b_topk_indices = s_q * self.k_value * h_k if cumulative_s_q is None else 0

        Q = cute.make_tensor(
            Q.iterator,
            cute.make_layout(
                (s_q, head_dim, (h_r, h_k, b)),
                stride=(
                    head_dim * h_r * h_k,
                    1,
                    (
                        head_dim * h_k,
                        head_dim,
                        stride_b_q,
                    ),  # TODO: head ordering is diff?
                ),
            ),
        )

        K = cute.make_tensor(
            K.iterator,
            cute.make_layout(
                (s_k, head_dim, (1, h_k, b)),
                stride=(
                    head_dim * h_k,
                    1,
                    (0, head_dim, stride_b_k),
                ),
            ),
        )

        # reshape LSE to (s_q, h_r, h_k, b)
        LSE = cute.make_tensor(
            LSE.iterator,
            cute.make_layout(
                (s_q, h_r, h_k, b),
                stride=(1, s_q * h_k, s_q, stride_b_lse),
            ),
        )

        Topk_scores = cute.make_tensor(
            Topk_scores.iterator,
            cute.make_layout(
                (s_q, self.k_value, (1, h_k, b)),
                stride=(self.k_value * h_k, 1, (0, self.k_value, stride_b_topk_scores)),
            ),
        )

        Topk_indices = cute.make_tensor(
            Topk_indices.iterator,
            cute.make_layout(
                (s_q, self.k_value, (1, h_k, b)),  # (s_q, k, h_k, b)
                stride=(
                    self.k_value * h_k,
                    1,
                    (0, self.k_value, stride_b_topk_indices),
                ),
            ),
        )

        self.q_major_mode = utils.LayoutEnum.from_tensor(Q).mma_major_mode()
        self.k_major_mode = utils.LayoutEnum.from_tensor(K).mma_major_mode()
        self.load_mma_Q_stage = 2
        self.load_mma_K_stage = 1
        self.load_compute_LSE_stage = 1
        self.mma_compute_S_stage = 1
        cta_group = tcgen05.CtaGroup.ONE

        QK_tiled_mma = sm100_utils.make_trivial_tiled_mma(
            self.element_dtype,
            self.element_dtype,
            self.q_major_mode,
            self.k_major_mode,
            self.acc_dtype,
            cta_group,
            self.mma_tiler[:2],
        )
        Q_smem_layout_staged = sm100_utils.make_smem_layout_a(
            QK_tiled_mma,
            self.mma_tiler,
            self.element_dtype,
            self.load_mma_Q_stage,
        )
        K_smem_layout_staged = sm100_utils.make_smem_layout_b(
            QK_tiled_mma,
            self.mma_tiler,
            self.element_dtype,
            self.load_mma_K_stage,
        )
        LSE_smem_layout = cute.make_layout((self.mma_tiler[0], 1))

        self.epi_tile = (self.mma_tiler[0], self.k_value)
        self.cluster_shape_mnk = (*self.cluster_shape_mn, 1)
        self.cluster_layout_vmnk = cute.tiled_divide(
            cute.make_layout(self.cluster_shape_mnk),
            (QK_tiled_mma.thr_id.shape,),
        )

        tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group)
        tma_store_op = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp()

        Q_smem_layout = cute.select(Q_smem_layout_staged, mode=[0, 1, 2])
        K_smem_layout = cute.select(K_smem_layout_staged, mode=[0, 1, 2])

        tma_atom_Q, tma_tensor_Q = cute.nvgpu.make_tiled_tma_atom_A(
            tma_load_op,
            Q,
            Q_smem_layout,
            self.mma_tiler,
            QK_tiled_mma,
            self.cluster_layout_vmnk,
        )
        tma_atom_K, tma_tensor_K = cute.nvgpu.make_tiled_tma_atom_B(
            tma_load_op,
            K,
            K_smem_layout,
            self.mma_tiler,
            QK_tiled_mma,
            self.cluster_layout_vmnk,
        )

        self.tma_copy_Q_bytes = cute.size_in_bytes(self.element_dtype, Q_smem_layout)
        self.tma_copy_K_bytes = cute.size_in_bytes(self.element_dtype, K_smem_layout)

        @cute.struct
        class SharedStorage:
            load_mma_Q_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_mma_Q_stage * 2]
            load_mma_K_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_mma_K_stage * 2]
            load_compute_LSE_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_compute_LSE_stage * 2]
            mma_compute_S_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.mma_compute_S_stage * 2]

            sQ: cute.struct.Align[
                cute.struct.MemRange[self.element_dtype, cute.cosize(Q_smem_layout_staged)],
                1024,
            ]
            sK: cute.struct.Align[
                cute.struct.MemRange[self.element_dtype, cute.cosize(K_smem_layout_staged)],
                1024,
            ]
            sLSE: cute.struct.Align[
                cute.struct.MemRange[self.acc_dtype, cute.cosize(LSE_smem_layout)],
                1024,
            ]

            tmem_holding_buf: Int32

        self.shared_storage = SharedStorage

        grid_shape = (cute.ceil_div(s_q_max, self.mma_tiler[0]), h_k, b)
        block_shape = (self.threads_per_cta, 1, 1)

        self.kernel(
            problem_size,
            QK_tiled_mma,
            tma_atom_Q,
            tma_tensor_Q,
            tma_atom_K,
            tma_tensor_K,
            LSE,
            cumulative_s_q,
            cumulative_s_k,
            Topk_scores,
            Topk_indices,
            Q_smem_layout_staged,
            K_smem_layout_staged,
            LSE_smem_layout,
            softmax_scale_log2_e,
        ).launch(
            grid=grid_shape,
            block=block_shape,
            smem=SharedStorage.size_in_bytes(),
            stream=stream,
            min_blocks_per_mp=1,
        )

    def make_and_init_load_mma_Q_pipeline(self, load_mma_Q_mbar_ptr):
        load_mma_Q_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, len([self.load_warp_id]))
        load_mma_Q_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, len([self.mma_warp_id]))
        return pipeline.PipelineTmaUmma.create(
            barrier_storage=load_mma_Q_mbar_ptr,
            num_stages=self.load_mma_Q_stage,
            producer_group=load_mma_Q_producer_group,
            consumer_group=load_mma_Q_consumer_group,
            tx_count=self.tma_copy_Q_bytes,
        )

    def make_and_init_load_mma_K_pipeline(self, load_mma_K_mbar_ptr):
        load_mma_K_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, len([self.load_warp_id]))
        load_mma_K_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, len([self.mma_warp_id]))
        return pipeline.PipelineTmaUmma.create(
            barrier_storage=load_mma_K_mbar_ptr,
            num_stages=self.load_mma_K_stage,
            producer_group=load_mma_K_producer_group,
            consumer_group=load_mma_K_consumer_group,
            tx_count=self.tma_copy_K_bytes,
        )

    def make_and_init_load_compute_LSE_pipeline(self, load_compute_lse_mbar_ptr):
        load_compute_lse_producer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread,
            self.threads_per_warp,
        )
        load_compute_lse_consumer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread,
            self.threads_per_warp * self.num_compute_warps,
        )
        return pipeline.PipelineCpAsync.create(
            barrier_storage=load_compute_lse_mbar_ptr,
            num_stages=self.load_compute_LSE_stage,
            producer_group=load_compute_lse_producer_group,
            consumer_group=load_compute_lse_consumer_group,
        )

    def make_and_init_mma_compute_S_pipeline(self, mma_compute_S_mbar_ptr):
        mma_compute_S_producer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread,
            len([self.mma_warp_id]),
        )
        mma_compute_S_consumer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread,
            self.num_compute_warps * self.threads_per_warp,
        )
        return pipeline.PipelineUmmaAsync.create(
            barrier_storage=mma_compute_S_mbar_ptr,
            num_stages=self.mma_compute_S_stage,
            producer_group=mma_compute_S_producer_group,
            consumer_group=mma_compute_S_consumer_group,
        )

    @cute.kernel
    def kernel(
        self,
        problem_size: Tuple[Int32, Int32, Int32, Int32, Int32, Int32],
        QK_tiled_mma: cute.TiledMma,
        tma_atom_Q: cute.CopyAtom,
        tma_tensor_Q: cute.Tensor,
        tma_atom_K: cute.CopyAtom,
        tma_tensor_K: cute.Tensor,
        LSE: cute.Tensor,
        cumulative_s_q: Optional[cute.Tensor],
        cumulative_s_k: Optional[cute.Tensor],
        Topk_scores: cute.Tensor,
        Topk_indices: cute.Tensor,
        Q_smem_layout_staged: cute.ComposedLayout,
        K_smem_layout_staged: cute.ComposedLayout,
        LSE_smem_layout: cute.Layout,
        softmax_scale_log2_e: Float32,
    ):
        tidx, tidy, _ = cute.arch.thread_idx()
        bidx, bidy, bidz = cute.arch.block_idx()

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

        b, s_q_max, s_k_max, h_q, h_k, head_dim = problem_size
        h_r = h_q // h_k

        smem = utils.SmemAllocator()
        storage = smem.allocate(self.shared_storage)

        load_mma_Q_pipeline = self.make_and_init_load_mma_Q_pipeline(storage.load_mma_Q_mbar_ptr.data_ptr())
        load_mma_K_pipeline = self.make_and_init_load_mma_K_pipeline(storage.load_mma_K_mbar_ptr.data_ptr())
        load_compute_LSE_pipeline = self.make_and_init_load_compute_LSE_pipeline(storage.load_compute_LSE_mbar_ptr.data_ptr())
        mma_compute_S_pipeline = self.make_and_init_mma_compute_S_pipeline(storage.mma_compute_S_mbar_ptr.data_ptr())

        cute.arch.barrier(barrier_id=self.cta_sync_bar_id, number_of_threads=self.threads_per_cta)

        sQ = storage.sQ.get_tensor(Q_smem_layout_staged.outer, swizzle=Q_smem_layout_staged.inner)
        sK = storage.sK.get_tensor(K_smem_layout_staged.outer, swizzle=K_smem_layout_staged.inner)
        sLSE = storage.sLSE.get_tensor(LSE_smem_layout)

        block_offset = (Int32(0), Int32(0), Int32(0), (Int32(0), Int32(0), Int32(0)))
        cur_s_q = s_q_max
        cur_s_k = s_k_max
        if cutlass.const_expr(cumulative_s_q is not None):
            cur_s_q = cumulative_s_q[bidz + 1] - cumulative_s_q[bidz]
            cur_s_k = cumulative_s_k[bidz + 1] - cumulative_s_k[bidz]
            block_offset = (
                cumulative_s_q[bidz],
                cumulative_s_k[bidz],
                Int32(0),
                (Int32(0), Int32(0), Int32(0)),
            )

        k_tile_count = cute.ceil_div(cur_s_k, self.mma_tiler[1])
        k_tile_idx = 0
        if cutlass.const_expr(self.is_causal):
            q_max = min(cur_s_q, (bidx + 1) * self.mma_tiler[0])
            k_cols = min(cur_s_k, q_max // self.compress_block_sliding_stride)
            k_tile_count = cute.ceil_div(k_cols, self.mma_tiler[1])

        mQ = cute.domain_offset(cute.select(block_offset, mode=[0, 2, 3]), tma_tensor_Q)
        mK = cute.domain_offset(cute.select(block_offset, mode=[1, 2, 3]), tma_tensor_K)
        mTopk_scores = cute.make_tensor(
            Topk_scores.iterator + cute.assume(block_offset[0] * Topk_scores.stride[0], divby=self.k_value),
            Topk_scores.layout,
        )
        mTopk_indices = cute.make_tensor(
            Topk_indices.iterator + cute.assume(block_offset[0] * Topk_indices.stride[0], divby=self.k_value),
            Topk_indices.layout,
        )

        # (MMA_M, MMA_K, REST_M, REST_K, (H_r, H_k, B))
        gQ = cute.local_tile(mQ, cute.select(self.mma_tiler, mode=[0, 2]), (None, None, None))
        # (MMA_N, MMA_K, REST_N, REST_K, (1, H_k, B))
        gK = cute.local_tile(mK, cute.select(self.mma_tiler, mode=[0, 2]), (None, None, None))

        # (MMA_M, MMA_K, H_r)
        gQ = gQ[None, None, bidx, 0, (None, bidy, bidz)]
        # (MMA_N, MMA_K, REST_N)
        gK = gK[None, None, None, 0, (0, bidy, bidz)]

        thr_mma = QK_tiled_mma.get_slice(0)

        tSgQ = thr_mma.partition_A(gQ)
        tSgK = thr_mma.partition_B(gK)

        # tQgQ: (MMA, H_r)
        # tQsQ: (MMA, PIPE)
        tQsQ, tQgQ = cute.nvgpu.cpasync.tma_partition(
            tma_atom_Q,
            0,
            cute.make_layout(1),
            cute.group_modes(sQ, 0, 3),
            cute.group_modes(tSgQ, 0, 3),
        )

        # tKgK: (MMA, REST_N)
        # tKsK: (MMA, PIPE)
        tKsK, tKgK = cute.nvgpu.cpasync.tma_partition(
            tma_atom_K,
            0,
            cute.make_layout(1),
            cute.group_modes(sK, 0, 3),
            cute.group_modes(tSgK, 0, 3),
        )

        tSrQ = QK_tiled_mma.make_fragment_A(sQ)
        tSrK = QK_tiled_mma.make_fragment_B(sK)
        tStS_shape = QK_tiled_mma.partition_shape_C(cute.select(self.mma_tiler, mode=[0, 1]))
        # ((MMA_M, MMA_N), REST_M, REST_N)
        tStS = QK_tiled_mma.make_fragment_C(tStS_shape)
        # another tmem for reduction
        tStS_reduce = cute.make_tensor(tStS.iterator + self.mma_tiler[1], tStS.layout)

        if bidx * self.mma_tiler[0] < cur_s_q and k_tile_count > 0:

            load_iter_count = k_tile_count
            load_iter_index = k_tile_idx
            mma_iter_count = k_tile_count
            mma_iter_index = k_tile_idx
            compute_iter_count = k_tile_count
            compute_iter_index = k_tile_idx

            # LOAD Q K WARP
            if warp_idx == self.load_warp_id:
                # TODO: reconfig regs
                cute.arch.setmaxregister_decrease(self.num_regs_other)

                load_mma_Q_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.load_mma_Q_stage)
                load_mma_K_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.load_mma_K_stage)
                load_compute_LSE_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.load_compute_LSE_stage)

                while load_iter_count > 0:
                    # Wait for K to be empty
                    load_mma_K_pipeline.producer_acquire(load_mma_K_producer_state)
                    K_tma_barrier = load_mma_K_pipeline.producer_get_barrier(load_mma_K_producer_state)

                    # Load K tile
                    cute.copy(
                        tma_atom_K,
                        tKgK[None, load_iter_index],
                        tKsK[None, load_mma_K_producer_state.index],
                        tma_bar_ptr=K_tma_barrier,
                    )

                    load_mma_K_producer_state.advance()

                    # Load Q and LSE
                    for h_r_idx in cutlass.range(cute.size(tQgQ, mode=[1])):

                        load_compute_LSE_pipeline.producer_acquire(load_compute_LSE_producer_state)

                        # Load LSE
                        thread_idx = tidx % self.threads_per_warp

                        async_copy_num_elts = sLSE.shape[0] // self.threads_per_warp
                        atom_async_copy = cute.make_copy_atom(
                            cpasync.CopyG2SOp(cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS),
                            self.acc_dtype,
                            num_bits_per_copy=self.acc_dtype.width,
                        )

                        sLSE_for_copy = cute.flat_divide(sLSE, (1,))
                        LSE_for_copy = cute.flat_divide(LSE, (1,))
                        LSE_idx_offset = block_offset[0] * LSE.stride[0]
                        for i in cutlass.range_constexpr(async_copy_num_elts):
                            LSE_idx = self.mma_tiler[0] * bidx + thread_idx * async_copy_num_elts
                            if cute.elem_less(LSE_idx + i, cur_s_q):
                                cute.copy(
                                    atom_async_copy,
                                    LSE_for_copy[
                                        None,
                                        LSE_idx_offset + LSE_idx + i,
                                        h_r_idx,
                                        bidy,
                                        bidz,
                                    ],
                                    sLSE_for_copy[
                                        None,
                                        thread_idx * async_copy_num_elts + i,
                                        load_compute_LSE_producer_state.index,
                                    ],
                                )
                            else:
                                sLSE_for_copy[
                                    None,
                                    thread_idx * async_copy_num_elts + i,
                                    load_compute_LSE_producer_state.index,
                                ].fill(0.0)

                        load_compute_LSE_pipeline.producer_commit(load_compute_LSE_producer_state)
                        load_compute_LSE_producer_state.advance()

                        # Wait for Q to be empty
                        load_mma_Q_pipeline.producer_acquire(load_mma_Q_producer_state)
                        Q_tma_barrier = load_mma_Q_pipeline.producer_get_barrier(load_mma_Q_producer_state)

                        # Load Q tile
                        cute.copy(
                            tma_atom_Q,
                            tQgQ[None, h_r_idx],
                            tQsQ[None, load_mma_Q_producer_state.index],
                            tma_bar_ptr=Q_tma_barrier,
                        )

                        load_mma_Q_producer_state.advance()

                    load_iter_count -= 1
                    load_iter_index += 1

            # MMA WARP
            if warp_idx == self.mma_warp_id:
                # TODO: reconfig regs
                cute.arch.setmaxregister_decrease(self.num_regs_other)

                num_tmem_cols = 512
                cute.arch.alloc_tmem(num_tmem_cols, storage.tmem_holding_buf.ptr)
                cute.arch.barrier(
                    barrier_id=self.tmem_alloc_sync_bar_id,
                    number_of_threads=self.threads_per_warp,
                )

                load_mma_Q_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.load_mma_Q_stage)
                load_mma_K_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.load_mma_K_stage)
                mma_compute_S_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.mma_compute_S_stage)

                while mma_iter_count > 0:
                    # Wait for K to be full
                    load_mma_K_pipeline.consumer_wait(load_mma_K_consumer_state)

                    for h_r_idx in cutlass.range(cute.size(tQgQ, mode=[1])):
                        # Wait for Q to be full
                        load_mma_Q_pipeline.consumer_wait(load_mma_Q_consumer_state)
                        mma_compute_S_pipeline.producer_acquire(mma_compute_S_producer_state)

                        QK_tiled_mma.set(cute.nvgpu.tcgen05.Field.ACCUMULATE, False)
                        for k_block_idx in cutlass.range_constexpr(cute.size(tSrQ, mode=[2])):
                            cute.gemm(
                                QK_tiled_mma,
                                tStS,
                                tSrQ[
                                    None,
                                    None,
                                    k_block_idx,
                                    load_mma_Q_consumer_state.index,
                                ],
                                tSrK[
                                    None,
                                    None,
                                    k_block_idx,
                                    load_mma_K_consumer_state.index,
                                ],
                                tStS,
                            )
                            QK_tiled_mma.set(cute.nvgpu.tcgen05.Field.ACCUMULATE, True)

                        mma_compute_S_pipeline.producer_commit(mma_compute_S_producer_state)
                        mma_compute_S_producer_state.advance()

                        load_mma_Q_pipeline.consumer_release(load_mma_Q_consumer_state)
                        load_mma_Q_consumer_state.advance()

                    mma_iter_count -= 1
                    mma_iter_index += 1

                    load_mma_K_pipeline.consumer_release(load_mma_K_consumer_state)
                    load_mma_K_consumer_state.advance()

            # COMPUTE WARP
            if warp_idx in self.compute_warp_id:
                cute.arch.setmaxregister_increase(self.num_regs_compute)

                mma_compute_S_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.mma_compute_S_stage)
                load_compute_LSE_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.load_compute_LSE_stage)

                load_compute_LSE_pipeline.consumer_wait(load_compute_LSE_consumer_state)
                thread_idx = tidx % (self.threads_per_warp * self.num_compute_warps)

                heap_size_ref = cute.make_rmem_tensor((1,), Int32)
                heap_size_ref[0] = 0
                # # Create temporary register heaps for computation
                scores_heap_rf = cute.make_rmem_tensor(((4, self.k_value // 4), 1, 1), Float32)
                idx_heap_rf = cute.make_rmem_tensor(((4, self.k_value // 4), 1, 1), Int32)

                tmem_load_atom = cute.make_copy_atom(
                    tcgen05.Ld32x32bOp(tcgen05.copy.Repetition(32)),
                    self.acc_dtype,
                )
                tmem_store_atom = cute.make_copy_atom(
                    tcgen05.St32x32bOp(tcgen05.copy.Repetition(32)),
                    self.acc_dtype,
                )

                cS = cute.make_identity_tensor((self.mma_tiler[0], self.mma_tiler[1]))

                comp_tile_size = 32
                tStS_tiled = cute.logical_divide(tStS, cute.make_layout((self.mma_tiler[0], comp_tile_size)))
                tStS_compute_tiled = cute.logical_divide(tStS_reduce, cute.make_layout((self.mma_tiler[0], comp_tile_size)))
                cS_tiled = cute.logical_divide(cS, cute.make_layout((self.mma_tiler[0], comp_tile_size)))

                tStS_slice = tStS_tiled[None, 0]  # ((128, 16), 8)
                tStS_compute_slice = tStS_compute_tiled[None, 0]

                # (MMA_M, MMA_N)
                tiled_t2r = tcgen05.make_tmem_copy(tmem_load_atom, tStS_slice)
                thr_t2r = tiled_t2r.get_slice(thread_idx)

                tTR_cS = thr_t2r.partition_D(cS_tiled)
                tTR_tS = thr_t2r.partition_S(tStS_tiled)
                tTR_tS_compute = thr_t2r.partition_S(tStS_compute_tiled)
                tTR_rS = cute.make_rmem_tensor(tTR_cS[None, None, 0].shape, self.acc_dtype)
                tTR_rS_compute = cute.make_rmem_tensor(tTR_cS[None, None, 0].shape, self.acc_dtype)

                tiled_r2t = tcgen05.make_tmem_copy(tmem_store_atom, tStS_compute_slice)
                thr_r2t = tiled_r2t.get_slice(thread_idx)
                tRT_tS_compute = thr_r2t.partition_D(tStS_compute_tiled)

                tiled_t2r_reduce = tcgen05.make_tmem_copy(tmem_load_atom, tStS[(None, None), 0, 0])
                thr_t2r_reduce = tiled_t2r_reduce.get_slice(thread_idx)
                tTR_tS_reduce = thr_t2r_reduce.partition_S(tStS_reduce[(None, None), 0, 0])
                tTR_cS_reduce = thr_t2r_reduce.partition_D(cS)
                tTR_rS_reduce = cute.make_rmem_tensor(tTR_cS_reduce.shape, self.acc_dtype)

                tmp = cute.make_rmem_tensor((self.mma_tiler[1] // self.num_elem_for_reduction), self.acc_dtype)

                while compute_iter_count > 0:

                    for h_r_idx in range(cute.size(tQgQ, mode=[1])):
                        mma_compute_S_pipeline.consumer_wait(mma_compute_S_consumer_state)

                        # TODO: Added this as we should wait for the producer to load
                        load_compute_LSE_pipeline.consumer_wait(load_compute_LSE_consumer_state)

                        for sub_tile in cutlass.range(self.mma_tiler[1] // comp_tile_size):
                            tTR_tS_sub_tile = tTR_tS[None, None, sub_tile]
                            tTR_tS_compute_sub_tile = tTR_tS_compute[None, None, sub_tile]
                            tRT_tS_compute_sub_tile = tRT_tS_compute[None, None, sub_tile]
                            tTR_cS_sub_tile = tTR_cS[None, None, sub_tile]

                            # Copy S from tmem to rmem
                            cute.copy(tiled_t2r, tTR_tS_sub_tile, tTR_rS)

                            is_residual_k = compute_iter_index * self.mma_tiler[1] + self.mma_tiler[1] > cur_s_k

                            leading_causal_masking = cutlass.Boolean(False)
                            if cutlass.const_expr(self.is_causal):
                                leading_causal_masking = (
                                    (compute_iter_index + 1) * self.mma_tiler[1] + 1
                                ) * self.compress_block_sliding_stride - 1 > bidx * self.mma_tiler[0]
                                leading_causal_masking = cute.arch.shuffle_sync(leading_causal_masking, 0)
                            trailing_residual_masking = cutlass.Boolean(False)
                            trailing_residual_masking = is_residual_k
                            trailing_residual_masking = cute.arch.shuffle_sync(trailing_residual_masking, 0)

                            is_masked_tile = leading_causal_masking or trailing_residual_masking

                            # Apply mask
                            if is_masked_tile:
                                for i in cutlass.range(cute.size(tTR_rS), unroll_full=True):
                                    q_idx = cute.get(tTR_cS_sub_tile[i], mode=[0]) + bidx * self.mma_tiler[0]
                                    k_block_idx = cute.get(tTR_cS_sub_tile[i], mode=[1]) + compute_iter_index * self.mma_tiler[1]

                                    if is_masked_tile:
                                        if cutlass.const_expr(self.is_causal):
                                            k_idx = (k_block_idx + 1) * self.compress_block_sliding_stride - 1
                                            if k_idx > q_idx:
                                                tTR_rS[i] = -cutlass.Float32.inf
                                            if q_idx > cur_s_q or k_block_idx > cur_s_k:
                                                tTR_rS[i] = -cutlass.Float32.inf
                                        else:
                                            if q_idx > cur_s_q or k_block_idx > cur_s_k:
                                                tTR_rS[i] = -cutlass.Float32.inf

                            # P = exp2(S * softmax_scale_log2_e - LSE),
                            # LSE should be set negative before and has be already multiplied by log2_e

                            # Copy S_reduce from tmem to rmem
                            cute.copy(tiled_t2r, tTR_tS_compute_sub_tile, tTR_rS_compute)

                            for i in cutlass.range(0, cute.size(tTR_rS, mode=[0]), 2):
                                lse = (
                                    sLSE[
                                        cute.get(tTR_cS_sub_tile[i], mode=[0]),
                                        load_compute_LSE_consumer_state.index,
                                    ],
                                    sLSE[
                                        cute.get(tTR_cS_sub_tile[i + 1], mode=[0]),
                                        load_compute_LSE_consumer_state.index,
                                    ],
                                )

                                tTR_rS[i], tTR_rS[i + 1] = cute.arch.fma_packed_f32x2(
                                    (tTR_rS[i], tTR_rS[i + 1]),
                                    (
                                        softmax_scale_log2_e,
                                        softmax_scale_log2_e,
                                    ),
                                    lse,
                                )
                                tTR_rS[i] = cute.math.exp2(tTR_rS[i], fastmath=True)
                                tTR_rS[i + 1] = cute.math.exp2(tTR_rS[i + 1], fastmath=True)

                                if h_r_idx == 0:
                                    tTR_rS_compute[i], tTR_rS_compute[i + 1] = cute.arch.add_packed_f32x2(
                                        (0.0, 0.0),
                                        (tTR_rS[i], tTR_rS[i + 1]),
                                    )
                                else:
                                    tTR_rS_compute[i], tTR_rS_compute[i + 1] = cute.arch.add_packed_f32x2(
                                        (tTR_rS_compute[i], tTR_rS_compute[i + 1]),
                                        (tTR_rS[i], tTR_rS[i + 1]),
                                    )

                            cute.arch.fence_view_async_tmem_load()
                            cute.arch.barrier(
                                barrier_id=self.compute_sync_bar_id,
                                number_of_threads=self.num_compute_warps * self.threads_per_warp,
                            )

                            # Copy tS_reduce back to tmem
                            cute.copy(tiled_r2t, tTR_rS_compute, tRT_tS_compute_sub_tile)

                            cute.arch.fence_view_async_tmem_store()

                        load_compute_LSE_pipeline.consumer_release(load_compute_LSE_consumer_state)
                        load_compute_LSE_consumer_state.advance()

                        mma_compute_S_pipeline.consumer_release(mma_compute_S_consumer_state)
                        mma_compute_S_consumer_state.advance()

                    # Reduce
                    cute.copy(tiled_t2r_reduce, tTR_tS_reduce, tTR_rS_reduce)

                    # Contraint: self.mma_tiler[1] % self.num_elem_for_reduction == 0
                    tTR_rS_reduce_reshape = cute.composition(
                        tTR_rS_reduce,
                        cute.make_layout(
                            (
                                self.mma_tiler[1] // self.num_elem_for_reduction,
                                self.num_elem_for_reduction,
                            ),
                            stride=(self.num_elem_for_reduction, 1),
                        ),
                    )
                    tTR_rS_reduce_vec = tTR_rS_reduce_reshape.load()
                    tTR_rS_reduce_sum_vec = tTR_rS_reduce_vec.reduce(
                        cute.ReductionOp.ADD,
                        self.acc_dtype(0.0),
                        reduction_profile=(None, 1),
                    )

                    tmp.store(tTR_rS_reduce_sum_vec)

                    q_row_idx = bidx * self.mma_tiler[0] + thread_idx
                    # TODO: make sure the self.mma_tile[1] // self.num_elem_for_reduction is always larger than self.k_value
                    if compute_iter_index == k_tile_idx:
                        self.topk_step(
                            tmp,
                            scores_heap_rf,
                            idx_heap_rf,
                            q_row_idx,
                            compute_iter_index,
                            self.k_value,
                        )
                    else:
                        self.topk_step(
                            tmp,
                            scores_heap_rf,
                            idx_heap_rf,
                            q_row_idx,
                            compute_iter_index,
                            0,
                        )

                    compute_iter_count -= 1
                    compute_iter_index += 1

                # (s_q, k_value, (1, h_k, b))
                gTopk_scores = cute.flat_divide(mTopk_scores, (self.epi_tile[0], self.k_value))
                gTopk_indices = cute.flat_divide(mTopk_indices, (self.epi_tile[0], self.k_value))
                gTopk_scores = gTopk_scores[None, None, bidx, 0, (0, bidy, bidz)]
                gTopk_indices = gTopk_indices[None, None, bidx, 0, (0, bidy, bidz)]
                cTopk = cute.make_identity_tensor((self.epi_tile[0], self.k_value))
                cTopk = cute.domain_offset((bidx * self.epi_tile[0], 0), cTopk)

                copy_atom = cute.make_copy_atom(
                    cute.nvgpu.CopyUniversalOp(),
                    self.acc_dtype,
                    num_bits_per_copy=128,
                )

                thr_layout = cute.make_layout(
                    (self.num_compute_warps * self.threads_per_warp, 1),
                    stride=(1, 1),
                )
                val_layout = cute.make_layout(
                    (1, self.k_value),
                    stride=(self.k_value, 1),
                )
                copy_op = cute.make_tiled_copy_tv(
                    copy_atom,
                    thr_layout,
                    val_layout,
                )

                thr_copy = copy_op.get_slice(thread_idx)

                tgTopk_scores = thr_copy.partition_D(gTopk_scores)
                tgTopk_indices = thr_copy.partition_D(gTopk_indices)
                tcTopk = thr_copy.partition_D(cTopk)

                pred_shape = (tcTopk.shape[0][1], tcTopk.shape[1], tcTopk.shape[2])
                preds = cute.make_rmem_tensor(pred_shape, cutlass.Boolean)
                for v in cutlass.range_constexpr(preds.shape[0]):
                    for m in cutlass.range_constexpr(preds.shape[1]):
                        for n in cutlass.range_constexpr(preds.shape[2]):
                            lhs = tcTopk[v, m, n]
                            val = cute.elem_less(lhs, (cur_s_q, self.k_value))
                            preds[v, m, n] = val

                cute.copy(copy_atom, scores_heap_rf, tgTopk_scores, pred=preds)
                cute.copy(copy_atom, idx_heap_rf, tgTopk_indices, pred=preds)

                cute.arch.barrier(
                    barrier_id=self.compute_sync_bar_id,
                    number_of_threads=self.num_compute_warps * self.threads_per_warp,
                )

                if warp_idx % self.num_compute_warps == 0:
                    tmem_ptr = cute.arch.retrieve_tmem_ptr(
                        self.acc_dtype,
                        alignment=16,
                        ptr_to_buffer_holding_addr=storage.tmem_holding_buf.ptr,
                    )
                    cute.arch.dealloc_tmem(tmem_ptr, 512)

    @cute.jit
    def topk_step(
        self,
        tiled_scores: cute.Tensor,  # (reduced_tile) probabilities
        scores_heap_rf: cute.Tensor,
        idx_heap_rf: cute.Tensor,
        query_index: Int32,
        k_tile_idx: Int32,
        heap_size: int,
    ):

        q_selection_block_idx = (query_index + 1) // self.selection_block_size
        blocks_per_tile = self.mma_tiler[1] // self.num_elem_for_reduction

        score = 0.0
        gmem_idx = 0
        is_valid = cutlass.Boolean(False)

        for i in cutlass.range_constexpr(heap_size):
            score = tiled_scores[i]
            gmem_idx = i + k_tile_idx * blocks_per_tile

            is_valid = (gmem_idx > 0) and (gmem_idx < q_selection_block_idx - 2)

            val_score = score if is_valid else -cutlass.Float32.inf
            val_index = gmem_idx if is_valid else -1
            scores_heap_rf[i] = val_score
            idx_heap_rf[i] = val_index

        for i in cutlass.range(heap_size, cute.size(tiled_scores)):
            score = tiled_scores[i]
            gmem_idx = i + k_tile_idx * blocks_per_tile

            # Decide if candidate is valid for consideration
            is_valid = (gmem_idx > 0) and (gmem_idx < q_selection_block_idx - 2)

            # Full: branchless min-scan to find current minimum (with tie-break)
            min_score = scores_heap_rf[0]
            min_index = idx_heap_rf[0]
            for j in cutlass.range(1, self.k_value, unroll_full=True):
                s = scores_heap_rf[j]
                t = idx_heap_rf[j]
                is_smaller_score = s < min_score
                is_tie_break = (s == min_score) and (t > min_index)
                if is_smaller_score or is_tie_break:
                    # swap to make sure the min_slot is always 0
                    scores_heap_rf[0], scores_heap_rf[j] = s, min_score
                    idx_heap_rf[0], idx_heap_rf[j] = t, min_index
                    min_score = s
                    min_index = t

            # Decide replacement under tie-break rules
            if is_valid:
                is_larger_score = score > min_score
                is_tie_better = (score == min_score) and (gmem_idx < min_index)
                if is_larger_score or is_tie_better:
                    scores_heap_rf[0] = score
                    idx_heap_rf[0] = gmem_idx
