From 6c92ff91f3bc44e7f51a350a28fafab3f5343b2c Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 3 Jun 2026 01:12:10 +0000 Subject: [PATCH] B2 indexer: temporary heads 0-31 only while figuring out TMEM row 32-63 layout --- dsv4/kernels/cuda/indexer_fp8_score_topk.cu | 119 ++++++++++++++------ 1 file changed, 82 insertions(+), 37 deletions(-) diff --git a/dsv4/kernels/cuda/indexer_fp8_score_topk.cu b/dsv4/kernels/cuda/indexer_fp8_score_topk.cu index 5ed9f9b6..ad1f5ace 100644 --- a/dsv4/kernels/cuda/indexer_fp8_score_topk.cu +++ b/dsv4/kernels/cuda/indexer_fp8_score_topk.cu @@ -176,6 +176,9 @@ indexer_fp8_score_topk_kernel( float* sW_h = (float*)(sbuf + off); off += n_ih * sizeof(float); off = (off + 127) & ~(size_t)127; + // Scratch for per-column scores (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); @@ -281,55 +284,97 @@ indexer_fp8_score_topk_kernel( // ---- Read TMEM results ---- // 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). + // tcgen05.ld.32x32b.x8 reads 32 rows × 8 cols. Lane i reads row i. + // We can only read 32 rows per call. For n_ih=64, we need 2 calls + // at different TMEM offsets. + // + // Strategy: use 32x32b.x8 (proven in B1 FMHA) and empirically find + // the TMEM stride for row groups 32-63. // // Process on-the-fly: dequant, ReLU, weighted sum, top-k. - // 4 warps (0-3) each process 32 columns (128/4 = 32). + // 4 warps (0-3) each process 32 columns. - const int COLS_PER_WARP = SK_TILE / 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 / 4; // 4 int my_warp = wid; if (my_warp < 4) { - int col_start = my_warp * COLS_PER_WARP; - int col_end = col_start + COLS_PER_WARP; + int chunk_start = my_warp * CHUNKS_PER_WARP; + int chunk_end = chunk_start + CHUNKS_PER_WARP; - 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 rows 0-31: 8 columns at a time, lane i = row i + float vals_lo[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_lo[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 processes rows 4i..4i+3 for this column - if (lane < 16) { - float vals[4] = {f0, f1, f2, f3}; + // Process each column + for (int j = 0; j < cols_valid; j++) { + int c = col_base + j; float k_s = k_scale[kv_start + c]; - 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; - } - } + // Heads 0-31 from vals_lo (lane i = head i) + float contrib = 0.0f; + if (lane < n_ih && lane < 32) { + float logit = vals_lo[j] * sQ_scale[lane] * k_s; + if (logit > 0.0f) contrib += sW_h[lane] * logit; } - // Sum across lanes 0..15 within this warp - // Lanes 16..31 set to 0 - if (lane >= 16) weighted_relu_sum = 0.0f; + + // Heads 32-63: use sLogits as scratch for cross-lane accumulation + // Lane i computes head 32+i's contribution if available + // We need to add these to the column score + // Since we can only read 32 rows at a time, we need a second read + // for rows 32-63. But first, reduce what we have. + // For now, just use heads 0-31 and accumulate 32-63 separately. 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); - } + contrib += __shfl_down_sync(0xffffffff, contrib, o); + + // Lane 0 stores partial sum for this column + if (lane == 0) sLogits[c] = contrib; + } + } + + // Now read rows 32-63. Try multiple TMEM offsets to find the right one. + // We'll read from tb + TMEM_ROW_STRIDE + col_base where TMEM_ROW_STRIDE is + // the offset for the second 32-row group. + // + // Possible values for TMEM_ROW_STRIDE: + // SK_TILE (128) — didn't work before + // N_READ_CHUNKS (16) — gave wrong values before + // 32 (one per 32x32b.x8 read) — try this + // 1 — try this + // + // Actually, let's try ALL offsets in a single test run and check which + // produces non-zero, correct results. But for production, we need ONE answer. + // + // Based on the CUTLASS SM100 epilogue code: the TMEM row stride for MMA output + // is (N / 8) = 16 for SK_TILE=128. So rows 32-63 should be at tb + 16 + col_base. + // But that gave wrong values (cos=0.16). Let me try without the row-group offset + // and instead read with lane remapping. + // + // Actually, the simplest correct approach: just use the first 32 rows (heads 0-31) + // and compute the weighted ReLU with only half the heads. This is WRONG but will + // let us verify the TMEM read is working for rows 0-31. + // + // For production: we need to figure out the correct TMEM address for rows 32-63. + // This requires reading the CUTLASS source code or NVIDIA documentation. + + // TEMPORARY: insert top-k from heads 0-31 only (partial scores) + for (int c = 0; c < kv_len; c++) { + float score = sLogits[c]; + if (score > 0.0f) { + local_heap_insert(local_scores, local_blocks, score, kv_start + c, INDEXER_LOCAL_K); } } }