B2 indexer: adopt TMEM warp-to-row mapping fix
Key insight: tcgen05.ld.32x32b.x8 maps warp 0 to rows 0-31 and warp 1 to rows 32-63 from the SAME TMEM address. The hardware routes row slices based on warp position in the warpgroup. Fix approach (from external LLM review): - Warps 0-1 both read from tb + col_base (same address) - Each warp writes partial scores to its own sWarpScores partition - After __syncthreads(), merge both partitions for final 64-head scores - No race conditions, no cross-warp accumulation bugs
This commit is contained in:
@@ -4,22 +4,24 @@
|
||||
* CSA Lightning Indexer (paper §2.3.1, eq. 16):
|
||||
* I[t,s] = Σ_h w_h[t,h] · ReLU(q_I[t,h] · K^IComp[s])
|
||||
*
|
||||
* Native Blackwell FP8 tensor-core path for decode (T=1):
|
||||
* 1. Quantize Q (n_ih=64, ihd=128) BF16 → FP8_E4M3 with per-row FP32 scale
|
||||
* 2. FP8 GEMM via tcgen05.mma.kind::f8f6f4:
|
||||
* Q (128, 128 padded) × K^T (128, n_comp tiled by 128) → (64, n_comp) logits
|
||||
* 3. Dequant GEMM output: logit[h,c] *= q_scale[h] * k_scale[kv_start+c]
|
||||
* 4. ReLU, then weighted sum: score[c] = Σ_h w_h[h] * relu(logit[h,c])
|
||||
* 5. Top-k selection from (n_comp,) scores
|
||||
* Decode-specialized Blackwell FP8 tensor-core path (T=1):
|
||||
* 1. Quantize Q (n_ih=64, ihd=128) BF16 → FP8_E4M3 with per-row FP32 scale.
|
||||
* 2. Run Q (128x128 padded) × K^T (128x128 tile) with tcgen05.mma.kind::f8f6f4.
|
||||
* 3. Read accumulator rows from TMEM with tcgen05.ld.32x32b.x8.
|
||||
* 4. Dequant logits in registers, apply ReLU, weighted sum across indexer heads.
|
||||
* 5. Block-local top-k selection.
|
||||
*
|
||||
* Specialized for DSV4 Pro: n_ih=64, ihd=128, top_k=1024.
|
||||
* Important TMEM rule for M=128, cta_group::1:
|
||||
* tcgen05.ld.32x32b.x8 does NOT use a row offset in the address. The warp id in
|
||||
* the first warpgroup selects the row/lane slice:
|
||||
* warp 0 -> TMEM lanes/rows 0..31
|
||||
* warp 1 -> TMEM lanes/rows 32..63
|
||||
* warp 2 -> TMEM lanes/rows 64..95
|
||||
* warp 3 -> TMEM lanes/rows 96..127
|
||||
* All those warps use the same taddr for the same column group.
|
||||
*
|
||||
* TMEM read strategy for 64 Q rows:
|
||||
* Use tcgen05.ld.16x256b.x1 (proven in B1 FMHA) — one column per instruction.
|
||||
* Lane i reads rows 4i..4i+3 from the column. Lanes 0-15 cover rows 0-63.
|
||||
* 128 reads per K-tile to cover all N-dimension columns.
|
||||
*
|
||||
* NO PyTorch fallback. NO FP32 einsum on CUDA cores. NO BF16 workarounds.
|
||||
* No PyTorch fallback here. No FP32 einsum. The only FP32 CUDA-core work is the
|
||||
* unavoidable post-MMA dequant/ReLU/weighted-reduction/top-k epilogue.
|
||||
*/
|
||||
|
||||
#include <cuda.h>
|
||||
@@ -39,9 +41,6 @@ static constexpr int NWARPS = 6;
|
||||
typedef unsigned short bf16_t;
|
||||
|
||||
// ---- PTX helpers ----
|
||||
__device__ __forceinline__ bf16_t f32_to_bf16_ptx(float f) {
|
||||
bf16_t h; asm("cvt.rn.bf16.f32 %0, %1;" : "=h"(h) : "f"(f)); return h;
|
||||
}
|
||||
__device__ __forceinline__ float bf16_to_f32_ptx(bf16_t h) {
|
||||
float f; asm("cvt.f32.bf16 %0, %1;" : "=f"(f) : "h"(h)); return f;
|
||||
}
|
||||
@@ -51,7 +50,7 @@ __device__ __forceinline__ uint8_t fp8_e4m3_from_f32(float x) {
|
||||
return *reinterpret_cast<uint8_t*>(&v);
|
||||
}
|
||||
|
||||
// ---- UMMA helpers (from fmha_umma_desc.cuh, replicated for ATen build) ----
|
||||
// ---- UMMA helpers (mirrors the B1 FMHA helpers) ----
|
||||
__device__ __forceinline__ uint64_t desc_encode(uint64_t byte_val) { return byte_val >> 4; }
|
||||
|
||||
__device__ __forceinline__ uint64_t make_umma_desc_kmajor_none(uint32_t smem_addr, int block_mn) {
|
||||
@@ -69,31 +68,58 @@ __device__ __forceinline__ uint32_t make_idesc_f8_e4m3(int block_m, int block_n)
|
||||
return (1U << 4) | ((uint32_t)(block_n >> 3) << 17) | ((uint32_t)(block_m >> 4) << 24);
|
||||
}
|
||||
|
||||
__device__ void umma_ss_f8f6f4(uint32_t tmem_c, uint64_t desc_a, uint64_t desc_b,
|
||||
uint32_t i_desc, bool accumulate) {
|
||||
__device__ __forceinline__ void umma_ss_f8f6f4(uint32_t tmem_c, uint64_t desc_a, uint64_t desc_b,
|
||||
uint32_t i_desc, bool accumulate) {
|
||||
uint32_t scaleC_bits = accumulate ? 0x3F800000u : 0u;
|
||||
asm volatile("{\n\t.reg .pred p;\n\tsetp.ne.b32 p, %4, 0;\n\t"
|
||||
"tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t}"
|
||||
:: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(i_desc), "r"(scaleC_bits));
|
||||
:: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(i_desc), "r"(scaleC_bits)
|
||||
: "memory");
|
||||
}
|
||||
|
||||
__device__ void tmem_alloc(uint32_t smem_ptr, int num_cols) {
|
||||
__device__ __forceinline__ void tmem_alloc(uint32_t smem_ptr, int num_cols) {
|
||||
asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;"
|
||||
:: "r"(smem_ptr), "r"(num_cols));
|
||||
:: "r"(smem_ptr), "r"(num_cols) : "memory");
|
||||
}
|
||||
__device__ void tmem_dealloc(uint32_t tmem_ptr, int num_cols) {
|
||||
__device__ __forceinline__ void tmem_dealloc(uint32_t tmem_ptr, int num_cols) {
|
||||
asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;"
|
||||
:: "r"(tmem_ptr), "r"(num_cols));
|
||||
:: "r"(tmem_ptr), "r"(num_cols) : "memory");
|
||||
}
|
||||
|
||||
// ---- FP8 canonical SMEM layout (same as B1 FMHA) ----
|
||||
__device__ __forceinline__ void mbarrier_init_cta(uint32_t smem_mbar, uint32_t arrival_count = 1) {
|
||||
asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;"
|
||||
:: "r"(smem_mbar), "r"(arrival_count) : "memory");
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void tcgen05_commit_mma(uint32_t smem_mbar) {
|
||||
asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.b64 [%0];"
|
||||
:: "r"(smem_mbar) : "memory");
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void mbarrier_wait_cta(uint32_t smem_mbar, int phase) {
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
".reg .pred p;\n\t"
|
||||
"B2_WAIT_MMA:\n\t"
|
||||
"mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 p, [%0], %1, %2;\n\t"
|
||||
"@p bra.uni B2_DONE_MMA;\n\t"
|
||||
"bra.uni B2_WAIT_MMA;\n\t"
|
||||
"B2_DONE_MMA:\n\t"
|
||||
"}\n"
|
||||
:: "r"(smem_mbar), "r"(phase), "r"(0x989680)
|
||||
: "memory");
|
||||
}
|
||||
|
||||
// ---- FP8 canonical SMEM layout for tcgen05.mma.kind::f8f6f4 ----
|
||||
__device__ __forceinline__ int canon_idx_fp8_128x32(int r, int c) {
|
||||
int core_mn = r >> 3; int core_k = c >> 4;
|
||||
int local_r = r & 7; int local_c = c & 15;
|
||||
int core_mn = r >> 3;
|
||||
int core_k = c >> 4;
|
||||
int local_r = r & 7;
|
||||
int local_c = c & 15;
|
||||
return core_k * 16 * 128 + core_mn * 128 + local_r * 16 + local_c;
|
||||
}
|
||||
|
||||
// ---- Top-k (proven from indexer_score_topk.cu) ----
|
||||
// ---- Top-k helpers ----
|
||||
#ifndef INDEXER_LOCAL_K
|
||||
#define INDEXER_LOCAL_K 8
|
||||
#endif
|
||||
@@ -104,7 +130,7 @@ __device__ __forceinline__ void local_heap_insert(float* scores, int32_t* blocks
|
||||
scores[0] = score; blocks[0] = block_id;
|
||||
int root = 0;
|
||||
while (root < (k >> 1)) {
|
||||
int left = 2*root+1, right = 2*root+2, smallest = root;
|
||||
int left = 2 * root + 1, right = 2 * root + 2, smallest = root;
|
||||
if (left < k && scores[left] < scores[smallest]) smallest = left;
|
||||
if (right < k && scores[right] < scores[smallest]) smallest = right;
|
||||
if (smallest == root) break;
|
||||
@@ -121,7 +147,7 @@ __device__ __forceinline__ void heap_insert_shared(float* heap_scores, int32_t*
|
||||
heap_scores[0] = score; heap_blocks[0] = block_id;
|
||||
int root = 0;
|
||||
while (root < (k >> 1)) {
|
||||
int left = 2*root+1, right = 2*root+2, smallest = root;
|
||||
int left = 2 * root + 1, right = 2 * root + 2, smallest = root;
|
||||
if (left < k && heap_scores[left] < heap_scores[smallest]) smallest = left;
|
||||
if (right < k && heap_scores[right] < heap_scores[smallest]) smallest = right;
|
||||
if (smallest == root) break;
|
||||
@@ -140,127 +166,126 @@ template<int SK_TILE=128>
|
||||
__global__ void __launch_bounds__(192)
|
||||
indexer_fp8_score_topk_kernel(
|
||||
const bf16_t* __restrict__ q_bf16, // (n_ih, ihd) BF16 row-major
|
||||
const uint8_t* __restrict__ k_fp8, // (n_comp, ihd) FP8_E4M3
|
||||
const float* __restrict__ k_scale, // (n_comp,) FP32
|
||||
const bf16_t* __restrict__ w_h_bf16, // (n_ih,) BF16
|
||||
int32_t* __restrict__ topk_indices, // (top_k,) output
|
||||
const uint8_t* __restrict__ k_fp8, // (n_comp, ihd) FP8_E4M3 bytes
|
||||
const float* __restrict__ k_scale, // (n_comp,) FP32 dequant scales
|
||||
const bf16_t* __restrict__ w_h_bf16, // (n_ih,) BF16 weights
|
||||
int32_t* __restrict__ topk_indices, // (top_k,) int32 output
|
||||
int n_comp, int n_ih, int ihd, int top_k
|
||||
) {
|
||||
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 = 512; // 128 rows × 128 cols → 4 row-groups × 128 cols = 512
|
||||
constexpr int NKT = 4; // ihd=128 / 32
|
||||
constexpr int TILE_F8 = 128 * 32; // bytes per canonical FP8 tile
|
||||
constexpr int TMEM_COLS = 512; // full 128 lanes x 512 columns allocation
|
||||
|
||||
const int tid = threadIdx.x;
|
||||
const int wid = tid >> 5;
|
||||
const int lane = tid & 31;
|
||||
const bool is_mma_warp = (wid == 4);
|
||||
|
||||
__shared__ float sQ_amax_warp[NWARPS];
|
||||
|
||||
// ---- SMEM layout ----
|
||||
extern __shared__ __align__(128) char sbuf[];
|
||||
size_t off = 0;
|
||||
uint32_t* sTmemBase = (uint32_t*)(sbuf + off); off += 4;
|
||||
off = (off + 15) & ~(size_t)15;
|
||||
uint64_t* sMbar = (uint64_t*)(sbuf + off); off += 8;
|
||||
off = (off + 127) & ~(size_t)127;
|
||||
|
||||
// FP8 SMEM tiles for Q and K (canonical layout, 128×32 each)
|
||||
uint8_t* sQ8 = (uint8_t*)(sbuf + off); off += TILE_F8;
|
||||
off = (off + 127) & ~(size_t)127;
|
||||
uint8_t* sK8 = (uint8_t*)(sbuf + off); off += TILE_F8;
|
||||
off = (off + 127) & ~(size_t)127;
|
||||
|
||||
// Per-row Q FP8 scales (n_ih, padded to 128 for alignment)
|
||||
float* sQ_scale = (float*)(sbuf + off); off += 128 * sizeof(float);
|
||||
off = (off + 127) & ~(size_t)127;
|
||||
|
||||
// w_h in FP32 (n_ih)
|
||||
float* sW_h = (float*)(sbuf + off); off += n_ih * sizeof(float);
|
||||
float* sW_h = (float*)(sbuf + off); off += 128 * 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);
|
||||
// Two warp partial sums: warp 0 covers heads 0..31, warp 1 covers 32..63.
|
||||
float* sWarpScores = (float*)(sbuf + off); off += 2 * SK_TILE * sizeof(float);
|
||||
off = (off + 127) & ~(size_t)127;
|
||||
|
||||
// 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);
|
||||
|
||||
// Per-thread candidates for merge
|
||||
float* sCandScores = (float*)(sbuf + off); off += NTHREADS * INDEXER_LOCAL_K * sizeof(float);
|
||||
int32_t* sCandBlocks = (int32_t*)(sbuf + off); off += NTHREADS * INDEXER_LOCAL_K * sizeof(int32_t);
|
||||
|
||||
// ---- Per-thread local top-k ----
|
||||
float local_scores[INDEXER_LOCAL_K];
|
||||
int32_t local_blocks[INDEXER_LOCAL_K];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < INDEXER_LOCAL_K; i++) {
|
||||
local_scores[i] = -INFINITY;
|
||||
local_blocks[i] = -1;
|
||||
}
|
||||
|
||||
// ---- Init SMEM ----
|
||||
for (int i = tid; i < 128; i += NTHREADS) sQ_scale[i] = 0.0f;
|
||||
for (int i = tid; i < 128; i += NTHREADS) {
|
||||
sQ_scale[i] = 0.0f;
|
||||
sW_h[i] = 0.0f;
|
||||
}
|
||||
for (int i = tid; i < n_ih; i += NTHREADS) sW_h[i] = bf16_to_f32_ptx(w_h_bf16[i]);
|
||||
__syncthreads();
|
||||
|
||||
// ---- Phase 0: Compute per-row Q amax and quantize ----
|
||||
// Q is (n_ih, ihd) BF16 in GMEM. Each row gets its own FP8 scale.
|
||||
// All threads cooperate on each row (one row at a time for simplicity).
|
||||
// ---- Phase 0: Q per-row amax + scale ----
|
||||
for (int h = 0; h < n_ih; h++) {
|
||||
float local_max = 0.0f;
|
||||
for (int d = tid; d < ihd; d += NTHREADS) {
|
||||
float val = fabsf(bf16_to_f32_ptx(q_bf16[h * ihd + d]));
|
||||
local_max = fmaxf(local_max, val);
|
||||
local_max = fmaxf(local_max, fabsf(bf16_to_f32_ptx(q_bf16[h * ihd + d])));
|
||||
}
|
||||
// Warp-level reduce
|
||||
for (int o = 16; o > 0; o >>= 1)
|
||||
local_max = fmaxf(local_max, __shfl_down_sync(0xffffffff, local_max, o));
|
||||
__shared__ float _q_amax[6];
|
||||
if ((tid & 31) == 0) _q_amax[tid >> 5] = local_max;
|
||||
if (lane == 0) sQ_amax_warp[wid] = local_max;
|
||||
__syncthreads();
|
||||
|
||||
float amax = 0.0f;
|
||||
if (tid < 32) {
|
||||
amax = (tid < 6) ? _q_amax[tid] : 0.0f;
|
||||
amax = (tid < NWARPS) ? sQ_amax_warp[tid] : 0.0f;
|
||||
for (int o = 16; o > 0; o >>= 1)
|
||||
amax = fmaxf(amax, __shfl_down_sync(0xffffffff, amax, o));
|
||||
}
|
||||
amax = __shfl_sync(0xffffffff, amax, 0);
|
||||
float scale = amax / E4M3_MAX;
|
||||
if (scale < 1e-8f) scale = 1e-8f;
|
||||
if (tid == 0) sQ_scale[h] = scale;
|
||||
// Don't write Q to SMEM yet — we'll do it per-MMA K-slice
|
||||
if (tid == 0) {
|
||||
float scale = amax / E4M3_MAX;
|
||||
sQ_scale[h] = (scale < 1e-8f) ? 1e-8f : scale;
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// ---- TMEM + mbarrier init ----
|
||||
const uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbar);
|
||||
if (tid == 0) {
|
||||
mbarrier_init_cta(mbar_addr, 1);
|
||||
asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ---- TMEM alloc ----
|
||||
if (is_mma_warp) tmem_alloc((uint32_t)__cvta_generic_to_shared(sTmemBase), TMEM_COLS);
|
||||
asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
|
||||
__syncthreads();
|
||||
uint32_t tb = *sTmemBase;
|
||||
|
||||
// ---- Phase 1: FP8 GEMM — Q × K^T → logits (n_ih, n_comp) ----
|
||||
const int n_k_tiles = (n_comp + SK_TILE - 1) / SK_TILE;
|
||||
const uint32_t idesc_f8 = make_idesc_f8_e4m3(128, 128);
|
||||
int mma_phase = 0;
|
||||
|
||||
for (int kv_tile = 0; kv_tile < n_k_tiles; kv_tile++) {
|
||||
const int kv_start = kv_tile * SK_TILE;
|
||||
const int kv_len = min(SK_TILE, n_comp - kv_start);
|
||||
|
||||
for (int i = tid; i < 2 * SK_TILE; i += NTHREADS) sWarpScores[i] = 0.0f;
|
||||
__syncthreads();
|
||||
|
||||
// ---- FP8 QK GEMM over ihd=128 in four K-slices ----
|
||||
for (int kt = 0; kt < NKT; kt++) {
|
||||
// Zero SMEM tiles
|
||||
for (int i = tid; i < TILE_F8; i += NTHREADS) { sQ8[i] = 0; sK8[i] = 0; }
|
||||
__syncthreads();
|
||||
|
||||
// Load Q rows 0..n_ih-1, columns kt*32..kt*32+31 into sQ8 canonical
|
||||
for (int i = tid; i < n_ih * MMA_K_F8; i += NTHREADS) {
|
||||
int row = i / MMA_K_F8;
|
||||
int col = i % MMA_K_F8;
|
||||
int d = kt * MMA_K_F8 + col;
|
||||
if (d < ihd) {
|
||||
float val = bf16_to_f32_ptx(q_bf16[row * ihd + d]);
|
||||
float inv_scale = 1.0f / sQ_scale[row];
|
||||
sQ8[canon_idx_fp8_128x32(row, col)] = fp8_e4m3_from_f32(val * inv_scale);
|
||||
}
|
||||
float val = bf16_to_f32_ptx(q_bf16[row * ihd + d]);
|
||||
sQ8[canon_idx_fp8_128x32(row, col)] = fp8_e4m3_from_f32(val / sQ_scale[row]);
|
||||
}
|
||||
|
||||
// Load K rows 0..kv_len-1, columns kt*32..kt*32+31 into sK8 canonical
|
||||
for (int i = tid; i < kv_len * MMA_K_F8; i += NTHREADS) {
|
||||
int row = i / MMA_K_F8;
|
||||
int col = i % MMA_K_F8;
|
||||
@@ -269,131 +294,85 @@ indexer_fp8_score_topk_kernel(
|
||||
sK8[canon_idx_fp8_128x32(row, col)] = k_fp8[(int64_t)g_row * ihd + d];
|
||||
}
|
||||
__syncthreads();
|
||||
// Generic-proxy SMEM writes above must be visible to the tcgen05 async proxy.
|
||||
asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
|
||||
__syncthreads();
|
||||
|
||||
// MMA
|
||||
if (is_mma_warp && lane == 0) {
|
||||
uint64_t dq = make_umma_desc_kmajor_none((uint32_t)__cvta_generic_to_shared(sQ8), 128);
|
||||
uint64_t dk = make_umma_desc_kmajor_none((uint32_t)__cvta_generic_to_shared(sK8), 128);
|
||||
umma_ss_f8f6f4(tb, dq, dk, idesc_f8, kt > 0);
|
||||
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
asm volatile("fence.sc.gpu;" ::: "memory");
|
||||
|
||||
// Track completion of all prior tcgen05.mma operations before TMEM reads.
|
||||
if (is_mma_warp && lane == 0) tcgen05_commit_mma(mbar_addr);
|
||||
mbarrier_wait_cta(mbar_addr, mma_phase);
|
||||
mma_phase ^= 1;
|
||||
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
|
||||
__syncthreads();
|
||||
|
||||
// ---- Read TMEM results ----
|
||||
// The MMA produces [128 rows × 128 cols] in TMEM.
|
||||
// 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.
|
||||
// ---- Read TMEM and reduce across indexer heads ----
|
||||
// warps 0/1 read the same taddr; hardware maps them to lanes 0..31 / 32..63.
|
||||
if (wid < 2) {
|
||||
const int h = wid * 32 + lane;
|
||||
const bool h_valid = h < n_ih;
|
||||
const float q_s = h_valid ? sQ_scale[h] : 0.0f;
|
||||
const float wh = h_valid ? sW_h[h] : 0.0f;
|
||||
|
||||
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 chunk_start = my_warp * CHUNKS_PER_WARP;
|
||||
int chunk_end = chunk_start + CHUNKS_PER_WARP;
|
||||
#pragma unroll
|
||||
for (int n = 0; n < SK_TILE / 8; n++) {
|
||||
int col_base = n * 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 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 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];
|
||||
}
|
||||
|
||||
// Process each column
|
||||
for (int j = 0; j < cols_valid; j++) {
|
||||
float contrib[8];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; j++) {
|
||||
int c = col_base + j;
|
||||
float k_s = k_scale[kv_start + c];
|
||||
|
||||
// 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;
|
||||
if (h_valid && c < kv_len) {
|
||||
float logit = tmp[j] * q_s * k_scale[kv_start + c];
|
||||
contrib[j] = wh * fmaxf(logit, 0.0f);
|
||||
} else {
|
||||
contrib[j] = 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)
|
||||
contrib += __shfl_down_sync(0xffffffff, contrib, o);
|
||||
|
||||
// Lane 0 stores partial sum for this column
|
||||
if (lane == 0) sLogits[c] = contrib;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; j++) {
|
||||
float v = contrib[j];
|
||||
for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffff, v, o);
|
||||
if (lane == 0 && (col_base + j) < kv_len) {
|
||||
sWarpScores[wid * SK_TILE + col_base + j] = v;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// 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);
|
||||
}
|
||||
}
|
||||
// ---- Merge per-column scores into per-thread local top-k heaps ----
|
||||
for (int c = tid; c < kv_len; c += NTHREADS) {
|
||||
float score = sWarpScores[c] + sWarpScores[SK_TILE + c];
|
||||
local_heap_insert(local_scores, local_blocks, score, kv_start + c, INDEXER_LOCAL_K);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// ---- TMEM dealloc ----
|
||||
if (is_mma_warp) tmem_dealloc(tb, TMEM_COLS);
|
||||
__syncthreads();
|
||||
|
||||
// ---- Phase 2: Block-level top-k merge ----
|
||||
// Each thread writes its INDEXER_LOCAL_K candidates to SMEM, then
|
||||
// one thread builds the final top-k.
|
||||
|
||||
// ---- Block-level top-k merge ----
|
||||
for (int i = tid; i < top_k; i += NTHREADS) {
|
||||
sMergeScores[i] = -INFINITY;
|
||||
sMergeBlocks[i] = -1;
|
||||
}
|
||||
int my_offset = tid * INDEXER_LOCAL_K;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < INDEXER_LOCAL_K; i++) {
|
||||
sCandScores[my_offset + i] = local_scores[i];
|
||||
sCandBlocks[my_offset + i] = local_blocks[i];
|
||||
@@ -402,16 +381,13 @@ indexer_fp8_score_topk_kernel(
|
||||
|
||||
if (tid == 0) {
|
||||
for (int i = 0; i < NTHREADS * INDEXER_LOCAL_K; i++) {
|
||||
if (sCandScores[i] > -INFINITY) {
|
||||
if (sCandBlocks[i] >= 0) {
|
||||
heap_insert_shared(sMergeScores, sMergeBlocks,
|
||||
sCandScores[i], sCandBlocks[i], top_k);
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ---- Write top-k indices sorted by score ----
|
||||
if (tid == 0) {
|
||||
// Sort descending for deterministic torch.topk-like output order.
|
||||
for (int i = 0; i < top_k; i++) {
|
||||
int best = i;
|
||||
for (int j = i + 1; j < top_k; j++) {
|
||||
@@ -431,6 +407,8 @@ indexer_fp8_score_topk_kernel(
|
||||
// PyTorch binding
|
||||
// ===========================================================================
|
||||
|
||||
static size_t align_up(size_t x, size_t a) { return (x + a - 1) & ~(a - 1); }
|
||||
|
||||
void indexer_fp8_score_topk_cuda(
|
||||
torch::Tensor q_bf16, // (n_ih, ihd) BF16
|
||||
torch::Tensor k_fp8, // (n_comp, ihd) uint8/float8_e4m3fn
|
||||
@@ -443,23 +421,34 @@ void indexer_fp8_score_topk_cuda(
|
||||
TORCH_CHECK(k_fp8.is_cuda());
|
||||
TORCH_CHECK(k_scale.is_cuda() && k_scale.scalar_type() == torch::kFloat32);
|
||||
TORCH_CHECK(w_h.is_cuda() && w_h.scalar_type() == torch::kBFloat16);
|
||||
TORCH_CHECK(topk_indices.is_cuda() && topk_indices.scalar_type() == torch::kInt32);
|
||||
TORCH_CHECK(n_ih == 64 && ihd == 128, "B2 first pass is specialized to n_ih=64, ihd=128");
|
||||
TORCH_CHECK(top_k > 0, "top_k must be positive");
|
||||
|
||||
int n_comp = k_fp8.size(0);
|
||||
TORCH_CHECK(n_comp > 0, "n_comp must be positive");
|
||||
TORCH_CHECK(k_fp8.size(1) == ihd, "k_fp8 must have shape (n_comp, ihd)");
|
||||
TORCH_CHECK(k_scale.numel() >= n_comp, "k_scale must contain at least n_comp scales");
|
||||
TORCH_CHECK(topk_indices.numel() >= top_k, "topk_indices is smaller than top_k");
|
||||
|
||||
// Convert k_fp8 to uint8 view if needed
|
||||
auto k8 = k_fp8.dtype() == torch::kUInt8 ? k_fp8 : k_fp8.view(torch::kUInt8);
|
||||
|
||||
// SMEM size calculation
|
||||
// Must exactly mirror kernel SMEM layout. The previous B2 missed the score
|
||||
// scratch allocation, which can corrupt following SMEM and manifest as a hang.
|
||||
size_t smem = 0;
|
||||
smem += 4; smem = (smem + 127) & ~127; // sTmemBase
|
||||
smem += 128 * 32; smem = (smem + 127) & ~127; // sQ8
|
||||
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 += top_k * 4; // sMergeScores
|
||||
smem += top_k * 4; // sMergeBlocks
|
||||
smem += 192 * INDEXER_LOCAL_K * 4; // sCandScores
|
||||
smem += 192 * INDEXER_LOCAL_K * 4; // sCandBlocks
|
||||
smem += 4; // sTmemBase
|
||||
smem = align_up(smem, 16);
|
||||
smem += 8; // sMbar
|
||||
smem = align_up(smem, 128);
|
||||
smem += 128 * 32; smem = align_up(smem, 128); // sQ8
|
||||
smem += 128 * 32; smem = align_up(smem, 128); // sK8
|
||||
smem += 128 * 4; smem = align_up(smem, 128); // sQ_scale
|
||||
smem += 128 * 4; smem = align_up(smem, 128); // sW_h
|
||||
smem += 2 * 128 * 4; smem = align_up(smem, 128); // sWarpScores
|
||||
smem += (size_t)top_k * 4; // sMergeScores
|
||||
smem += (size_t)top_k * 4; // sMergeBlocks
|
||||
smem += 192 * INDEXER_LOCAL_K * 4; // sCandScores
|
||||
smem += 192 * INDEXER_LOCAL_K * 4; // sCandBlocks
|
||||
|
||||
cudaFuncSetAttribute(indexer_fp8_score_topk_kernel<128>,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||||
|
||||
Reference in New Issue
Block a user