"""
Stage A: Bare Q@K^T via tcgen05.mma → TMEM → GMEM

Validates the tcgen05 MMA path with zero attention logic.
Following dense_blockscaled_gemm_persistent.py for all Blackwell idioms.

Shape: Q(M=128, K=512) @ K^T(512, N=16) → S(128, 16) fp32 output to GMEM.
One CTA. 6 warps: 4 epilogue, 1 MMA, 1 TMA load.

Pipeline: TMA loads Q and K into SMEM, MMA warp issues tcgen05.mma with TMEM accumulator,
epilogue warps TMEM→reg→SMEM→TMA→GMEM.
"""

import torch
import math

try:
    import cutlass
    import cutlass.cute as cute
    import cutlass.torch as cutlass_torch
    import cutlass.utils as utils
    import cutlass.pipeline as pipeline
    import cutlass.utils.blackwell_helpers as sm100_utils
    from cutlass.cute.nvgpu import tcgen05, cpasync
    from cutlass import BFloat16, Float32
    from cutlass.cute.runtime import make_ptr
    from typing import Tuple
    import cuda.bindings.driver as cuda
    HAS_CUTEDSL = True
except ImportError:
    HAS_CUTEDSL = False
    print("WARNING: CuTeDSL not available")

# ── Problem dimensions ────────────────────────────────────────────────
HG = 128   # query heads per CTA (M dimension)
KT = 64    # KV positions per tile (N dimension) — minimum for tcgen05 is 64
HD = 512   # head dim (K dimension)

# ── Warp specialization (mirrors the dense GEMM) ─────────────────────
EPILOGUE_WARP_IDS = (0, 1, 2, 3)
MMA_WARP_ID = 4
TMA_WARP_ID = 5
THREADS_PER_WARP = 32
NUM_WARPS = 6
NUM_THREADS = THREADS_PER_WARP * NUM_WARPS  # 192


