diff --git a/csrc/tests/gqa_decode_test.cu b/csrc/tests/gqa_decode_test.cu
index a81caf0..0dfb910 100644
--- a/csrc/tests/gqa_decode_test.cu
+++ b/csrc/tests/gqa_decode_test.cu
@@ -10,6 +10,9 @@ nvcc -I csrc -arch=sm_89 -O3 \
 #include <cmath>
 #include <sys/time.h>
 #include "../kernels/gqa_decode_attn.cuh"
+#ifndef ASTRAI_NO_MMA
+#include "../kernels/gqa_decode_attn_mma.cuh"
+#endif
 
 static double now_ms() {
     struct timeval tv;
@@ -17,6 +20,80 @@ static double now_ms() {
     return tv.tv_sec * 1000.0 + tv.tv_usec / 1000.0;
 }
 
+// Split-K scratch (torch-free): the production launcher allocates these from
+// torch; here we pass pre-allocated device buffers so the bench loop doesn't
+// pay a cudaMalloc per iteration. Size for the maximum split count (32).
+struct DecodeScratch {
+    float* o_part = nullptr;
+    float* ml_part = nullptr;
+};
+
+// Launch the production decode path (tensor-core head-packing MMA on sm_80+,
+// scalar fallback otherwise), mirroring dispatch_decode() in gqa_decode_attn.cu.
+#ifndef ASTRAI_NO_MMA
+static bool decode_use_mma(const GQAParams& p) {
+    int G = p.q_head / p.kv_head;
+    return !p.use_mask && G > 1 && G <= 16;
+}
+
+static int decode_num_splits(const GQAParams& p, int tiles_total) {
+    int sm_count = 0;
+    cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
+    int base_blocks = p.kv_head * p.batch;
+    int desired = 2 * (sm_count > 0 ? sm_count : 64);
+    int n = (desired + base_blocks - 1) / base_blocks;
+    return max(1, min(n, min(tiles_total, 32)));
+}
+
+template <int HEAD_DIM, int BC>
+static void launch_mma_decode(GQAParams& p, DecodeScratch& sc) {
+    constexpr int BR = 16, LD = HEAD_DIM;
+    int smem = (2 * BC * LD + BR * LD) * (int)sizeof(bf16);
+    int tiles_total = (p.kv_len + BC - 1) / BC;
+    int num_splits = decode_num_splits(p, tiles_total);
+
+    if (num_splits <= 1) {
+        cudaFuncSetAttribute(gqa_decode_attn_mma_kernel<HEAD_DIM, BC>,
+                             cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
+        gqa_decode_attn_mma_kernel<HEAD_DIM, BC>
+            <<<dim3(p.kv_head, p.batch), 32, smem>>>(p);
+        return;
+    }
+    cudaFuncSetAttribute(gqa_decode_attn_mma_splitk_kernel<HEAD_DIM, BC>,
+                         cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
+    gqa_decode_attn_mma_splitk_kernel<HEAD_DIM, BC>
+        <<<dim3(p.kv_head, p.batch, num_splits), 32, smem>>>(
+            p, sc.o_part, sc.ml_part, num_splits);
+    gqa_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(
+        sc.o_part, sc.ml_part, p.o, num_splits, p.head_dim);
+}
+#endif
+
+static void launch_scalar_decode(const GQAParams& p) {
+    int gs = p.q_head / p.kv_head;
+    size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
+    gqa_decode_attn_kernel<<<p.batch * p.kv_head, dim3(32, gs), smem>>>(p);
+}
+
+template <int HEAD_DIM>
+static void dispatch_decode_t(GQAParams& p, DecodeScratch& sc) {
+#ifndef ASTRAI_NO_MMA
+    if (decode_use_mma(p)) { launch_mma_decode<HEAD_DIM, 32>(p, sc); return; }
+#endif
+    (void)sc;
+    launch_scalar_decode(p);
+}
+
+static void dispatch_decode(GQAParams& p, DecodeScratch& sc) {
+    switch (p.head_dim) {
+        case 32:  dispatch_decode_t<32>(p, sc);  break;
+        case 64:  dispatch_decode_t<64>(p, sc);  break;
+        case 128: dispatch_decode_t<128>(p, sc); break;
+        case 256: dispatch_decode_t<256>(p, sc); break;
+        default:  printf("bench: unsupported D=%d\n", p.head_dim);
+    }
+}
+
 static void cpu_decode(const float* Q, const float* K, const float* V,
                         const bool* mask, float* O,
                         int B, int Hq, int Hk, int seq_len, int D) {
@@ -54,6 +131,84 @@ static bf16 f2bf(float x) { return __float2bfloat16(x); }
 static float bf2f(bf16 x) { return __bfloat162float(x); }
 static float randf() { return (float)rand() / (float)RAND_MAX - 0.5f; }
 
+// Warmed-up, CUDA-event timed sweep over the production decode MMA path.
+// Decode (q_len==1) is memory-bound: the two matmuls are GEMV-shaped, so we
+// report both effective K/V read bandwidth and the (small) attention FLOP/s.
+// FLOP/s = 2 matmuls (q@K^T, P@V), each 2*B*Hq*kv*D flops.
+// Bytes    = K + V read = 2 * B*Hk*kv*D * sizeof(bf16).
+static void bench() {
+    const int cfgs[][5] = {
+        {1, 32, 4, 512, 128},    // B,Hq,Hk,seq,D
+        {1, 32, 4, 1024, 128},
+        {1, 32, 4, 2048, 128},
+        {1, 32, 4, 4096, 128},
+        {16, 32, 4, 2048, 128},
+        {32, 32, 4, 1024, 128},
+    };
+    int n = sizeof(cfgs)/sizeof(cfgs[0]);
+    const int WARMUP = 10, ITERS = 100;
+    printf("\n===== DECODE BENCH (warmup=%d iters=%d) =====\n", WARMUP, ITERS);
+    printf("%-46s | %10s | %10s | %10s\n",
+           "config", "latency", "bandwidth", "throughput");
+    printf("---------------------------------------------------------------"
+           "----------------------------\n");
+
+    for (int ci = 0; ci < n; ci++) {
+        int B=cfgs[ci][0], Hq=cfgs[ci][1], Hk=cfgs[ci][2];
+        int sl=cfgs[ci][3], D=cfgs[ci][4];
+        size_t nQ=(size_t)B*Hq*D, nKV=(size_t)B*Hk*sl*D;
+
+        bf16 *dQ,*dK,*dV,*dO,*tmp;
+        cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
+        cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
+        size_t big = nQ>nKV?nQ:nKV; tmp=new bf16[big];
+        for (size_t i=0;i<nQ;i++)  tmp[i]=f2bf(randf());
+        cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
+        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
+        cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
+        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
+        cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
+
+        GQAParams p;
+        p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D;
+        p.use_mask=0; p.is_causal=0; p.causal_offset=0;
+        p.scale=1.0f/sqrtf((float)D);
+        p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
+
+        DecodeScratch sc;
+        cudaMalloc(&sc.o_part, (size_t)B*Hq*32*D*sizeof(float));
+        cudaMalloc(&sc.ml_part, (size_t)B*Hq*32*2*sizeof(float));
+
+        for (int i=0;i<WARMUP;i++) dispatch_decode(p, sc);
+        cudaDeviceSynchronize();
+        cudaError_t err=cudaGetLastError();
+        if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return;}
+
+        cudaEvent_t s,e; cudaEventCreate(&s); cudaEventCreate(&e);
+        cudaEventRecord(s);
+        for (int i=0;i<ITERS;i++) dispatch_decode(p, sc);
+        cudaEventRecord(e); cudaEventSynchronize(e);
+        float ms=0; cudaEventElapsedTime(&ms,s,e); ms/=ITERS;
+
+        double flops = 4.0*B*Hq*(double)sl*D;
+        double tflops = flops/(ms*1e-3)/1e12;
+        // HBM traffic: K + V read (B*Hk*sl*D each), bf16; Q/O negligible.
+        double bytes = 2.0 * (2.0*nKV);
+        double gbps = bytes/(ms*1e-3)/1e9;
+
+        char cfg[64];
+        snprintf(cfg, sizeof(cfg),
+                 "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d causal=%d",
+                 B,Hq,Hk,1,sl,D,0);
+        printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
+               cfg, ms, gbps, tflops);
+
+        cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
+        cudaFree(sc.o_part);cudaFree(sc.ml_part);
+        delete[]tmp; cudaEventDestroy(s); cudaEventDestroy(e);
+    }
+}
+
 int main() {
     const int configs[][5] = {
         {1, 2, 1, 64, 32},    // B,Hq,Hk,seq_len,D
@@ -92,18 +247,17 @@ int main() {
 
         GQAParams p;
         p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D;
-        p.use_mask=1; p.is_causal=0; p.causal_offset=0;
+        p.use_mask=0; p.is_causal=0; p.causal_offset=0;
         p.scale=1.0f/sqrtf((float)D);
-        p.q=dQ; p.k=dK; p.v=dV; p.mask=dMask; p.o=dO;
+        p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
 
-        size_t smem=DC_CHUNK*D*sizeof(bf16);
-        dim3 block(32, gs);
-        dim3 grid(B*Hk);
-        printf("grid=(%d,1,1) block=(%d,%d,1) smem=%zu\n",
-               grid.x, block.x, block.y, smem);
+        // Split-K scratch (max 32 splits), sized for the production MMA path.
+        DecodeScratch sc;
+        cudaMalloc(&sc.o_part, (size_t)B*Hq*32*D*sizeof(float));
+        cudaMalloc(&sc.ml_part, (size_t)B*Hq*32*2*sizeof(float));
 
         double t0=now_ms();
-        gqa_decode_attn_kernel<<<grid,block,smem>>>(p);
+        dispatch_decode(p, sc);
         cudaDeviceSynchronize();
         double kms=now_ms()-t0;
         cudaError_t err=cudaGetLastError();
@@ -123,8 +277,10 @@ int main() {
         printf("kernel: %.3f ms  max_err: %.6e\n\n",kms,max_err);
 
         cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);cudaFree(dMask);
+        cudaFree(sc.o_part);cudaFree(sc.ml_part);
         delete[]hQ;delete[]hK;delete[]hV;delete[]hMask;delete[]hOut;delete[]ref;delete[]tmp;
     }
     printf("All tests passed!\n");
+    bench();
     return 0;
 }
diff --git a/csrc/tests/gqa_prefill_test.cu b/csrc/tests/gqa_prefill_test.cu
index 3acebec..6d2a753 100644
--- a/csrc/tests/gqa_prefill_test.cu
+++ b/csrc/tests/gqa_prefill_test.cu
@@ -10,6 +10,9 @@ nvcc -I csrc -arch=sm_89 -O3 \
 #include <cmath>
 #include <sys/time.h>
 #include "../kernels/gqa_prefill_attn.cuh"
+#ifndef ASTRAI_NO_MMA
+#include "../kernels/gqa_prefill_attn_mma.cuh"
+#endif
 
 static double now_ms() {
     struct timeval tv;
@@ -17,6 +20,32 @@ static double now_ms() {
     return tv.tv_sec * 1000.0 + tv.tv_usec / 1000.0;
 }
 
+// Launch the production prefill path (tensor-core MMA on sm_80+, else the
+// scalar fallback), mirroring dispatch_prefill() in gqa_prefill_attn.cu.
+template <int HEAD_DIM>
+static void launch_prefill(GQAParams& p) {
+#ifndef ASTRAI_NO_MMA
+    constexpr int WARPS = 4, BC = 16, BR = 16;
+    constexpr int MIN_BLOCKS = (HEAD_DIM <= 64) ? 6 : (HEAD_DIM <= 128) ? 3 : 2;
+    dim3 grid((p.q_len + BR * WARPS - 1) / (BR * WARPS), p.q_head, p.batch);
+    dim3 block(WARPS * 32, 1, 1);
+    gqa_prefill_attn_mma_kernel<HEAD_DIM, WARPS, BC, MIN_BLOCKS><<<grid, block>>>(p);
+#else
+    constexpr int G = 8, ROWS = 32, P_BC = 32;
+    dim3 grid((p.q_len + ROWS - 1) / ROWS, p.q_head, p.batch);
+    dim3 block(G, ROWS, 1);
+    gqa_prefill_attn_kernel_t<HEAD_DIM, G, ROWS, P_BC><<<grid, block>>>(p);
+#endif
+}
+
+static void dispatch_prefill(GQAParams& p) {
+    switch (p.head_dim) {
+        case 64:  launch_prefill<64>(p);  break;
+        case 128: launch_prefill<128>(p); break;
+        default:  printf("bench: unsupported D=%d\n", p.head_dim);
+    }
+}
+
 static void cpu_attention(const float* Q, const float* K, const float* V, float* O,
                           int B, int Hq, int Hk, int q_len, int kv_len, int D,
                           int is_causal, int causal_off) {
@@ -56,6 +85,78 @@ static __nv_bfloat16 f2bf(float x) { return __float2bfloat16(x); }
 static float bf2f(__nv_bfloat16 x) { return __bfloat162float(x); }
 static float randf() { return (float)rand() / (float)RAND_MAX - 0.5f; }
 
+// Warmed-up, CUDA-event timed throughput sweep over the production MMA path.
+// Reports per-call latency and effective tensor-core TFLOP/s (2 matmuls:
+// QK^T and P@V, each 2*B*Hq*ql*kl*D flops; halved for causal).
+static void bench() {
+    const int cfgs[][7] = {
+        {1,32,4,512,512,128,0},
+        {1,32,4,1024,1024,128,0},
+        {1,32,4,2048,2048,128,0},
+        {1,32,4,2048,2048,128,1},
+        {4,32,4,2048,2048,128,1},
+        {1,32,4,4096,4096,128,1},
+    };
+    int n = sizeof(cfgs)/sizeof(cfgs[0]);
+    const int WARMUP = 10, ITERS = 50;
+    printf("\n===== PREFILL BENCH (warmup=%d iters=%d) =====\n", WARMUP, ITERS);
+    printf("%-46s | %10s | %10s | %10s\n",
+           "config", "latency", "bandwidth", "throughput");
+    printf("---------------------------------------------------------------"
+           "----------------------------\n");
+
+    for (int ci = 0; ci < n; ci++) {
+        int B=cfgs[ci][0], Hq=cfgs[ci][1], Hk=cfgs[ci][2];
+        int ql=cfgs[ci][3], kl=cfgs[ci][4], D=cfgs[ci][5], causal=cfgs[ci][6];
+        size_t nQ=(size_t)B*Hq*ql*D, nKV=(size_t)B*Hk*kl*D;
+
+        bf16 *dQ,*dK,*dV,*dO,*tmp;
+        cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
+        cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
+        size_t big = nQ>nKV?nQ:nKV; tmp=new bf16[big];
+        for (size_t i=0;i<nQ;i++)  tmp[i]=f2bf(randf());
+        cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
+        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
+        cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
+        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
+        cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
+
+        GQAParams p;
+        p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
+        p.use_mask=0; p.is_causal=causal; p.causal_offset=0;
+        p.scale=1.0f/sqrtf((float)D);
+        p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
+
+        for (int i=0;i<WARMUP;i++) dispatch_prefill(p);
+        cudaDeviceSynchronize();
+        cudaError_t err=cudaGetLastError();
+        if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return;}
+
+        cudaEvent_t s,e; cudaEventCreate(&s); cudaEventCreate(&e);
+        cudaEventRecord(s);
+        for (int i=0;i<ITERS;i++) dispatch_prefill(p);
+        cudaEventRecord(e); cudaEventSynchronize(e);
+        float ms=0; cudaEventElapsedTime(&ms,s,e); ms/=ITERS;
+
+        double flops = 4.0*B*Hq*(double)ql*kl*D;
+        if (causal) flops *= 0.5;
+        double tflops = flops/(ms*1e-3)/1e12;
+        // HBM traffic: Q + O (B*Hq*ql*D each) + K + V (B*Hk*kl*D each), bf16.
+        double bytes = 2.0 * (2.0*nQ + 2.0*nKV);
+        double gbps = bytes/(ms*1e-3)/1e9;
+
+        char cfg[64];
+        snprintf(cfg, sizeof(cfg),
+                 "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d causal=%d",
+                 B,Hq,Hk,ql,kl,D,causal);
+        printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
+               cfg, ms, gbps, tflops);
+
+        cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
+        delete[]tmp; cudaEventDestroy(s); cudaEventDestroy(e);
+    }
+}
+
 int main() {
     const int configs[][7] = {
         {1,2,1,64,128,64,0},     // tiny: B,Hq,Hk,q,kv,D,causal
@@ -94,18 +195,8 @@ int main() {
         p.scale=1.0f/sqrtf((float)D);
         p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
 
-        constexpr int G=8, ROWS=32, P_BC=32;
-        dim3 grid((ql+ROWS-1)/ROWS, Hq, B);
-        dim3 block(G, ROWS, 1);
-        printf("grid=(%d,%d,%d) block=(%d,%d,%d)\n",
-               grid.x,grid.y,grid.z, block.x,block.y,block.z);
-
         double t0=now_ms();
-        switch (D) {
-            case 64:  gqa_prefill_attn_kernel_t<64, G,ROWS,P_BC><<<grid,block>>>(p); break;
-            case 128: gqa_prefill_attn_kernel_t<128,G,ROWS,P_BC><<<grid,block>>>(p); break;
-            default: printf("unsupported D=%d\n",D); return 1;
-        }
+        dispatch_prefill(p);
         cudaDeviceSynchronize();
         double kms=now_ms()-t0;
         cudaError_t err=cudaGetLastError();
@@ -128,5 +219,6 @@ int main() {
         delete[]hQ;delete[]hK;delete[]hV;delete[]hOut;delete[]ref;delete[]tmp;
     }
     printf("All tests passed!\n");
+    bench();
     return 0;
 }