refactor: extract shared MMA utils into gqa_mma_utils.cuh

- Move mma16816, ld2, pk2, pkb, ldmatrix_x4/x2/x2_trans to shared header
- gqa_prefill_attn_mma.cuh and gqa_decode_attn_mma.cuh both include it

csrc/kernels/gqa_decode_attn_mma.cuh

@@ -1,6 +1,6 @@
 #pragma once
 #include "gqa_common.cuh"
-#include "gqa_prefill_attn_mma.cuh"  // mma16816, ldmatrix_*, pk2
+#include "gqa_mma_utils.cuh"
 
 // Tensor-core decode via GQA head-packing.
 //

csrc/kernels/gqa_mma_utils.cuh

@@ -0,0 +1,57 @@
+#pragma once
+
+// Shared MMA utilities for tensor-core GQA kernels.
+// mma.sync.m16n8k16 PTX wrappers, ldmatrix helpers, and bf16 packing.
+
+// mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32
+__device__ __forceinline__ void mma16816(float* d, const unsigned* a,
+                                          const unsigned* b, const float* c) {
+    asm volatile(
+        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};"
+        : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3])
+        : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
+          "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
+}
+
+// read two adjacent bf16 from smem as one packed .b32 (elem0 low, elem1 high)
+__device__ __forceinline__ unsigned ld2(const bf16* p) {
+    return *reinterpret_cast<const unsigned*>(p);
+}
+
+// pack two floats into one bf16x2 as .b32
+__device__ __forceinline__ unsigned pk2(float a, float b) {
+    __nv_bfloat162 v = __floats2bfloat162_rn(a, b);
+    return *reinterpret_cast<unsigned*>(&v);
+}
+
+// pack two (non-contiguous) bf16 into one .b32
+__device__ __forceinline__ unsigned pkb(bf16 a, bf16 b) {
+    __nv_bfloat162 v;
+    v.x = a;
+    v.y = b;
+    return *reinterpret_cast<unsigned*>(&v);
+}
+
+// ldmatrix: cooperatively load mma fragments from smem (one instruction per
+// 16x16 / 16x8 tile) with the exact register layout mma expects — replaces the
+// scalar per-thread fragment packing, cutting shared-load instructions and bank
+// conflicts. Each lane supplies the shared address of one 8-wide row.
+__device__ __forceinline__ void ldmatrix_x4(unsigned* r, const bf16* p) {
+    unsigned a = __cvta_generic_to_shared(p);
+    asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
+                 : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
+                 : "r"(a));
+}
+__device__ __forceinline__ void ldmatrix_x2(unsigned* r, const bf16* p) {
+    unsigned a = __cvta_generic_to_shared(p);
+    asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];"
+                 : "=r"(r[0]), "=r"(r[1])
+                 : "r"(a));
+}
+__device__ __forceinline__ void ldmatrix_x2_trans(unsigned* r, const bf16* p) {
+    unsigned a = __cvta_generic_to_shared(p);
+    asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0,%1}, [%2];"
+                 : "=r"(r[0]), "=r"(r[1])
+                 : "r"(a));
+}

csrc/kernels/gqa_prefill_attn_mma.cuh

@@ -1,5 +1,6 @@
 #pragma once
 #include "gqa_common.cuh"
+#include "gqa_mma_utils.cuh"
 
 // Tensor-core prefill, register-resident flash attention (raw mma.sync PTX).
 // One warp owns BR=16 query rows. S = Q@K^T and O = P@V run on bf16 tensor
@@ -11,57 +12,6 @@
 // repack; row reductions fold across the 4-lane thread group. Templated on
 // <HEAD_DIM, WARPS, BC> with BC a multiple of 16.
 
-// mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32
-// (only compiled when ASTRAI_HAS_MMA is set, i.e. built for sm_80+)
-__device__ __forceinline__ void mma16816(float* d, const unsigned* a,
-                                          const unsigned* b, const float* c) {
-    asm volatile(
-        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
-        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};"
-        : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3])
-        : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
-          "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
-}
-
-// read two adjacent bf16 from smem as one packed .b32 (elem0 low, elem1 high)
-__device__ __forceinline__ unsigned ld2(const bf16* p) {
-    return *reinterpret_cast<const unsigned*>(p);
-}
-__device__ __forceinline__ unsigned pk2(float a, float b) {
-    __nv_bfloat162 v = __floats2bfloat162_rn(a, b);
-    return *reinterpret_cast<unsigned*>(&v);
-}
-// pack two (non-contiguous) bf16 into one .b32
-__device__ __forceinline__ unsigned pkb(bf16 a, bf16 b) {
-    __nv_bfloat162 v;
-    v.x = a;
-    v.y = b;
-    return *reinterpret_cast<unsigned*>(&v);
-}
-
-// ldmatrix: cooperatively load mma fragments from smem (one instruction per
-// 16x16 / 16x8 tile) with the exact register layout mma expects — replaces the
-// scalar per-thread fragment packing, cutting shared-load instructions and bank
-// conflicts. Each lane supplies the shared address of one 8-wide row.
-__device__ __forceinline__ void ldmatrix_x4(unsigned* r, const bf16* p) {
-    unsigned a = __cvta_generic_to_shared(p);
-    asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
-                 : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
-                 : "r"(a));
-}
-__device__ __forceinline__ void ldmatrix_x2(unsigned* r, const bf16* p) {
-    unsigned a = __cvta_generic_to_shared(p);
-    asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];"
-                 : "=r"(r[0]), "=r"(r[1])
-                 : "r"(a));
-}
-__device__ __forceinline__ void ldmatrix_x2_trans(unsigned* r, const bf16* p) {
-    unsigned a = __cvta_generic_to_shared(p);
-    asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0,%1}, [%2];"
-                 : "=r"(r[0]), "=r"(r[1])
-                 : "r"(a));
-}
-
 template <int HEAD_DIM, int WARPS, int BC>
 __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
     constexpr int BR = 16;