refactor: streamline Q block mapping

- bypass shared mapping for contiguous attention - centralize paged Q tile broadcast in KV policy helpers
2026-08-10 08:40:18 +08:00
parent 9b58fef222
commit 9f48cb8928
3 changed files with 23 additions and 22 deletions
@@ -36,13 +36,11 @@ 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__ QTileMapper<ROWS, KV> qmap;
-    if (!qmap.init(p, blockIdx.x, blockIdx.z))
+    int batch, q_tile;
+    if (!map_q_block<ROWS, KV>(p, batch, q_tile))
        return;

-    int q_tile = qmap.q_tile;
    int q_head = blockIdx.y;
-    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;