From 7732c93f625a9f89f89d1b3e45163815791babb9 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 3 Jun 2026 01:08:48 +0000 Subject: [PATCH] Fix B2 indexer: use 16x256b.x1 TMEM read with TMEM_COLS=512 Revert to 16x256b.x1 approach (reads 64 rows from single column). Previous hang was likely due to TMEM_COLS=128 (too small). With TMEM_COLS=512, the full 128-row MMA output fits in TMEM. Lane i reads rows 4i..4i+3. Lanes 0-15 cover rows 0-63. 4 warps (0-3) each process 32 columns, computing weighted ReLU scores. --- dsv4/kernels/cuda/indexer_fp8_score_topk.cu | 102 ++++++++------------ 1 file changed, 39 insertions(+), 63 deletions(-) diff --git a/dsv4/kernels/cuda/indexer_fp8_score_topk.cu b/dsv4/kernels/cuda/indexer_fp8_score_topk.cu index e68fb21b..5ed9f9b6 100644 --- a/dsv4/kernels/cuda/indexer_fp8_score_topk.cu +++ b/dsv4/kernels/cuda/indexer_fp8_score_topk.cu @@ -176,9 +176,6 @@ indexer_fp8_score_topk_kernel( float* sW_h = (float*)(sbuf + off); off += n_ih * sizeof(float); off = (off + 127) & ~(size_t)127; - // Scratch buffer for cross-warp score accumulation (SK_TILE floats) - float* sLogits = (float*)(sbuf + off); off += SK_TILE * sizeof(float); - // Merge buffer for top-k: scores (top_k floats) + indices (top_k ints) float* sMergeScores = (float*)(sbuf + off); off += top_k * sizeof(float); int32_t* sMergeBlocks = (int32_t*)(sbuf + off); off += top_k * sizeof(int32_t); @@ -283,76 +280,56 @@ indexer_fp8_score_topk_kernel( __syncthreads(); // ---- Read TMEM results ---- - // tcgen05.ld.32x32b.x8 reads 32 rows × 8 cols per instruction. - // Per P7 docs: warp 0 reads rows 0-31, warp 1 reads rows 32-63 from the - // SAME TMEM address. Different warps see different row slices. + // The MMA produces [128 rows × 128 cols] in TMEM. + // tcgen05.ld.16x256b.x1 reads 16 lanes × 4 rows = 64 rows from ONE column. + // Lane i reads rows 4i..4i+3. Lanes 0-15 cover rows 0-63. + // Lanes 16-31 cover rows 64-127 (padding, ignored). // - // Warp 0 and 1 each handle 8 chunks of 8 columns (64 columns each). - // After per-warp reduce, the two warps exchange partial sums via SMEM - // to compute the full 64-head weighted ReLU score per column. + // Process on-the-fly: dequant, ReLU, weighted sum, top-k. + // 4 warps (0-3) each process 32 columns (128/4 = 32). - const int COLS_PER_READ = 8; - const int N_READ_CHUNKS = SK_TILE / COLS_PER_READ; // 16 - const int CHUNKS_PER_WARP = N_READ_CHUNKS / 2; // 8 chunks per read warp + const int COLS_PER_WARP = SK_TILE / 4; // 32 int my_warp = wid; + if (my_warp < 4) { + int col_start = my_warp * COLS_PER_WARP; + int col_end = col_start + COLS_PER_WARP; - if (my_warp < 2) { - int chunk_start = my_warp * CHUNKS_PER_WARP; - int chunk_end = chunk_start + CHUNKS_PER_WARP; - int h_base = my_warp * 32; // warp 0 → heads 0-31, warp 1 → heads 32-63 + for (int c = col_start; c < col_end; c++) { + if (c >= kv_len) break; - for (int ch = chunk_start; ch < chunk_end; ch++) { - int col_base = ch * COLS_PER_READ; - if (col_base >= kv_len) break; - int cols_valid = min(COLS_PER_READ, kv_len - col_base); + // Read column c from TMEM + uint32_t r0, r1, r2, r3; + asm volatile("tcgen05.ld.sync.aligned.16x256b.x1.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) : "r"(tb + c)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - // Read 8 columns. Warp 0 gets rows 0-31, warp 1 gets rows 32-63. - float vals[8] = {}; - { - float tmp[8]; - asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];" - : "=f"(tmp[0]),"=f"(tmp[1]),"=f"(tmp[2]),"=f"(tmp[3]), - "=f"(tmp[4]),"=f"(tmp[5]),"=f"(tmp[6]),"=f"(tmp[7]) - : "r"(tb + col_base)); - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - for (int j = 0; j < 8; j++) vals[j] = tmp[j]; - } + float f0, f1, f2, f3; + memcpy(&f0, &r0, 4); memcpy(&f1, &r1, 4); + memcpy(&f2, &r2, 4); memcpy(&f3, &r3, 4); - // Lane i has row h_base+i. Compute per-column weighted ReLU contributions. - for (int j = 0; j < cols_valid; j++) { - int c = col_base + j; - int h = h_base + lane; + // Lane i processes rows 4i..4i+3 for this column + if (lane < 16) { + float vals[4] = {f0, f1, f2, f3}; float k_s = k_scale[kv_start + c]; - float contrib = 0.0f; - if (h < n_ih) { - float logit = vals[j] * sQ_scale[h] * k_s; - contrib = (logit > 0.0f) ? sW_h[h] * logit : 0.0f; - } - // Warp-level reduce - for (int o = 16; o > 0; o >>= 1) - contrib += __shfl_down_sync(0xffffffff, contrib, o); - // Lane 0 has this warp's partial sum for this column - if (lane == 0) { - // Use sLogits as scratch for cross-warp accumulation - // First warp to arrive writes, second adds - if (my_warp == 0) { - sLogits[c] = contrib; - } else { - sLogits[c] += contrib; - } - } - } - __syncthreads(); - // Both warps can now insert into local top-k - if (my_warp < 2 && lane == 0) { - for (int j = 0; j < cols_valid; j++) { - int c = col_base + j; - float score = sLogits[c]; - if (score > 0.0f) { - local_heap_insert(local_scores, local_blocks, score, kv_start + c, INDEXER_LOCAL_K); + float weighted_relu_sum = 0.0f; + for (int j = 0; j < 4; j++) { + int h = lane * 4 + j; + if (h < n_ih) { + float logit = vals[j] * sQ_scale[h] * k_s; + if (logit > 0.0f) { + weighted_relu_sum += sW_h[h] * logit; + } } } + // Sum across lanes 0..15 within this warp + // Lanes 16..31 set to 0 + if (lane >= 16) weighted_relu_sum = 0.0f; + for (int o = 16; o > 0; o >>= 1) + weighted_relu_sum += __shfl_down_sync(0xffffffff, weighted_relu_sum, o); + if (lane == 0 && weighted_relu_sum > 0.0f) { + local_heap_insert(local_scores, local_blocks, weighted_relu_sum, kv_start + c, INDEXER_LOCAL_K); + } } } } @@ -434,7 +411,6 @@ void indexer_fp8_score_topk_cuda( smem += 128 * 32; smem = (smem + 127) & ~127; // sK8 smem += 128 * 4; smem = (smem + 127) & ~127; // sQ_scale smem += n_ih * 4; smem = (smem + 127) & ~127; // sW_h - smem += 128 * 4; // sLogits scratch (SK_TILE=128) smem += top_k * 4; // sMergeScores smem += top_k * 4; // sMergeBlocks smem += 192 * INDEXER_LOCAL_K * 4; // sCandScores