class StageAQKTKernel:
    """Stage A: Q @ K^T → TMEM → GMEM, no softmax, no PV GEMM.
    
    This is dense_blockscaled_gemm_persistent.py stripped to attention shapes:
    - BF16 inputs, FP32 accumulator
    - No block scaling (SFA/SFB) — plain bf16 MMA
    - No persistent scheduler — one tile per CTA
    - Single AB stage (one KV tile, no double-buffer needed for Stage A)
    - TMA load for Q and K, tcgen05.mma, TMEM→reg→SMEM→GMEM epilogue
    """

    def __init__(self, mma_tiler_mn=(HG, KT)):
        self.mma_tiler_mn = mma_tiler_mn
        self.use_2cta_instrs = mma_tiler_mn[0] == 256
        self.cta_group = tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE
        
        # Warp IDs and thread count
        self.epilog_warp_id = EPILOGUE_WARP_IDS
        self.mma_warp_id = MMA_WARP_ID
        self.tma_warp_id = TMA_WARP_ID
        self.threads_per_warp = THREADS_PER_WARP
        self.threads_per_cta = NUM_THREADS
        
        # Named barriers
        self.epilog_sync_barrier = pipeline.NamedBarrier(
            barrier_id=1,
            num_threads=THREADS_PER_WARP * len(EPILOGUE_WARP_IDS),
        )
        self.tmem_alloc_barrier = pipeline.NamedBarrier(
            barrier_id=2,
            num_threads=THREADS_PER_WARP * (1 + len(EPILOGUE_WARP_IDS)),
        )
        
        self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")

    def _setup_attributes(self, a_dtype, b_dtype, c_dtype, a_major_mode, b_major_mode, c_layout):
        """Setup attributes that depend on input types."""
        self.a_dtype = a_dtype
        self.b_dtype = b_dtype
        self.c_dtype = c_dtype
        self.acc_dtype = Float32
        self.a_major_mode = a_major_mode
        self.b_major_mode = b_major_mode
        self.c_layout = c_layout
        
        # Create tiled MMA — plain BF16 (no block scaling)
        self.tiled_mma = sm100_utils.make_trivial_tiled_mma(
            a_dtype, b_dtype,
            a_major_mode, b_major_mode,
            Float32,  # acc_dtype
            self.cta_group,
            self.mma_tiler_mn,
        )
        
        atom_thr_size = cute.size(self.tiled_mma.thr_id.shape)
        mma_inst_shape_k = cute.size(self.tiled_mma.shape_mnk, mode=[2])
        mma_inst_tile_k = 4
        self.mma_tiler = (
            self.mma_tiler_mn[0],
            self.mma_tiler_mn[1],
            mma_inst_shape_k * mma_inst_tile_k,
        )
        
        self.cta_tile_shape_mnk = (
            self.mma_tiler[0] // atom_thr_size,
            self.mma_tiler[1],
            self.mma_tiler[2],
        )
        
        # Cluster shape (1,1) — no clustering for attention
        self.cluster_shape_mn = (1, 1)
        self.cluster_layout_vmnk = cute.tiled_divide(
            cute.make_layout((1, 1, 1)),
            (self.tiled_mma.thr_id.shape,),
        )
        self.num_mcast_ctas_a = 1
        self.num_mcast_ctas_b = 1
        self.is_a_mcast = False
        self.is_b_mcast = False
        
        # Epilogue tile
        self.epi_tile = sm100_utils.compute_epilogue_tile_shape(
            self.cta_tile_shape_mnk,
            self.use_2cta_instrs,
            self.c_layout,
            self.c_dtype,
        )
        self.epi_tile_n = cute.size(self.epi_tile[1])
        
        # Stage counts
        self.num_ab_stage = 1
        self.num_acc_stage = 1
        self.num_c_stage = 2
        
        # SMEM layouts
        self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
            self.tiled_mma, self.mma_tiler, a_dtype, self.num_ab_stage,
        )
        self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
            self.tiled_mma, self.mma_tiler, b_dtype, self.num_ab_stage,
        )
        self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi(
            c_dtype, self.c_layout, self.epi_tile, self.num_c_stage,
        )
        
        # TMEM columns for accumulator
        self.num_accumulator_tmem_cols = self.cta_tile_shape_mnk[1]
        self.overlapping_accum = False

    @cute.jit
    def __call__(
        self,
        a_ptr: cute.Pointer,
        b_ptr: cute.Pointer,
        c_ptr: cute.Pointer,
        problem_m: cutlass.Int32,
        problem_n: cutlass.Int32,
        problem_k: cutlass.Int32,
        stream: cuda.CUstream,
    ):
        a_dtype = a_ptr.value_type
        b_dtype = b_ptr.value_type
        c_dtype = c_ptr.value_type
        
        a_major_mode, b_major_mode, c_layout = self._layouts
        self._setup_attributes(a_dtype, b_dtype, c_dtype, a_major_mode, b_major_mode, c_layout)
        
        m, n, k = problem_m, problem_n, problem_k
        
        # Make GMEM tensors — include batch dim (l=1) for local_tile compatibility
        a_layout = cute.make_ordered_layout((m, cute.assume(k, 32), 1), order=(1, 0, 2))
        b_layout = cute.make_ordered_layout((n, cute.assume(k, 32), 1), order=(1, 0, 2))
        c_layout_obj = cute.make_ordered_layout((m, cute.assume(n, 32), 1), order=(1, 0, 2))
        
        mA = cute.make_tensor(a_ptr, a_layout)
        mB = cute.make_tensor(b_ptr, b_layout)
        mC = cute.make_tensor(c_ptr, c_layout_obj)
        
        # TMA descriptors
        a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
        b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
        
        tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
            sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, self.tiled_mma.thr_id),
            mA, a_smem_layout, self.mma_tiler, self.tiled_mma,
            self.cluster_layout_vmnk.shape,
        )
        tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
            sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, self.tiled_mma.thr_id),
            mB, b_smem_layout, self.mma_tiler, self.tiled_mma,
            self.cluster_layout_vmnk.shape,
        )
        
        a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout)
        b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout)
        self.num_tma_load_bytes = (a_copy_size + b_copy_size) * cute.size(self.tiled_mma.thr_id.shape)
        
        epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0))
        tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom(
            cpasync.CopyBulkTensorTileS2GOp(),
            mC, epi_smem_layout, self.epi_tile,
        )
        
        @cute.struct
        class SharedStorage:
            ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage]
            ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage]
            acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
            acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
            tmem_dealloc_mbar: cutlass.Int64
            tmem_holding_buf: cutlass.Int32
            sA: cute.struct.Align[
                cute.struct.MemRange[self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)],
                1024,
            ]
            sB: cute.struct.Align[
                cute.struct.MemRange[self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)],
                1024,
            ]
            sC: cute.struct.Align[
                cute.struct.MemRange[self.c_dtype, cute.cosize(self.c_smem_layout_staged.outer)],
                1024,
            ]
        
        self.shared_storage = SharedStorage
        
        self._kernel(
            self.tiled_mma,
            tma_atom_a, tma_tensor_a,
            tma_atom_b, tma_tensor_b,
            tma_atom_c, tma_tensor_c,
            self.cluster_layout_vmnk,
            self.a_smem_layout_staged,
            self.b_smem_layout_staged,
            self.c_smem_layout_staged,
            self.epi_tile,
        ).launch(
            grid=(1, 1, 1),
            block=[self.threads_per_cta, 1, 1],
            stream=stream,
            min_blocks_per_mp=1,
        )

    @cute.kernel
    def _kernel(
        self,
        tiled_mma,
        tma_atom_a, mA_mkl,
        tma_atom_b, mB_nkl,
        tma_atom_c, mC_mnl,
        cluster_layout_vmnk,
        a_smem_layout_staged,
        b_smem_layout_staged,
        c_smem_layout_staged,
        epi_tile,
    ):
        warp_idx = cute.arch.warp_idx()
        warp_idx = cute.arch.make_warp_uniform(warp_idx)
        tidx, _, _ = cute.arch.thread_idx()
        
        use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2
        is_leader_cta = True
        
        # ── Shared memory ───────────────────────────────────────
        smem = utils.SmemAllocator()
        storage = smem.allocate(self.shared_storage)
        
        # Init AB pipeline
        ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
        num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
        ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_tma_producer)
        ab_pipeline = pipeline.PipelineTmaUmma.create(
            barrier_storage=storage.ab_full_mbar_ptr.data_ptr(),
            num_stages=self.num_ab_stage,
            producer_group=ab_pipeline_producer_group,
            consumer_group=ab_pipeline_consumer_group,
            tx_count=self.num_tma_load_bytes,
            cta_layout_vmnk=cluster_layout_vmnk,
            defer_sync=True,
        )
        
        # Init accumulator pipeline
        acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
        num_acc_consumer_threads = self.threads_per_warp * len(self.epilog_warp_id) * (2 if use_2cta_instrs else 1)
        acc_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_acc_consumer_threads)
        acc_pipeline = pipeline.PipelineUmmaAsync.create(
            barrier_storage=storage.acc_full_mbar_ptr.data_ptr(),
            num_stages=self.num_acc_stage,
            producer_group=acc_pipeline_producer_group,
            consumer_group=acc_pipeline_consumer_group,
            cta_layout_vmnk=cluster_layout_vmnk,
            defer_sync=True,
        )
        
        # TMEM allocator
        tmem = utils.TmemAllocator(
            storage.tmem_holding_buf.ptr,
            barrier_for_retrieve=self.tmem_alloc_barrier,
            allocator_warp_id=self.epilog_warp_id[0],
            is_two_cta=use_2cta_instrs,
            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
        )
        
        pipeline.pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
        
        # ── SMEM tensors ────────────────────────────────────────
        sA = storage.sA.get_tensor(
            a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner
        )
        sB = storage.sB.get_tensor(
            b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner
        )
        sC = storage.sC.get_tensor(
            c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner
        )
        
        # ── Partition global tensors ────────────────────────────
        gA_mkl = cute.local_tile(
            mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
        )
        gB_nkl = cute.local_tile(
            mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
        )
        gC_mnl = cute.local_tile(
            mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
        )
        k_tile_cnt = cute.size(gA_mkl, mode=[3])
        
        mma_tile_coord_v = 0
        thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
        tCgA = thr_mma.partition_A(gA_mkl)
        tCgB = thr_mma.partition_B(gB_nkl)
        tCgC = thr_mma.partition_C(gC_mnl)
        
        # TMA partition for A and B
        a_cta_layout = cute.make_layout(1)
        tAsA, tAgA = cpasync.tma_partition(
            tma_atom_a, 0, a_cta_layout,
            cute.group_modes(sA, 0, 3),
            cute.group_modes(tCgA, 0, 3),
        )
        b_cta_layout = cute.make_layout(1)
        tBsB, tBgB = cpasync.tma_partition(
            tma_atom_b, 0, b_cta_layout,
            cute.group_modes(sB, 0, 3),
            cute.group_modes(tCgB, 0, 3),
        )
        
        # Slice to the single tile coordinate
        # (mma_tile_coord_mnl = (0,0,0) for single-tile)
        tAgA_slice = tAgA[(None, 0, None, 0)]
        tBgB_slice = tBgB[(None, 0, None, 0)]
        
        # MMA fragments
        tCrA = tiled_mma.make_fragment_A(sA)
        tCrB = tiled_mma.make_fragment_B(sB)
        
        # TMEM accumulator
        acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
        tCtAcc_fake = tiled_mma.make_fragment_C(
            cute.append(acc_shape, self.num_acc_stage)
        )
        
        pipeline.pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
        
        # ══════════════════════════════════════════════════════════
        # TMA LOAD WARP (warp 5)
        # ══════════════════════════════════════════════════════════
        if warp_idx == self.tma_warp_id:
            cpasync.prefetch_descriptor(tma_atom_a)
            cpasync.prefetch_descriptor(tma_atom_b)
            cpasync.prefetch_descriptor(tma_atom_c)
            
            ab_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.num_ab_stage
            )
            
            for k_tile in cutlass.range(k_tile_cnt, unroll=1):
                ab_pipeline.producer_acquire(ab_producer_state)
                
                cute.copy(
                    tma_atom_a,
                    tAgA_slice[(None, ab_producer_state.count)],
                    tAsA[(None, ab_producer_state.index)],
                    tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
                )
                cute.copy(
                    tma_atom_b,
                    tBgB_slice[(None, ab_producer_state.count)],
                    tBsB[(None, ab_producer_state.index)],
                    tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
                )
                ab_producer_state.advance()
            
            ab_pipeline.producer_tail(ab_producer_state)
        
        # ══════════════════════════════════════════════════════════
        # MMA WARP (warp 4)
        # ══════════════════════════════════════════════════════════
        if warp_idx == self.mma_warp_id:
            tmem.wait_for_alloc()
            acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
            tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
            tCtAcc = tCtAcc_base[(None, None, None, 0)]
            
            ab_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.num_ab_stage
            )
            acc_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.num_acc_stage
            )
            
            acc_pipeline.producer_acquire(acc_producer_state)
            tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
            
            for k_tile in range(k_tile_cnt):
                if is_leader_cta:
                    ab_pipeline.consumer_wait(ab_consumer_state, cutlass.Boolean(1))
                
                num_kblocks = cute.size(tCrA, mode=[2])
                for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
                    kblock_coord = (None, None, kblock_idx, ab_consumer_state.index)
                    cute.gemm(
                        tiled_mma,
                        tCtAcc,
                        tCrA[kblock_coord],
                        tCrB[kblock_coord],
                        tCtAcc,
                    )
                    tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
                
                if is_leader_cta:
                    ab_pipeline.consumer_release(ab_consumer_state)
                ab_consumer_state.advance()
            
            if is_leader_cta:
                acc_pipeline.producer_commit(acc_producer_state)
            acc_producer_state.advance()
            acc_pipeline.producer_tail(acc_producer_state)
        
        # ══════════════════════════════════════════════════════════
        # EPILOGUE WARPS (0..3)
        # ══════════════════════════════════════════════════════════
        if warp_idx < self.mma_warp_id:
            tmem.allocate(self.num_accumulator_tmem_cols)
            tmem.wait_for_alloc()
            
            acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
            tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
            
            epi_tidx = tidx
            copy_atom_t2r = sm100_utils.get_tmem_load_op(
                self.cta_tile_shape_mnk,
                self.c_layout,
                self.c_dtype,
                self.acc_dtype,
                epi_tile,
                use_2cta_instrs,
            )
            tAcc_epi = cute.flat_divide(
                tCtAcc_base[((None, None), 0, 0, None)],
                epi_tile,
            )
            tiled_copy_t2r = tcgen05.make_tmem_copy(
                copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)]
            )
            thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx)
            tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi)
            
            tTR_rAcc = cute.make_rmem_tensor(
                thr_copy_t2r.partition_D(
                    cute.flat_divide(
                        tCgC[((None, None), 0, 0, None, None, None)], epi_tile
                    )
                )[(None, None, None, 0, 0, 0, 0, 0)].shape,
                self.acc_dtype,
            )
            tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype)
            
            copy_atom_r2s = sm100_utils.get_smem_store_op(
                self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r
            )
            tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r)
            thr_copy_r2s = tiled_copy_r2s.get_slice(epi_tidx)
            tRS_sC = thr_copy_r2s.partition_D(sC)
            tRS_rC = tiled_copy_r2s.retile(tTR_rC)
            
            gC_epi = cute.flat_divide(
                tCgC[((None, None), 0, 0, None, None, None)], epi_tile
            )
            sC_for_tma = cute.group_modes(sC, 0, 2)
            gC_for_tma = cute.group_modes(gC_epi, 0, 2)
            bSG_sC, bSG_gC = cpasync.tma_partition(
                tma_atom_c, 0, cute.make_layout(1),
                sC_for_tma, gC_for_tma,
            )
            
            acc_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.num_acc_stage
            )
            c_producer_group = pipeline.CooperativeGroup(
                pipeline.Agent.Thread,
                self.threads_per_warp * len(self.epilog_warp_id),
            )
            c_pipeline = pipeline.PipelineTmaStore.create(
                num_stages=self.num_c_stage,
                producer_group=c_producer_group,
            )
            
            acc_pipeline.consumer_wait(acc_consumer_state)
            
            tTR_tAcc_g = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
            bSG_gC_g = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))
            
            subtile_cnt = cute.size(tTR_tAcc_g.shape, mode=[3])
            for subtile_idx in cutlass.range(subtile_cnt):
                tTR_tAcc_mn = tTR_tAcc_g[(None, None, None, subtile_idx)]
                cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc)
                
                acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load()
                tRS_rC.store(acc_vec.to(self.c_dtype))
                
                c_buffer = subtile_idx % self.num_c_stage
                cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)])
                cute.arch.fence_proxy("async.shared", space="cta")
                self.epilog_sync_barrier.arrive_and_wait()
                
                if warp_idx == self.epilog_warp_id[0]:
                    cute.copy(
                        tma_atom_c,
                        bSG_sC[(None, c_buffer)],
                        bSG_gC_g[(None, subtile_idx)],
                    )
                    c_pipeline.producer_commit()
                    c_pipeline.producer_acquire()
                self.epilog_sync_barrier.arrive_and_wait()
            
            acc_pipeline.consumer_release(acc_consumer_state)
            acc_consumer_state.advance()
            
            tmem.relinquish_alloc_permit()
            self.epilog_sync_barrier.arrive_and_wait()
            tmem.free(acc_tmem_ptr)
            c_pipeline.producer_tail()


