/**
 * Debug test for B1 prefill kernel T>1 path.
 * 
 * Tests T=2 N=128 step by step:
 * 1. Compute QK (noPE + RoPE) for 2 query rows
 * 2. Verify QK logits against CPU reference
 * 3. Compute softmax
 * 4. Compute PV and verify against CPU reference
 * 5. Full T=2 prefill vs CPU reference
 */
#include <cuda_runtime.h>
#include <cuda_fp8.h>
#include <cuda_fp8.hpp>
#include <cstdio>
#include <cmath>
#include <cstdlib>
#include <cassert>

// Include kernel headers
#include "dsv4/kernels/attention/fmha_common.cuh"
#include "dsv4/kernels/attention/fmha_umma_desc.cuh"
#include "dsv4/kernels/attention/fmha_mixed_fp8_prefill.cuh"

using namespace dsv4::kernels::attention;

// ---- CPU reference functions ----

static void cpu_fp8_e4m3_quantize(const float* src, uint8_t* dst, float* scale,
                                    int rows, int cols) {
    for (int r = 0; r < rows; r++) {
        float amax = 0.0f;
        for (int c = 0; c < cols; c++) amax = fmaxf(amax, fabsf(src[r * cols + c]));
        float s = amax / 448.0f;
        if (s < 1e-12f) s = 1.0f;
        scale[r] = s;
        for (int c = 0; c < cols; c++) {
            float v = src[r * cols + c] / s;
            v = fmaxf(-448.0f, fminf(448.0f, v));
            __nv_fp8_e4m3 fp8; fp8.__x = 0;
            // Simplest quantize: round to FP8
            memcpy(&fp8, &v, 1); // This won't work, use proper conversion
            dst[r * cols + c] = 0; // placeholder
        }
    }
}

static float fp8_to_f32(uint8_t b) {
    __nv_fp8_e4m3 v; v.__x = b;
    return (float)v;
}

static bf16_t f32_to_bf16_host(float f) {
    uint32_t u; memcpy(&u, &f, 4);
    uint16_t h = (u + 0x8000) >> 16;
    return h;
}

static float bf16_to_f32_host(bf16_t h) {
    uint32_t u = (uint32_t)h << 16;
    float f; memcpy(&f, &u, 4);
    return f;
}

// ---- Minimal T=2 kernel that prints intermediate values ----

