perf: load prefill Q fragments directly from global (drop sQ staging)

- read the 8 Q elements each lane needs straight from global into the mma
  A-operand layout, pre-scaled, instead of staging through shared sQ
- removes the sQ smem area (20KB->16KB) and the serialized per-warp prologue
  with its WARPS __syncthreads barriers
- result vs torch SDPA: prefill 0.70-0.82x -> 0.85-1.00x (matches torch at
  seq=128; 2048 1.251->1.159ms), correctness unchanged across head dims

csrc/kernels/gqa_prefill_attn_mma.cuh

@@ -5,8 +5,9 @@
 // Tensor-core prefill 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
 // cores via mma.sync.m16n8k16 (f32 accumulate). Q fragments are loaded once
-// per warp via shared sQ staging and kept resident in registers across the
+// straight from global into the mma A-operand layout (no smem staging) and
-// tile loop. S, O, and the online-softmax stats (m, l) also live in registers.
+// kept resident in registers across the tile loop. S, O, and the online-softmax
+// stats (m, l) also live in registers.
 // Shared memory is statically sized via template parameters — no dynamic
 // allocation. The mma fragment layout is used directly: the S accumulator
 // (f32) maps element-for-element onto the P matrix_a (bf16) operand, so
@@ -26,8 +27,8 @@
 // so no second barrier is needed. Predicated cp.async (cp_async_16_pred)
 // zero-fills rows past kv_len, unifying full and partial tiles on one path.
 //
-// Optimizations: shared sQ staging (single area, serialized per-warp load);
+// Optimizations: load Q fragments directly from global in mma A-operand layout
-// pre-scale Q by attention scale during Q load; packed bf16x2 output stores;
+// (no sQ staging, no prologue barriers); pre-scale Q by attention scale during Q load; packed bf16x2 output stores;
 // causal tile skipping (block-level prefetch bound + warp-level compute skip);
 // XOR swizzle (swiz_col) → eliminates ldmatrix bank conflicts without LD
 // padding (LD=HEAD_DIM).
@@ -61,33 +62,32 @@ void gqa_prefill_attn_mma_kernel(GQAParams p) {
     constexpr int STAGES = 2;
     __shared__ __align__(16) bf16 sK[STAGES * BC * LD];
     __shared__ __align__(16) bf16 sV[STAGES * BC * LD];
-    __shared__ __align__(16) bf16 sQ[BR * LD];
 
-    // Load Q into sQ with pre-scaling (staged per-warp to avoid smem conflicts).
+    // Load the Q fragments straight from global into the mma A-operand layout
-    // Pre-scale by attention scale so softmax doesn't need to multiply later.
+    // (m16n8k16, row-major): no sQ staging area and no serialized per-warp
-    // Q fragments are loaded via ldmatrix during this phase and kept resident
+    // prologue barriers. Each lane reads exactly the 8 Q elements ldmatrix
-    // in registers across the tile loop — sQ is only valid for the current warp
+    // would have produced, pre-scaled by the attention scale. Kept resident in
-    // during this loop, so Qa must be loaded here, not in the tile loop.
+    // registers across the tile loop.
+    //   frag[0]/[2]: row = qrow0 + gid ;  frag[1]/[3]: row = qrow0 + gid + 8
+    //   frag[0]/[1]: cols kt*16 + tid4*2 + {0,1} ; frag[2]/[3]: + 8
     const int q_base = ((batch * p.q_head + q_head) * p.q_len) * HEAD_DIM;
+    const int qra = qrow0 + gid;
+    const int qrb = qrow0 + gid + 8;
+    const bool va = qra < p.q_len, vb = qrb < p.q_len;
     unsigned Qa[KD][4];
-    bf16 scale_bf16 = __float2bfloat16(p.scale);
+#pragma unroll
-    int qrow_l = (lane & 7) + (lane & 8);       // 0..15
+    for (int kt = 0; kt < KD; kt++) {
-    int qcol_l = (lane & 16) ? 8 : 0;
+        int c = kt * 16 + tid4 * 2;
-    for (int w = 0; w < WARPS; w++) {
+        const bf16* pa = &p.q[q_base + qra * HEAD_DIM + c];
-        if (warp == w) {
+        const bf16* pb = &p.q[q_base + qrb * HEAD_DIM + c];
-            for (int i = lane; i < BR * HEAD_DIM; i += 32) {
+        Qa[kt][0] = va ? pk2(__bfloat162float(pa[0]) * p.scale,
-                int r = i / HEAD_DIM, d = i % HEAD_DIM;
+                             __bfloat162float(pa[1]) * p.scale) : 0u;
-                int qr = qrow0 + r;
+        Qa[kt][1] = vb ? pk2(__bfloat162float(pb[0]) * p.scale,
-                bf16 qv = (qr < p.q_len) ? p.q[q_base + qr * HEAD_DIM + d]
+                             __bfloat162float(pb[1]) * p.scale) : 0u;
-                                         : __float2bfloat16(0.0f);
+        Qa[kt][2] = va ? pk2(__bfloat162float(pa[8]) * p.scale,
-                sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = __hmul(qv, scale_bf16);
+                             __bfloat162float(pa[9]) * p.scale) : 0u;
-            }
+        Qa[kt][3] = vb ? pk2(__bfloat162float(pb[8]) * p.scale,
-            __syncwarp();
+                             __bfloat162float(pb[9]) * p.scale) : 0u;
-        #pragma unroll
-            for (int kt = 0; kt < KD; kt++)
-                ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + swiz_col(kt * 16 + qcol_l, qrow_l, SWIZ_MASK)]);
-        }
-        __syncthreads();  // prevent next warp from overwriting sQ prematurely
     }
 
     float Oacc[DN8][4];