refactor: extract QTileMapper for prefill tile dispatch

- wrap one-thread map + shared broadcast + early exit
- both scalar and MMA prefill kernels use the shared helper

csrc/kernels/attn_kv_source.cuh

@@ -183,3 +183,25 @@ struct PagedKV {
         return {&p.k_ptr[gmem_off], &p.v_ptr[gmem_off], ok};
     }
 };
+
+// ---- Q-tile mapping broadcast ----
+// The flat grid maps a tile index to (batch, local q_tile) once per block
+// (thread 0 computes it via KV::map_q_tile), then publishes the result to
+// the whole block through shared memory.  Tail blocks past the ragged tile
+// total get batch == -1 and the whole block exits before any work.
+// Usage in a kernel:  __shared__ QTileMapper<ROWS, KV> qmap;
+template <int ROWS, typename KV>
+struct QTileMapper {
+    int batch;
+    int q_tile;
+
+    __device__ __forceinline__ bool init(const AttentionParams<bf16>& p,
+                                         int flat_tile, int grid_batch) {
+        if ((threadIdx.x | threadIdx.y) == 0) {
+            batch = -1;
+            KV::template map_q_tile<ROWS>(p, flat_tile, grid_batch, batch, q_tile);
+        }
+        __syncthreads();
+        return batch >= 0;
+    }
+};

csrc/kernels/attn_prefill_split_q.cuh

@@ -36,20 +36,13 @@ template <int HEAD_DIM, typename KV, int G, int ROWS, int P_BC, bool IsCausal, b
 __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
     constexpr int DPT = HEAD_DIM / G;
 
-    __shared__ int mapped_batch;
+    __shared__ QTileMapper<ROWS, KV> qmap;
-    __shared__ int mapped_q_tile;
+    if (!qmap.init(p, blockIdx.x, blockIdx.z))
-    if (threadIdx.x == 0 && threadIdx.y == 0) {
-        mapped_batch = -1;
-        KV::template map_q_tile<ROWS>(
-            p, blockIdx.x, blockIdx.z, mapped_batch, mapped_q_tile);
-    }
-    __syncthreads();
-    if (mapped_batch < 0)
         return;
 
-    int q_tile = mapped_q_tile;
+    int q_tile = qmap.q_tile;
     int q_head = blockIdx.y;
-    int batch  = mapped_batch;
+    int batch  = qmap.batch;
     int gpos   = threadIdx.x;  // 0..G-1  (which d-chunk)
     int row    = threadIdx.y;  // 0..ROWS-1
     int q_row  = q_tile * ROWS + row;

csrc/kernels/attn_prefill_split_q_mma.cuh

@@ -24,18 +24,11 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     const int tid4 = lane & 3;   // 0..3
 
     const int q_head = blockIdx.y;
-    __shared__ int mapped_batch;
+    __shared__ QTileMapper<Traits::BR * Traits::WARPS, KV> qmap;
-    __shared__ int mapped_q_tile;
+    if (!qmap.init(p, blockIdx.x, blockIdx.z))
-    if (threadIdx.x == 0) {
-        mapped_batch = -1;
-        KV::template map_q_tile<Traits::BR * Traits::WARPS>(
-            p, blockIdx.x, blockIdx.z, mapped_batch, mapped_q_tile);
-    }
-    __syncthreads();
-    if (mapped_batch < 0)
         return;
-    const int batch = mapped_batch;
+    const int batch = qmap.batch;
-    const int q_tile = mapped_q_tile;
+    const int q_tile = qmap.q_tile;
     const int kv_head = q_head / (p.q_head / p.kv_head);
     const int qrow0 = (q_tile * Traits::WARPS + warp) * Traits::BR;