diff --git a/dsv4/kernels/cuda/indexer_fp8_score_topk.cu b/dsv4/kernels/cuda/indexer_fp8_score_topk.cu index 0bca3807..69975dd2 100644 --- a/dsv4/kernels/cuda/indexer_fp8_score_topk.cu +++ b/dsv4/kernels/cuda/indexer_fp8_score_topk.cu @@ -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); } } }