diff --git a/dsv4/kernels/attention/fmha_6warp_multirow.cuh b/dsv4/kernels/attention/fmha_6warp_multirow.cuh index 6a6ff54f..4c95504c 100644 --- a/dsv4/kernels/attention/fmha_6warp_multirow.cuh +++ b/dsv4/kernels/attention/fmha_6warp_multirow.cuh @@ -9,27 +9,27 @@ * Grid: (1, n_h, batch) — each CTA processes one head of one batch item. * * TMEM operations ALL use 32x32b.x8 format exclusively. - * The 16x256b.x1 format crashes on the 2nd call — verified on B200. + * 16x256b.x1 stores crash on 2nd call (verified on B200). * - * TMEM ADDRESSING for 32x32b.x8: - * addr = tmem_base + (row_page << 16) + col_group * 8 - * - col_group selects 8 consecutive TMEM columns - * - row_page selects a 32-row "page" (row_page=0 → rows 0-31, =1 → rows 32-63, etc.) - * - Lane l within the warp reads row (row_page*32 + l) across the 8 columns - * - ALL 32 lanes must call with the same address (warp-collective) + * KEY DISCOVERY (verified on B200): + * After UMMA (QK GEMM), 4 warps reading TMEM with 32x32b.x8 each + * see a DIFFERENT 32-row partition of the S matrix: + * Warp 0 → rows 0-31, Warp 1 → rows 32-63, + * Warp 2 → rows 64-95, Warp 3 → rows 96-127. + * Lane l in warp w reads row w*32 + l across 8 columns. + * This is because the MMA writes Layout D to TMEM, and each warp's + * 32x32b.x8 read naturally accesses its partition. * * SOFTMAX (two regimes): - * - T <= 32: warp 0 only. row_page=0. Lane l processes row l. 32x32b.x8 reads. - * - T > 32: 4 softmax warps. Warp w uses row_page=w. - * Lane l in warp w processes row w*32+l. Two-pass: - * (1) row_max, (2) exp+sum, stream P to sPk per K-tile. + * - T <= 32: only warp 0 active. Lane l handles row l. + * - T > 32: 4 softmax warps. Warp w handles rows [w*32, min((w+1)*32, T)). + * Lane l in warp w handles row w*32 + l. + * Two-pass: (1) row_max over all S columns, (2) exp/sum, stream P to sPk. * * P STAGING FOR PV (no s_p_vals buffer): - * P is written one K-tile at a time to sPk (128, 16) SMEM in canonical - * layout. Softmax warps compute P values in registers, write to sPk. - * Then MMA warp does PV with sPk × sV. - * The softmax and PV are interleaved K-tile-by-K-tile (required because - * we can't hold the full 128×128 P matrix in SMEM). + * P written one K-tile at a time to sPk (128, 16) SMEM in canonical layout. + * Softmax warps compute P in registers, write to sPk. Then MMA does PV. + * Softmax and PV interleaved per K-tile (can't hold full 128×128 P in SMEM). * * PV GEMM: * N=16 sub-tiles (avoids Layout D N=64 bug). SS MMA: sPk(128,16) × sV(16,16). @@ -37,20 +37,8 @@ * * EPILOGUE: * 32x32b.x8 TMEM read → normalize (1/row_sum) → BF16 → GMEM. + * Each softmax warp reads its own 32-row partition of O from TMEM. * LSE = ln(row_sum) + row_max for D5 KV merge. - * - * ================================================================== - * SMEM BUDGET (HD=64) - * ================================================================== - * sTmemBase: 8 B - * sRowMax: 128 × 4 = 512 B - * sRowSum: 128 × 4 = 512 B - * padding: 128 B (128B alignment for UMMA) - * sQ0: 128 × 16 × 2 = 4096 B (one K-tile) - * sK0: 128 × 16 × 2 = 4096 B (one K-tile) - * sPk: 128 × 16 × 2 = 4096 B (one K-tile P buffer) - * sV: 16 × 16 × 2 = 512 B (one N-sub-tile, one K-tile) - * Total: ~14 KB — well within 232 KB budget * ================================================================== */ @@ -80,14 +68,14 @@ struct FmhaMultiRowParams { template __global__ void __launch_bounds__(192) fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { - static constexpr int NKT_QK = HD / MMA_K_BF16; // K-tiles for QK GEMM - static constexpr int NKT_PV = SK_TILE / MMA_K_BF16; // K-tiles for PV GEMM (= 8) - static constexpr int N_NSUB = HD / 16; // N-sub-tiles for PV (avoids N=64 bug) - static constexpr int TILE_SZ = 128 * MMA_K_BF16; // 2048 BF16 elements - static constexpr int V_SUB_SZ = 16 * MMA_K_BF16; // 256 BF16 elements (16, 16) canonical + static constexpr int NKT_QK = HD / MMA_K_BF16; + static constexpr int NKT_PV = SK_TILE / MMA_K_BF16; + static constexpr int N_NSUB = HD / 16; + static constexpr int TILE_SZ = 128 * MMA_K_BF16; + static constexpr int V_SUB_SZ = 16 * MMA_K_BF16; static constexpr int TMEM_N = (HD <= 128) ? 128 : 256; static constexpr int MAX_ROWS = 128; - static constexpr int CORES_MN = 128 / 8; // 16 core matrices along M + static constexpr int CORES_MN = 128 / 8; const int head_idx = blockIdx.y; const int batch_idx = blockIdx.z; @@ -101,7 +89,6 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { const int s_k = params.s_k; const float scale = params.scale; - // Per-head GMEM pointers const bf16_t* __restrict__ q_head = params.q + head_idx * params.q_head_stride + batch_idx * params.q_batch_stride; const bf16_t* __restrict__ k_head = params.k + head_idx * params.k_head_stride + batch_idx * params.k_batch_stride; const bf16_t* __restrict__ v_head = params.v + head_idx * params.v_head_stride + batch_idx * params.v_batch_stride; @@ -112,18 +99,15 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { // SMEM allocation — NO s_p_vals buffer // ================================================================ extern __shared__ char sbuf[]; - uint32_t* sTmemBase = (uint32_t*)sbuf; // 4 B (8B-aligned) - float* sRowMax = (float*)(sbuf + 8); // [MAX_ROWS] = 512 B - float* sRowSum = sRowMax + MAX_ROWS; // [MAX_ROWS] = 512 B - // Align sQ0 to 128B for UMMA descriptors + uint32_t* sTmemBase = (uint32_t*)sbuf; + float* sRowMax = (float*)(sbuf + 8); + float* sRowSum = sRowMax + MAX_ROWS; bf16_t* sQ0 = (bf16_t*)(((uintptr_t)(sRowSum + MAX_ROWS) + 127) & ~(uintptr_t)127); - bf16_t* sK0 = sQ0 + TILE_SZ; // (128, 16) BF16 canonical - bf16_t* sPk = (bf16_t*)(((uintptr_t)(sK0 + TILE_SZ) + 127) & ~(uintptr_t)127); // (128, 16) BF16 canonical - bf16_t* sV = (bf16_t*)(((uintptr_t)(sPk + TILE_SZ) + 127) & ~(uintptr_t)127); // (16, 16) BF16 canonical + bf16_t* sK0 = sQ0 + TILE_SZ; + bf16_t* sPk = (bf16_t*)(((uintptr_t)(sK0 + TILE_SZ) + 127) & ~(uintptr_t)127); + bf16_t* sV = (bf16_t*)(((uintptr_t)(sPk + TILE_SZ) + 127) & ~(uintptr_t)127); - // ================================================================ // TMEM allocation (warp 4) - // ================================================================ if (is_mma_warp) { uint32_t smem_ptr = __cvta_generic_to_shared(sTmemBase); tmem_alloc(smem_ptr, TMEM_N); @@ -134,21 +118,17 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { // ================================================================ // Row assignment for softmax // ================================================================ - // T <= 32: only warp 0, row_page = 0, lane l = row l - // T > 32: 4 softmax warps, warp w uses row_page = w - // lane l in warp w handles row w*32 + l + // After UMMA, warp w's 32x32b.x8 read gives rows [w*32, w*32+31]. + // Lane l in warp w reads row w*32 + l. const bool my_warp_active = (T <= 32) ? (wid == 0) : is_softmax_warp; - const int my_row_page = (T <= 32) ? 0 : wid; - const int my_row = my_row_page * 32 + lane; + const int my_row = my_warp_active ? (wid * 32 + lane) : 0; const bool my_row_active = my_warp_active && (my_row < T); - const uint32_t my_tmem_row_offset = (uint32_t)(my_row_page * 32) << 16; // ================================================================ // QK GEMM: accumulate all K-tiles into TMEM // ================================================================ for (int kt = 0; kt < NKT_QK; kt++) { if (is_load_warp) { - // Load Q K-tile: Q is (T, HD), columns [kt*16, kt*16+16) for (int i = lane; i < TILE_SZ; i += 32) sQ0[i] = 0; for (int r = 0; r < T; r++) { for (int d = lane; d < MMA_K_BF16; d += 32) { @@ -160,7 +140,6 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { } } } - // Load K K-tile: K is (s_k, HD), columns [kt*16, kt*16+16) for (int i = lane; i < TILE_SZ; i += 32) sK0[i] = 0; for (int r = 0; r < s_k; r++) { for (int d = lane; d < MMA_K_BF16; d += 32) { @@ -185,26 +164,19 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { __syncthreads(); } - // TMEM now contains S[128, s_k] — the raw QK dot products. - // S[row, col] needs scaling by `scale` before softmax. - // ================================================================ // SOFTMAX PASS 1: compute row_max for each row // ================================================================ - // Read all SK_TILE columns (16 groups of 8 columns each). - // 32x32b.x8: addr = tb + (row_page*32 << 16) + n*8 - // Each lane reads its assigned row across 8 columns per read. float my_row_max = -INFINITY; - const int NUM_READS = SK_TILE / 8; // = 16 + const int NUM_READS = SK_TILE / 8; if (my_warp_active) { for (int n = 0; n < NUM_READS; n++) { float tmp[8]; - uint32_t addr = tb + my_tmem_row_offset + n * 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"(addr)); + : "r"(tb + n * 8)); asm volatile("tcgen05.wait::ld.sync.aligned;"); if (my_row_active) { @@ -219,47 +191,32 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { } } - // Store row_max to SMEM if (my_row_active) sRowMax[my_row] = my_row_max; __syncthreads(); // ================================================================ // SOFTMAX PASS 2 + PV GEMM (interleaved per PV K-tile) // ================================================================ - // For each PV K-tile (16 columns of S): - // 1. Softmax warps re-read TMEM, compute exp(S*scale - row_max) - // 2. Write P to sPk in canonical layout - // 3. Load V sub-tile to sV - // 4. MMA does PV: sPk(128,16) × sV(16,16) → TMEM - // - // We accumulate row_sum across K-tiles in registers. - // P written to sPk is UN-normalized (exp(S*scale - row_max)). - // Normalization happens in the epilogue (divide by row_sum). - // ================================================================ float my_row_sum = 0.0f; for (int pv_kt = 0; pv_kt < NKT_PV; pv_kt++) { const int col_start = pv_kt * MMA_K_BF16; const int col_end = col_start + MMA_K_BF16; - // TMEM column groups for this K-tile: [col_start/8, col_end/8) const int rg_start = col_start / 8; const int rg_end = col_end / 8; - // P values for this K-tile (un-normalized) - float p_vals[16]; // MMA_K_BF16 values + float p_vals[16]; for (int c = 0; c < MMA_K_BF16; c++) p_vals[c] = 0.0f; if (my_warp_active) { float rm = my_row_active ? sRowMax[my_row] : 0.0f; - // Read TMEM for this K-tile's columns for (int n = rg_start; n < rg_end; n++) { float tmp[8]; - uint32_t addr = tb + my_tmem_row_offset + n * 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"(addr)); + : "r"(tb + n * 8)); asm volatile("tcgen05.wait::ld.sync.aligned;"); if (my_row_active) { @@ -275,13 +232,13 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { } } - // Zero sPk first (load warp does this for all rows) + // Zero sPk if (is_load_warp) { for (int i = lane; i < TILE_SZ; i += 32) sPk[i] = 0; } __syncthreads(); - // Write P to sPk in canonical layout (softmax warps) + // Write P to sPk in canonical layout if (my_row_active) { for (int c = 0; c < MMA_K_BF16; c++) { int ck = c / 8, lc = c % 8; @@ -296,10 +253,6 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { int d_base = n_sub * 16; if (is_load_warp) { - // Load V sub-tile: V is (HD, s_k) in GMEM - // Need V[d_base..d_base+15, col_start..col_end-1] - // Canonical layout for (16, 16): - // g_mn = d/8, g_k = lr/8, llr = d%8, lc = lr%8 for (int i = lane; i < V_SUB_SZ; i += 32) sV[i] = 0; for (int dd = lane; dd < 16; dd += 32) { for (int lr = 0; lr < MMA_K_BF16; lr++) { @@ -318,9 +271,6 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { uint32_t idesc_pv = make_idesc(128, 16); uint64_t dp = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sPk), 128); uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sV), 16); - // accumulate = true for pv_kt > 0 (K-tile accumulation) - // For n_sub > 0, the PV result goes to different TMEM columns (n_sub*16) - // so first call for each n_sub should NOT accumulate if (tid == 128) umma_ss_f16(tb + n_sub * 16, dp, dv, idesc_pv, pv_kt > 0); asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); } @@ -328,33 +278,27 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { } } - // Store row_sum to SMEM if (my_row_active) sRowSum[my_row] = my_row_sum; __syncthreads(); // ================================================================ // EPILOGUE: TMEM → regs → normalize → BF16 → GMEM // ================================================================ + // After PV GEMM, O[128, HD] is in TMEM. Each softmax warp reads + // its own 32-row partition (same as after QK GEMM). + // HD/8 = N_NSUB*2 read groups. if (my_warp_active) { if (my_row_active) { float rm = sRowMax[my_row]; float rs = sRowSum[my_row]; float inv_rs = 1.0f / rs; - // Read O from TMEM: HD values in HD/8 groups of 8 - // O is in TMEM at columns 0..N_NSUB*16-1 - // After PV GEMM, TMEM layout is: - // N-sub-tile n writes to columns n*16..n*16+15 - // 32x32b.x8 reads 8 columns at a time - // So we read (N_NSUB*16)/8 = N_NSUB*2 groups - // But we need to read ALL N_NSUB*2 groups with correct row for (int n = 0; n < N_NSUB * 2; n++) { float tmp[8]; - uint32_t addr = tb + my_tmem_row_offset + n * 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"(addr)); + : "r"(tb + n * 8)); asm volatile("tcgen05.wait::ld.sync.aligned;"); for (int c = 0; c < 8; c++) { @@ -365,7 +309,6 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { } } - // LSE if (lse_head) { lse_head[my_row] = logf(rs) + rm; } @@ -373,10 +316,7 @@ fmha_6warp_multirow_kernel(FmhaMultiRowParams params) { } __syncthreads(); - // TMEM dealloc (warp 4) - if (is_mma_warp) { - tmem_dealloc(tb, TMEM_N); - } + if (is_mma_warp) tmem_dealloc(tb, TMEM_N); } } // namespace dsv4::kernels::attention