nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[kernel] sm_100a backward: keep each softmax warp's bf16 overlay inside its own fp32 columns

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
root
2026-09-05 11:29:04 -07:00
co-authored by Claude Fable 5.1
parent 749a9c5d97
commit 88319b75d5
1 changed files with 11 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh
+11 -8
View File
@@ -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<const float2*>(sLSE + qblock_in_quad * BLOCK);
const float2* delta2 = reinterpret_cast<const float2*>(sDelta + qblock_in_quad * BLOCK);