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

Fix B2 indexer: replace broken 16x256b TMEM read with proven 32x32b.x8

Browse Source 
ROOT CAUSES:
1. tcgen05.ld.16x256b.x1 was hanging — either invalid instruction or unaligned
2. TMEM_COLS=128 was too small for 64-row MMA output (needs 256 for 2 row-groups)
3. TMEM row-group addressing: rows 32-63 are at offset SK_TILE (128) in TMEM

Fixes:
- Use tcgen05.ld.32x32b.x8 (proven in B1 FMHA) instead of 16x256b.x1
- Increase TMEM_COLS from 128 to 256
- Read both row-groups (0-31 and 32-63) per 8-column chunk
- Each lane handles head i (from row-group 0) and head 32+i (from row-group 1)
- Warp-level reduce sums contributions from all 64 heads per column
This commit is contained in:
biondizzle
2026-06-03 00:39:49 +00:00
parent 99e50fcb58
commit afb82b9c89
1 changed files with 60 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -149,7 +149,8 @@ indexer_fp8_score_topk_kernel(
constexpr int MMA_K_F8 = 32;
constexpr int NKT = 4; // ihd=128 / MMA_K_F8=32
constexpr int TILE_F8 = 128 * 32; // 4096 bytes per SMEM tile
constexpr int TMEM_COLS = 128;
constexpr int TMEM_COLS = 256; // 128 rows × 128 cols needs 4×128 = 512,
// but only 64 rows used (2×128 = 256)
const int tid = threadIdx.x;
const int wid = tid >> 5;
@@ -280,58 +281,76 @@ indexer_fp8_score_topk_kernel(
__syncthreads();
// ---- Read TMEM results ----
// We need rows 0..n_ih-1 (64 rows) × SK_TILE (128 columns) from TMEM.
// Using tcgen05.ld.16x256b.x1: lane i reads rows 4i..4i+3 from one column.
// Lanes 0..15 cover rows 0..63. Lanes 16..31 cover rows 64..127 (ignored).
// TMEM layout: MMA produces [128 rows × 128 cols].
// tcgen05.ld.32x32b.x8 reads 32 rows × 8 cols per instruction.
// Row groups: rows 0-31 at [tb+0..tb+127], rows 32-63 at [tb+128..tb+255].
//
// Process logits on-the-fly: dequant, ReLU, weighted sum, top-k update.
// No SMEM staging of the full logits matrix needed.
//
// Parallel read: warps 0-3 each read 32 columns (128/4=32), processing
// independently. Each warp computes the weighted ReLU sum for its columns
// and updates per-thread local top-k.
// Each warp processes 4 chunks of 8 columns (32 columns total).
// For each chunk, read both row-groups (0-31 and 32-63), then
// compute per-column weighted ReLU scores.
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 row group 0-31: lane i = row i, 8 column values per lane
float vals_lo[8] = {};
// All 32 lanes must participate in the TMEM read
{
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);
// Read row group 32-63: lane i = row 32+i, 8 column values per lane
float vals_hi[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 + SK_TILE + col_base));
asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory");
for (int j = 0; j < 8; j++) vals_hi[j] = tmp[j];
}
// Lane i processes rows 4i..4i+3 for this column
if (lane < (n_ih + 3) / 4) {
float vals[4] = {f0, f1, f2, f3};
// Process each column: compute score[c] = sum_h w_h[h] * relu(q_s[h]*k_s[c]*logit[h,c])
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;
}
}
// Lane i contributes head i (from vals_lo) and head 32+i (from vals_hi)
float contrib = 0.0f;
// Head i (row i, 0-31)
if (lane < n_ih && lane < 32) {
float logit0 = vals_lo[j] * sQ_scale[lane] * k_s;
if (logit0 > 0.0f) contrib += sW_h[lane] * logit0;
}
// Sum across lanes 0..15 within this warp
if (lane >= 16) weighted_relu_sum = 0.0f;
// Head 32+i (row 32+i)
int h1 = lane + 32;
if (h1 < n_ih) {
float logit1 = vals_hi[j] * sQ_scale[h1] * k_s;
if (logit1 > 0.0f) contrib += sW_h[h1] * logit1;
}
// Sum contributions across all 32 lanes (64 heads total)
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) {
int c_global = kv_start + c;
local_heap_insert(local_scores, local_blocks, weighted_relu_sum, c_global, INDEXER_LOCAL_K);
contrib += __shfl_down_sync(0xffffffff, contrib, o);
if (lane == 0 && contrib > 0.0f) {
local_heap_insert(local_scores, local_blocks, contrib, kv_start + c, INDEXER_LOCAL_K);
}
}
}