From 88319b75d559e08cbe1cf10ede8746f6377f080d Mon Sep 17 00:00:00 2001 From: root Date: Sat, 5 Sep 2026 11:29:04 -0700 Subject: [PATCH] [kernel] sm_100a backward: keep each softmax warp's bf16 overlay inside its own fp32 columns Review finding: the two softmax warps of a TMEM lane group split the 128-column S^T (and dP^T) tile into fp32 halves, but the col_half == 1 warp packed its bf16x2 P^T / dS^T into columns 32-63, inside the fp32 range the col_half == 0 warp was still loading. Nothing but the shared full_bar_st / full_bar_dpt wait ordered the two, so a delayed col_half == 0 warp could read bf16 pairs where it expected fp32 scores. Fix, no extra synchronization: a warp's bf16x2 overlay now starts at the same column as the 64 fp32 columns it loads (ST_QBLOCK_COLS = 64 per q64 block), so it only overwrites data it has itself consumed; the dV / dK TS MMAs read slot p's bf16 atoms from column p * ST_QBLOCK_COLS instead of 32 * p. The dQ GEMM already overwrote the dP^T tile only after the dK GEMM consumed dS^T (MMA issue order), unchanged. Verified: CPU-reference checks incl. odd top-k, ragged variable_block_sizes and zero-count kv blocks; tests/test_block_sparse_bwd_sm100a.py 15 passed, tests/test_block_sparse_sm100a_dispatch.py 10 passed, tests/test_block_sparse_sm100a.py 37 passed (extension rebuilt from this tree). Interleaved A/B vs the previous build: neutral (4k..131k within noise). Co-Authored-By: Claude Fable 5.1 --- .../block_sparse_bwd_kernel_sm100a.cuh | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh b/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh index bdac3bdd5..ce3f673fa 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh @@ -68,10 +68,12 @@ constexpr int SMEM_TOTAL = 2 * KV_TILE_BYTES + NUM_Q_STAGES * Q_RING_SLOT_BYTES DQConfig<>::DQ_STAGE_BUFFERS * DQConfig<>::DQ_STAGE_BYTES + 2 * Q_QUAD * (int)sizeof(float) + NUM_BARS * 8 + CLC_STAGES * 16 + 48; -constexpr int ST_COLS = Q_QUAD / 2; -constexpr int DV_COLS = HEAD_DIM; -constexpr int DK_COLS = HEAD_DIM; -constexpr int TMEM_TOTAL = ST_COLS + DV_COLS + ST_COLS + DK_COLS; +constexpr int ST_COLS = Q_QUAD / 2; +constexpr int ST_QBLOCK_COLS = BLOCK; +constexpr int DV_COLS = HEAD_DIM; +constexpr int DK_COLS = HEAD_DIM; +constexpr int TMEM_TOTAL = ST_COLS + DV_COLS + ST_COLS + DK_COLS; +static_assert(ST_COLS == 2 * ST_QBLOCK_COLS, "two q blocks per lane-half"); static_assert(TMEM_TOTAL == 512, "TMEM map must fill exactly 512 columns"); extern __shared__ __align__(1024) uint8_t bwd_smem[]; @@ -395,8 +397,9 @@ __global__ void __cluster_dims__(1, 1, 1) __launch_bounds__(N_WARPS * 32, 1) vsa const uint64_t db = desc_ring + (uint64_t)slot * RING_DELTA; #pragma unroll for (int ki = 0; ki < K_ATOMS_PER_QBLOCK; ++ki) { - const int a = p * K_ATOMS_PER_QBLOCK + ki; - const uint32_t tmem_a = tmem_a_base + (uint32_t)(a * BF16X2_COLS_PER_K16); + const int a = p * K_ATOMS_PER_QBLOCK + ki; + const uint32_t tmem_a = + tmem_a_base + (uint32_t)(p * ST_QBLOCK_COLS + ki * BF16X2_COLS_PER_K16); const bool accumulate = (!first) || (a != 0); tcgen05_mma_ws_f16_ts_1sm_predicated(lead, tmem_acc, tmem_a, db + ki * K16_COLS_DELTA, idesc_dv_dk, accumulate); @@ -529,8 +532,8 @@ __global__ void __cluster_dims__(1, 1, 1) __launch_bounds__(N_WARPS * 32, 1) vsa const int qblock_in_quad = 2 * col_half + q_half; const uint32_t tmem_lane_base = (uint32_t)(lane_group * 32) << 16; - const uint32_t tmem_f32_offset = tmem_lane_base + (uint32_t)(col_half * HALF_COLS); - const uint32_t tmem_bf16x2_offset = tmem_lane_base + (uint32_t)(col_half * HALF_COLS / 2); + const uint32_t tmem_f32_offset = tmem_lane_base + (uint32_t)(col_half * ST_QBLOCK_COLS); + const uint32_t tmem_bf16x2_offset = tmem_f32_offset; const float2* lse2 = reinterpret_cast(sLSE + qblock_in_quad * BLOCK); const float2* delta2 = reinterpret_cast(sDelta + qblock_in_quad * BLOCK);