def test_stage_a():
    """Test Stage A: Q @ K^T via tcgen05.mma → TMEM → GMEM."""
    if not HAS_CUTEDSL:
        print("CuTeDSL not available, skipping")
        return
    
    device = torch.device("cuda")
    torch.manual_seed(42)
    
    prob_m, prob_n, prob_k = HG, KT, HD
    
    tQ = torch.randn(prob_m, prob_k, dtype=torch.bfloat16, device=device)
    tK = torch.randn(prob_n, prob_k, dtype=torch.bfloat16, device=device)
    ref = torch.matmul(tQ.to(torch.float32), tK.to(torch.float32).T)
    tC = torch.zeros(prob_m, prob_n, dtype=torch.bfloat16, device=device)
    
    # Compile using make_ptr pattern (like dense GEMM)
    a_ptr = make_ptr(BFloat16, 0, cute.AddressSpace.gmem, assumed_align=16)
    b_ptr = make_ptr(BFloat16, 0, cute.AddressSpace.gmem, assumed_align=16)
    c_ptr = make_ptr(BFloat16, 0, cute.AddressSpace.gmem, assumed_align=16)
    
    a_major_mode = tcgen05.OperandMajorMode.K
    b_major_mode = tcgen05.OperandMajorMode.K
    c_layout = utils.LayoutEnum.ROW_MAJOR
    
    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
    
    # Compile using make_ptr pattern (like dense GEMM)
    a_ptr_fake = make_ptr(BFloat16, 0, cute.AddressSpace.gmem, assumed_align=16)
    b_ptr_fake = make_ptr(BFloat16, 0, cute.AddressSpace.gmem, assumed_align=16)
    c_ptr_fake = make_ptr(BFloat16, 0, cute.AddressSpace.gmem, assumed_align=16)
    
    a_major_mode = tcgen05.OperandMajorMode.K
    b_major_mode = tcgen05.OperandMajorMode.K
    c_layout = utils.LayoutEnum.ROW_MAJOR
    
    kernel = StageAQKTKernel()
    kernel._layouts = (a_major_mode, b_major_mode, c_layout)
    compiled = cute.compile(
        kernel,
        a_ptr_fake, b_ptr_fake, c_ptr_fake,
        cutlass.Int32(prob_m), cutlass.Int32(prob_n), cutlass.Int32(prob_k),
        stream,
    )
    
    # Create runtime pointers from torch tensor data
    a_ptr = make_ptr(BFloat16, tQ.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    b_ptr = make_ptr(BFloat16, tK.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    c_ptr = make_ptr(BFloat16, tC.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    
    compiled(
        a_ptr, b_ptr, c_ptr,
        cutlass.Int32(prob_m), cutlass.Int32(prob_n), cutlass.Int32(prob_k),
        stream,
    )
    torch.cuda.synchronize()
    
    output = tC.to(torch.float32)
    cos = torch.nn.functional.cosine_similarity(
        output.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)
    ).item()
    max_err = (output - ref).abs().max().item()
    mean_err = (output - ref).abs().mean().item()
    
    print(f"Stage A: Q({prob_m},{prob_k}) @ K^T({prob_k},{prob_n}) → S({prob_m},{prob_n})")
    print(f"  Cosine similarity: {cos:.6f}")
    print(f"  Max absolute error: {max_err:.6f}")
    print(f"  Mean absolute error: {mean_err:.6f}")
    
    if cos >= 0.99:
        print("  ✅ PASS")
    else:
        print("  ❌ FAIL — cosine < 0.99")
    
    return cos


if __name__ == "__main__":
    test_stage_a()
