diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py new file mode 100644 index 00000000000..fb1e782fcbe --- /dev/null +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -0,0 +1,963 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Gluon (gfx950) DeepSeek-V4 sparse-MLA decode. Adapted from the vLLM DSv4 sparse +attention kernels. + +K == V, so each tile is gathered once into one [BLOCK_K, HEAD] bf16 LDS buffer and +read permuted for QK, direct for PV. fp8 dequants to bf16 on the way in. One kernel +serves three KV formats via the UNIFORM constexpr on the fp8 gather: + + UNIFORM=False packed fp8_ds_mla: NoPE fp8 + embedded UE8M0 + RoPE bf16 + UNIFORM=True uniform pool: whole head fp8 + a separate fp32 kv_scales + IS_FP8=False bf16 pool + +Two-loop (SWA + top-k) or a single merged segment; 2D and 3D (split-K + reduce) +share one kernel. Launchers: aiter/ops/triton/attention/pa_decode_sparse.py. +""" + +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 + + +@gluon.jit +def _cache_load(ptr, off, USE_BUFFER_LOAD: gl.constexpr, mask=None, other=None): + # gfx950 buffer_load carries a 32-bit offset (2 GB cap), a cache past that gathers + # via 64-bit gl.load instead. + if USE_BUFFER_LOAD: + return gl.amd.cdna4.buffer_load( + ptr=ptr, offsets=off.to(gl.int32), mask=mask, other=other, cache=".cg" + ) + return gl.load(ptr + off.to(gl.int64), mask=mask, other=other, cache_modifier=".cg") + + +@gluon.jit +def _fp8_to_f32(x_u8, FP8_FNUZ: gl.constexpr): + # gfx950's native e4m3 cvt is OCP (float8e4nv). fnuz (float8e4b8) -> f32 has no + # native cvt and lowers to a ~5x software unpack that spills; but fnuz -> bf16 is + # cheap and fp8 -> bf16 is exact (3 mantissa bits), so route fnuz through bf16. + if FP8_FNUZ: + return x_u8.to(gl.float8e4b8, bitcast=True).to(gl.bfloat16).to(gl.float32) + return x_u8.to(gl.float8e4nv, bitcast=True).to(gl.float32) + + +@gluon.jit +def _decode_tile( + q_dot, + cache_ptr, + cache_bf16_ptr, + indices_ptr, + seg_start, + k_start, + hi, + cs0, + num_rows, + m_i, + l_i, + acc, + head_mask, + qk_scale, + kv_smem, + offs_full, + offs_rope, + k_rng_slot, + qk_layout: gl.constexpr, + pv_layout: gl.constexpr, + k_layout: gl.constexpr, + v_layout: gl.constexpr, + p_layout: gl.constexpr, + gather_l: gl.constexpr, + gather_rope_l: gl.constexpr, + IS_FP8: gl.constexpr, + BLOCK_SIZE: gl.constexpr, + NOPE_DIM: gl.constexpr, + ROPE_DIM: gl.constexpr, + HEAD_SIZE: gl.constexpr, + BLOCK_M: gl.constexpr, + BLOCK_K: gl.constexpr, + HEAD_ALIGNED: gl.constexpr, + MASKED: gl.constexpr, + UNIFORM: gl.constexpr, + USE_BUFFER_LOAD: gl.constexpr, + HAS_INVALID: gl.constexpr, + FP8_FNUZ: gl.constexpr, +): + """One KV tile -> online-softmax update. MASKED=False (peeled full tiles) drops + the in-range / gather / score masking; MASKED=True (the tail) keeps it. When + HAS_INVALID, full tiles also clamp -1 sentinels in-bounds for the gather and + mask their scores to -inf (matching the tail's slot-validity handling).""" + neg_inf = float("-inf") + if not USE_BUFFER_LOAD: + cs0 = cs0.to(gl.int64) # >2 GB cache: 64-bit gather offsets (see _cache_load) + k_pos = k_start + k_rng_slot + if MASKED: + in_range = k_pos < hi + slot = gl.load(indices_ptr + seg_start + k_pos, mask=in_range, other=-1) + valid1d = in_range & (slot >= 0) & (slot < num_rows) + safe_slot = gl.where(valid1d, slot, 0) + elif HAS_INVALID: + slot = gl.load(indices_ptr + seg_start + k_pos) + valid1d = slot >= 0 # -1 sentinels: clamp in-bounds, mask score below + safe_slot = gl.where(valid1d, slot, 0) + else: + slot = gl.load(indices_ptr + seg_start + k_pos) + safe_slot = slot + block_idx = (safe_slot // BLOCK_SIZE).to(gl.int32) + pos = (safe_slot % BLOCK_SIZE).to(gl.int32) + block_idx_g = gl.convert_layout(block_idx, gl.SliceLayout(1, gather_l)) + pos_g = gl.convert_layout(pos, gl.SliceLayout(1, gather_l)) + if MASKED: + valid_g = gl.convert_layout(valid1d, gl.SliceLayout(1, gather_l)) + + if IS_FP8 and UNIFORM: + # uniform pool: one fp8 gather over the whole head + separate fp32 scales. + NGRP: gl.constexpr = HEAD_SIZE // 64 + kv_off = (block_idx_g * cs0 + pos_g * HEAD_SIZE)[:, None] + offs_full[None, :] + scl_off = (block_idx_g * NGRP)[:, None] + (offs_full[None, :] // 64) + if MASKED: + x_u8 = _cache_load( + cache_ptr, kv_off, USE_BUFFER_LOAD, mask=valid_g[:, None], other=0 + ) + sc = _cache_load( + cache_bf16_ptr, + scl_off, + USE_BUFFER_LOAD, + mask=valid_g[:, None], + other=0.0, + ) + else: + x_u8 = _cache_load(cache_ptr, kv_off, USE_BUFFER_LOAD) + sc = _cache_load(cache_bf16_ptr, scl_off, USE_BUFFER_LOAD) + kv_smem.store((_fp8_to_f32(x_u8, FP8_FNUZ) * sc).to(gl.bfloat16)) + elif IS_FP8: + # DSv4 packed fp8_ds_mla: NoPE fp8 + embedded UE8M0 + separate RoPE-bf16. + nope_off = (block_idx_g * cs0 + pos_g * 576)[:, None] + offs_full[None, :] + scl_off = (block_idx_g * cs0 + BLOCK_SIZE * 576 + pos_g * 8)[:, None] + ( + offs_full[None, :] // 64 + ) + if MASKED: + x_u8 = _cache_load( + cache_ptr, nope_off, USE_BUFFER_LOAD, mask=valid_g[:, None], other=0 + ) + exps = _cache_load( + cache_ptr, scl_off, USE_BUFFER_LOAD, mask=valid_g[:, None], other=127 + ) + else: + x_u8 = _cache_load(cache_ptr, nope_off, USE_BUFFER_LOAD) + exps = _cache_load(cache_ptr, scl_off, USE_BUFFER_LOAD) + x_fp8 = x_u8.to(gl.float8e4nv, bitcast=True) + scales = gl.exp2(exps.to(gl.float32) - 127.0) + k_nope = (x_fp8.to(gl.float32) * scales).to(gl.bfloat16) + kv_smem.store(k_nope) + block_idx_gr = gl.convert_layout(block_idx, gl.SliceLayout(1, gather_rope_l)) + pos_gr = gl.convert_layout(pos, gl.SliceLayout(1, gather_rope_l)) + rope_off = (block_idx_gr * (cs0 // 2) + pos_gr * 288 + 224)[ + :, None + ] + offs_rope[None, :] + if MASKED: + valid_gr = gl.convert_layout(valid1d, gl.SliceLayout(1, gather_rope_l)) + k_rope = _cache_load( + cache_bf16_ptr, + rope_off, + USE_BUFFER_LOAD, + mask=valid_gr[:, None], + other=0.0, + ) + else: + k_rope = _cache_load(cache_bf16_ptr, rope_off, USE_BUFFER_LOAD) + kv_smem.slice(NOPE_DIM, ROPE_DIM, dim=1).store(k_rope) + else: + off = (block_idx_g * cs0 + pos_g * HEAD_SIZE)[:, None] + offs_full[None, :] + if MASKED: + kv = _cache_load( + cache_bf16_ptr, off, USE_BUFFER_LOAD, mask=valid_g[:, None], other=0.0 + ) + else: + kv = _cache_load(cache_bf16_ptr, off, USE_BUFFER_LOAD) + kv_smem.store(kv) + + k = kv_smem.permute([1, 0]).load(k_layout) # [HEAD_SIZE, BLOCK_K] + S = gl.amd.cdna4.mfma( + q_dot, k, gl.zeros([BLOCK_M, BLOCK_K], gl.float32, layout=qk_layout) + ) + # exp2 softmax: qk_scale folds in log2(e) so we hit the HW exp2 directly. + # Running max stays in raw-score space; masked cols (-inf) give exp2=0. + COL_VALID: gl.constexpr = MASKED or HAS_INVALID # valid1d defined in both cases + NEED_MASK: gl.constexpr = COL_VALID or (not HEAD_ALIGNED) + if NEED_MASK: + if COL_VALID: + col_mask = gl.convert_layout(valid1d, gl.SliceLayout(0, qk_layout))[None, :] + if not HEAD_ALIGNED: + col_mask = ( + gl.convert_layout(head_mask, gl.SliceLayout(1, qk_layout))[:, None] + & col_mask + ) + else: # not HEAD_ALIGNED and no col invalidity -> head mask only + col_mask = gl.convert_layout(head_mask, gl.SliceLayout(1, qk_layout))[ + :, None + ] + S = gl.where(col_mask, S, neg_inf) + + m_block = gl.max(S, axis=1) + m_new = gl.maximum(m_i, m_block) + m_new = gl.where(m_new > neg_inf, m_new, 0.0) # guard all-masked rows + m_new_s = m_new * qk_scale + p = gl.exp2(S * qk_scale - m_new_s[:, None]) + alpha = gl.exp2(m_i * qk_scale - m_new_s) + l_new = l_i * alpha + gl.sum(p, axis=1) + + v = kv_smem.load(v_layout) # [BLOCK_K, HEAD_SIZE] + p_dot = gl.convert_layout(p.to(gl.bfloat16), p_layout) + alpha_pv = gl.convert_layout(alpha, gl.SliceLayout(1, pv_layout)) + acc = acc * alpha_pv[:, None] + acc = gl.amd.cdna4.mfma(p_dot, v, acc) + return m_new, l_new, acc + + +@gluon.jit +def _gd_fp8( + cache_ptr, + cache_bf16_ptr, + indices_ptr, + seg_start, + k_start, + cs0, + offs_full, + offs_rope, + k_rng_slot, + gather_l: gl.constexpr, + gather_rope_l: gl.constexpr, + BLOCK_SIZE: gl.constexpr, + HEAD_SIZE: gl.constexpr, + UNIFORM: gl.constexpr, + USE_BUFFER_LOAD: gl.constexpr, + HAS_INVALID: gl.constexpr, + FP8_FNUZ: gl.constexpr, +): + """Gather + dequant one full fp8 tile to bf16 regs (k_nope, k_rope). Split from + the LDS-write/MFMA so the pipeline runs this tile's gather+dequant while the + previous tile's MFMAs are on the matrix core -> the dequant hides behind them. + UNIFORM: one fp8 gather + separate fp32 scales, no RoPE (k_rope is a dead + alias). Else: packed NoPE fp8 (over-read) + embedded UE8M0 + RoPE bf16. + Returns (k_nope, k_rope, valid); valid is the per-slot >=0 mask (all-True and + DCE'd when HAS_INVALID is False), consumed by _qkpv_fp8 for the score mask.""" + slot = gl.load(indices_ptr + seg_start + k_start + k_rng_slot) + if not USE_BUFFER_LOAD: + cs0 = cs0.to(gl.int64) # >2 GB cache: 64-bit gather offsets (see _cache_load) + valid = slot >= 0 + if HAS_INVALID: + slot = gl.where(valid, slot, 0) # clamp -1 sentinels in-bounds for the gather + block_idx = (slot // BLOCK_SIZE).to(gl.int32) + pos = (slot % BLOCK_SIZE).to(gl.int32) + bg = gl.convert_layout(block_idx, gl.SliceLayout(1, gather_l)) + pg = gl.convert_layout(pos, gl.SliceLayout(1, gather_l)) + if UNIFORM: + NGRP: gl.constexpr = HEAD_SIZE // 64 + kv_off = (bg * cs0 + pg * HEAD_SIZE)[:, None] + offs_full[None, :] + scl_off = (bg * NGRP)[:, None] + (offs_full[None, :] // 64) + x_u8 = _cache_load(cache_ptr, kv_off, USE_BUFFER_LOAD) + sc = _cache_load(cache_bf16_ptr, scl_off, USE_BUFFER_LOAD) + k_nope = (_fp8_to_f32(x_u8, FP8_FNUZ) * sc).to(gl.bfloat16) + k_rope = k_nope # unused for UNIFORM (rope slice-store skipped) -> DCE'd + else: + nope_off = (bg * cs0 + pg * 576)[:, None] + offs_full[None, :] + scl_off = (bg * cs0 + BLOCK_SIZE * 576 + pg * 8)[:, None] + ( + offs_full[None, :] // 64 + ) + x_u8 = _cache_load(cache_ptr, nope_off, USE_BUFFER_LOAD) + exps = _cache_load(cache_ptr, scl_off, USE_BUFFER_LOAD) + k_nope = ( + x_u8.to(gl.float8e4nv, bitcast=True).to(gl.float32) + * gl.exp2(exps.to(gl.float32) - 127.0) + ).to(gl.bfloat16) + bgr = gl.convert_layout(block_idx, gl.SliceLayout(1, gather_rope_l)) + pgr = gl.convert_layout(pos, gl.SliceLayout(1, gather_rope_l)) + rope_off = (bgr * (cs0 // 2) + pgr * 288 + 224)[:, None] + offs_rope[None, :] + k_rope = _cache_load(cache_bf16_ptr, rope_off, USE_BUFFER_LOAD) + return k_nope, k_rope, valid + + +@gluon.jit +def _qkpv_fp8( + k_nope, + k_rope, + valid, + q_dot, + m_i, + l_i, + acc, + head_mask, + qk_scale, + kv_smem, + qk_layout: gl.constexpr, + pv_layout: gl.constexpr, + k_layout: gl.constexpr, + v_layout: gl.constexpr, + p_layout: gl.constexpr, + NOPE_DIM: gl.constexpr, + ROPE_DIM: gl.constexpr, + HEAD_SIZE: gl.constexpr, + BLOCK_M: gl.constexpr, + BLOCK_K: gl.constexpr, + HEAD_ALIGNED: gl.constexpr, + UNIFORM: gl.constexpr, + HAS_INVALID: gl.constexpr, +): + """Write a prefetched bf16 tile to LDS, then QK -> softmax -> PV. UNIFORM skips + the RoPE slice-store (the whole head is one fp8 tile). When HAS_INVALID, mask the + columns of -1-sentinel slots (``valid`` from _gd_fp8) to -inf.""" + neg_inf = float("-inf") + kv_smem.store(k_nope) + if not UNIFORM: + kv_smem.slice(NOPE_DIM, ROPE_DIM, dim=1).store(k_rope) + k = kv_smem.permute([1, 0]).load(k_layout) + S = gl.amd.cdna4.mfma( + q_dot, k, gl.zeros([BLOCK_M, BLOCK_K], gl.float32, layout=qk_layout) + ) + NEED_MASK: gl.constexpr = HAS_INVALID or (not HEAD_ALIGNED) + if NEED_MASK: + if HAS_INVALID: + col_mask = gl.convert_layout(valid, gl.SliceLayout(0, qk_layout))[None, :] + if not HEAD_ALIGNED: + col_mask = ( + gl.convert_layout(head_mask, gl.SliceLayout(1, qk_layout))[:, None] + & col_mask + ) + else: + col_mask = gl.convert_layout(head_mask, gl.SliceLayout(1, qk_layout))[ + :, None + ] + S = gl.where(col_mask, S, neg_inf) + m_block = gl.max(S, axis=1) + m_new = gl.maximum(m_i, m_block) + m_new = gl.where(m_new > neg_inf, m_new, 0.0) + m_new_s = m_new * qk_scale + p = gl.exp2(S * qk_scale - m_new_s[:, None]) + alpha = gl.exp2(m_i * qk_scale - m_new_s) + l_new = l_i * alpha + gl.sum(p, axis=1) + v = kv_smem.load(v_layout) + p_dot = gl.convert_layout(p.to(gl.bfloat16), p_layout) + alpha_pv = gl.convert_layout(alpha, gl.SliceLayout(1, pv_layout)) + acc = acc * alpha_pv[:, None] + acc = gl.amd.cdna4.mfma(p_dot, v, acc) + return m_new, l_new, acc + + +@gluon.jit +def _process_segment( + q_dot, + cache_ptr, + cache_bf16_ptr, + indices_ptr, + seg_start, + lo, + hi, + cs0, + num_rows, + m_i, + l_i, + acc, + head_mask, + qk_scale, + kv_smem, + qk_layout: gl.constexpr, + pv_layout: gl.constexpr, + k_layout: gl.constexpr, + v_layout: gl.constexpr, + p_layout: gl.constexpr, + gather_l: gl.constexpr, + gather_rope_l: gl.constexpr, + slot_l: gl.constexpr, + IS_FP8: gl.constexpr, + BLOCK_SIZE: gl.constexpr, + NOPE_DIM: gl.constexpr, + ROPE_DIM: gl.constexpr, + HEAD_SIZE: gl.constexpr, + BLOCK_M: gl.constexpr, + BLOCK_K: gl.constexpr, + HEAD_ALIGNED: gl.constexpr, + UNIFORM: gl.constexpr, + USE_BUFFER_LOAD: gl.constexpr, + HAS_INVALID: gl.constexpr, + FP8_FNUZ: gl.constexpr, +): + offs_full = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, gather_l)) + offs_rope = gl.arange(0, ROPE_DIM, layout=gl.SliceLayout(0, gather_rope_l)) + k_rng_slot = gl.arange(0, BLOCK_K, layout=slot_l) + + # Peel the (possibly partial) last tile: [lo, hi_full) are full BLOCK_K tiles + # whose slots are all valid -> mask-free. Only the peeled tail carries masking. + hi_full = lo + ((hi - lo) // BLOCK_K) * BLOCK_K + + if IS_FP8: + # Prologue/epilogue software pipeline: gather + DEQUANT tile i into bf16 + # registers while the previous tile's MFMAs run (VALU+matrix-core co-issue + # hides the quant cost). One LDS buffer, so the LDS write of tile i still + # waits on tile i-1's reads, but the gather+dequant is off the critical path. + n_full = (hi_full - lo) // BLOCK_K + if n_full > 0: + kn, kr, vld = _gd_fp8( + cache_ptr, + cache_bf16_ptr, + indices_ptr, + seg_start, + lo, + cs0, + offs_full, + offs_rope, + k_rng_slot, + gather_l, + gather_rope_l, + BLOCK_SIZE, + HEAD_SIZE, + UNIFORM, + USE_BUFFER_LOAD, + HAS_INVALID, + FP8_FNUZ, + ) + for i in range(1, n_full): + kn2, kr2, vld2 = _gd_fp8( + cache_ptr, + cache_bf16_ptr, + indices_ptr, + seg_start, + lo + i * BLOCK_K, + cs0, + offs_full, + offs_rope, + k_rng_slot, + gather_l, + gather_rope_l, + BLOCK_SIZE, + HEAD_SIZE, + UNIFORM, + USE_BUFFER_LOAD, + HAS_INVALID, + FP8_FNUZ, + ) + m_i, l_i, acc = _qkpv_fp8( + kn, + kr, + vld, + q_dot, + m_i, + l_i, + acc, + head_mask, + qk_scale, + kv_smem, + qk_layout, + pv_layout, + k_layout, + v_layout, + p_layout, + NOPE_DIM, + ROPE_DIM, + HEAD_SIZE, + BLOCK_M, + BLOCK_K, + HEAD_ALIGNED, + UNIFORM, + HAS_INVALID, + ) + kn, kr, vld = kn2, kr2, vld2 + m_i, l_i, acc = _qkpv_fp8( + kn, + kr, + vld, + q_dot, + m_i, + l_i, + acc, + head_mask, + qk_scale, + kv_smem, + qk_layout, + pv_layout, + k_layout, + v_layout, + p_layout, + NOPE_DIM, + ROPE_DIM, + HEAD_SIZE, + BLOCK_M, + BLOCK_K, + HEAD_ALIGNED, + UNIFORM, + HAS_INVALID, + ) + else: + for k_start in range(lo, hi_full, BLOCK_K): + m_i, l_i, acc = _decode_tile( + q_dot, + cache_ptr, + cache_bf16_ptr, + indices_ptr, + seg_start, + k_start, + hi, + cs0, + num_rows, + m_i, + l_i, + acc, + head_mask, + qk_scale, + kv_smem, + offs_full, + offs_rope, + k_rng_slot, + qk_layout, + pv_layout, + k_layout, + v_layout, + p_layout, + gather_l, + gather_rope_l, + IS_FP8, + BLOCK_SIZE, + NOPE_DIM, + ROPE_DIM, + HEAD_SIZE, + BLOCK_M, + BLOCK_K, + HEAD_ALIGNED, + False, + UNIFORM, + USE_BUFFER_LOAD, + HAS_INVALID, + FP8_FNUZ, + ) + + if hi_full < hi: + m_i, l_i, acc = _decode_tile( + q_dot, + cache_ptr, + cache_bf16_ptr, + indices_ptr, + seg_start, + hi_full, + hi, + cs0, + num_rows, + m_i, + l_i, + acc, + head_mask, + qk_scale, + kv_smem, + offs_full, + offs_rope, + k_rng_slot, + qk_layout, + pv_layout, + k_layout, + v_layout, + p_layout, + gather_l, + gather_rope_l, + IS_FP8, + BLOCK_SIZE, + NOPE_DIM, + ROPE_DIM, + HEAD_SIZE, + BLOCK_M, + BLOCK_K, + HEAD_ALIGNED, + True, + UNIFORM, + USE_BUFFER_LOAD, + HAS_INVALID, + FP8_FNUZ, + ) + return m_i, l_i, acc + + +_pa_decode_sparse_repr = make_kernel_repr( + "_pa_decode_sparse", + ["BLOCK_M", "BLOCK_K", "HEAD_SIZE", "NUM_SPLITS", "UNIFORM", "MAIN_IS_FP8"], +) + + +@gluon.jit(repr=_pa_decode_sparse_repr) +def _pa_decode_sparse( + q_ptr, + main_cache_ptr, + main_cache_bf16_ptr, + main_indices_ptr, + main_indptr_ptr, + extra_cache_ptr, + extra_cache_bf16_ptr, + extra_indices_ptr, + extra_indptr_ptr, + attn_sink_ptr, + out_ptr, + part_m_ptr, + part_l_ptr, + part_acc_ptr, + scale: gl.constexpr, + q_stride0: gl.constexpr, + q_stride1: gl.constexpr, + out_stride0: gl.constexpr, + out_stride1: gl.constexpr, + main_cs0, + extra_cs0, + main_num_rows, + extra_num_rows, + pm_stride0: gl.constexpr, + pm_stride_s: gl.constexpr, + pa_stride0: gl.constexpr, + pa_stride_s: gl.constexpr, + pa_stride_h: gl.constexpr, + num_heads: gl.constexpr, + HAS_EXTRA: gl.constexpr, + HAS_SINK: gl.constexpr, + MAIN_IS_FP8: gl.constexpr, + EXTRA_IS_FP8: gl.constexpr, + MAIN_BLOCK_SIZE: gl.constexpr, + EXTRA_BLOCK_SIZE: gl.constexpr, + NOPE_DIM: gl.constexpr, + ROPE_DIM: gl.constexpr, + HEAD_SIZE: gl.constexpr, + BLOCK_M: gl.constexpr, + BLOCK_K: gl.constexpr, + NUM_SPLITS: gl.constexpr, + HEAD_ALIGNED: gl.constexpr, + MFMA_K: gl.constexpr, + UNIFORM: gl.constexpr, + USE_BUFFER_LOAD: gl.constexpr, + HAS_INVALID: gl.constexpr, + FP8_FNUZ: gl.constexpr, +): + """One program = (query t, split, head-block). Two-loop: main (SWA) then + extra (top-k). NUM_SPLITS==1 writes the output directly; NUM_SPLITS>1 stores + un-normalized partials for the reduce kernel. HAS_INVALID gates -1-sentinel + handling (clamp + score mask) on the full-tile fast paths.""" + NUM_WARPS: gl.constexpr = gl.num_warps() + query_idx = gl.program_id(0) + split_id = gl.program_id(1) + pid_h = gl.program_id(2) + + qk_layout: gl.constexpr = gl.amd.AMDMFMALayout( + version=4, + instr_shape=[16, 16, MFMA_K], + transposed=True, + warps_per_cta=[1, NUM_WARPS], + ) + pv_layout: gl.constexpr = gl.amd.AMDMFMALayout( + version=4, + instr_shape=[16, 16, MFMA_K], + transposed=True, + warps_per_cta=[1, NUM_WARPS], + ) + KW: gl.constexpr = MFMA_K // 2 + q_layout: gl.constexpr = gl.DotOperandLayout(0, qk_layout, KW) + k_layout: gl.constexpr = gl.DotOperandLayout(1, qk_layout, KW) + p_layout: gl.constexpr = gl.DotOperandLayout(0, pv_layout, KW) + v_layout: gl.constexpr = gl.DotOperandLayout(1, pv_layout, KW) + + # 16 uint8 = 128-bit fp8 gather loads + GSPT: gl.constexpr = 16 + gather_l: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, GSPT], + threads_per_warp=[8, 8], + warps_per_cta=[1, NUM_WARPS], + order=[1, 0], + ) + gather_rope_l: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[8, 8], + warps_per_cta=[1, NUM_WARPS], + order=[1, 0], + ) + slot_l: gl.constexpr = gl.SliceLayout(1, gather_l) + blocked_q: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[8, 8], + warps_per_cta=[1, NUM_WARPS], + order=[1, 0], + ) + kv_shared: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( + [[HEAD_SIZE, 8]], [BLOCK_K, HEAD_SIZE], [1, 0] + ) + + h_off = pid_h * BLOCK_M + + # ---- load Q [BLOCK_M, HEAD_SIZE] ---- + offs_m_q = gl.arange(0, BLOCK_M, layout=gl.SliceLayout(1, blocked_q)) + offs_d_q = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, blocked_q)) + h_q = h_off + offs_m_q + h_mask_q = h_q < num_heads + q_off = (query_idx * q_stride0 + h_q[:, None] * q_stride1 + offs_d_q[None, :]).to( + gl.int32 + ) + q = gl.amd.cdna4.buffer_load( + ptr=q_ptr, offsets=q_off, mask=h_mask_q[:, None], other=0.0 + ) + q_dot = gl.convert_layout(q, q_layout) + + # head mask in pv-slice layout (for output / partial masking) + offs_m_pv = gl.arange(0, BLOCK_M, layout=gl.SliceLayout(1, pv_layout)) + h_pv = h_off + offs_m_pv + head_mask_pv = h_pv < num_heads + + # ---- online-softmax state ---- + m_i = gl.full( + [BLOCK_M], float("-inf"), gl.float32, layout=gl.SliceLayout(1, qk_layout) + ) + l_i = gl.zeros([BLOCK_M], gl.float32, layout=gl.SliceLayout(1, qk_layout)) + acc = gl.zeros([BLOCK_M, HEAD_SIZE], gl.float32, layout=pv_layout) + + kv_smem = gl.allocate_shared_memory(gl.bfloat16, [BLOCK_K, HEAD_SIZE], kv_shared) + + # exp2 softmax: fold scale*log2(e) into the loop exponent; keep raw `scale` + # for the sink/normalization (sink is a scaled-score-space logit). + RCP_LN2: gl.constexpr = 1.4426950408889634 + qk_scale = scale * RCP_LN2 + + # ---- main (SWA) segment ---- + main_start = gl.load(main_indptr_ptr + query_idx) + main_end = gl.load(main_indptr_ptr + query_idx + 1) + main_len = main_end - main_start + main_chunk = (main_len + NUM_SPLITS - 1) // NUM_SPLITS + main_lo = split_id * main_chunk + main_hi = gl.minimum(main_lo + main_chunk, main_len) + m_i, l_i, acc = _process_segment( + q_dot, + main_cache_ptr, + main_cache_bf16_ptr, + main_indices_ptr, + main_start, + main_lo, + main_hi, + main_cs0, + main_num_rows, + m_i, + l_i, + acc, + head_mask_pv, + qk_scale, + kv_smem, + qk_layout, + pv_layout, + k_layout, + v_layout, + p_layout, + gather_l, + gather_rope_l, + slot_l, + MAIN_IS_FP8, + MAIN_BLOCK_SIZE, + NOPE_DIM, + ROPE_DIM, + HEAD_SIZE, + BLOCK_M, + BLOCK_K, + HEAD_ALIGNED, + UNIFORM, + USE_BUFFER_LOAD, + HAS_INVALID, + FP8_FNUZ, + ) + + if HAS_EXTRA: + extra_start = gl.load(extra_indptr_ptr + query_idx) + extra_end = gl.load(extra_indptr_ptr + query_idx + 1) + extra_len = extra_end - extra_start + extra_chunk = (extra_len + NUM_SPLITS - 1) // NUM_SPLITS + extra_lo = split_id * extra_chunk + extra_hi = gl.minimum(extra_lo + extra_chunk, extra_len) + m_i, l_i, acc = _process_segment( + q_dot, + extra_cache_ptr, + extra_cache_bf16_ptr, + extra_indices_ptr, + extra_start, + extra_lo, + extra_hi, + extra_cs0, + extra_num_rows, + m_i, + l_i, + acc, + head_mask_pv, + qk_scale, + kv_smem, + qk_layout, + pv_layout, + k_layout, + v_layout, + p_layout, + gather_l, + gather_rope_l, + slot_l, + EXTRA_IS_FP8, + EXTRA_BLOCK_SIZE, + NOPE_DIM, + ROPE_DIM, + HEAD_SIZE, + BLOCK_M, + BLOCK_K, + HEAD_ALIGNED, + UNIFORM, + USE_BUFFER_LOAD, + HAS_INVALID, + FP8_FNUZ, + ) + + # m_i/l_i are in SliceLayout(1, qk_layout); acc in pv_layout. Move the row + # reductions into pv-slice space for output/partials. + m_pv = gl.convert_layout(m_i, gl.SliceLayout(1, pv_layout)) + l_pv = gl.convert_layout(l_i, gl.SliceLayout(1, pv_layout)) + + if NUM_SPLITS == 1: + if HAS_SINK: + # m_pv is the RAW row-max; the sink is a scaled-score logit. Combine + # in scaled space (Ms = scale*m) using exp2. + Ms = m_pv * scale + sink = gl.amd.cdna4.buffer_load( + ptr=attn_sink_ptr, offsets=h_pv, mask=head_mask_pv, other=float("-inf") + ).to(gl.float32) + m_final = gl.maximum(Ms, sink) + alpha = gl.exp2((Ms - m_final) * RCP_LN2) + l_final = l_pv * alpha + gl.exp2((sink - m_final) * RCP_LN2) + acc = acc * alpha[:, None] + else: + l_final = l_pv + out = acc / l_final[:, None] + # store output + offs_d_o = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, pv_layout)) + o_off = ( + query_idx * out_stride0 + h_pv[:, None] * out_stride1 + offs_d_o[None, :] + ).to(gl.int32) + gl.amd.cdna4.buffer_store( + out.to(out_ptr.dtype.element_ty), + ptr=out_ptr, + offsets=o_off, + mask=head_mask_pv[:, None], + ) + else: + # store un-normalized partials for the reduce kernel. m is stored in the + # base-2 exponent domain (row-max * softmax_scale * log2e) so the reduce/ + # skip_reduce partials match the triton convention. + pm_base = query_idx * pm_stride0 + split_id * pm_stride_s + gl.amd.cdna4.buffer_store( + m_pv * (scale * RCP_LN2), + ptr=part_m_ptr + pm_base, + offsets=h_pv.to(gl.int32), + mask=head_mask_pv, + ) + gl.amd.cdna4.buffer_store( + l_pv, ptr=part_l_ptr + pm_base, offsets=h_pv.to(gl.int32), mask=head_mask_pv + ) + offs_d_a = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, pv_layout)) + a_base = query_idx * pa_stride0 + split_id * pa_stride_s + a_off = (a_base + h_pv[:, None] * pa_stride_h + offs_d_a[None, :]).to(gl.int32) + gl.amd.cdna4.buffer_store( + acc, ptr=part_acc_ptr, offsets=a_off, mask=head_mask_pv[:, None] + ) + + +_pa_decode_sparse_reduce_repr = make_kernel_repr( + "_pa_decode_sparse_reduce", + ["BLOCK_M", "HEAD_SIZE", "NUM_SPLITS"], +) + + +@gluon.jit(repr=_pa_decode_sparse_reduce_repr) +def _pa_decode_sparse_reduce( + part_m_ptr, + part_l_ptr, + part_acc_ptr, + attn_sink_ptr, + out_ptr, + out_stride0: gl.constexpr, + out_stride1: gl.constexpr, + pm_stride0: gl.constexpr, + pm_stride_s: gl.constexpr, + pa_stride0: gl.constexpr, + pa_stride_s: gl.constexpr, + pa_stride_h: gl.constexpr, + num_heads: gl.constexpr, + HAS_SINK: gl.constexpr, + HEAD_SIZE: gl.constexpr, + BLOCK_M: gl.constexpr, + NUM_SPLITS: gl.constexpr, + HEAD_ALIGNED: gl.constexpr, +): + """Split-KV combine: merge the per-split partials, fold the attn sink, and write + the final output. Partials store m in the base-2 exponent domain (row-max * + softmax_scale * log2e), matching the triton reduce. Grid: (num_queries, heads_blocks). + """ + NUM_WARPS: gl.constexpr = gl.num_warps() + RCP_LN2: gl.constexpr = 1.4426950408889634 + query_idx = gl.program_id(0) + pid_h = gl.program_id(1) + + BLK: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[8, 8], + warps_per_cta=[1, NUM_WARPS], + order=[1, 0], + ) + row_l: gl.constexpr = gl.SliceLayout(1, BLK) # [BLOCK_M] + + h_off = pid_h * BLOCK_M + offs_m = gl.arange(0, BLOCK_M, layout=row_l) + h = h_off + offs_m + head_mask = h < num_heads + offs_d = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, BLK)) + + neg_inf = float("-inf") + m_final = gl.full([BLOCK_M], neg_inf, gl.float32, layout=row_l) + # pass 1: global max over splits + for s in range(NUM_SPLITS): + base = query_idx * pm_stride0 + s * pm_stride_s + m_s = gl.amd.cdna4.buffer_load( + ptr=part_m_ptr + base, offsets=h, mask=head_mask, other=neg_inf + ) + m_final = gl.maximum(m_final, m_s) # m_s already in base-2 exponent domain + if HAS_SINK: + sink = gl.amd.cdna4.buffer_load( + ptr=attn_sink_ptr, offsets=h, mask=head_mask, other=neg_inf + ).to(gl.float32) + m_final = gl.maximum(m_final, sink * RCP_LN2) # lift sink to base-2 + + # pass 2: weighted sums + l_final = gl.zeros([BLOCK_M], gl.float32, layout=row_l) + acc = gl.zeros([BLOCK_M, HEAD_SIZE], gl.float32, layout=BLK) + for s in range(NUM_SPLITS): + base = query_idx * pm_stride0 + s * pm_stride_s + m_s = gl.amd.cdna4.buffer_load( + ptr=part_m_ptr + base, offsets=h, mask=head_mask, other=neg_inf + ) + l_s = gl.amd.cdna4.buffer_load( + ptr=part_l_ptr + base, offsets=h, mask=head_mask, other=0.0 + ) + w = gl.exp2(m_s - m_final) + l_final = l_final + w * l_s + a_base = query_idx * pa_stride0 + s * pa_stride_s + a_off = (a_base + h[:, None] * pa_stride_h + offs_d[None, :]).to(gl.int32) + acc_s = gl.amd.cdna4.buffer_load( + ptr=part_acc_ptr, offsets=a_off, mask=head_mask[:, None], other=0.0 + ) + acc = acc + w[:, None] * acc_s + + if HAS_SINK: + sink = gl.amd.cdna4.buffer_load( + ptr=attn_sink_ptr, offsets=h, mask=head_mask, other=neg_inf + ).to(gl.float32) + l_final = l_final + gl.exp2(sink * RCP_LN2 - m_final) + + out = acc / l_final[:, None] + o_off = (query_idx * out_stride0 + h[:, None] * out_stride1 + offs_d[None, :]).to( + gl.int32 + ) + gl.amd.cdna4.buffer_store( + out.to(out_ptr.dtype.element_ty), + ptr=out_ptr, + offsets=o_off, + mask=head_mask[:, None], + ) diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 49f10d351b9..5585ab0cc68 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -8,23 +8,40 @@ This module exposes ``pa_decode_sparse`` — a 3D split-K + widened-BLOCK_H + pipelined-K-loop variant suitable for sparse decode (e.g. V4 top-k gather) where each token's K range is an unordered subset of a unified KV pool. + +On gfx950 (CDNA4) DeepSeek-V4 sparse-MLA decode has a dedicated gluon +implementation (bottom of this module): ``pa_decode_sparse`` routes all formats +to the merged ``_pa_decode_sparse_gfx950_gluon`` driver -- packed fp8_ds_mla / +bf16 block cache (3D; optional SWA+top-k two-loop via ``extra_*``) and the +uniform fp8 / bf16 pool (2D). """ -from typing import Optional +import math import torch import triton +from aiter.ops.triton._gluon_kernels.gfx1250.attention.pa_decode_sparse import ( + _pa_decode_sparse as gluon_pa_decode_sparse, +) +from aiter.ops.triton._gluon_kernels.gfx1250.attention.pa_decode_sparse import ( + _pa_decode_sparse_reduce as gluon_pa_decode_sparse_reduce, +) +from aiter.ops.triton._gluon_kernels.gfx950.attention.pa_decode_sparse import ( + _pa_decode_sparse as _pa_decode_sparse_gfx950, +) +from aiter.ops.triton._gluon_kernels.gfx950.attention.pa_decode_sparse import ( + _pa_decode_sparse_reduce as _pa_decode_sparse_reduce_gfx950, +) from aiter.ops.triton._triton_kernels.attention.pa_decode_sparse import ( _pa_decode_sparse as triton_pa_decode_sparse, +) +from aiter.ops.triton._triton_kernels.attention.pa_decode_sparse import ( _pa_decode_sparse_reduce as triton_pa_decode_sparse_reduce, ) from aiter.ops.triton.utils._triton import arch_info +from aiter.ops.triton.utils.device_info import get_num_sms from aiter.ops.triton.utils.logger import AiterTritonLogger -from aiter.ops.triton._gluon_kernels.gfx1250.attention.pa_decode_sparse import ( - _pa_decode_sparse as gluon_pa_decode_sparse, - _pa_decode_sparse_reduce as gluon_pa_decode_sparse_reduce, -) DEVICE_ARCH = arch_info.get_arch() @@ -42,12 +59,16 @@ def pa_decode_sparse( kv_indptr: torch.Tensor, attn_sink: torch.Tensor, softmax_scale: float, - kv_scales: Optional[torch.Tensor] = None, - block_h: Optional[int] = None, - kv_splits: Optional[int] = None, - has_invalid: Optional[bool] = True, - skip_reduce: Optional[bool] = False, - USE_EXP2: Optional[bool] = None, + kv_scales: torch.Tensor | None = None, + block_h: int | None = None, + kv_splits: int | None = None, + has_invalid: bool | None = True, + skip_reduce: bool | None = False, + USE_EXP2: bool | None = None, + *, + extra_cache: torch.Tensor | None = None, + extra_indices: torch.Tensor | None = None, + extra_indptr: torch.Tensor | None = None, ) -> torch.Tensor: """Sparse paged-decode attention with split-K + widened BLOCK_H. @@ -72,6 +93,13 @@ def pa_decode_sparse( ``kv_splits == 1`` (the single-CTA path already produces the final ``out`` directly). Useful for profiling the main kernel in isolation and for callers that fold the reduce into a downstream op. + extra_cache/extra_indices/extra_indptr: gfx950 packed-only — the SWA+top-k + two-loop's second (top-k) cache + index set; must be None otherwise. + + On gfx950 the DSv4 gluon driver handles this: a 3D ``unified_kv`` selects the + packed fp8_ds_mla / bf16 block cache (``extra_*`` = the two-loop), a 2D one the + uniform pool (``kv_scales`` present = fp8). ``kv_splits``/``skip_reduce`` are + honored; ``block_h`` and fp16 ``q`` fall through to the triton path. Returns: ``[N, H, D]`` attention output, same dtype as ``q``. When @@ -92,6 +120,54 @@ def pa_decode_sparse( if q.dtype not in (torch.bfloat16, torch.float16): raise RuntimeError(f"pa_decode_sparse expects fp16/bf16 q, got {q.dtype}") + # gfx950: route to the merged DSv4 sparse-MLA gluon driver. Format is inferred + # from the cache: 3D -> packed fp8_ds_mla / bf16 block cache (optional SWA+top-k + # two-loop via extra_*); 2D -> uniform pool (OCP fp8 + fp32 kv_scales, or bf16). + # kv_splits and skip_reduce are honored here; block_h and fp16 q fall through to + # the triton path below (the gluon kernel is bf16-only: bf16 LDS + bf16 MFMA). + if DEVICE_ARCH == "gfx950" and block_h is None and q.dtype == torch.bfloat16: + if unified_kv.ndim == 3: + _ok = kv_scales is None and ( + unified_kv.dtype == torch.uint8 or unified_kv.dtype == q.dtype + ) + else: + _fp8 = unified_kv.dtype in ( + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + torch.uint8, + ) + _ok = (kv_scales is not None and _fp8) or ( + kv_scales is None and unified_kv.dtype == q.dtype + ) + # fnuz vs OCP e4m3 (2D fp8 only) selects the in-kernel dequant bias. + fp8_fnuz = unified_kv.ndim == 2 and unified_kv.dtype == torch.float8_e4m3fnuz + if _ok: + cache = ( + unified_kv.view(torch.uint8) + if (unified_kv.ndim == 2 and kv_scales is not None) + else unified_kv + ) + return _pa_decode_sparse_gfx950_gluon( + q, + cache, + kv_scales, + kv_indices, + kv_indptr, + softmax_scale, + attn_sink, + extra_cache=extra_cache, + extra_indices=extra_indices, + extra_indptr=extra_indptr, + kv_splits=kv_splits, + skip_reduce=skip_reduce, + has_invalid=bool(has_invalid), + fp8_fnuz=fp8_fnuz, + ) + + assert ( + extra_cache is None and extra_indices is None and extra_indptr is None + ), "extra_cache/extra_indices/extra_indptr are gfx950 packed-only" + quant_kv = kv_scales is not None if quant_kv: assert unified_kv.dtype == _FP8_DTYPE, ( @@ -310,3 +386,259 @@ def pa_decode_sparse( waves_per_eu=reduce_waves_per_eu, ) return out + + +def _as_int32_contiguous_1d(x: torch.Tensor) -> torch.Tensor: + if x.dtype == torch.int32 and x.ndim == 1 and x.is_contiguous(): + return x + return x.to(torch.int32).contiguous() + + +def _decode_num_splits( + num_queries, heads_blocks, avg_main=0.0, avg_extra=0.0, block_k=64 +): + """Pick the split-K count by minimizing a cost model of the decode work: + + cost(s) = waves(s) * iters(s) + GAMMA * s + DELTA * fill(s) + s = # splits + Tuned on gfx950 for DSv4 decode (H=16, D=512, BLOCK_K=64); split count is + capped at 16. + """ + cu = max(1, get_num_sms()) + base = max(1, num_queries * heads_blocks) + GAMMA, DELTA, FILL_CU = 0.32, 2.0, 0.75 + thr = FILL_CU * cu + best_splits, best_cost = 1, None + for splits in range(1, 17): + m_it = math.ceil(math.ceil(avg_main / splits) / block_k) if avg_main > 0 else 0 + e_it = ( + math.ceil(math.ceil(avg_extra / splits) / block_k) if avg_extra > 0 else 0 + ) + waves = (base * splits + cu - 1) // cu + fill = max(0.0, 1.0 - base * splits / thr) / splits + cost = waves * (m_it + e_it) + GAMMA * splits + DELTA * fill + if best_cost is None or cost < best_cost - 1e-9: + best_splits, best_cost = splits, cost + return best_splits + + +def _pa_decode_sparse_gfx950_gluon( + q, + cache, + cache_scales, + indices, + indptr, + scale, + attn_sink, + extra_cache=None, + extra_indices=None, + extra_indptr=None, + kv_splits=None, + skip_reduce=False, + has_invalid=False, + fp8_fnuz=False, +): + """Merged gfx950 gluon DSv4 sparse-MLA decode driver. Format from ``cache.ndim``: + 3D [nb, block, ...] -> packed fp8_ds_mla (uint8: 448 NoPE fp8 e4m3 OCP + + embedded UE8M0 per-64 scale + 64 RoPE bf16) or a bf16 + block cache; pass ``extra_*`` for the SWA+top-k two-loop, + else a single segment. + 2D [pages, D] -> uniform pool: fp8 (uint8) + ``cache_scales`` + [pages, D//64] fp32, or bf16 (``cache_scales`` None). + ``kv_splits`` overrides the split-K count. ``skip_reduce`` (only takes effect when + the chosen split count > 1) returns the pre-reduce partials + ``(part_acc, part_m, part_l)`` -- shapes ``([N, S, H, D], [N, S, H], [N, S, H])`` + fp32; ``m`` is the row-max in the base-2 exponent domain (row-max * softmax_scale + * log2e) and ``l``/``acc`` are per-split un-normalized -- same convention as the + triton skip_reduce partials -- instead of the final ``[N, H, D]`` output. + ``has_invalid`` (default False): when True, -1 sentinels anywhere in a token's + index range are clamped in-bounds for the gather and masked out of the softmax. + ``fp8_fnuz`` (uniform-pool fp8 only): fp8 e4m3 flavor -- False = OCP (bias 7), + True = fnuz (bias 8); selects the in-kernel dequant. Packed fp8_ds_mla is OCP. + """ + assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}" + assert DEVICE_ARCH == "gfx950", "gluon DSv4 decode kernel is gfx950-only" + + # Tuned launch config (gfx950 / MI355), inlined. BLOCK_M = heads per MFMA M-tile; + # BLOCK_K = KV tile; num_warps = BLOCK_K // 16 (warps tile the dot-N, MFMA N=16). + BLOCK_M, BLOCK_K, MFMA_K, waves_per_eu = 16, 64, 16, 0 + num_warps = BLOCK_K // 16 + NOPE_DIM, ROPE_DIM = 448, 64 + MAX_BYTES = 2**31 - 1 # buffer_load 32-bit offset cap; larger -> gl.load int64 + + num_queries, num_heads, head_dim = q.shape + indices = _as_int32_contiguous_1d(indices) + indptr = _as_int32_contiguous_1d(indptr) + has_sink = attn_sink is not None + attn_sink = ( + attn_sink.contiguous().to(torch.float32) + if has_sink + else torch.empty(1, device=q.device, dtype=torch.float32) + ) + + if cache.ndim == 2: + # uniform pool: one fp8 gather over the whole head + separate fp32 scales, + # or bf16. page_size=1 -> block_idx=slot, pos=0; scales ride the bf16 ptr. + UNIFORM = True + main_is_fp8 = cache.dtype == torch.uint8 + if main_is_fp8: + assert cache_scales is not None and cache_scales.dtype == torch.float32 + main_bf16 = cache_scales.contiguous() + else: + main_bf16 = cache + # if HAS_EXTRA=False, reuse main tensors as unread placeholders. + extra_cache, extra_bf16, extra_indices, extra_indptr = ( + cache, + main_bf16, + indices, + indptr, + ) + extra_is_fp8 = main_is_fp8 + has_extra = False + main_block, extra_block = 1, 1 + nope_dim = head_dim + main_num_rows = extra_num_rows = cache.shape[0] + cache_bytes = cache.nelement() * cache.element_size() + avg_main = indices.numel() / max(1, num_queries) # one segment; no extra + avg_extra = 0.0 + else: + # packed fp8_ds_mla [nb, block, 584] (embedded scale) or bf16 block cache. + UNIFORM = False + main_is_fp8 = cache.dtype == torch.uint8 + main_bf16 = cache.view(torch.bfloat16) if main_is_fp8 else cache + has_extra = ( + extra_cache is not None + and extra_indices is not None + and extra_indptr is not None + ) + if has_extra: + extra_indices = _as_int32_contiguous_1d(extra_indices) + extra_indptr = _as_int32_contiguous_1d(extra_indptr) + else: + extra_cache, extra_indices, extra_indptr = cache, indices, indptr + extra_is_fp8 = extra_cache.dtype == torch.uint8 + extra_bf16 = extra_cache.view(torch.bfloat16) if extra_is_fp8 else extra_cache + main_block, extra_block = cache.shape[1], extra_cache.shape[1] + nope_dim = NOPE_DIM + main_num_rows = cache.shape[0] * cache.shape[1] + extra_num_rows = extra_cache.shape[0] * extra_cache.shape[1] + cache_bytes = max( + cache.nelement() * cache.element_size(), + extra_cache.nelement() * extra_cache.element_size(), + ) + avg_main = indices.numel() / max(1, num_queries) + avg_extra = extra_indices.numel() / max(1, num_queries) if has_extra else 0.0 + + use_buffer_load = cache_bytes <= MAX_BYTES + HEAD_ALIGNED = num_heads % BLOCK_M == 0 + heads_blocks = (num_heads + BLOCK_M - 1) // BLOCK_M + out = torch.empty_like(q, dtype=torch.bfloat16) + + if kv_splits is not None: + num_splits = max(1, int(kv_splits)) + else: + num_splits = _decode_num_splits( + num_queries, heads_blocks, avg_main, avg_extra, BLOCK_K + ) + + if num_splits > 1: + part_m = torch.empty( + (num_queries, num_splits, num_heads), dtype=torch.float32, device=q.device + ) + part_l = torch.empty_like(part_m) + part_acc = torch.empty( + (num_queries, num_splits, num_heads, head_dim), + dtype=torch.float32, + device=q.device, + ) + pm_stride0, pm_stride_s = part_m.stride(0), part_m.stride(1) + pa_stride0, pa_stride_s, pa_stride_h = ( + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), + ) + else: + part_m = part_l = part_acc = out # unused placeholders (never dereferenced) + pm_stride0 = pm_stride_s = pa_stride0 = pa_stride_s = pa_stride_h = 0 + + grid = (num_queries, num_splits, heads_blocks) + _pa_decode_sparse_gfx950[grid]( + q, + cache, + main_bf16, + indices, + indptr, + extra_cache, + extra_bf16, + extra_indices, + extra_indptr, + attn_sink, + out, + part_m, + part_l, + part_acc, + scale, + q.stride(0), + q.stride(1), + out.stride(0), + out.stride(1), + cache.stride(0), + extra_cache.stride(0), + main_num_rows, + extra_num_rows, + pm_stride0, + pm_stride_s, + pa_stride0, + pa_stride_s, + pa_stride_h, + num_heads, + HAS_EXTRA=has_extra, + HAS_SINK=has_sink, + MAIN_IS_FP8=main_is_fp8, + EXTRA_IS_FP8=extra_is_fp8, + MAIN_BLOCK_SIZE=main_block, + EXTRA_BLOCK_SIZE=extra_block, + NOPE_DIM=nope_dim, + ROPE_DIM=ROPE_DIM, + HEAD_SIZE=head_dim, + BLOCK_M=BLOCK_M, + BLOCK_K=BLOCK_K, + NUM_SPLITS=num_splits, + HEAD_ALIGNED=HEAD_ALIGNED, + MFMA_K=MFMA_K, + UNIFORM=UNIFORM, + USE_BUFFER_LOAD=use_buffer_load, + HAS_INVALID=has_invalid, + FP8_FNUZ=fp8_fnuz, + num_warps=num_warps, + waves_per_eu=waves_per_eu, + ) + + if num_splits == 1: + return out + if skip_reduce: + return part_acc, part_m, part_l + + rgrid = (num_queries, heads_blocks) + _pa_decode_sparse_reduce_gfx950[rgrid]( + part_m, + part_l, + part_acc, + attn_sink, + out, + out.stride(0), + out.stride(1), + pm_stride0, + pm_stride_s, + pa_stride0, + pa_stride_s, + pa_stride_h, + num_heads, + HAS_SINK=has_sink, + HEAD_SIZE=head_dim, + BLOCK_M=BLOCK_M, + NUM_SPLITS=num_splits, + HEAD_ALIGNED=HEAD_ALIGNED, + num_warps=4, + ) + return out diff --git a/op_tests/triton_tests/attention/test_pa_decode_sparse.py b/op_tests/triton_tests/attention/test_pa_decode_sparse.py index a751ce6b806..d8734ce8c00 100644 --- a/op_tests/triton_tests/attention/test_pa_decode_sparse.py +++ b/op_tests/triton_tests/attention/test_pa_decode_sparse.py @@ -351,3 +351,139 @@ def test_pa_decode_sparse_fp8_vs_reference(T, H, D, kv_len, var_len): ) torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2) + + +def make_packed_cache(num_tokens, D, dtype): + device = "cuda" + rope = 64 # DSv4 RoPE dim, stored bf16 + block = 256 # packed cache page size + nope = D - rope # NoPE dim, stored fp8 e4m3 OCP + nb = triton.cdiv(num_tokens, block) + if dtype == "bf16": + cache = (torch.randn(nb, block, D, device=device) * 0.4).to(torch.bfloat16) + return cache, cache.reshape(nb * block, D).float() + + # per token: [nope fp8 (1B) | rope bf16 (2B) | 8 UE8M0 scale bytes] + data_bytes = nope + rope * 2 + scale_bytes = 8 + row_bytes = data_bytes + scale_bytes + cache = torch.zeros(nb, block, row_bytes, dtype=torch.uint8, device=device) + flat = cache.view(nb, block * row_bytes) + data = flat[:, : block * data_bytes].view(nb, block, data_bytes) + scales_region = flat[:, block * data_bytes :].view(nb, block, scale_bytes) + nope_fp8 = (torch.randn(nb, block, nope, device=device) * 0.4).to( + torch.float8_e4m3fn + ) + data[:, :, :nope] = nope_fp8.view(torch.uint8) + rope_bf16 = (torch.randn(nb, block, rope, device=device) * 0.4).to(torch.bfloat16) + data[:, :, nope:data_bytes] = rope_bf16.view(torch.uint8).view(nb, block, rope * 2) + num_groups = nope // 64 + exps = torch.randint( + 124, 130, (nb, block, num_groups), device=device, dtype=torch.uint8 + ) + scales_region[:, :, :num_groups] = exps + scales = torch.exp2(exps.float() - 127.0).repeat_interleave(64, dim=2) + kv_deq = torch.cat([nope_fp8.float() * scales, rope_bf16.float()], dim=2) + return cache, kv_deq.reshape(nb * block, D) + + +def two_loop_reference( + q, + main_deq, + main_idx, + main_indptr, + extra_deq, + extra_idx, + extra_indptr, + attn_sink, + softmax_scale, +): + """Reference for the SWA(main) + top-k(extra) two-loop: concatenate the two + dequantized pools, merge the two ragged index sets (extra slots shifted past + the main pool), then reuse ``pa_decode_sparse_reference``. + """ + main_pages = main_deq.shape[0] + combined = torch.cat([main_deq, extra_deq], dim=0).to(q.dtype) + T = main_indptr.numel() - 1 + mi, mp = main_idx.long(), main_indptr.long() + ei, ep = extra_idx.long(), extra_indptr.long() + rows, lens = [], [] + for tok in range(T): + row = torch.cat( + [mi[mp[tok] : mp[tok + 1]], ei[ep[tok] : ep[tok + 1]] + main_pages] + ) + rows.append(row) + lens.append(row.numel()) + combined_idx = torch.cat(rows).to(torch.int32) + combined_indptr = torch.zeros(T + 1, dtype=torch.int32, device=q.device) + combined_indptr[1:] = torch.tensor(lens, device=q.device).cumsum(0) + return pa_decode_sparse_reference( + q, combined, combined_idx, combined_indptr, attn_sink, softmax_scale + ) + + +@pytest.mark.parametrize("T", [1, 32, 128]) +@pytest.mark.parametrize("H", [16]) +@pytest.mark.parametrize("D", [512]) +@pytest.mark.parametrize("main_len", [128]) +@pytest.mark.parametrize("extra_len", [8, 256]) +@pytest.mark.parametrize("dtype", ["bf16", "fp8"]) +def test_pa_decode_sparse_two_loop(T, H, D, main_len, extra_len, dtype): + """gfx950 vLLM DSv4 decode path: SWA (main) + top-k (extra) two-loop over + packed caches. fp8 (fp8_ds_mla) is the vLLM production format; bf16 is also + exercised. Skipped off gfx950 (extra_* is a packed-only gluon path).""" + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + if arch_info.get_arch() != "gfx950": + pytest.skip("two-loop (extra_*) is a gfx950 packed-cache-only path") + + device = "cuda" + torch.manual_seed(0) + q = torch.randn(T, H, D, dtype=torch.bfloat16, device=device) * 0.125 + attn_sink = torch.randn(H, dtype=torch.float32, device=device) * 0.1 + softmax_scale = float(D) ** -0.5 + + # main = contiguous SWA window per query + main_cache, main_deq = make_packed_cache(T * main_len, D, dtype) + query_base = (torch.arange(T, device=device) * main_len)[:, None] + main_idx = ( + (query_base + torch.arange(main_len, device=device)).to(torch.int32).reshape(-1) + ) + main_indptr = torch.arange( + 0, T * main_len + 1, main_len, dtype=torch.int32, device=device + ) + # extra = scattered top-k over a pool + extra_pool = T * extra_len + extra_cache, extra_deq = make_packed_cache(extra_pool, D, dtype) + extra_idx = torch.randint( + 0, extra_pool, (T, extra_len), device=device, dtype=torch.int32 + ).reshape(-1) + extra_indptr = torch.arange( + 0, T * extra_len + 1, extra_len, dtype=torch.int32, device=device + ) + + ref = two_loop_reference( + q, + main_deq, + main_idx, + main_indptr, + extra_deq, + extra_idx, + extra_indptr, + attn_sink, + softmax_scale, + ) + out = pa_decode_sparse( + q, + main_cache, + main_idx, + main_indptr, + attn_sink, + softmax_scale, + extra_cache=extra_cache, + extra_indices=extra_idx, + extra_indptr=extra_indptr, + ) + + tol = 1e-2 if dtype == "fp8" else 5e-3 + torch.testing.assert_close(out, ref, atol=tol, rtol=tol)