__global__ void prefill_t2_debug_kernel(
    const uint8_t* __restrict__ q_nope_fp8,
    const float* __restrict__ q_nope_scale,
    const bf16_t* __restrict__ q_rope_bf16,
    const uint8_t* __restrict__ k_nope_fp8,
    const float* __restrict__ k_nope_scale,
    const bf16_t* __restrict__ k_rope_bf16,
    int T, int N, int HD, int NOPE, int ROPE,
    float scale) 
{
    // Only one CTA for debug
    if (blockIdx.x > 0 || blockIdx.y > 0 || blockIdx.z > 0) return;
    
    constexpr int SK_TILE = 128;
    constexpr int MMA_K_F8 = 32;
    constexpr int MMA_K_F16 = 16;
    constexpr int NKT_NOPE = 448 / MMA_K_F8;  // 14
    constexpr int NKT_ROPE = 64 / MMA_K_F16;   // 4
    constexpr int N_SUB = 512 / 16;  // 32
    constexpr int NKT_PV = SK_TILE / MMA_K_F16;  // 8
    constexpr int TILE_F8 = 128 * MMA_K_F8;   // 4096
    constexpr int TILE_F16 = 128 * MMA_K_F16;  // 2048
    constexpr int V_SUB_SZ = 16 * MMA_K_F16;   // 256
    constexpr int TMEM_COLS = 512;
    constexpr int T_ACT = 2;

    const int tid = threadIdx.x;
    const int wid = tid >> 5;
    const int lane = tid & 31;
    const bool is_mma_warp = (wid == 4);

    extern __shared__ __align__(128) char sbuf[];
    size_t off = 0;
    uint32_t* sTmemBase = (uint32_t*)(sbuf + off); off += 4;
    off = (off + 127) & ~(size_t)127;
    uint8_t* sQ8 = (uint8_t*)(sbuf + off); off += TILE_F8;
    off = (off + 127) & ~(size_t)127;
    uint8_t* sK8 = (uint8_t*)(sbuf + off); off += TILE_F8;
    off = (off + 127) & ~(size_t)127;
    bf16_t* sQ16 = (bf16_t*)(sbuf + off); off += TILE_F16 * sizeof(bf16_t);
    off = (off + 127) & ~(size_t)127;
    bf16_t* sK16 = (bf16_t*)(sbuf + off); off += TILE_F16 * sizeof(bf16_t);
    off = (off + 127) & ~(size_t)127;
    bf16_t* sPk = (bf16_t*)(sbuf + off); off += TILE_F16 * sizeof(bf16_t);
    off = (off + 127) & ~(size_t)127;
    bf16_t* sV = (bf16_t*)(sbuf + off); off += V_SUB_SZ * sizeof(bf16_t);
    off = (off + 127) & ~(size_t)127;
    float* sLogits = (float*)(sbuf + off); off += T_ACT * SK_TILE * sizeof(float);
    float* sP = (float*)(sbuf + off); off += T_ACT * SK_TILE * sizeof(float);
    float* sOacc = (float*)(sbuf + off); off += T_ACT * HD * sizeof(float);
    float* sRunningMax = (float*)(sbuf + off); off += T_ACT * sizeof(float);
    float* sRunningSum = (float*)(sbuf + off); off += T_ACT * sizeof(float);

    // TMEM alloc
    if (is_mma_warp) tmem_alloc((uint32_t)__cvta_generic_to_shared(sTmemBase), TMEM_COLS);
    asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
    __syncthreads();
    uint32_t tb = *sTmemBase;

    const uint32_t idesc_f8_qk = make_idesc_f8_e4m3(128, 128);
    const uint32_t idesc_f16_qk = make_idesc(128, 128);
    const uint32_t idesc_pv = make_idesc(128, 16);

    // Init accumulators
    for (int i = tid; i < T_ACT * HD; i += blockDim.x) sOacc[i] = 0.0f;
    for (int t = tid; t < T_ACT; t += blockDim.x) {
        sRunningMax[t] = -INFINITY;
        sRunningSum[t] = 0.0f;
    }
    __syncthreads();

    // Single KV tile (N=128)
    const int kv_len = min(SK_TILE, N);
    
    // ---- QK noPE: FP8 ----
    for (int kt = 0; kt < NKT_NOPE; kt++) {
        for (int i = tid; i < TILE_F8; i += blockDim.x) { sQ8[i] = 0; sK8[i] = 0; }
        __syncthreads();
        for (int r = tid; r < T_ACT; r += blockDim.x) {
            for (int c = 0; c < MMA_K_F8; c++) {
                int d = kt * MMA_K_F8 + c;
                if (d < NOPE) sQ8[_pfill_cidx_f8(r, c)] = q_nope_fp8[r * NOPE + d];
            }
        }
        for (int i = tid; i < kv_len * MMA_K_F8; i += blockDim.x) {
            int r = i / MMA_K_F8, c = i % MMA_K_F8;
            int d = kt * MMA_K_F8 + c;
            if (d < NOPE) sK8[_pfill_cidx_f8(r, c)] = k_nope_fp8[r * NOPE + d];
        }
        __syncthreads();
        if (is_mma_warp && lane == 0) {
            uint64_t dq = make_umma_desc_kmajor_none((uint32_t)__cvta_generic_to_shared(sQ8), 128);
            uint64_t dk = make_umma_desc_kmajor_none((uint32_t)__cvta_generic_to_shared(sK8), 128);
            umma_ss_f8f6f4(tb, dq, dk, idesc_f8_qk, kt > 0);
            asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
        }
        __syncthreads();
    }
    asm volatile("fence.sc.gpu;" ::: "memory");
    __syncthreads();

    // Read QK noPE
    prefill_read_qk_rows<SK_TILE>(tb, sLogits, T_ACT, kv_len);
    __syncthreads();

    // Print QK noPE logits for rows 0,1 (first 8 values)
    if (tid == 0) {
        printf("QK noPE (row 0, first 8): ");
        for (int c = 0; c < 8; c++) printf("%.4f ", sLogits[0 * SK_TILE + c]);
        printf("\n");
        printf("QK noPE (row 1, first 8): ");
        for (int c = 0; c < 8; c++) printf("%.4f ", sLogits[1 * SK_TILE + c]);
        printf("\n");
    }
    __syncthreads();

    // Apply scales
    for (int r = tid; r < T_ACT; r += blockDim.x) {
        float q_s = q_nope_scale[r];
        for (int c = 0; c < kv_len; c++) {
            sLogits[r * SK_TILE + c] *= q_s * k_nope_scale[c];
        }
    }
    __syncthreads();

    if (tid == 0) {
        printf("QK noPE scaled (row 0, first 8): ");
        for (int c = 0; c < 8; c++) printf("%.4f ", sLogits[0 * SK_TILE + c]);
        printf("\n");
        printf("QK noPE scaled (row 1, first 8): ");
        for (int c = 0; c < 8; c++) printf("%.4f ", sLogits[1 * SK_TILE + c]);
        printf("\n");
    }
    __syncthreads();

    // ---- QK RoPE: BF16 ----
    for (int kt = 0; kt < NKT_ROPE; kt++) {
        for (int i = tid; i < TILE_F16; i += blockDim.x) { sQ16[i] = 0; sK16[i] = 0; }
        __syncthreads();
        for (int r = tid; r < T_ACT; r += blockDim.x) {
            for (int c = 0; c < MMA_K_F16; c++) {
                int d = kt * MMA_K_F16 + c;
                if (d < ROPE) sQ16[_pfill_cidx_bf16_128(r, c)] = q_rope_bf16[r * ROPE + d];
            }
        }
        for (int i = tid; i < kv_len * MMA_K_F16; i += blockDim.x) {
            int r = i / MMA_K_F16, c = i % MMA_K_F16;
            int d = kt * MMA_K_F16 + c;
            if (d < ROPE) sK16[_pfill_cidx_bf16_128(r, c)] = k_rope_bf16[(int64_t)r * ROPE + d];
        }
        __syncthreads();
        if (is_mma_warp && lane == 0) {
            uint64_t dq = make_umma_desc_kmajor_none((uint32_t)__cvta_generic_to_shared(sQ16), 128);
            uint64_t dk = make_umma_desc_kmajor_none((uint32_t)__cvta_generic_to_shared(sK16), 128);
            umma_ss_f16(tb, dq, dk, idesc_f16_qk, kt > 0);
            asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
        }
        __syncthreads();
    }
    asm volatile("fence.sc.gpu;" ::: "memory");
    __syncthreads();

    // Add RoPE to noPE
    prefill_read_qk_rows<SK_TILE>(tb, sP, T_ACT, kv_len);
    __syncthreads();
    for (int i = tid; i < T_ACT * kv_len; i += blockDim.x) {
        sLogits[i] += sP[i];
    }
    __syncthreads();

    if (tid == 0) {
        printf("QK total (row 0, first 8): ");
        for (int c = 0; c < 8; c++) printf("%.4f ", sLogits[0 * SK_TILE + c] * scale);
        printf("\n");
        printf("QK total (row 1, first 8): ");
        for (int c = 0; c < 8; c++) printf("%.4f ", sLogits[1 * SK_TILE + c] * scale);
        printf("\n");
    }
    __syncthreads();

    // ---- Softmax ----
    for (int r = tid; r < T_ACT; r += blockDim.x) {
        float tile_max = -INFINITY;
        for (int c = 0; c < kv_len; c++)
            tile_max = fmaxf(tile_max, sLogits[r * SK_TILE + c] * scale);
        
        float tile_sum = 0.0f;
        for (int c = 0; c < kv_len; c++) {
            float pv = expf(sLogits[r * SK_TILE + c] * scale - tile_max);
            sP[r * SK_TILE + c] = pv;
            tile_sum += pv;
        }
        for (int c = kv_len; c < SK_TILE; c++) sP[r * SK_TILE + c] = 0.0f;

        float old_max = sRunningMax[r];
        float new_max = fmaxf(old_max, tile_max);
        float rescale_old = (old_max > -INFINITY) ? expf(old_max - new_max) : 0.0f;
        for (int d = 0; d < HD; d++) sOacc[r * HD + d] *= rescale_old;
        float rescale_new = expf(tile_max - new_max);
        sRunningSum[r] = sRunningSum[r] * rescale_old + tile_sum * rescale_new;
        sRunningMax[r] = new_max;

        sLogits[r * SK_TILE] = rescale_new;
    }
    __syncthreads();

    if (tid == 0) {
        printf("Softmax P (row 0, first 8): ");
        for (int c = 0; c < 8; c++) printf("%.6f ", sP[0 * SK_TILE + c]);
        printf(" sum=%.6f\n", sRunningSum[0]);
        printf("Softmax P (row 1, first 8): ");
        for (int c = 0; c < 8; c++) printf("%.6f ", sP[1 * SK_TILE + c]);
        printf(" sum=%.6f\n", sRunningSum[1]);
        printf("Rescale: row0=%.6f row1=%.6f\n", sLogits[0 * SK_TILE], sLogits[1 * SK_TILE]);
    }
    __syncthreads();

    // ---- PV: per query row ----
    for (int qr = 0; qr < T_ACT; qr++) {
        float p_rescale = sLogits[qr * SK_TILE];
        
        if (tid == 0) printf("PV for qr=%d: p_rescale=%.6f\n", qr, p_rescale);

        for (int n_sub = 0; n_sub < N_SUB; n_sub++) {
            int d_base = n_sub * 16;
            for (int pv_kt = 0; pv_kt < NKT_PV; pv_kt++) {
                const int col_start = pv_kt * MMA_K_F16;
                for (int i = tid; i < TILE_F16; i += blockDim.x) sPk[i] = 0;
                for (int i = tid; i < V_SUB_SZ; i += blockDim.x) sV[i] = 0;
                __syncthreads();

                for (int c = tid; c < MMA_K_F16; c += blockDim.x) {
                    int gc = col_start + c;
                    sPk[_pfill_cidx_bf16_128(qr, c)] = f32_to_bf16(sP[qr * SK_TILE + gc]);
                }

                for (int i = tid; i < 16 * MMA_K_F16; i += blockDim.x) {
                    int dd = i / MMA_K_F16, kk = i % MMA_K_F16;
                    int row = col_start + kk;
                    int g_row = row;
                    int d = d_base + dd;
                    bf16_t vbits = 0;
                    if (row < kv_len) {
                        if (d < NOPE) {
                            uint8_t b = k_nope_fp8[(int64_t)g_row * NOPE + d];
                            float v = _prefill_fp8_to_f32(b) * k_nope_scale[g_row];
                            vbits = f32_to_bf16(v);
                        } else {
                            vbits = k_rope_bf16[(int64_t)g_row * ROPE + (d - NOPE)];
                        }
                    }
                    sV[_pfill_cidx_bf16_16(dd, kk)] = vbits;
                }
                __syncthreads();

                bool first = (pv_kt == 0);
                if (is_mma_warp && lane == 0) {
                    uint64_t dp = make_umma_desc_kmajor_none((uint32_t)__cvta_generic_to_shared(sPk), 128);
                    uint64_t dv = make_umma_desc_kmajor_none((uint32_t)__cvta_generic_to_shared(sV), 16);
                    umma_ss_f16(tb + n_sub * 16, dp, dv, idesc_pv, !first);
                    asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
                }
                __syncthreads();
            }
        }

        // Read PV result for row qr
        asm volatile("fence.sc.gpu;" ::: "memory");
        __syncthreads();
        prefill_read_pv_all_subs<512, 32>(tb, qr, sOacc, p_rescale);
        __syncthreads();

        // Print first few accumulated values
        if (tid == 0 && qr == 0) {
            printf("sOacc qr=0 (first 8): ");
            for (int d = 0; d < 8; d++) printf("%.6f ", sOacc[0 * HD + d]);
            printf("\n");
        }
        if (tid == 0 && qr == 1) {
            printf("sOacc qr=1 (first 8): ");
            for (int d = 0; d < 8; d++) printf("%.6f ", sOacc[1 * HD + d]);
            printf("\n");
        }
        __syncthreads();
    }

    // Normalize and print final output
    if (tid == 0) {
        printf("sRunningSum: row0=%.6f row1=%.6f\n", sRunningSum[0], sRunningSum[1]);
        printf("sRunningMax: row0=%.6f row1=%.6f\n", sRunningMax[0], sRunningMax[1]);
        printf("Final output row0 (first 8): ");
        for (int d = 0; d < 8; d++) printf("%.6f ", sOacc[0 * HD + d] / sRunningSum[0]);
        printf("\n");
        printf("Final output row1 (first 8): ");
        for (int d = 0; d < 8; d++) printf("%.6f ", sOacc[1 * HD + d] / sRunningSum[1]);
        printf("\n");
        
        // Check for NaN
        bool has_nan0 = false, has_nan1 = false;
        for (int d = 0; d < HD; d++) {
            if (isnan(sOacc[0 * HD + d])) has_nan0 = true;
            if (isnan(sOacc[1 * HD + d])) has_nan1 = true;
        }
        printf("NaN check: row0=%s row1=%s\n", has_nan0 ? "YES" : "no", has_nan1 ? "YES" : "no");
    }

    if (is_mma_warp) tmem_dealloc(tb, TMEM_COLS);
}

