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

Fix B2 indexer: use 16x256b.x1 TMEM read with TMEM_COLS=512

Browse Source 
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.
This commit is contained in:
biondizzle
2026-06-03 01:08:48 +00:00
parent a75a9843af
commit 7732c93f62
1 changed files with 39 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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