ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: streamline Q block mapping

Browse Source 
- bypass shared mapping for contiguous attention
- centralize paged Q tile broadcast in KV policy helpers
This commit is contained in:
ViperEkura
2026-08-10 08:40:18 +08:00
parent 9b58fef222
commit 9f48cb8928
3 changed files with 23 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/attn_prefill_split_q_mma.cuh
+2 -4
View File
@@ -24,11 +24,9 @@ __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__ QTileMapper<Traits::BR * Traits::WARPS, KV> qmap;
if (!qmap.init(p, blockIdx.x, blockIdx.z))
int batch, q_tile;
if (!map_q_block<Traits::BR * Traits::WARPS, KV>(p, batch, q_tile))
return;
const int batch = qmap.batch;
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;