int main() {
    constexpr int T = 2;
    constexpr int N = 128;
    constexpr int HD = 512;
    constexpr int NOPE = 448;
    constexpr int ROPE = 64;
    const float scale = 1.0f / sqrtf((float)HD);

    printf("=== Prefill T=2 Debug Test ===\n");
    printf("T=%d N=%d HD=%d NOPE=%d ROPE=%d scale=%.6f\n", T, N, HD, NOPE, ROPE, scale);

    // Generate random data on CPU, then upload
    srand(42);
    
    // Q: (T, HD) FP32 → quantize noPE to FP8, keep RoPE as BF16
    float* h_q = (float*)malloc(T * HD * sizeof(float));
    for (int i = 0; i < T * HD; i++) h_q[i] = (float)rand() / RAND_MAX * 0.5f - 0.25f;

    // K: (N, HD) FP32 → quantize noPE to FP8, keep RoPE as BF16
    float* h_k = (float*)malloc(N * HD * sizeof(float));
    for (int i = 0; i < N * HD; i++) h_k[i] = (float)rand() / RAND_MAX * 0.5f - 0.25f;

    // Q noPE FP8 quantization (per-row scale)
    uint8_t* h_q_nope_fp8 = (uint8_t*)malloc(T * NOPE);
    float* h_q_nope_scale = (float*)malloc(T * sizeof(float));
    for (int r = 0; r < T; r++) {
        float amax = 0.0f;
        for (int c = 0; c < NOPE; c++) amax = fmaxf(amax, fabsf(h_q[r * HD + c]));
        float s = amax / 448.0f;
        if (s < 1e-12f) s = 1.0f;
        h_q_nope_scale[r] = s;
        for (int c = 0; c < NOPE; c++) {
            float v = h_q[r * HD + c] / s;
            v = fmaxf(-448.0f, fminf(448.0f, v));
            __nv_fp8_e4m3 fp8 = __nv_fp8_e4m3(v);
            h_q_nope_fp8[r * NOPE + c] = fp8.__x;
        }
    }

    // Q RoPE BF16
    bf16_t* h_q_rope_bf16 = (bf16_t*)malloc(T * ROPE * sizeof(bf16_t));
    for (int r = 0; r < T; r++)
        for (int c = 0; c < ROPE; c++)
            h_q_rope_bf16[r * ROPE + c] = f32_to_bf16_host(h_q[r * HD + NOPE + c]);

    // K noPE FP8 quantization
    uint8_t* h_k_nope_fp8 = (uint8_t*)malloc(N * NOPE);
    float* h_k_nope_scale = (float*)malloc(N * sizeof(float));
    for (int r = 0; r < N; r++) {
        float amax = 0.0f;
        for (int c = 0; c < NOPE; c++) amax = fmaxf(amax, fabsf(h_k[r * HD + c]));
        float s = amax / 448.0f;
        if (s < 1e-12f) s = 1.0f;
        h_k_nope_scale[r] = s;
        for (int c = 0; c < NOPE; c++) {
            float v = h_k[r * HD + c] / s;
            v = fmaxf(-448.0f, fminf(448.0f, v));
            __nv_fp8_e4m3 fp8 = __nv_fp8_e4m3(v);
            h_k_nope_fp8[r * NOPE + c] = fp8.__x;
        }
    }

    // K RoPE BF16
    bf16_t* h_k_rope_bf16 = (bf16_t*)malloc(N * ROPE * sizeof(bf16_t));
    for (int r = 0; r < N; r++)
        for (int c = 0; c < ROPE; c++)
            h_k_rope_bf16[r * ROPE + c] = f32_to_bf16_host(h_k[r * HD + NOPE + c]);

    // Upload to GPU
    uint8_t *d_q_nope_fp8, *d_k_nope_fp8;
    float *d_q_nope_scale, *d_k_nope_scale;
    bf16_t *d_q_rope_bf16, *d_k_rope_bf16;
    
    cudaMalloc(&d_q_nope_fp8, T * NOPE);
    cudaMalloc(&d_q_nope_scale, T * sizeof(float));
    cudaMalloc(&d_q_rope_bf16, T * ROPE * sizeof(bf16_t));
    cudaMalloc(&d_k_nope_fp8, N * NOPE);
    cudaMalloc(&d_k_nope_scale, N * sizeof(float));
    cudaMalloc(&d_k_rope_bf16, N * ROPE * sizeof(bf16_t));

    cudaMemcpy(d_q_nope_fp8, h_q_nope_fp8, T * NOPE, cudaMemcpyHostToDevice);
    cudaMemcpy(d_q_nope_scale, h_q_nope_scale, T * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(d_q_rope_bf16, h_q_rope_bf16, T * ROPE * sizeof(bf16_t), cudaMemcpyHostToDevice);
    cudaMemcpy(d_k_nope_fp8, h_k_nope_fp8, N * NOPE, cudaMemcpyHostToDevice);
    cudaMemcpy(d_k_nope_scale, h_k_nope_scale, N * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(d_k_rope_bf16, h_k_rope_bf16, N * ROPE * sizeof(bf16_t), cudaMemcpyHostToDevice);

    // Compute CPU reference QK
    printf("\n=== CPU Reference QK ===\n");
    float ref_qk[2][128] = {};
    for (int r = 0; r < T; r++) {
        for (int c = 0; c < N; c++) {
            float dot = 0.0f;
            // noPE: FP8 dequant dot product
            for (int d = 0; d < NOPE; d++) {
                float qv = fp8_to_f32(h_q_nope_fp8[r * NOPE + d]) * h_q_nope_scale[r];
                float kv = fp8_to_f32(h_k_nope_fp8[c * NOPE + d]) * h_k_nope_scale[c];
                dot += qv * kv;
            }
            // RoPE: BF16 dot product
            for (int d = 0; d < ROPE; d++) {
                float qv = bf16_to_f32_host(h_q_rope_bf16[r * ROPE + d]);
                float kv = bf16_to_f32_host(h_k_rope_bf16[c * ROPE + d]);
                dot += qv * kv;
            }
            ref_qk[r][c] = dot * scale;
        }
    }
    printf("CPU ref QK (row 0, first 8): ");
    for (int c = 0; c < 8; c++) printf("%.4f ", ref_qk[0][c]);
    printf("\n");
    printf("CPU ref QK (row 1, first 8): ");
    for (int c = 0; c < 8; c++) printf("%.4f ", ref_qk[1][c]);
    printf("\n");

    // Compute CPU reference softmax
    printf("\n=== CPU Reference Softmax + Attention ===\n");
    float ref_softmax[2][128] = {};
    for (int r = 0; r < T; r++) {
        float mx = ref_qk[r][0];
        for (int c = 1; c < N; c++) mx = fmaxf(mx, ref_qk[r][c]);
        float sm = 0.0f;
        for (int c = 0; c < N; c++) {
            ref_softmax[r][c] = expf(ref_qk[r][c] - mx);
            sm += ref_softmax[r][c];
        }
        for (int c = 0; c < N; c++) ref_softmax[r][c] /= sm;
    }
    printf("CPU ref softmax (row 0, first 8): ");
    for (int c = 0; c < 8; c++) printf("%.6f ", ref_softmax[0][c]);
    printf("\n");

    // Compute CPU reference attention output
    float ref_out[2][512] = {};
    for (int r = 0; r < T; r++) {
        for (int d = 0; d < HD; d++) {
            float val = 0.0f;
            for (int c = 0; c < N; c++) {
                float kv;
                if (d < NOPE) {
                    kv = fp8_to_f32(h_k_nope_fp8[c * NOPE + d]) * h_k_nope_scale[c];
                } else {
                    kv = bf16_to_f32_host(h_k_rope_bf16[c * ROPE + (d - NOPE)]);
                }
                val += ref_softmax[r][c] * kv;
            }
            ref_out[r][d] = val;
        }
    }
    printf("CPU ref output (row 0, first 8): ");
    for (int d = 0; d < 8; d++) printf("%.6f ", ref_out[0][d]);
    printf("\n");
    printf("CPU ref output (row 1, first 8): ");
    for (int d = 0; d < 8; d++) printf("%.6f ", ref_out[1][d]);
    printf("\n");

    // Launch debug kernel
    printf("\n=== GPU Kernel Execution ===\n");
    int smem_size = 200 * 1024;  // ~149KB needed, stay under 232KB limit
    cudaFuncSetAttribute(prefill_t2_debug_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
    prefill_t2_debug_kernel<<<dim3(1,1,1), 192, smem_size>>>(
        d_q_nope_fp8, d_q_nope_scale, d_q_rope_bf16,
        d_k_nope_fp8, d_k_nope_scale, d_k_rope_bf16,
        T, N, HD, NOPE, ROPE, scale);
    cudaDeviceSynchronize();
    
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        printf("Kernel launch FAILED: %s\n", cudaGetErrorString(err));
    } else {
        printf("Kernel completed successfully.\n");
    }

    // Cleanup
    cudaFree(d_q_nope_fp8); cudaFree(d_q_nope_scale); cudaFree(d_q_rope_bf16);
    cudaFree(d_k_nope_fp8); cudaFree(d_k_nope_scale); cudaFree(d_k_rope_bf16);
    free(h_q); free(h_k);
    free(h_q_nope_fp8); free(h_q_nope_scale); free(h_q_rope_bf16);
    free(h_k_nope_fp8); free(h_k_nope_scale); free(h_k_rope_bf16);

    printf("\n=== Done ===\n");
    return 0;
}