biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

B2 indexer: temporary heads 0-31 only while figuring out TMEM row 32-63 layout

Browse Source
This commit is contained in:
biondizzle
2026-06-03 01:12:10 +00:00
parent 7732c93f62
commit 6c92ff91f3
1 changed files with 82 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

119
dsv4/kernels/cuda/indexer_fp8_score_topk.cu
View File

@@ -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);
}
}
}