diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_attention_dsv4_bwd.py new file mode 100644 index 0000000000..7284a98cfc --- /dev/null +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_attention_dsv4_bwd.py @@ -0,0 +1,533 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Gluon kernels for the DeepSeek-V4 sparse-MLA training BACKWARD (gfx950 / CDNA4). + +The two MFMA phases. Both operate on the official V4 form: ``K == V == kv`` as one dense +512-wide tensor, RoPE already applied in place caller-side, scale ``1/sqrt(512)``, ``attn_sink`` +in the softmax denominator only, ``topk == -1`` masked. + +``_dq_v4_kernel`` + Per (query token, head block): ``S = Q@kv^T``, ``P = exp(S - lse)``, ``dP = dO@kv^T``, + ``dS = P*(dP - delta)*scale``, ``dQ += dS@kv``. Three MFMAs per tile; the gathered KV tile + is read from LDS once and feeds both the ``S`` and ``dP`` MFMAs. Also emits the ``dS`` / ``P`` + chunks the dKV-interm kernel consumes. + +``_dkv_interm_v4_kernel`` + ``interm[t, slot, d] = sum_h ( dS[t,h,slot]*Q[t,h,d] + P[t,h,slot]*dO[t,h,d] )``, contracting + over ALL heads inside one MFMA pair so nothing accumulates across a loop over heads. Q and dO + are transposed once into registers and D is split across ``grid.y``, which is what keeps them + read once instead of ``topk/TILE_K`` times. + +Launchers live in ``aiter.ops.triton.attention.sparse_attention_dsv4_bwd``; this module stays +free of torch so the kernels can be called without it. +""" + +import triton.language as tl +from triton.experimental import gluon +from triton.experimental.gluon import language as gl + +from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr + +_dq_v4_kernel_repr = make_kernel_repr( + "_dq_v4_kernel", + [ + "R_CHUNK", + "BLOCK_H", + "TILE_K", + "D", + "IS_FIRST_CHUNK", + ], +) + + +@gluon.jit(repr=_dq_v4_kernel_repr) +def _dq_v4_kernel( + Q_ptr, # [T, H, D] bf16 + KV_ptr, # [T, D] bf16 (K == V) + dO_ptr, # [T, H, D] bf16 + TopK_ptr, # [T, TOPK_padded] int32 + LSE_ptr, # [T, H] fp32 (sink-inclusive) + Delta_ptr, # [T, H] fp32 + dQ_ptr, # [T, H, D] bf16 (RMW across chunks) + dS_ptr, # [T, H, R_CHUNK] bf16 + P_ptr, # [T, H, R_CHUNK] bf16 + stride_q_t: tl.int64, + stride_q_h: tl.int64, + stride_kv_t: tl.int64, + stride_do_t: tl.int64, + stride_do_h: tl.int64, + stride_dq_t: tl.int64, + stride_dq_h: tl.int64, + stride_topk_t: tl.int64, + stride_ds_t: tl.int64, + stride_ds_h: tl.int64, + scale: tl.float32, + num_heads: tl.int32, + R_START: tl.int32, + R_CHUNK: gl.constexpr, + BLOCK_H: gl.constexpr, + TILE_K: gl.constexpr, + D: gl.constexpr, + IS_FIRST_CHUNK: gl.constexpr, +): + gl.static_assert(TILE_K % 32 == 0, "16x16x32 needs TILE_K multiple of 32") + gl.static_assert(D % 32 == 0, "16x16x32 needs D multiple of 32") + + # ---- single 16x16x32 MFMA parent (score contracts D; accumulate contracts TILE_K) ---- + mma: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 32], + transposed=True, + warps_per_cta=[4, 1], + ) + qa: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mma, k_width=8) + qb: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mma, k_width=8) + + # ---- blocked layouts for global loads ---- + _q_tpw_k: gl.constexpr = min(64, D // 8) + _q_tpw_m: gl.constexpr = 64 // _q_tpw_k + blk_q: gl.constexpr = gl.BlockedLayout( # [BLOCK_H, D] (Q, dO, dQ) + size_per_thread=[1, 8], + threads_per_warp=[_q_tpw_m, _q_tpw_k], + warps_per_cta=[4, 1], + order=[1, 0], + ) + _kv_tpw_m: gl.constexpr = min(64, D // 8) + _kv_tpw_n: gl.constexpr = 64 // _kv_tpw_m + blk_kv: gl.constexpr = gl.BlockedLayout( # [D, TILE_K] + size_per_thread=[8, 1], + threads_per_warp=[_kv_tpw_m, _kv_tpw_n], + warps_per_cta=[1, 4], + order=[0, 1], + ) + sh_kv: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( + [[512, 16]], [D, TILE_K], [0, 1] + ) + + # ---- program ids ---- + token_idx = gl.program_id(axis=0) + hg_idx = gl.program_id(axis=1) + hg_offset = hg_idx * BLOCK_H + NUM_TILES: gl.constexpr = R_CHUNK // TILE_K + + # ---- Q / dO offsets + load (register, convert to dot operand) ---- + offs_h_q = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_q)) + offs_d_q = gl.arange(0, D, layout=gl.SliceLayout(0, blk_q)) + mask_h_q = offs_h_q < num_heads + q_base = token_idx.to(tl.int64) * stride_q_t + q_offs = ( + q_base + + offs_h_q[:, None].to(tl.int64) * stride_q_h + + offs_d_q[None, :].to(tl.int64) + ) + do_base = token_idx.to(tl.int64) * stride_do_t + do_offs = ( + do_base + + offs_h_q[:, None].to(tl.int64) * stride_do_h + + offs_d_q[None, :].to(tl.int64) + ) + + q_blk = gl.amd.cdna4.buffer_load( + ptr=Q_ptr, offsets=q_offs.to(tl.int32), mask=mask_h_q[:, None], other=0.0 + ) + do_blk = gl.amd.cdna4.buffer_load( + ptr=dO_ptr, offsets=do_offs.to(tl.int32), mask=mask_h_q[:, None], other=0.0 + ) + Q_dot = gl.convert_layout(q_blk, qa) + dO_dot = gl.convert_layout(do_blk, qa) + + # ---- topk / KV offsets ---- + topk_base = token_idx.to(tl.int64) * stride_topk_t + R_START + stride_kv_t_i32: tl.int32 = stride_kv_t.to(tl.int32) + offs_tile_kv = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_kv)) + offs_tile_mma = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mma)) + offs_d_kv = gl.arange(0, D, layout=gl.SliceLayout(1, blk_kv)) + + smem_kv = gl.allocate_shared_memory( + KV_ptr.dtype.element_ty, [2, D, TILE_K], layout=sh_kv + ) + + dQ_acc = gl.zeros([BLOCK_H, D], dtype=gl.float32, layout=mma) + + offs_h_s = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mma)) + mask_h_s = offs_h_s < num_heads + lse = gl.amd.cdna4.buffer_load( + ptr=LSE_ptr, + offsets=(token_idx * num_heads + offs_h_s).to(tl.int32), + mask=mask_h_s, + other=0.0, + ) + delta = gl.amd.cdna4.buffer_load( + ptr=Delta_ptr, + offsets=(token_idx * num_heads + offs_h_s).to(tl.int32), + mask=mask_h_s, + other=0.0, + ) + + # ---- prologue: gather kv tile 0 ---- + topk_pos_kv = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + offs_tile_kv).to(tl.int32), + mask=offs_tile_kv < R_CHUNK, + other=-1, + ) + topk_pos_mma = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + offs_tile_mma).to(tl.int32), + mask=offs_tile_mma < R_CHUNK, + other=-1, + ) + valid_kv = topk_pos_kv != -1 + valid_mma = topk_pos_mma != -1 + safe_kv = gl.where(valid_kv, topk_pos_kv, 0) + kv_offs = safe_kv[None, :] * stride_kv_t_i32 + offs_d_kv[:, None] + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_kv.index(0), ptr=KV_ptr, offsets=kv_offs, mask=valid_kv[None, :] + ) + gl.amd.cdna4.async_copy.commit_group() + + ds_base = ( + token_idx.to(tl.int64) * stride_ds_t + + hg_idx.to(tl.int64) * BLOCK_H * stride_ds_h + ) + offs_h_dsp = gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mma)) + offs_tile_dsp = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mma)) + mask_h_dsp = (hg_offset + offs_h_dsp) < num_heads + + cur_buf = 0 + for t in range(NUM_TILES - 1): + next_offs_kv = (t + 1) * TILE_K + offs_tile_kv + next_offs_mma = (t + 1) * TILE_K + offs_tile_mma + topk_pos_kv_next = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + next_offs_kv).to(tl.int32), + mask=next_offs_kv < R_CHUNK, + other=-1, + ) + topk_pos_mma_next = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + next_offs_mma).to(tl.int32), + mask=next_offs_mma < R_CHUNK, + other=-1, + ) + valid_kv_next = (next_offs_kv < R_CHUNK) & (topk_pos_kv_next != -1) + valid_mma_next = (next_offs_mma < R_CHUNK) & (topk_pos_mma_next != -1) + safe_kv_next = gl.where(valid_kv_next, topk_pos_kv_next, 0) + + next_buf = 1 - cur_buf + kv_offs_next = safe_kv_next[None, :] * stride_kv_t_i32 + offs_d_kv[:, None] + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_kv.index(next_buf), + ptr=KV_ptr, + offsets=kv_offs_next, + mask=valid_kv_next[None, :], + ) + gl.amd.cdna4.async_copy.commit_group() + + gl.amd.cdna4.async_copy.wait_group(1) + + kv_smem_cur = smem_kv.index(cur_buf) + # score K (direct); V (permuted) read LATE, before the accumulate + K_T_dot = gl.amd.cdna4.async_copy.load_shared_relaxed(kv_smem_cur, qb) + + S = gl.amd.cdna4.mfma( + Q_dot, K_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mma) + ) + S = S * scale + offs_h_mma = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mma)) + valid_mask = valid_mma[None, :] & (offs_h_mma < num_heads)[:, None] + S = gl.where(valid_mask, S, float("-inf")) + + P = gl.exp(S - lse[:, None]) + P = gl.where(valid_mask, P, 0.0) + dP = gl.amd.cdna4.mfma( + dO_dot, K_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mma) + ) + dS = P * (dP - delta[:, None]) * scale + dS = gl.where(valid_mask, dS, 0.0) + + dS_bf = dS.to(KV_ptr.dtype.element_ty) + dS_dot = gl.convert_layout(dS_bf, qa) + K_v_dot = gl.amd.cdna4.async_copy.load_shared_relaxed( + kv_smem_cur.permute([1, 0]), qb + ) # load V LATE + dQ_acc = gl.amd.cdna4.mfma(dS_dot, K_v_dot, dQ_acc) + + col = t * TILE_K + offs_tile_dsp + dsp_offs = ( + ds_base + + offs_h_dsp[:, None].to(tl.int64) * stride_ds_h + + col[None, :].to(tl.int64) + ) + gl.amd.cdna4.buffer_store( + stored_value=dS_bf, + ptr=dS_ptr, + offsets=dsp_offs.to(tl.int32), + mask=mask_h_dsp[:, None], + ) + gl.amd.cdna4.buffer_store( + stored_value=P.to(KV_ptr.dtype.element_ty), + ptr=P_ptr, + offsets=dsp_offs.to(tl.int32), + mask=mask_h_dsp[:, None], + ) + + cur_buf = next_buf + valid_mma = valid_mma_next + + # ---- epilogue: last tile ---- + gl.amd.cdna4.async_copy.wait_group(0) + t = NUM_TILES - 1 + kv_smem_cur = smem_kv.index(cur_buf) + K_T_dot = gl.amd.cdna4.async_copy.load_shared_relaxed(kv_smem_cur, qb) + + S = gl.amd.cdna4.mfma( + Q_dot, K_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mma) + ) + S = S * scale + offs_h_mma = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mma)) + valid_mask = valid_mma[None, :] & (offs_h_mma < num_heads)[:, None] + S = gl.where(valid_mask, S, float("-inf")) + + P = gl.exp(S - lse[:, None]) + P = gl.where(valid_mask, P, 0.0) + dP = gl.amd.cdna4.mfma( + dO_dot, K_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mma) + ) + dS = P * (dP - delta[:, None]) * scale + dS = gl.where(valid_mask, dS, 0.0) + + dS_bf = dS.to(KV_ptr.dtype.element_ty) + dS_dot = gl.convert_layout(dS_bf, qa) + K_v_dot = gl.amd.cdna4.async_copy.load_shared_relaxed( + kv_smem_cur.permute([1, 0]), qb + ) + dQ_acc = gl.amd.cdna4.mfma(dS_dot, K_v_dot, dQ_acc) + + col = t * TILE_K + offs_tile_dsp + dsp_offs = ( + ds_base + + offs_h_dsp[:, None].to(tl.int64) * stride_ds_h + + col[None, :].to(tl.int64) + ) + gl.amd.cdna4.buffer_store( + stored_value=dS_bf, + ptr=dS_ptr, + offsets=dsp_offs.to(tl.int32), + mask=mask_h_dsp[:, None], + ) + gl.amd.cdna4.buffer_store( + stored_value=P.to(KV_ptr.dtype.element_ty), + ptr=P_ptr, + offsets=dsp_offs.to(tl.int32), + mask=mask_h_dsp[:, None], + ) + + # ---- store dQ (RMW across chunks) ---- + dq_base = token_idx.to(tl.int64) * stride_dq_t + offs_h_o = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_q)) + offs_d_o = gl.arange(0, D, layout=gl.SliceLayout(0, blk_q)) + mask_h_o = offs_h_o < num_heads + dq_offs = ( + dq_base + + offs_h_o[:, None].to(tl.int64) * stride_dq_h + + offs_d_o[None, :].to(tl.int64) + ) + dq_blk = gl.convert_layout(dQ_acc.to(dQ_ptr.dtype.element_ty), blk_q) + if not IS_FIRST_CHUNK: + prev = gl.amd.cdna4.buffer_load( + ptr=dQ_ptr, offsets=dq_offs.to(tl.int32), mask=mask_h_o[:, None], other=0.0 + ) + dq_blk = (dq_blk.to(gl.float32) + prev.to(gl.float32)).to( + dQ_ptr.dtype.element_ty + ) + gl.amd.cdna4.buffer_store( + stored_value=dq_blk, + ptr=dQ_ptr, + offsets=dq_offs.to(tl.int32), + mask=mask_h_o[:, None], + ) + + +_dkv_interm_v4_kernel_repr = make_kernel_repr( + "_dkv_interm_v4_kernel", + [ + "R_CHUNK", + "TILE_K", + "NH", + "BD", + "D", + "MFMA_K", + "DUAL_STAGE", + ], +) + + +@gluon.jit(repr=_dkv_interm_v4_kernel_repr) +def _dkv_interm_v4_kernel( + Q_ptr, # [T, H, D] bf16 + dO_ptr, # [T, H, D] bf16 + dS_ptr, # [T, H, R_CHUNK] bf16 + P_ptr, # [T, H, R_CHUNK] bf16 + Interm_ptr, # [T, R_CHUNK, D] bf16 + stride_q_t: tl.int64, + stride_q_h: tl.int64, + stride_do_t: tl.int64, + stride_do_h: tl.int64, + stride_ds_t: tl.int64, + stride_ds_h: tl.int64, + stride_interm_t: tl.int64, + stride_interm_r: tl.int64, + num_heads: tl.int32, + R_CHUNK: gl.constexpr, + TILE_K: gl.constexpr, + NH: gl.constexpr, + BD: gl.constexpr, + D: gl.constexpr, + MFMA_K: gl.constexpr, + DUAL_STAGE: gl.constexpr, +): + """Grid (T, D//BD). NH is the padded head count and the mfma contraction dim.""" + # instr_shape[2]=32: on gfx950 v_mfma_f32_16x16x32_bf16 does 2x the FLOPs of the 16-deep + # form in the same 16 cycles. The old kernel used 16 and it did not matter there because it + # was bandwidth-saturated at 7.3 TB/s; once the traffic is halved the matrix rate binds. + mfma: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, MFMA_K], + transposed=True, + warps_per_cta=[4, 1], + ) + _q_tpw_k: gl.constexpr = min(64, BD // 8) + _q_tpw_m: gl.constexpr = 64 // _q_tpw_k + blk_q: gl.constexpr = gl.BlockedLayout( # [H, BD] global load + size_per_thread=[1, 8], + threads_per_warp=[_q_tpw_m, _q_tpw_k], + warps_per_cta=[4, 1], + order=[1, 0], + ) + blk_ds: gl.constexpr = gl.BlockedLayout( # [H, TILE_K] global load + size_per_thread=[1, 4], + threads_per_warp=[16, 4], + warps_per_cta=[4, 1], + order=[1, 0], + ) + sh_q: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( + [[512, 16]], [NH, BD], [1, 0] + ) + + dot_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8) + dot_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma, k_width=8) + + token_idx = gl.program_id(axis=0) + dblk = gl.program_id(axis=1) + d_off = dblk * BD + NUM_TILES: gl.constexpr = R_CHUNK // TILE_K + + q_base = token_idx.to(tl.int64) * stride_q_t + do_base = token_idx.to(tl.int64) * stride_do_t + ds_base = token_idx.to(tl.int64) * stride_ds_t + interm_base = token_idx.to(tl.int64) * stride_interm_t + + # ---- prologue: stage Q, transpose to [BD, H] registers; reuse the buffer for dO ---- + offs_h_q = gl.arange(0, NH, layout=gl.SliceLayout(1, blk_q)) + offs_d_q = d_off + gl.arange(0, BD, layout=gl.SliceLayout(0, blk_q)) + mask_h_q = offs_h_q < num_heads + q_offs = ( + q_base + + offs_h_q[:, None].to(tl.int64) * stride_q_h + + offs_d_q[None, :].to(tl.int64) + ) + do_offs = ( + do_base + + offs_h_q[:, None].to(tl.int64) * stride_do_h + + offs_d_q[None, :].to(tl.int64) + ) + + if DUAL_STAGE: + # two buffers -> BOTH HBM->LDS copies in flight behind ONE drain. Costs 2x LDS + # (128 KB at BD=256, so occ-1) but halves the exposed prologue latency. + smem_q = gl.allocate_shared_memory( + Q_ptr.dtype.element_ty, [NH, BD], layout=sh_q + ) + smem_do = gl.allocate_shared_memory( + dO_ptr.dtype.element_ty, [NH, BD], layout=sh_q + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_q, ptr=Q_ptr, offsets=q_offs.to(tl.int32), mask=mask_h_q[:, None] + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_do, + ptr=dO_ptr, + offsets=do_offs.to(tl.int32), + mask=mask_h_q[:, None], + ) + gl.amd.cdna4.async_copy.commit_group() + gl.amd.cdna4.async_copy.wait_group(0) + Q_T = smem_q.permute([1, 0]).load(dot_a) # [BD, H], register-resident + dO_T = smem_do.permute([1, 0]).load(dot_a) + else: + # one buffer, re-used after a barrier: half the LDS (occ-2 at BD=256) but the two + # HBM round trips serialize. + smem_stage = gl.allocate_shared_memory( + Q_ptr.dtype.element_ty, [NH, BD], layout=sh_q + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_stage, + ptr=Q_ptr, + offsets=q_offs.to(tl.int32), + mask=mask_h_q[:, None], + ) + gl.amd.cdna4.async_copy.commit_group() + gl.amd.cdna4.async_copy.wait_group(0) + Q_T = smem_stage.permute([1, 0]).load(dot_a) + gl.barrier() # all warps done reading before reuse + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_stage, + ptr=dO_ptr, + offsets=do_offs.to(tl.int32), + mask=mask_h_q[:, None], + ) + gl.amd.cdna4.async_copy.commit_group() + gl.amd.cdna4.async_copy.wait_group(0) + dO_T = smem_stage.permute([1, 0]).load(dot_a) + + # ---- main loop: rank tile inner, head contraction folded into the mfma ---- + offs_h_ds = gl.arange(0, NH, layout=gl.SliceLayout(1, blk_ds)) + offs_k_ds = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_ds)) + mask_h_ds = offs_h_ds < num_heads + offs_d_st = d_off + gl.arange(0, BD, layout=gl.SliceLayout(1, mfma)) + offs_col_st = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma)) + + for t in range(NUM_TILES): + col_c = t * TILE_K + offs_k_ds + offs_c = ( + ds_base + + offs_h_ds[:, None].to(tl.int64) * stride_ds_h + + col_c[None, :].to(tl.int64) + ) + dS_blk = gl.amd.cdna4.buffer_load( + ptr=dS_ptr, offsets=offs_c.to(tl.int32), mask=mask_h_ds[:, None], other=0.0 + ) + P_blk = gl.amd.cdna4.buffer_load( + ptr=P_ptr, offsets=offs_c.to(tl.int32), mask=mask_h_ds[:, None], other=0.0 + ) + + dS_dot = gl.convert_layout(dS_blk, dot_b) + P_dot = gl.convert_layout(P_blk, dot_b) + + dKV = gl.zeros([BD, TILE_K], dtype=gl.float32, layout=mfma) + dKV = gl.amd.cdna4.mfma(Q_T, dS_dot, dKV) + dKV = gl.amd.cdna4.mfma(dO_T, P_dot, dKV) + + col_st = t * TILE_K + offs_col_st + interm_offs = ( + interm_base + + col_st[None, :].to(tl.int64) * stride_interm_r + + offs_d_st[:, None].to(tl.int64) + ) + gl.amd.cdna4.buffer_store( + stored_value=dKV.to(Interm_ptr.dtype.element_ty), + ptr=Interm_ptr, + offsets=interm_offs.to(tl.int32), + ) diff --git a/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py new file mode 100644 index 0000000000..b2295c8066 --- /dev/null +++ b/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Triton kernels for the DeepSeek-V4 sparse-MLA training BACKWARD (gfx950 / CDNA4). + +The two memory-bound phases; the MFMA phases are Gluon and live in +``aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_attention_dsv4_bwd``. + +``_delta_v4_kernel`` + ``delta = rowsum(O * dO)`` -- the standard flash-attention "o_dot_do" preamble. Streams the + bf16 inputs and accumulates in fp32, so it moves exactly the working set. + +``_bwd_dkv_gather_acc_v4`` + Reduces ``interm[t, slot, :]`` into ``dkv[kv_row, :]`` over the top-k mapping. The scatter is + inverted into a CSR gather (each output KV row collects its own contributors), so no atomics + are needed. + +Launchers live in ``aiter.ops.triton.attention.sparse_attention_dsv4_bwd``; this module stays +free of torch so the kernels can be called without it. +""" + +import triton +import triton.language as tl + +from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr + +_delta_v4_kernel_repr = make_kernel_repr( + "_delta_v4_kernel", + [ + "D", + "BLOCK_R", + ], +) + + +@triton.jit(repr=_delta_v4_kernel_repr) +def _delta_v4_kernel( + O_ptr, # [n_rows, D] bf16 (rows = T*H, contiguous) + dO_ptr, # [n_rows, D] bf16 + Delta_ptr, # [n_rows] fp32 + n_rows, + D: tl.constexpr, + BLOCK_R: tl.constexpr, +): + """Grid (cdiv(n_rows, BLOCK_R),) — each program reduces BLOCK_R rows of width D.""" + pid = tl.program_id(0) + rows = pid * BLOCK_R + tl.arange(0, BLOCK_R) + mask = rows < n_rows + offs = rows.to(tl.int64)[:, None] * D + tl.arange(0, D)[None, :] + o = tl.load(O_ptr + offs, mask=mask[:, None], other=0.0).to(tl.float32) + d = tl.load(dO_ptr + offs, mask=mask[:, None], other=0.0).to(tl.float32) + tl.store(Delta_ptr + rows, tl.sum(o * d, axis=1), mask=mask) + + +_bwd_dkv_gather_acc_v4_repr = make_kernel_repr( + "_bwd_dkv_gather_acc_v4", + [ + "D", + "BLOCK_E", + "ACCUMULATE", + ], +) + + +@triton.jit(repr=_bwd_dkv_gather_acc_v4_repr) +def _bwd_dkv_gather_acc_v4( + Interm_ptr, # [T, R_CHUNK, D] bf16, flat [T*R_CHUNK, D] + InvPtr_ptr, # [num_kv+1] int32 — CSR row pointers + InvData_ptr, # [valid] int32 — encoded q*R_CHUNK+local_r, sorted by KV token + dKV_acc_ptr, # [num_kv, D] fp32 — accumulator + stride_interm_r: tl.int64, + stride_acc_t: tl.int64, + D: tl.constexpr, + BLOCK_E: tl.constexpr, + ACCUMULATE: tl.constexpr, +): + """Grid (num_kv,) — one CTA per KV token, BLOCK_E CSR entries in flight. + + ``BLOCK_E`` entries are carried per iteration, which the gather needs for two reasons: + + * **load width.** A bare ``tl.arange(0, D)`` block over 256 threads is 2 bf16 = 4 B per + lane -- a dword. The [BLOCK_E, D] block gives ``BLOCK_E*D/threads`` elements per lane, + so the loads become dwordx4. The gather is issue-bound, so this dominates. + * **trip count.** ``tl.sum`` folds the entry axis, so the run is consumed BLOCK_E at a + time. A realistic top-k gives run lengths up to ~3000 on the pool rows. + + ``ACCUMULATE=False`` writes the destination instead of reading it back first. The caller + uses it for the first chunk, where the accumulator is still zero. + """ + k = tl.program_id(0) + offs_d = tl.arange(0, D) + offs_e = tl.arange(0, BLOCK_E) + start = tl.load(InvPtr_ptr + k) + end = tl.load(InvPtr_ptr + k + 1) + acc_base = k.to(tl.int64) * stride_acc_t + + if ACCUMULATE: + acc = tl.load(dKV_acc_ptr + acc_base + offs_d).to(tl.float32) + else: + acc = tl.zeros([D], dtype=tl.float32) + + for i0 in range(start, end, BLOCK_E): + idx = i0 + offs_e + m = idx < end + entry = tl.load(InvData_ptr + idx, mask=m, other=0).to(tl.int64) + vals = tl.load( + Interm_ptr + entry[:, None] * stride_interm_r + offs_d[None, :], + mask=m[:, None], + other=0.0, + ) + acc += tl.sum(vals.to(tl.float32), axis=0) + + tl.store(dKV_acc_ptr + acc_base + offs_d, acc) diff --git a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py new file mode 100644 index 0000000000..bdec379db7 --- /dev/null +++ b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py @@ -0,0 +1,522 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""DeepSeek-V4 sparse-MLA training BACKWARD (gfx950 / CDNA4). + +Counterpart to the DSv4 sparse prefill forward. The op is the official V4 form: shared-KV GQA +where ``K == V == kv`` is a single dense 512-wide tensor, RoPE already applied in place +caller-side, scale ``1/sqrt(512)``, ``attn_sink`` folded into the softmax denominator only, and +``topk_indices == -1`` masked out. + + P = exp(Q@kv^T * scale - lse) + dP = dO@kv^T + delta = rowsum(O * dO) + dS = P * (dP - delta) * scale + dQ = dS @ kv + dKV = scatter_add over top-k of sum_h ( dS*Q + P*dO ) + +Pipeline (five kernels + one torch reduction), per rank chunk: + + delta triton rowsum(O*dO) + dQ gluon also emits this chunk's dS / P + dKV-interm gluon interm[t, slot, d] = sum_h (dS*Q + P*dO) + CSR build torch inverted top-k index (sort + searchsorted) + dKV gather triton reduce interm over the top-k mapping, atomic-free + d_sink torch 26 us, not worth a kernel + +``lse`` and ``o`` come from the forward. The merged Gluon prefill kernel produces both:: + + from aiter.ops.triton.gluon.mla_gluon import mla_gluon + o, lse = mla_gluon(..., has_pe=False, attn_sink=sink, return_lse=True) + +Its ``lse`` is sink-inclusive (the sink is folded into ``e_max``/``e_sum`` before +``lse = e_max + log(e_sum)``), which is the convention this backward expects. + +Measured on gfx950 (MI355X) at ``T=4096 H=128 topk=512`` with a realistic SWA(128)+pool top-k: +delta 0.178 / dQ 1.391 / interm 1.152 / CSR build 0.130 / gather 0.503 / d_sink 0.026 ms, +3.380 ms total = 407 TFLOPS. +""" + +from dataclasses import dataclass + +import torch +import triton + +from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_attention_dsv4_bwd import ( + _dkv_interm_v4_kernel, + _dq_v4_kernel, +) +from aiter.ops.triton._triton_kernels.attention.sparse_attention_dsv4_bwd import ( + _bwd_dkv_gather_acc_v4, + _delta_v4_kernel, +) +from aiter.ops.triton.utils._triton import arch_info + +_BLOCK_H_DQ = 64 +_TILE_K_DQ = 32 +_BD_DKV = 256 +_TILE_K_DKV = 128 + + +def sparse_mla_bwd_dq( + q, + kv, + do, + topk, + lse, + delta, + dq, + chunk_dS, + chunk_P, + scale, + r_start, + R_CHUNK, + BLOCK_H=64, + TILE_K=32, + is_first_chunk=True, +): + """Launch the dQ kernel for one rank chunk. Writes ``dq`` (RMW when not the first chunk) + plus this chunk's ``chunk_dS`` / ``chunk_P``.""" + T, H, D = q.shape + _dq_v4_kernel[(T, triton.cdiv(H, BLOCK_H))]( + q, + kv, + do, + topk, + lse, + delta, + dq, + chunk_dS, + chunk_P, + q.stride(0), + q.stride(1), + kv.stride(0), + do.stride(0), + do.stride(1), + dq.stride(0), + dq.stride(1), + topk.stride(0), + chunk_dS.stride(0), + chunk_dS.stride(1), + scale, + H, + r_start, + R_CHUNK=R_CHUNK, + BLOCK_H=BLOCK_H, + TILE_K=TILE_K, + D=D, + IS_FIRST_CHUNK=is_first_chunk, + num_warps=4, + waves_per_eu=1, + ) + + +def sparse_mla_bwd_dkv_interm_v4( + q, + do, + chunk_dS, + chunk_P, + R_CHUNK, + BD=256, + TILE_K=128, + MFMA_K=32, + DUAL_STAGE=1, + H_POW2=None, + num_warps=4, + interm=None, +): + """V4 dKV-interm, Q/dO read once. Returns interm [T, R_CHUNK, D] bf16. + + ``BD`` splits D across ``grid.y``; dS/P are re-read once per D block, so a larger BD moves + less of them. ``MFMA_K=32`` is the CDNA4 16x16x32 depth. + """ + T, H, D = q.shape + assert R_CHUNK % TILE_K == 0 + assert D % BD == 0 + h_pow2 = H_POW2 or triton.next_power_of_2(H) + if interm is None: + interm = torch.empty(T, R_CHUNK, D, dtype=torch.bfloat16, device=q.device) + _dkv_interm_v4_kernel[(T, D // BD)]( + q, + do, + chunk_dS, + chunk_P, + interm, + q.stride(0), + q.stride(1), + do.stride(0), + do.stride(1), + chunk_dS.stride(0), + chunk_dS.stride(1), + interm.stride(0), + interm.stride(1), + H, + R_CHUNK=R_CHUNK, + TILE_K=TILE_K, + NH=h_pow2, + BD=BD, + D=D, + MFMA_K=MFMA_K, + DUAL_STAGE=DUAL_STAGE, + num_warps=num_warps, + ) + return interm + + +def delta_v4(o, do, out=None, BLOCK_R=8, num_warps=8): + # BLOCK_R=8 keeps each lane loading >= 8 bf16, i.e. a dwordx4; narrower blocks drop to a + # dword and the kernel loses most of its bandwidth. + """o[T,H,D] bf16, do[T,H,D] bf16 -> delta[T,H] fp32 = sum_d o*do. + + ``do`` must already be the D-wide (lora) slice, contiguous — same contract as the dQ kernel. + """ + assert o.shape == do.shape and o.is_contiguous() and do.is_contiguous() + T, H, D = o.shape + n_rows = T * H + if out is None: + out = torch.empty(T, H, dtype=torch.float32, device=o.device) + _delta_v4_kernel[(triton.cdiv(n_rows, BLOCK_R),)]( + o, + do, + out, + n_rows, + D=D, + BLOCK_R=BLOCK_R, + num_warps=num_warps, + ) + return out + + +def build_inverted_topk(topk_indices_slice, num_kv): + """CSR inverted index over ``num_kv`` KV rows. + + One stable sort yields both the permutation (``inv_data``) and the sorted keys; + ``inv_ptr[k] = searchsorted(sorted, k, 'left')`` = the number of entries with value < k. + Invalid (-1) entries sort to the front, so ``inv_ptr[0]`` starts past them and they are + never visited. + + The sort key is narrowed to int16 when ``num_kv`` fits, which is what keeps the radix sort + to two byte-passes. + + Returns ``inv_ptr[num_kv+1]`` int32, ``inv_data[T*R]`` int32. + """ + # row_ids is the searchsorted query: [0 .. num_kv], one per KV row plus the end sentinel. + # Its dtype must match `keys` -- searchsorted is built per branch for that reason, not by + # accident. + flat_kv = topk_indices_slice.reshape(-1) # [T*R] int32; -1 = invalid + # row_ids reaches num_kv, so int16 holds it as long as num_kv is within the type; + # the -1 sentinel is fine either way. + if num_kv <= torch.iinfo(torch.int16).max: + keys = flat_kv.to(torch.int16) + row_ids = torch.arange(num_kv + 1, device=flat_kv.device, dtype=torch.int16) + else: + keys = flat_kv.to(torch.int32) + row_ids = torch.arange(num_kv + 1, device=flat_kv.device, dtype=torch.int32) + sorted_vals, inv_data = torch.sort(keys, stable=True) + inv_ptr = torch.searchsorted(sorted_vals, row_ids).to(torch.int32) + return inv_ptr, inv_data.to(torch.int32) + + +def dkv_gather_acc( + interm, inv_ptr, inv_data, dkv_acc, BLOCK_E=64, num_warps=8, accumulate=True +): + """interm[T,R,D] bf16 -> dkv_acc[num_kv,D] fp32 via the entry-blocked CSR gather. + + Grid is ``num_kv`` (from ``dkv_acc``), not ``T``, so a compressed-pool KV works. + """ + _, _, D = interm.shape + num_kv = dkv_acc.shape[0] + _bwd_dkv_gather_acc_v4[(num_kv,)]( + interm, + inv_ptr, + inv_data, + dkv_acc, + interm.stride(1), + dkv_acc.stride(0), + D=D, + BLOCK_E=BLOCK_E, + ACCUMULATE=accumulate, + num_warps=num_warps, + ) + + +@dataclass +class _BwdPlan: + """Validated inputs, tile choices and workspace for one backward call. + + Holds the values that flow between phases (``delta``, the CSR pair, the accumulator) so + ``_bwd_phases`` can hand out independent thunks that still compose into the real pipeline. + """ + + q: torch.Tensor + kv: torch.Tensor + do: torch.Tensor + o: torch.Tensor + lse: torch.Tensor + topk_indices: torch.Tensor + attn_sink: object + scale: float + T: int + H: int + D: int + TOPK: int + num_kv: int + R_CHUNK: int + tk_dq: int + tk_dkv: int + delta: torch.Tensor + dq: torch.Tensor + dkv_acc: torch.Tensor + chunk_dS: torch.Tensor + chunk_P: torch.Tensor + interm: torch.Tensor + inv_ptr: object = None + inv_data: object = None + d_sink: object = None + + def build_csr(self, r): + self.inv_ptr, self.inv_data = build_inverted_topk( + self.topk_indices[:, r : r + self.R_CHUNK], self.num_kv + ) + + def compute_d_sink(self): + # d_sink[h] = -sum_t exp(sink[h] - lse[t,h]) * delta[t,h] + self.d_sink = -( + torch.exp(self.attn_sink[None, :].float() - self.lse) * self.delta + ).sum(dim=0) + + def result(self): + return self.dq, self.dkv_acc.to(self.kv.dtype), self.d_sink + + +def _plan_bwd( + q, kv, do, o, lse, topk_indices, attn_sink=None, scale=None, R_CHUNK=None +): + """Validate the inputs, pick the tile widths and allocate the workspace for one call. + + Split out from ``sparse_mla_bwd_dsv4`` so the op benchmark can build the same plan and then + time ``_bwd_phases`` against it. Argument meanings are documented on the public entry. + """ + assert ( + arch_info.get_arch() == "gfx950" + ), f"sparse_mla_bwd_dsv4 requires gfx950 (CDNA4), got {arch_info.get_arch()}" + + for name, t in ( + ("q", q), + ("kv", kv), + ("do", do), + ("o", o), + ("lse", lse), + ("topk_indices", topk_indices), + ): + if not t.is_cuda: + raise RuntimeError( + f"sparse_mla_bwd_dsv4 requires CUDA/HIP tensors, {name} is on {t.device}" + ) + + if q.dtype != torch.bfloat16: + raise RuntimeError(f"sparse_mla_bwd_dsv4 expects bf16 q, got {q.dtype}") + for name, t in (("kv", kv), ("do", do), ("o", o)): + if t.dtype != q.dtype: + raise RuntimeError(f"{name} dtype mismatch: {name}={t.dtype}, q={q.dtype}") + + T, H, D = q.shape + assert topk_indices.ndim == 2 and topk_indices.shape[0] == T, ( + f"topk_indices must be [T, TOPK] with T={T}, got {tuple(topk_indices.shape)} -- the dQ " + "grid is sized from q, so a shorter index tensor is read past its end" + ) + TOPK = topk_indices.shape[1] + num_kv = kv.shape[0] + assert D == 512, f"DSv4 sparse-MLA backward is fixed to head_dim 512, got {D}" + assert kv.shape[-1] == D, f"kv must be [num_kv, {D}], got {tuple(kv.shape)}" + assert ( + do.shape == q.shape + ), f"do must match q {tuple(q.shape)}, got {tuple(do.shape)}" + assert o.shape == q.shape, f"o must match q {tuple(q.shape)}, got {tuple(o.shape)}" + assert lse.shape == (T, H), f"lse must be [{T}, {H}], got {tuple(lse.shape)}" + assert num_kv >= T, f"num_kv ({num_kv}) must be >= T ({T})" + if attn_sink is not None: + # Optional, so it misses the device sweep above; d_sink does GPU math with it. + if not attn_sink.is_cuda: + raise RuntimeError( + f"sparse_mla_bwd_dsv4 requires CUDA/HIP tensors, attn_sink is on " + f"{attn_sink.device}" + ) + if attn_sink.dtype != torch.float32: + raise RuntimeError( + f"sparse_mla_bwd_dsv4 expects fp32 attn_sink, got {attn_sink.dtype}" + ) + assert attn_sink.shape == ( + H, + ), f"attn_sink must be [{H}], got {tuple(attn_sink.shape)}" + assert ( + topk_indices.dtype == torch.int32 + ), f"topk_indices must be int32, got {topk_indices.dtype}" + for name, t in ( + ("q", q), + ("kv", kv), + ("do", do), + ("o", o), + ("topk_indices", topk_indices), + ): + assert t.is_contiguous(), f"{name} must be contiguous" + + if scale is None: + scale = 1.0 / (D**0.5) + if R_CHUNK is None: + R_CHUNK = TOPK + lse = lse.float().contiguous() + + # Both mfma tiles must divide the chunk width; step down rather than making it the + # caller's problem, so a small R_CHUNK still works. + tk_dkv = next( + (t for t in (_TILE_K_DKV, 64, 32) if t <= R_CHUNK and R_CHUNK % t == 0), None + ) + tk_dq = next( + (t for t in (_TILE_K_DQ, 32) if t <= R_CHUNK and R_CHUNK % t == 0), None + ) + assert ( + tk_dkv is not None and tk_dq is not None + ), f"R_CHUNK={R_CHUNK} must be a multiple of 32 (it is the mfma tile width)" + + # A tail chunk narrower than R_CHUNK is wrong in three places at once. The kernels take + # R_CHUNK as a constexpr and mask the top-k load against it rather than against TOPK, so + # they read past the end of each row -- into the next token's indices for tokens 0..T-2, + # and past the tensor entirely for the last one (measured: that faults the GPU). The CSR + # build meanwhile gets a torch-clamped, narrower slice, and the gather then indexes interm, + # which is still R_CHUNK wide, with entries encoded against that narrower width. Require + # the divisor rather than paying a second compile for a narrower tail. + assert TOPK % R_CHUNK == 0, ( + f"R_CHUNK={R_CHUNK} must divide TOPK={TOPK} -- a partial tail chunk reads past the end " + "of each top-k row and desynchronizes the CSR index space from interm" + ) + + return _BwdPlan( + q=q, + kv=kv, + do=do, + lse=lse, + topk_indices=topk_indices, + attn_sink=attn_sink, + scale=scale, + T=T, + H=H, + D=D, + TOPK=TOPK, + num_kv=num_kv, + R_CHUNK=R_CHUNK, + tk_dq=tk_dq, + tk_dkv=tk_dkv, + delta=torch.empty(T, H, dtype=torch.float32, device=q.device), + dq=torch.empty_like(q), + dkv_acc=torch.zeros(num_kv, D, dtype=torch.float32, device=q.device), + chunk_dS=torch.empty(T, H, R_CHUNK, dtype=torch.bfloat16, device=q.device), + chunk_P=torch.empty(T, H, R_CHUNK, dtype=torch.bfloat16, device=q.device), + interm=torch.empty(T, R_CHUNK, D, dtype=torch.bfloat16, device=q.device), + o=o, + ) + + +def _bwd_phases(plan): + """``(name, thunk)`` for every phase of one backward call, in execution order. + + ``sparse_mla_bwd_dsv4`` runs these; the op benchmark times them one at a time. Sharing the + sequence is the point -- a per-phase measurement is only meaningful if it is measuring the + phases the op actually runs, with the same tile configuration and the same workspace. + """ + yield "delta", lambda: delta_v4(plan.o, plan.do, out=plan.delta) + + for r in range(0, plan.TOPK, plan.R_CHUNK): + yield ( + "dq", + lambda r=r: sparse_mla_bwd_dq( + plan.q, + plan.kv, + plan.do, + plan.topk_indices, + plan.lse, + plan.delta, + plan.dq, + plan.chunk_dS, + plan.chunk_P, + plan.scale, + r, + plan.R_CHUNK, + BLOCK_H=_BLOCK_H_DQ, + TILE_K=plan.tk_dq, + is_first_chunk=(r == 0), + ), + ) + yield ( + "interm", + lambda: sparse_mla_bwd_dkv_interm_v4( + plan.q, + plan.do, + plan.chunk_dS, + plan.chunk_P, + plan.R_CHUNK, + BD=_BD_DKV, + TILE_K=plan.tk_dkv, + interm=plan.interm, + ), + ) + yield "csr_build", lambda r=r: plan.build_csr(r) + # dkv_acc was just zeroed, so the first chunk can write instead of read-modify-write. + # Unchunked -- the default -- that is the only chunk, and it saves reading back the + # whole [num_kv, 512] fp32 accumulator. + yield ( + "gather", + lambda r=r: dkv_gather_acc( + plan.interm, + plan.inv_ptr, + plan.inv_data, + plan.dkv_acc, + accumulate=(r > 0), + ), + ) + + if plan.attn_sink is not None: + yield "d_sink", plan.compute_d_sink + + +def sparse_mla_bwd_dsv4( + q, + kv, + do, + o, + lse, + topk_indices, + attn_sink=None, + scale=None, + R_CHUNK=None, +): + """Backward for the DSv4 sparse-MLA prefill attention. gfx950 (CDNA4) only. + + Args: + q: [T, H, 512] bf16 + kv: [num_kv, 512] bf16, K == V. ``num_kv >= T``; rows ``T..num_kv-1`` are + the compressed pool, which ``topk_indices`` may reference. + do: [T, H, 512] bf16, gradient of the attention output + o: [T, H, 512] bf16, the forward output + lse: [T, H] fp32, sink-inclusive log-sum-exp from the forward + topk_indices: [T, TOPK] int32, -1 marks an invalid slot + attn_sink: [H] fp32 per-head sink bias, or None + scale: softmax scale, defaults to ``1/sqrt(512)`` + R_CHUNK: split the rank dimension into chunks of this width. ``None`` (default) + runs unchunked, which is what you want. Chunking exists only to bound the + ``interm`` intermediate, which is ``T*TOPK*512`` bf16 (2.0 GiB at + T=4096, TOPK=512); it costs a dQ read-modify-write between chunks and one + CSR build per chunk. Must be a multiple of 32 (the mfma tile width) and + must divide ``TOPK`` -- a partial tail chunk is rejected rather than + handled, since the chunk width is a kernel constexpr. + + Returns: + dq [T, H, 512] bf16, dkv [num_kv, 512] bf16, d_sink [H] fp32 (None if no ``attn_sink``) + """ + plan = _plan_bwd(q, kv, do, o, lse, topk_indices, attn_sink, scale, R_CHUNK) + for _, run in _bwd_phases(plan): + run() + return plan.result() + + +__all__ = ["sparse_mla_bwd_dsv4"] diff --git a/aiter/test_mha_common.py b/aiter/test_mha_common.py index 936d79798b..97ac286ff2 100644 --- a/aiter/test_mha_common.py +++ b/aiter/test_mha_common.py @@ -724,3 +724,51 @@ def opus_check_lse(tag, lse, lse_ref): diff = (lse[finite] - lse_ref[finite]).abs().max().item() if finite.any() else 0.0 print(f"[{tag}] lse max diff: {diff}") assert diff <= 0.01, f"{tag}: lse diff {diff} > 0.01" + + +def sparse_mla_dsv4_ref(q, kv, topk_indices, attn_sink=None, scale=None): + """fp32 reference forward for the DeepSeek-V4 sparse MLA attention. + + The official V4 form: shared-KV GQA where ``K == V == kv`` is one dense ``head_dim``-wide + tensor, RoPE already applied in place caller-side, ``attn_sink`` folded into the softmax + denominator only, and ``topk_indices == -1`` masked out. + + Differentiable in ``q`` / ``kv`` / ``attn_sink``, so a backward reference is just autograd + through this. It is a per-token python loop, so keep the shapes small. + + Args: + q: [T, H, D] float32, requires_grad for a backward reference + kv: [num_kv, D] float32, ``num_kv >= T``; rows ``T..`` are the compressed pool + topk_indices: [T, TOPK] int, -1 marks an invalid slot + attn_sink: [H] float32 or None + scale: softmax scale, defaults to ``1/sqrt(D)`` + + Returns: + ``(o, lse)`` -- ``o`` [T, H, D] float32 carrying grad, and ``lse`` [T, H] float32 + detached and sink-inclusive, which is the convention the backward kernel expects. + A row whose top-k is entirely -1 gets ``o = 0`` and ``lse = -inf``. + """ + if scale is None: + scale = 1.0 / (q.shape[-1] ** 0.5) + outs, lses = [], [] + for t in range(q.shape[0]): + idx = topk_indices[t].long() + valid = idx != -1 + k = kv[idx.clamp(min=0)] + k = torch.where(valid[:, None], k, torch.zeros_like(k)) + s = (q[t] @ k.t()) * scale + s = torch.where(valid[None, :], s, torch.full_like(s, float("-inf"))) + row_max = s.max(dim=1).values + m = row_max if attn_sink is None else torch.maximum(row_max, attn_sink) + p = torch.where(valid[None, :], torch.exp(s - m[:, None]), torch.zeros_like(s)) + denom = p.sum(dim=1) + if attn_sink is not None: + denom = denom + torch.exp(attn_sink - m) + # A row whose top-k is entirely -1 has no contributors at all, and without a sink the + # denominator is then 0. Define that row's output as zero rather than letting 0/0 make + # it NaN; the lse stays -inf, which is the honest value for an empty row. + outs.append( + (p @ k) / torch.where(denom > 0, denom, torch.ones_like(denom))[:, None] + ) + lses.append((m + torch.log(denom)).detach()) + return torch.stack(outs, dim=0), torch.stack(lses, dim=0) diff --git a/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py b/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py new file mode 100644 index 0000000000..efb38e5581 --- /dev/null +++ b/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py @@ -0,0 +1,280 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Benchmark for the DSv4 sparse MLA training backward (gfx950 / CDNA4). + +Reproduces the numbers quoted for `sparse_mla_bwd_dsv4`. The top-k is not uniform random: +it is a sliding window plus a causally-visible compressed pool, the distribution the V4 +indexer actually produces. That matters here -- a uniform top-k gives every KV row about the +same number of contributors, while the real one gives the pool rows runs of a few thousand, +and the dKV gather is sized for those runs. Benchmarking on uniform indices flatters the +gather by roughly 2x. + +`--breakdown` additionally times each phase on its own, which is how the per-kernel table in +the PR description was produced. + +Usage: + python op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py + python op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py --breakdown + python op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py --cfgs 4096,128,512 +""" + +import argparse + +import torch +import triton + +from aiter.ops.triton.attention.sparse_attention_dsv4_bwd import ( + _bwd_phases, + _plan_bwd, + sparse_mla_bwd_dsv4, +) +from aiter.ops.triton.utils._triton import arch_info +from aiter.test_mha_common import sparse_mla_dsv4_ref + +D = 512 +SWA = 128 # sliding-window width +CR_POOL = ( + 4 # pool compression ratio: block b is visible to token t once (b+1)*CR_POOL-1 <= t +) + + +# --------------------------------------------------------------------------- +# Bench data builder +# --------------------------------------------------------------------------- +def _build_topk_swa_pool(T, topk, num_pool, device, generator): + """[T, topk] int32 -- SWA(128) window plus causally-visible pool ranks, -1 padded. + + KV row layout is ``[ per-token 0..T-1 | pool T..T+num_pool-1 ]``. Early tokens see fewer + than `topk - SWA` pool blocks, so their trailing slots stay -1, exactly as in production. + + ``topk`` must be at least the window width: the window is the floor of a V4 top-k, and a + narrower request has no meaning here. Returning the full window anyway would report one + ``topk`` in the results table while timing another. + """ + assert topk >= SWA, f"topk={topk} is narrower than the SWA({SWA}) window" + + idx = torch.arange(T, device=device) + off = torch.arange(SWA, device=device) + swa = idx[:, None] - (SWA - 1) + off[None, :] + swa = torch.where(swa >= 0, swa, torch.full_like(swa, -1)) + + n_pool = topk - SWA + if n_pool == 0 or num_pool == 0: + return swa.to(torch.int32).contiguous() + assert n_pool <= num_pool, f"need {n_pool} pool ranks but only {num_pool} blocks" + + n_visible = torch.clamp((idx + 1) // CR_POOL, min=0, max=num_pool) + blk = torch.arange(num_pool, device=device) + visible = blk[None, :] < n_visible[:, None] + + # Pick n_pool visible blocks at random: score the visible ones in [0,1), push the rest to + # 2.0, take the smallest. + score = torch.rand(T, num_pool, device=device, generator=generator) + score = torch.where(visible, score, torch.full_like(score, 2.0)) + sel = score.argsort(dim=1)[:, :n_pool] + pool = torch.where(torch.gather(visible, 1, sel), T + sel, torch.full_like(sel, -1)) + return torch.cat([swa, pool], dim=1).to(torch.int32).contiguous() + + +def _build_case(T, H, topk, device, num_pool=1024, seed=0): + """Inputs for one configuration. + + `o` / `lse` are synthetic. No kernel branches on their values, so timing is unaffected and + this avoids a reference forward that would dominate the run. + """ + gen = torch.Generator(device=device) + gen.manual_seed(seed) + indices = _build_topk_swa_pool(T, topk, num_pool, device, gen) + num_kv = T + num_pool if topk > SWA else T + + def _bf16(*shape): + return torch.randn(*shape, device=device, dtype=torch.bfloat16, generator=gen) + + return { + "q": _bf16(T, H, D), + "kv": _bf16(num_kv, D), + "do": _bf16(T, H, D), + "o": _bf16(T, H, D), + "lse": torch.randn(T, H, device=device, dtype=torch.float32, generator=gen), + "sink": ( + torch.randn(H, device=device, dtype=torch.float32, generator=gen) * 0.1 + ).contiguous(), + "indices": indices, + "num_kv": num_kv, + } + + +# --------------------------------------------------------------------------- +# Timing +# --------------------------------------------------------------------------- +def _flops(T, H, topk): + # dQ contributes S, dP and dS@kv; dKV-interm contributes dS*Q and P*dO. Five 2*D-flop + # products per (token, head, top-k slot). + return 10.0 * D * T * H * topk + + +def _time_phases(case, scale): + """Time each phase of the real pipeline. + + `_plan_bwd` and `_bwd_phases` are the same helpers `sparse_mla_bwd_dsv4` runs, so this cannot + drift from the op: the tile widths, the workspace and the phase order all come from the + wrapper rather than being restated here. Timing a phase repeatedly re-runs its side effects, + which is harmless -- only the duration is read. + """ + plan = _plan_bwd( + case["q"], + case["kv"], + case["do"], + case["o"], + case["lse"], + case["indices"], + case["sink"], + scale, + ) + return [(name, triton.testing.do_bench(run)) for name, run in _bwd_phases(plan)] + + +# --------------------------------------------------------------------------- +# Correctness gate +# --------------------------------------------------------------------------- +def check_correctness(device): + """Small-shape gate against autograd, so a broken kernel fails before it is timed.""" + print("\n========== CORRECTNESS ==========") + T, H, topk = 128, 64, 128 + scale = 1.0 / (D**0.5) + case = _build_case(T, H, topk, device, num_pool=64) + q, kv, do, sink, indices = ( + case["q"], + case["kv"], + case["do"], + case["sink"], + case["indices"], + ) + + qg = q.float().clone().requires_grad_(True) + kvg = kv.float().clone().requires_grad_(True) + sg = sink.clone().requires_grad_(True) + ref_o, lse = sparse_mla_dsv4_ref(qg, kvg, indices, sg, scale) + o = ref_o.detach().to(torch.bfloat16).contiguous() + + dq, dkv, d_sink = sparse_mla_bwd_dsv4( + q, kv, do, o, lse, indices, attn_sink=sink, scale=scale + ) + ref_o.backward(do.float()) + + def _cos(a, b): + return torch.nn.functional.cosine_similarity( + a.float().reshape(-1), b.float().reshape(-1), dim=0 + ).item() + + for name, got, want in ( + ("dq", dq, qg.grad), + ("dkv", dkv, kvg.grad), + ("d_sink", d_sink, sg.grad), + ): + cos = _cos(got, want) + assert cos > 0.999, f"{name} cos={cos:.6f}" + print(f" {name:7s}: OK (cos={cos:.6f})") + + +# --------------------------------------------------------------------------- +# Reporting +# --------------------------------------------------------------------------- +def _print_table(headers, rows): + def _fmt(x): + if isinstance(x, float): + return f"{x:.3f}" if x >= 1 or x == 0 else f"{x:.4f}" + return str(x) + + cells = [[_fmt(c) for c in r] for r in rows] + widths = [max(len(h), *(len(c[i]) for c in cells)) for i, h in enumerate(headers)] + print("| " + " | ".join(h.rjust(widths[i]) for i, h in enumerate(headers)) + " |") + print("| " + " | ".join("-" * widths[i] for i in range(len(headers))) + " |") + for c in cells: + print("| " + " | ".join(s.rjust(widths[i]) for i, s in enumerate(c)) + " |") + + +def run_bwd_bench(args, device): + print("\n========== BACKWARD ==========") + rows = [] + for T, H, topk in args.cfgs: + case = _build_case(T, H, topk, device) + scale = 1.0 / (D**0.5) + tflops = _flops(T, H, topk) / 1e12 + + ms = triton.testing.do_bench( + lambda c=case, s=scale: sparse_mla_bwd_dsv4( + c["q"], + c["kv"], + c["do"], + c["o"], + c["lse"], + c["indices"], + attn_sink=c["sink"], + scale=s, + ) + ) + rows.append([T, H, case["num_kv"], topk, ms, tflops / (ms * 1e-3)]) + + if args.breakdown: + print(f"\n per-kernel, T={T} H={H} topk={topk}:") + phases = _time_phases(case, scale) + total = sum(t for _, t in phases) + for name, t in phases: + print(f" {name:10s} {t:7.3f} ms ({100 * t / total:4.1f}%)") + print( + f" {'SUM':10s} {total:7.3f} ms -> {tflops / (total * 1e-3):.0f} TFLOPS" + ) + + _print_table(["T", "H", "Kv", "topk", "ms", "TFLOPS"], rows) + + +def _parse_args(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument( + "--cfgs", + nargs="+", + type=str, + default=[ # (T, H, topk) + "4096,128,512", + "4096,128,1024", + "8192,128,512", + ], + help="one or more shapes as T,H,topk (e.g. 4096,128,512). topk must be >= the sliding " + "window (128); the remainder is drawn from the compressed pool.", + ) + p.add_argument( + "--breakdown", + action="store_true", + help="also time each phase separately (the per-kernel table in the PR description)", + ) + p.add_argument( + "--skip-correctness", + action="store_true", + help="skip the small-shape autograd check that otherwise runs before any timing.", + ) + args = p.parse_args() + args.cfgs = [tuple(int(x) for x in s.split(",")) for s in args.cfgs] + return args + + +def main(): + args = _parse_args() + device = "cuda" + if arch_info.get_arch() != "gfx950": + print(f"sparse_mla_bwd_dsv4 is gfx950 only; this is {arch_info.get_arch()}") + return + print( + f"GPU: {torch.cuda.get_device_name(0)} " + f"({torch.cuda.get_device_properties(0).multi_processor_count} CUs)" + ) + print(f"Triton: {triton.__version__}") + + if not args.skip_correctness: + check_correctness(device) + run_bwd_bench(args, device) + + +if __name__ == "__main__": + main() diff --git a/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py b/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py new file mode 100644 index 0000000000..4493048e9b --- /dev/null +++ b/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Correctness for the DSv4 sparse-MLA training backward against torch autograd. + +The reference is ``aiter.test_mha_common.sparse_mla_dsv4_ref``, a differentiable fp32 +re-implementation of the V4 forward -- an independent path from the kernels, not a re-expression +of them. Backward values come from autograd through it. It is a per-token python loop, so the +shapes here are deliberately small; the op benchmark exercises the kernels at production shapes +and shares the same reference for its correctness gate. +""" + +import pytest +import torch + +from aiter.ops.triton.attention.sparse_attention_dsv4_bwd import sparse_mla_bwd_dsv4 +from aiter.ops.triton.utils._triton import arch_info +from aiter.test_mha_common import sparse_mla_dsv4_ref + +D = 512 +COS_TOL = 0.999 + +pytestmark = pytest.mark.skipif( + arch_info.get_arch() != "gfx950", + reason="DSv4 sparse-MLA backward is gfx950 (CDNA4) only", +) + + +def _cos(a, b): + return torch.nn.functional.cosine_similarity( + a.float().reshape(-1), b.float().reshape(-1), dim=0 + ).item() + + +@pytest.mark.parametrize( + "T, H, topk, npool, has_sink, r_chunk", + [ + (128, 64, 128, 0, True, None), # no pool, unchunked + (128, 128, 128, 0, True, None), # H=128 + (128, 128, 128, 64, True, None), # compressed pool (num_kv > T) + (128, 64, 128, 0, False, None), # no attn_sink + (128, 128, 128, 64, True, 64), # chunked: dQ RMW + per-chunk CSR build + ], +) +def test_sparse_mla_bwd_dsv4(T, H, topk, npool, has_sink, r_chunk): + torch.manual_seed(0) + dev = "cuda" + num_kv = T + npool + scale = 1.0 / (D**0.5) + + q = torch.randn(T, H, D, device=dev, dtype=torch.bfloat16) + kv = torch.randn(num_kv, D, device=dev, dtype=torch.bfloat16) + do = torch.randn(T, H, D, device=dev, dtype=torch.bfloat16) + sink = (torch.randn(H, device=dev, dtype=torch.float32) * 0.1) if has_sink else None + + indices = torch.randint(0, num_kv, (T, topk), dtype=torch.int32, device=dev) + invalid = torch.rand(T, topk, device=dev) < 0.1 + indices = torch.where(invalid, torch.full_like(indices, -1), indices).contiguous() + + qg = q.float().clone().requires_grad_(True) + kvg = kv.float().clone().requires_grad_(True) + sg = sink.clone().requires_grad_(True) if has_sink else None + ref_o, lse = sparse_mla_dsv4_ref(qg, kvg, indices, sg, scale) + o = ref_o.detach().to(torch.bfloat16).contiguous() + + dq, dkv, d_sink = sparse_mla_bwd_dsv4( + q, kv, do, o, lse, indices, attn_sink=sink, scale=scale, R_CHUNK=r_chunk + ) + + ref_o.backward(do.float()) + + assert _cos(dq, qg.grad) > COS_TOL, f"dq cos {_cos(dq, qg.grad)}" + assert _cos(dkv, kvg.grad) > COS_TOL, f"dkv cos {_cos(dkv, kvg.grad)}" + if has_sink: + assert _cos(d_sink, sg.grad) > COS_TOL, f"d_sink cos {_cos(d_sink, sg.grad)}" + else: + assert d_sink is None + + +def _dummy_inputs(T=64, H=64, topk=64, dev="cuda"): + return { + "q": torch.randn(T, H, D, device=dev, dtype=torch.bfloat16), + "kv": torch.randn(T, D, device=dev, dtype=torch.bfloat16), + "do": torch.randn(T, H, D, device=dev, dtype=torch.bfloat16), + "o": torch.randn(T, H, D, device=dev, dtype=torch.bfloat16), + "lse": torch.randn(T, H, device=dev, dtype=torch.float32), + "idx": torch.randint(0, T, (T, topk), dtype=torch.int32, device=dev), + } + + +def test_sparse_mla_bwd_dsv4_rejects_bad_chunk(): + """R_CHUNK must be a multiple of the mfma tile width.""" + t = _dummy_inputs(topk=64) + with pytest.raises(AssertionError, match="multiple of 32"): + sparse_mla_bwd_dsv4( + t["q"], t["kv"], t["do"], t["o"], t["lse"], t["idx"], R_CHUNK=48 + ) + + +def test_sparse_mla_bwd_dsv4_rejects_short_topk(): + """topk_indices must have one row per query token. + + The dQ grid is sized from `q`, so a shorter index tensor is read past its end rather than + producing a shape error anywhere downstream. + """ + t = _dummy_inputs(T=64, topk=64) + with pytest.raises(AssertionError, match="topk_indices must be"): + sparse_mla_bwd_dsv4(t["q"], t["kv"], t["do"], t["o"], t["lse"], t["idx"][:32]) + + +def test_sparse_mla_bwd_dsv4_rejects_indivisible_chunk(): + """A chunk width that is a valid tile multiple but does not divide TOPK is still rejected. + + 96 = 3*32 so it clears the tile-width check, but 128 % 96 == 32 would leave a 32-wide tail + chunk. The kernels take the chunk width as a constexpr and would read past the end of each + top-k row, silently, so this is an error rather than a handled case. + """ + t = _dummy_inputs(topk=128) + with pytest.raises(AssertionError, match="must divide TOPK"): + sparse_mla_bwd_dsv4( + t["q"], t["kv"], t["do"], t["o"], t["lse"], t["idx"], R_CHUNK=96 + )