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 deleted file mode 100644 index 54d526ca51b..00000000000 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ /dev/null @@ -1,1005 +0,0 @@ -# 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: -https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py - -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 _slots( - indices_ptr, - seg_start, - k_pos, - hi, - num_rows, - BLOCK_SIZE: gl.constexpr, - MASKED: gl.constexpr, - HAS_INVALID: gl.constexpr, -): - # Returns in whatever layout k_pos carries. Called once per gather layout: NoPE - # and RoPE tile their warps differently, and re-loading this tiny broadcast - # vector is cheaper than a cross-lane convert between the two. - if MASKED: - in_range = k_pos < hi - slot = gl.load(indices_ptr + seg_start + k_pos, mask=in_range, other=-1) - valid = in_range & (slot >= 0) & (slot < num_rows) - slot = gl.where(valid, slot, 0) - else: - slot = gl.load(indices_ptr + seg_start + k_pos) - valid = slot >= 0 # -1 sentinels: clamp in-bounds, mask score below - if HAS_INVALID: - slot = gl.where(valid, slot, 0) - return (slot // BLOCK_SIZE).to(gl.int32), (slot % BLOCK_SIZE).to(gl.int32), valid - - -@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, - k_rng_rope, - 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) - block_idx, pos, valid1d = _slots( - indices_ptr, - seg_start, - k_start + k_rng_slot, - hi, - num_rows, - BLOCK_SIZE, - MASKED, - HAS_INVALID, - ) - 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) - # Keep the dequant in f32: gfx950 has no bf16 multiply, so a bf16 path - # lowers to an emulated mul and doubles the loop. - 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, pos_gr, valid_gr = _slots( - indices_ptr, - seg_start, - k_start + k_rng_rope, - hi, - num_rows, - BLOCK_SIZE, - MASKED, - HAS_INVALID, - ) - rope_off = (block_idx_gr * (cs0 // 2) + pos_gr * 288 + 224)[ - :, None - ] + offs_rope[None, :] - if MASKED: - 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, - k_rng_rope, - 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. Split from the LDS-write/MFMA - so the gather issues an iteration early.""" - if not USE_BUFFER_LOAD: - cs0 = cs0.to(gl.int64) # >2 GB cache: 64-bit gather offsets (see _cache_load) - bg, pg, valid = _slots( - indices_ptr, - seg_start, - k_start + k_rng_slot, - 0, # hi/num_rows: unused at MASKED=False - 0, - BLOCK_SIZE, - False, - HAS_INVALID, - ) - 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, pgr, _ = _slots( - indices_ptr, - seg_start, - k_start + k_rng_rope, - 0, - 0, - BLOCK_SIZE, - False, - HAS_INVALID, - ) - 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) - k_rng_rope = gl.arange(0, BLOCK_K, layout=gl.SliceLayout(1, gather_rope_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: - 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, - k_rng_rope, - 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, - k_rng_rope, - 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, - k_rng_rope, - 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, - k_rng_rope, - 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], - ) - # Warps tile dim 0 here, one warp already covers all 64 RoPE columns, so putting - # them on dim 1 (as the NoPE gather does) overshoots 4x and every warp re-gathers - # the same tile. Worth ~10% on packed fp8. - gather_rope_l: gl.constexpr = gl.BlockedLayout( - size_per_thread=[1, 8], - threads_per_warp=[8, 8], - warps_per_cta=[NUM_WARPS, 1], - 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] - 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/_gluon_kernels/gfx950/attention/sparse_mla.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_mla.py new file mode 100644 index 00000000000..b5df196d086 --- /dev/null +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_mla.py @@ -0,0 +1,2340 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Gluon (gfx950) sparse decode for DeepSeek-V4 and MLA (GLM-5). + +One kernel serves two geometries, selected by the ROPE_SEPARATE constexpr: + + False: rope gets no LDS buffer of its own, either because it + lives inside the pow-2 row (DSv4) or because there is none at all + (GLM-5.3-Flash, ROPE_DIM=0). One LDS buffer, kv_smem + [BLOCK_K, HEAD_SIZE], one QK MFMA chain, V = that same buffer. + True (MLA): rope is appended (kv_lora_rank latent + rope = QK width; + V = the latent only). A second pow-2 LDS buffer, rope_smem, holds the + K-only rope and the QK contraction chains a second MFMA into the first. + The KV buffer is the entire V in both geometries. + +Cache formats (Fmt.KIND), per segment. vLLM calls both packed kinds +--kv-cache-dtype fp8_ds_mla, one name for two byte layouts told apart by model +generation, so the tags here carry the generation instead: + + "bf16" full-row bf16 + "fp8_scalar" whole-row fp8 + a single per-tensor f32 scale (k_scale). + The scale never touches the tile loop: K-side folds into + qk_scale, V-side into p. + "fp8_g64" whole-row fp8 + separate f32 per-64 kv_scales + "fp8_dsv4_mla" DeepSeek-V4, 584 B per token: 448 fp8 | 64 bf16 rope + (576 B data rows) + an 8 B UE8M0-per-64 trailer after the + block's rows (7 scales + 1 B pad). Rope lives inside the + 512-wide head dim, so V is the whole row. + "fp8_dsv32_mla" DeepSeek-V3.2 (also Kimi-K3), 656 B per token, token-major + records: 512 fp8 | 4 f32 per-128 scales | 64 bf16 rope. + Rope is appended, so this requires ROPE_SEPARATE. +""" + +from triton.experimental import gluon +from triton.experimental.gluon import language as gl +from triton.language.core import PropagateNan +from triton.language.core import _aggregate as aggregate + +from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr +from aiter.ops.triton.utils.common_utils import strip_annotate + +# Triton's default max ignores NaN, which on AMD costs a canonicalize per +# operand. Nothing here produces NaN (masked lanes are -inf and the all-masked +# row is guarded), so propagate instead. +_MAX_PROP_NAN: gl.constexpr = gl.constexpr(PropagateNan.ALL) + +_CG: gl.constexpr = gl.constexpr(".cg") +_NO_CACHE: gl.constexpr = gl.constexpr("") + + +@gluon.jit +def _max2(a, b): + return gl.maximum(a, b, propagate_nan=_MAX_PROP_NAN) + + +@gluon.jit +def _rmax(x, axis): + return gl.reduce(x, axis, _max2) + + +@gluon.jit +def _cache_load( + ptr, + row, + col, + USE_BUFFER_LOAD: gl.constexpr, + mask=None, + other=None, + CACHE: gl.constexpr = _CG, +): + """Gather rows[i] + col[j]. row is the per-token offset in ptr's element + units; col a small compile-time arange. Keeping them apart resolves one + pointer per token on the 64-bit path (the column offset folds into the + load's immediate) instead of a 64-bit add per element.""" + if USE_BUFFER_LOAD: + return gl.amd.cdna4.buffer_load( + ptr=ptr, + offsets=row.to(gl.int32)[:, None] + col.to(gl.int32)[None, :], + mask=mask, + other=other, + cache=CACHE, + ) + row_ptr = ptr + row.to(gl.int64) + return gl.load( + row_ptr[:, None] + col[None, :], + mask=mask, + other=other, + cache_modifier=CACHE, + ) + + +@gluon.jit +def _fp8_to_f32(x_u8): + return x_u8.to(gl.float8e4nv, bitcast=True).to(gl.float32) + + +@gluon.jit +def _fp8_to_bf16(x_u8): + # Exact: fp8's 3 mantissa bits fit bf16's 8. + return x_u8.to(gl.float8e4nv, bitcast=True).to(gl.bfloat16) + + +@gluon.jit +def _scale_load( + ptr, + row, + valid, + USE_BUFFER_LOAD: gl.constexpr, + gather_l: gl.constexpr, + scl_l: gl.constexpr, + NG: gl.constexpr, + W_FULL: gl.constexpr, + MASKED: gl.constexpr, + OTHER: gl.constexpr, +): + """Gather the NG per-group scales of each token and broadcast to W_FULL + columns. Indexing the full row with offs // GROUP would build a + [BLOCK_K, W_FULL] pointer tensor for NG distinct values, and only + buffer_load's 32-bit offsets CSE that redundancy away. scl_l is chosen so + the broadcast lands on gather_l as a register rename (assert_trivial + proves it). + Element type follows ptr: u8 E8M0 bytes (OTHER=127 -> 2^0) or f32 + (OTHER=0.0 -> masked lanes dequant to 0).""" + cols = gl.arange(0, NG, layout=gl.SliceLayout(0, scl_l)) + rows = gl.convert_layout(row, gl.SliceLayout(1, scl_l)) + if MASKED: + m = gl.convert_layout(valid, gl.SliceLayout(1, scl_l))[:, None] + if USE_BUFFER_LOAD: + sc = gl.amd.cdna4.buffer_load( + ptr=ptr, + offsets=rows.to(gl.int32)[:, None] + cols[None, :], + mask=m, + other=OTHER, + cache=".cg", + ) + else: + sc = gl.load( + (ptr + rows.to(gl.int64))[:, None] + cols[None, :], + mask=m, + other=OTHER, + cache_modifier=".cg", + ) + else: + if USE_BUFFER_LOAD: + sc = gl.amd.cdna4.buffer_load( + ptr=ptr, + offsets=rows.to(gl.int32)[:, None] + cols[None, :], + cache=".cg", + ) + else: + sc = gl.load( + (ptr + rows.to(gl.int64))[:, None] + cols[None, :], + cache_modifier=".cg", + ) + wide = gl.expand_dims(sc, 2).broadcast_to([sc.shape[0], NG, W_FULL // NG]) + return gl.convert_layout( + wide.reshape([sc.shape[0], W_FULL]), gather_l, assert_trivial=True + ) + + +@gluon.jit +def _split2(x): + """Register split along dim 1: [A, B] -> two [A, B//2] in the input's own + layout. Free only while the column direction is a per-lane register repeat; + assert_trivial turns anything else into a compile error.""" + layout: gl.constexpr = x.type.layout + x_r = x.reshape(x.shape[0], 2, x.shape[1] // 2).permute(0, 2, 1) + x0, x1 = gl.split(x_r) + x0 = gl.convert_layout(x0, layout, assert_trivial=True) + x1 = gl.convert_layout(x1, layout, assert_trivial=True) + return x0, x1 + + +@gluon.jit +def _split2_dim0(x): + """Dim-0 counterpart of _split2, for the row-wide gather layout where the + per-lane register repeats live on dim 0.""" + layout: gl.constexpr = x.type.layout + x_r = x.reshape(2, x.shape[0] // 2, x.shape[1]).permute(1, 2, 0) + x0, x1 = gl.split(x_r) + x0 = gl.convert_layout(x0, layout, assert_trivial=True) + x1 = gl.convert_layout(x1, layout, assert_trivial=True) + return x0, x1 + + +@gluon.jit +def _split_ax(x, AXIS: gl.constexpr): + if AXIS == 1: + a, b = _split2(x) + else: + a, b = _split2_dim0(x) + return a, b + + +@gluon.jit +def _deq_asm(x16, e_u8, W8: gl.constexpr, out_l: gl.constexpr): + """[BLOCK_K, W8/2] int16 (4 packed fp8) + raw E8M0 byte -> [BLOCK_K, W8] bf16. + + The hardware reads only bits [30:23] of the scale operand, i.e. bits [14:7] + of its high half, so a 16-bit e << 7 lands the exponent where an f32 + e << 23 would and costs one register instead of two. The scale is a power of + two, so this stays bit-identical to the unfused convert + multiply.""" + + _DEQ_LO: gl.constexpr = gl.constexpr("v_cvt_scalef32_pk_bf16_fp8 $0, $1, $2") + _DEQ_HI: gl.constexpr = gl.constexpr( + "v_cvt_scalef32_pk_bf16_fp8 $0, $1, $2 op_sel:[1,0,0]" + ) + _DEQ_CONS: gl.constexpr = gl.constexpr("=v,v,v") + sc16 = e_u8.to(gl.uint16) << 7 + lo = gl.inline_asm_elementwise( + _DEQ_LO, _DEQ_CONS, [x16, sc16], dtype=gl.bfloat16, is_pure=True, pack=2 + ) + hi = gl.inline_asm_elementwise( + _DEQ_HI, _DEQ_CONS, [x16, sc16], dtype=gl.bfloat16, is_pure=True, pack=2 + ) + # Flat order is 4i + 2*lohi + s; a lane's int16s are the same 16-byte run + # the gather used, so the interleave stays lane-local (assert_trivial). + W16: gl.constexpr = W8 // 2 + lo3 = lo.reshape(lo.shape[0], W16 // 2, 2) + hi3 = hi.reshape(hi.shape[0], W16 // 2, 2) + both = gl.join(lo3, hi3).permute(0, 1, 3, 2).reshape(lo.shape[0], W8) + # Back to the gather layout, or the kv_smem store lowers to narrow ds_writes. + return gl.convert_layout(both, out_l, assert_trivial=True) + + +@aggregate +@strip_annotate +class Cfg: + """Compile-time geometry, layouts, and behavior knobs shared by both segments.""" + + # geometry + BLOCK_M: gl.constexpr + BLOCK_K: gl.constexpr + KV_DIM: gl.constexpr # KV buffer width = LDS tile width = V width (pow-2) + ROPE_DIM: gl.constexpr # rope buffer width when ROPE_SEPARATE, else the + # bf16 tail width inside the KV buffer (DSv4 packed) + ROPE_L: gl.constexpr # rope buffer layout width. ROPE_DIM, or a stand-in + # when there is no rope at all. + ROPE_SEPARATE: gl.constexpr # False: rope inside the KV buffer. True: a + # K-only second buffer, QK contracts over KV_DIM + ROPE_DIM. + QK_DIM: gl.constexpr # q row width: KV_DIM (+ ROPE_DIM if separate) + MFMA_K: gl.constexpr + NUM_WARPS: gl.constexpr + GATHER_TW1: gl.constexpr + LDS_PAD: gl.constexpr + # behavior knobs + UNI_TILE: gl.constexpr + HAS_INVALID: gl.constexpr + HEAD_ALIGNED: gl.constexpr + IDX_BUFFER_LOAD: gl.constexpr + FP8_MFMA: gl.constexpr # "fp8_scalar" only: feed the matrix core the cache's + # own fp8 instead of dequantizing to bf16 + # Cache policy per load site + GATHER_CACHE: gl.constexpr + IDX_CACHE: gl.constexpr + # Gather straight into LDS (no register staging). FP8_MFMA only: the bf16 path + # stages dequantized values, which pass through registers by definition. + # Single-buffered, the launcher decides when it is on. + ASYNC_LDS: gl.constexpr + RELAXED_LOAD: gl.constexpr # read LDS with the syncedViaAsyncWait hint + ROPE_VEC: gl.constexpr # bytes per lane in the rope buffer's copy + # operator layouts + qk_layout: gl.constexpr + pv_layout: gl.constexpr + q_layout: gl.constexpr + k_layout: gl.constexpr + p_layout: gl.constexpr + v_layout: gl.constexpr + # memory layouts + gather_l: gl.constexpr + async_l: gl.constexpr + async_rope_l: gl.constexpr + slot_a_l: gl.constexpr + slot_rope_a_l: gl.constexpr + gather_rope_l: gl.constexpr + gather16_l: gl.constexpr + slot_l: gl.constexpr + blocked_q: gl.constexpr + kv_shared: gl.constexpr + rope_shared: gl.constexpr + + @gluon.constexpr_function + def __init__( + self, + BLOCK_M, + BLOCK_K, + KV_DIM, + ROPE_DIM, + ROPE_SEPARATE, + NUM_WARPS, + UNI_TILE, + HAS_INVALID, + HEAD_ALIGNED, + IDX_BUFFER_LOAD, + FP8_MFMA=False, + GATHER_CACHE=".cg", + IDX_CACHE="", + ASYNC_LDS=False, + RELAXED_LOAD=True, + PAD_INTERVAL=1024, + ): + self.BLOCK_M = gl.constexpr(BLOCK_M) + self.BLOCK_K = gl.constexpr(BLOCK_K) + self.KV_DIM = gl.constexpr(KV_DIM) + self.ROPE_DIM = gl.constexpr(ROPE_DIM) + self.ROPE_SEPARATE = gl.constexpr(ROPE_SEPARATE) + self.QK_DIM = gl.constexpr(KV_DIM + (ROPE_DIM if ROPE_SEPARATE else 0)) + self.NUM_WARPS = gl.constexpr(NUM_WARPS) + # Threads the row gather spends on the head dim + GATHER_TW1 = 32 + self.GATHER_TW1 = gl.constexpr(GATHER_TW1) + LDS_PAD = 16 if FP8_MFMA else 8 + self.LDS_PAD = gl.constexpr(LDS_PAD) + self.UNI_TILE = gl.constexpr(UNI_TILE) + self.HAS_INVALID = gl.constexpr(HAS_INVALID) + self.HEAD_ALIGNED = gl.constexpr(HEAD_ALIGNED) + self.IDX_BUFFER_LOAD = gl.constexpr(IDX_BUFFER_LOAD) + self.FP8_MFMA = gl.constexpr(FP8_MFMA) + self.GATHER_CACHE = gl.constexpr(GATHER_CACHE) + self.IDX_CACHE = gl.constexpr(IDX_CACHE) + self.ASYNC_LDS = gl.constexpr(ASYNC_LDS) + self.RELAXED_LOAD = gl.constexpr(RELAXED_LOAD) + ROPE_VEC = 16 + self.ROPE_VEC = gl.constexpr(ROPE_VEC) + MFMA_K = 32 if FP8_MFMA else 16 + self.MFMA_K = gl.constexpr(MFMA_K) + + self.qk_layout = gl.constexpr( + gl.amd.AMDMFMALayout( + version=4, + instr_shape=[16, 16, MFMA_K], + transposed=True, + warps_per_cta=[1, NUM_WARPS], + ) + ) + self.pv_layout = gl.constexpr( + gl.amd.AMDMFMALayout( + version=4, + instr_shape=[16, 16, MFMA_K], + transposed=True, + warps_per_cta=[1, NUM_WARPS], + ) + ) + KW = MFMA_K // 2 + self.q_layout = gl.constexpr(gl.DotOperandLayout(0, self.qk_layout, KW)) + self.k_layout = gl.constexpr(gl.DotOperandLayout(1, self.qk_layout, KW)) + self.p_layout = gl.constexpr(gl.DotOperandLayout(0, self.pv_layout, KW)) + self.v_layout = gl.constexpr(gl.DotOperandLayout(1, self.pv_layout, KW)) + + # 16 uint8 per thread = 128-bit gather loads. GATHER_TW1 lanes span a row, + # so one gather covers GATHER_TW1 * 16 = 512 B of it, at the cost of a + # longer slot vector. + GSPT = 16 + self.gather_l = gl.constexpr( + gl.BlockedLayout( + size_per_thread=[1, GSPT], + threads_per_warp=[64 // GATHER_TW1, GATHER_TW1], + warps_per_cta=[NUM_WARPS, 1], + order=[1, 0], + ) + ) + # One warp already covers all rope columns; warps on dim 1 would + # re-gather the same tile NUM_WARPS times. + self.gather_rope_l = gl.constexpr( + gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[8, 8], + warps_per_cta=[NUM_WARPS, 1], + order=[1, 0], + ) + ) + # gather_l with 2-byte elements: same 16-byte per-lane run, half the + # columns (for the fused-dequant int16 view). + self.gather16_l = gl.constexpr( + gl.BlockedLayout( + size_per_thread=[1, GSPT // 2], + threads_per_warp=[64 // GATHER_TW1, GATHER_TW1], + warps_per_cta=[NUM_WARPS, 1], + order=[1, 0], + ) + ) + self.slot_l = gl.constexpr(gl.SliceLayout(1, self.gather_l.value)) + self.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], + ) + ) + # Row pitch (KV_DIM + LDS_PAD) decides which banks the transposed K + # read (walks down a column) lands on. + self.kv_shared = gl.constexpr( + gl.PaddedSharedLayout.with_identity_for( + [[KV_DIM, LDS_PAD]], [BLOCK_K, KV_DIM], [1, 0] + ) + ) + # The rope buffer (K-only) exists when ROPE_SEPARATE, dead otherwise. + # compiler removes the dead code, so no side effect beyond eliminating errors + ROPE_L = ROPE_DIM if ROPE_DIM > 0 else 64 + self.ROPE_L = gl.constexpr(ROPE_L) + self.rope_shared = gl.constexpr( + gl.PaddedSharedLayout.with_identity_for( + [[ROPE_L, LDS_PAD]], [BLOCK_K, ROPE_L], [1, 0] + ) + ) + + AVEC = 16 # 128-bit LDS-DMA + ATW = KV_DIM // AVEC + self.async_l = gl.constexpr( + gl.BlockedLayout( + size_per_thread=[1, AVEC], + threads_per_warp=[64 // ATW, ATW], + warps_per_cta=[NUM_WARPS, 1], + order=[1, 0], + ) + ) + RTW = ROPE_L // ROPE_VEC + self.async_rope_l = gl.constexpr( + gl.BlockedLayout( + size_per_thread=[1, ROPE_VEC], + threads_per_warp=[64 // RTW, RTW], + warps_per_cta=[NUM_WARPS, 1], + order=[1, 0], + ) + ) + self.slot_a_l = gl.constexpr(gl.SliceLayout(1, self.async_l.value)) + self.slot_rope_a_l = gl.constexpr(gl.SliceLayout(1, self.async_rope_l.value)) + if ASYNC_LDS: + self.kv_shared = gl.constexpr( + gl.PaddedSharedLayout.with_identity_for( + [[PAD_INTERVAL, LDS_PAD]], [BLOCK_K, KV_DIM], [1, 0] + ) + ) + # 64 * ROPE_VEC is this buffer's warp run: the smallest legal interval. + self.rope_shared = gl.constexpr( + gl.PaddedSharedLayout.with_identity_for( + [[64 * ROPE_VEC, LDS_PAD]], [BLOCK_K, ROPE_L], [1, 0] + ) + ) + + +@aggregate +@strip_annotate +class Fmt: + """Compile-time description of one segment's cache format.""" + + # "bf16" | "fp8_scalar" | "fp8_g64" + # | "fp8_dsv4_mla" (584 B/token) | "fp8_dsv32_mla" (656 B/token) + KIND: gl.constexpr + IS_FP8: gl.constexpr # pipeline select: prefetched fp8 loop vs bf16 loop + BLOCK_SIZE: gl.constexpr + USE_BUFFER_LOAD: gl.constexpr + # dsv4 E8M0 dequant: "none" | "upcast" | "asm" (gathers the int16 view) + DEQ: gl.constexpr + NOPE_CHUNK: gl.constexpr + CHUNK_AXIS: gl.constexpr + NOPE_DIM: gl.constexpr # fp8 payload width (448 dsv4; KV_DIM elsewhere) + GROUP: gl.constexpr # scale group width (64 dsv4/uniform, 128 dsmla) + NG: gl.constexpr # scale groups per row (KV_DIM // GROUP) + NARROW_SCALE: gl.constexpr + scl_l: gl.constexpr + # packed-row addressing constants (element units named in the suffix) + TOK_U8: gl.constexpr # bytes per token row inside a block (576 / 656) + TOK_U16: gl.constexpr # the same row in bf16-view units + ROPE_U16_OFF: gl.constexpr # bf16-view offset of the rope tail in a row + SCL_TRAILER_U8: gl.constexpr # dsv4: scale bytes per token in the block trailer + TOK_F32: gl.constexpr # dsmla: f32-view stride per token row + SCL_F32_OFF: gl.constexpr # dsmla: f32-view offset of the group scales + TOK_EL: gl.constexpr # flat formats: cache elements per token row + + @gluon.constexpr_function + def __init__( + self, + cfg, + KIND, + BLOCK_SIZE, + USE_BUFFER_LOAD, + DEQ, + NOPE_DIM, + NOPE_CHUNK, + CHUNK_AXIS, + ): + self.KIND = gl.constexpr(KIND) + self.IS_FP8 = gl.constexpr(KIND != "bf16") + self.BLOCK_SIZE = gl.constexpr(BLOCK_SIZE) + self.USE_BUFFER_LOAD = gl.constexpr(USE_BUFFER_LOAD) + self.DEQ = gl.constexpr(DEQ if KIND == "fp8_dsv4_mla" else "none") + self.NOPE_CHUNK = gl.constexpr(NOPE_CHUNK) + self.CHUNK_AXIS = gl.constexpr(CHUNK_AXIS) + self.NOPE_DIM = gl.constexpr(NOPE_DIM) + + KV_DIM = cfg.KV_DIM.value + ROPE_DIM = cfg.ROPE_DIM.value + GROUP = 128 if KIND == "fp8_dsv32_mla" else 64 + NG = KV_DIM // GROUP + self.GROUP = gl.constexpr(GROUP) + self.NG = gl.constexpr(NG) + # 3-D companion of gather_l for _scale_load: dim 1 carries the NG scale + # groups, dim 2 the columns inside a group, so the [BLOCK_K, NG, GROUP] + # -> 2-D reshape reproduces gather_l exactly. Legal only when a group's + # columns fill whole threads. + GSPT = 16 + TW1 = cfg.GATHER_TW1.value + NARROW_SCALE = TW1 % NG == 0 and GSPT * (TW1 // NG) == GROUP + self.NARROW_SCALE = gl.constexpr(NARROW_SCALE) + scl_l3 = gl.BlockedLayout( + size_per_thread=[1, 1, GSPT], + threads_per_warp=[ + 64 // TW1, + NG if NARROW_SCALE else 1, + (TW1 // NG) if NARROW_SCALE else TW1, + ], + warps_per_cta=[cfg.NUM_WARPS.value, 1, 1], + order=[2, 1, 0], + ) + self.scl_l = gl.constexpr(gl.SliceLayout(2, scl_l3)) + + # dsv4 row: [NOPE_DIM fp8 | ROPE_DIM bf16] + 8 B UE8M0 per token after + # the block; dsmla row: [KV_DIM fp8 | NG f32 | ROPE_DIM bf16] inline. + if KIND == "fp8_dsv4_mla": + TOK_U8 = NOPE_DIM + 2 * ROPE_DIM # 448 + 128 = 576 + elif KIND == "fp8_dsv32_mla": + TOK_U8 = KV_DIM + 4 * NG + 2 * ROPE_DIM # 512 + 16 + 128 = 656 + else: + TOK_U8 = 0 + self.TOK_U8 = gl.constexpr(TOK_U8) + self.TOK_U16 = gl.constexpr(TOK_U8 // 2) + self.ROPE_U16_OFF = gl.constexpr( + (NOPE_DIM // 2) if KIND == "fp8_dsv4_mla" else ((KV_DIM + 4 * NG) // 2) + ) + self.SCL_TRAILER_U8 = gl.constexpr(8) + self.TOK_F32 = gl.constexpr(TOK_U8 // 4) + self.SCL_F32_OFF = gl.constexpr(KV_DIM // 4) + # Flat formats gather in cache elements: KV_DIM wide, plus the appended + # rope when the geometry separates it. + self.TOK_EL = gl.constexpr( + KV_DIM + (ROPE_DIM if cfg.ROPE_SEPARATE.value and KIND != "fp8_g64" else 0) + ) + + +@aggregate +@strip_annotate +class Seg: + """Runtime context of one segment. Pointer roles by format ("--" = unused + duplicate): + + KIND cache_ptr alt_ptr scl_ptr + bf16 -- bf16 cache -- + fp8_scalar u8 cache -- f32 scalar k_scale + fp8_g64 u8 cache f32 per-64 kv_scales -- + fp8_dsv4_mla u8 cache bf16 view of the cache -- + fp8_dsv32_mla u8 cache bf16 view (rope tail) f32 view (scales) + """ + + fmt: Fmt + cache_ptr: gl.tensor + alt_ptr: gl.tensor + scl_ptr: gl.tensor + indices_ptr: gl.tensor + seg_start: gl.tensor + cs0: gl.tensor + num_rows: gl.tensor + + @gluon.constexpr_function + def __init__( + self, fmt, cache_ptr, alt_ptr, scl_ptr, indices_ptr, seg_start, cs0, num_rows + ): + self.fmt = fmt + self.cache_ptr = cache_ptr + self.alt_ptr = alt_ptr + self.scl_ptr = scl_ptr + self.indices_ptr = indices_ptr + self.seg_start = seg_start + self.cs0 = cs0 + self.num_rows = num_rows + + +@gluon.jit +def _deq_store(x_u8, sc, kv_smem, off, cfg, fmt, AXIS: gl.constexpr): + """Dequant one fp8 slab into kv_smem[:, off:off+W]. Dequant stays in f32 + (gfx950 has no bf16 multiply). sc: raw UE8M0 byte (dsv4), f32 scale + (uniform/dsmla), or unused ("fp8_scalar": bare fp8 -> bf16 convert).""" + if fmt.DEQ == "asm": + # x_u8 is the int16 view here (see _gather_full), so its column count is W8/2. + W8: gl.constexpr = x_u8.shape[1] * 2 + # Adjacent fp8 columns share a scale (groups are even), so dropping + # every other broadcast column is exact. + s_even, _ = gl.split(sc.reshape(sc.shape[0], sc.shape[1] // 2, 2)) + s_even = gl.convert_layout(s_even, x_u8.type.layout) + val = _deq_asm(x_u8.to(gl.int16, bitcast=True), s_even, W8, cfg.gather_l) + if AXIS == 1: + kv_smem.slice(off, W8, dim=1).store(val) + else: + kv_smem.slice(off, x_u8.shape[0], dim=0).store(val) + else: + if fmt.DEQ == "upcast": + # Upstream fused fp8 x E8M0 -> bf16 upcast; the driver only asks for + # it when cdna4.scaled_upcast exists, else DEQ="asm" (_deq_asm). + # sc is the raw E8M0 byte, already in x_u8's shape and layout. + val = gl.amd.cdna4.scaled_upcast( + x_u8.to(gl.float8e4nv, bitcast=True), sc, gl.bfloat16 + ) + elif fmt.KIND == "fp8_scalar": + val = _fp8_to_bf16(x_u8) + else: + if fmt.KIND == "fp8_g64" or fmt.KIND == "fp8_dsv32_mla": + scale = sc + else: + scale = gl.exp2(sc.to(gl.float32) - 127.0) + val = (_fp8_to_f32(x_u8) * scale).to(gl.bfloat16) + if AXIS == 1: + kv_smem.slice(off, x_u8.shape[1], dim=1).store(val) + else: + kv_smem.slice(off, x_u8.shape[0], dim=0).store(val) + + +@gluon.jit +def _deq_store_tile(x_u8, sc, kv_smem, cfg, fmt): + """Dequant a gathered fp8 tile into kv_smem in NOPE_CHUNK-sized pieces + along CHUNK_AXIS (0 = rows, 1 = columns). The f32 expansion is 4x the fp8 + tile; chunking keeps only 1/pieces of it live (each piece's converts feed + its ds_writes and die). The splits are register renames (see _split2).""" + AXIS: gl.constexpr = fmt.CHUNK_AXIS + NOPE_CHUNK: gl.constexpr = fmt.NOPE_CHUNK + W: gl.constexpr = x_u8.shape[1] if AXIS == 1 else x_u8.shape[0] + if NOPE_CHUNK >= W: + _deq_store(x_u8, sc, kv_smem, 0, cfg, fmt, AXIS) + else: + x0, x1 = _split_ax(x_u8, AXIS) + s0, s1 = _split_ax(sc, AXIS) + W2: gl.constexpr = W // 2 + if NOPE_CHUNK >= W2: + _deq_store(x0, s0, kv_smem, 0, cfg, fmt, AXIS) + _deq_store(x1, s1, kv_smem, W2, cfg, fmt, AXIS) + else: + x00, x01 = _split_ax(x0, AXIS) + s00, s01 = _split_ax(s0, AXIS) + x10, x11 = _split_ax(x1, AXIS) + s10, s11 = _split_ax(s1, AXIS) + W4: gl.constexpr = W // 4 + if NOPE_CHUNK >= W4: + _deq_store(x00, s00, kv_smem, 0, cfg, fmt, AXIS) + _deq_store(x01, s01, kv_smem, W4, cfg, fmt, AXIS) + _deq_store(x10, s10, kv_smem, 2 * W4, cfg, fmt, AXIS) + _deq_store(x11, s11, kv_smem, 3 * W4, cfg, fmt, AXIS) + else: + W8: gl.constexpr = W // 8 + y0, y1 = _split_ax(x00, AXIS) + t0, t1 = _split_ax(s00, AXIS) + _deq_store(y0, t0, kv_smem, 0, cfg, fmt, AXIS) + _deq_store(y1, t1, kv_smem, W8, cfg, fmt, AXIS) + y0, y1 = _split_ax(x01, AXIS) + t0, t1 = _split_ax(s01, AXIS) + _deq_store(y0, t0, kv_smem, 2 * W8, cfg, fmt, AXIS) + _deq_store(y1, t1, kv_smem, 3 * W8, cfg, fmt, AXIS) + y0, y1 = _split_ax(x10, AXIS) + t0, t1 = _split_ax(s10, AXIS) + _deq_store(y0, t0, kv_smem, 4 * W8, cfg, fmt, AXIS) + _deq_store(y1, t1, kv_smem, 5 * W8, cfg, fmt, AXIS) + y0, y1 = _split_ax(x11, AXIS) + t0, t1 = _split_ax(s11, AXIS) + _deq_store(y0, t0, kv_smem, 6 * W8, cfg, fmt, AXIS) + _deq_store(y1, t1, kv_smem, 7 * W8, cfg, fmt, AXIS) + + +@gluon.jit +def _slots( + cfg, + seg, + k_pos, + hi, + num_rows, + MASKED: gl.constexpr, + UNI_TILE: gl.constexpr = False, +): + """Index-list read -> (block, pos, valid), in whatever layout k_pos carries. + A masked gl.load predicates on exec while a masked buffer_load folds the + mask into the offset, so the unmasked paths clamp the read in-range and + mask the score instead (UNI_TILE); -1 sentinels are handled the same way.""" + indices_ptr = seg.indices_ptr + seg_start = seg.seg_start + BLOCK_SIZE: gl.constexpr = seg.fmt.BLOCK_SIZE + HAS_INVALID: gl.constexpr = cfg.HAS_INVALID + IDX_BUFFER_LOAD: gl.constexpr = cfg.IDX_BUFFER_LOAD + if MASKED: + in_range = k_pos < hi + if IDX_BUFFER_LOAD: + slot = gl.amd.cdna4.buffer_load( + ptr=indices_ptr + seg_start, + offsets=k_pos, + mask=in_range, + other=-1, + cache=cfg.IDX_CACHE, + ) + else: + slot = gl.load( + indices_ptr + seg_start + k_pos, + mask=in_range, + other=-1, + cache_modifier=cfg.IDX_CACHE, + ) + valid = in_range & (slot >= 0) & (slot < num_rows) + slot = gl.where(valid, slot, 0) + else: + # hi >= 1 whenever UNI_TILE runs (guarded by n_full > 0). + off = gl.minimum(k_pos, hi - 1) if UNI_TILE else k_pos + if IDX_BUFFER_LOAD: + slot = gl.amd.cdna4.buffer_load( + ptr=indices_ptr + seg_start, offsets=off, cache=cfg.IDX_CACHE + ) + else: + slot = gl.load(indices_ptr + seg_start + off, cache_modifier=cfg.IDX_CACHE) + valid = (k_pos < hi) if UNI_TILE else (slot >= 0) + if HAS_INVALID: + if UNI_TILE: + valid = valid & (slot >= 0) + slot = gl.where(valid, slot, 0) # -1 sentinels: clamp, mask score below + return (slot // BLOCK_SIZE).to(gl.int32), (slot % BLOCK_SIZE).to(gl.int32), valid + + +@gluon.jit +def _qk_scores(cfg, q_dot, q_rope_dot, kv_smem, rope_smem): + """QK scores for one tile; ROPE_SEPARATE chains a second MFMA over the rope + buffer (MFMA accumulates natively, so KV_DIM + ROPE_DIM is two dots).""" + if cfg.ASYNC_LDS and cfg.RELAXED_LOAD: + # the async_wait already ordered the copy; stops the backend re-inserting + # a conservative vmcnt(0) before every LDS read + k = gl.amd.cdna4.async_copy.load_shared_relaxed( + kv_smem.permute([1, 0]), cfg.k_layout + ) + else: + k = kv_smem.permute([1, 0]).load(cfg.k_layout) # [KV_DIM, BLOCK_K] + if cfg.ASYNC_LDS: + k = k.to(gl.float8e4nv, bitcast=True) # raw cache bytes; layout-preserving + S = gl.amd.cdna4.mfma( + q_dot, + k, + gl.zeros([cfg.BLOCK_M, cfg.BLOCK_K], gl.float32, layout=cfg.qk_layout), + ) + if cfg.ROPE_SEPARATE: + if cfg.ASYNC_LDS and cfg.RELAXED_LOAD: + k_rope = gl.amd.cdna4.async_copy.load_shared_relaxed( + rope_smem.permute([1, 0]), cfg.k_layout + ) + else: + k_rope = rope_smem.permute([1, 0]).load(cfg.k_layout) + if cfg.ASYNC_LDS: + k_rope = k_rope.to(gl.float8e4nv, bitcast=True) + S = gl.amd.cdna4.mfma(q_rope_dot, k_rope, S) + return S + + +@gluon.jit +def _tile_rows(cfg, seg, k_start, seg_hi, k_rng, ROPE: gl.constexpr): + """Per-token row offsets for one tile, in k_rng's layout. Clamped like UNI_TILE + (the score mask kills the duplicate), so an over-fetched tile is addressable.""" + fmt = seg.fmt + cs0 = seg.cs0 + if not fmt.USE_BUFFER_LOAD: + cs0 = cs0.to(gl.int64) # >2 GB cache: 64-bit gather offsets + bg, pg, _ = _slots(cfg, seg, k_start + k_rng, seg_hi, 0, False, cfg.UNI_TILE) + row = bg * cs0 + pg * fmt.TOK_EL + if ROPE: + row = row + cfg.KV_DIM + return row + + +@gluon.jit +def _copy_lds_buffer( + seg, dst, row, col, USE_BUFFER_LOAD: gl.constexpr, CACHE: gl.constexpr +): + """One LDS buffer of one tile, global -> LDS, no register staging.""" + if USE_BUFFER_LOAD: + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dst, + seg.cache_ptr, + row.to(gl.int32)[:, None] + col.to(gl.int32)[None, :], + cache_modifier=CACHE, + ) + else: + row_ptr = seg.cache_ptr + row.to(gl.int64) + gl.amd.cdna4.async_copy.global_load_to_shared( + dst, row_ptr[:, None] + col[None, :], cache_modifier=CACHE + ) + + +@gluon.jit +def _copy_tile(cfg, seg, kv_smem, rope_smem, row_l, row_r, offs_l, offs_r): + """One commit group = one tile (both buffers), so wait_group counts tiles.""" + _copy_lds_buffer( + seg, kv_smem, row_l, offs_l, seg.fmt.USE_BUFFER_LOAD, cfg.GATHER_CACHE + ) + if cfg.ROPE_SEPARATE: + _copy_lds_buffer( + seg, rope_smem, row_r, offs_r, seg.fmt.USE_BUFFER_LOAD, cfg.GATHER_CACHE + ) + gl.amd.cdna4.async_copy.commit_group() + + +@gluon.jit +def _async_segment( + cfg, + seg, + q_dot, + q_rope_dot, + lo, + hi, + m_i, + l_i, + acc, + head_mask, + qk_scale, + v_scale, + kv_smem, + rope_smem, +): + """Direct-to-LDS tile walk, one LDS buffer. + + vmcnt is one in-order FIFO, so the wait goes before any plain load in the body: a + plain load outstanding past the newest copy group would drag wait_group into + waiting for that copy too. The body has no masked variant: the slot read is + clamped and the partial tail falls out of the score mask (UNI_TILE). + """ + BK: gl.constexpr = cfg.BLOCK_K + offs_l = gl.arange(0, cfg.KV_DIM, layout=gl.SliceLayout(0, cfg.async_l)) + offs_r = gl.arange(0, cfg.ROPE_L, layout=gl.SliceLayout(0, cfg.async_rope_l)) + rng_l = gl.arange(0, BK, layout=cfg.slot_a_l) + rng_r = gl.arange(0, BK, layout=cfg.slot_rope_a_l) + n_full = (hi - lo + BK - 1) // BK + for i in range(n_full): + _copy_tile( + cfg, + seg, + kv_smem, + rope_smem, + _tile_rows(cfg, seg, lo + i * BK, hi, rng_l, False), + _tile_rows(cfg, seg, lo + i * BK, hi, rng_r, True), + offs_l, + offs_r, + ) + gl.amd.cdna4.async_copy.wait_group(0) + m_i, l_i, acc = _qkpv_lds( + cfg, + seg, + None, + q_dot, + q_rope_dot, + m_i, + l_i, + acc, + head_mask, + qk_scale, + v_scale, + kv_smem, + rope_smem, + lo + i * BK, + hi, + ) + return m_i, l_i, acc + + +# --------------------------------------------------------------------------- +# Prefetched fp8 pipeline: _gather_full issues tile N+1's loads while _qkpv +# stages and dots tile N. +# --------------------------------------------------------------------------- + + +@gluon.jit +def _gather_full( + cfg, + seg, + k_start, + seg_hi, + offs_full, + offs_full16, + offs_rope, + k_rng_slot, + k_rng_rope, +): + """Gather one full fp8 tile, split from the LDS-write/MFMA so it issues an + iteration early. The prefetch stays in raw fp8: dequantizing here would + double the loop-carried registers, so the consumer dequants in chunks. + Returns (x, sc, k_rope, valid); unused slots carry a duplicate DCE removes.""" + fmt = seg.fmt + cs0 = seg.cs0 + if not fmt.USE_BUFFER_LOAD: + cs0 = cs0.to(gl.int64) # >2 GB cache: 64-bit gather offsets + bg, pg, valid = _slots( + cfg, + seg, + k_start + k_rng_slot, + seg_hi, # hi: unused unless UNI_TILE + 0, + False, + cfg.UNI_TILE, + ) + if fmt.KIND == "fp8_g64": + NGRP: gl.constexpr = cfg.KV_DIM // 64 + x_u8 = _cache_load( + seg.cache_ptr, + bg * cs0 + pg * cfg.KV_DIM, + offs_full, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + sc = _cache_load( + seg.alt_ptr, + bg * NGRP, + offs_full // 64, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + k_rope = x_u8 # no rope side-channel -> DCE'd + elif fmt.KIND == "fp8_scalar": + # Per-tensor scale is folded outside the loop (qk_scale / p), so this + # is a bare gather; the K-only rope tail follows when separated. + x_u8 = _cache_load( + seg.cache_ptr, + bg * cs0 + pg * fmt.TOK_EL, + offs_full, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + sc = x_u8 # no scale vector -> DCE'd + if cfg.ROPE_SEPARATE: + bgr, pgr, _ = _slots( + cfg, + seg, + k_start + k_rng_rope, + seg_hi, + 0, + False, + cfg.UNI_TILE, + ) + k_rope = _cache_load( + seg.cache_ptr, + bgr * cs0 + pgr * fmt.TOK_EL + cfg.KV_DIM, + offs_rope, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + else: + k_rope = x_u8 # rope lives inside the KV buffer -> DCE'd + elif fmt.KIND == "fp8_dsv32_mla": + nope_row = bg * cs0 + pg * fmt.TOK_U8 + scl_row = bg * (cs0 // 4) + pg * fmt.TOK_F32 + fmt.SCL_F32_OFF + # Scales issue before the bulk fp8: vmcnt is one in-order FIFO, so + # issued last they would stall the first dequant piece behind every + # data load as well. + if fmt.NARROW_SCALE and not fmt.USE_BUFFER_LOAD: + sc = _scale_load( + seg.scl_ptr, + scl_row, + scl_row, + fmt.USE_BUFFER_LOAD, + cfg.gather_l, + fmt.scl_l, + fmt.NG, + cfg.KV_DIM, + False, + 0.0, + ) + else: + sc = _cache_load( + seg.scl_ptr, + scl_row, + offs_full // fmt.GROUP, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + x_u8 = _cache_load( + seg.cache_ptr, + nope_row, + offs_full, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + bgr, pgr, _ = _slots( + cfg, + seg, + k_start + k_rng_rope, + seg_hi, + 0, + False, + cfg.UNI_TILE, + ) + k_rope = _cache_load( + seg.alt_ptr, + bgr * (cs0 // 2) + pgr * fmt.TOK_U16 + fmt.ROPE_U16_OFF, + offs_rope, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + else: # "fp8_dsv4_mla" + nope_row = bg * cs0 + pg * fmt.TOK_U8 + scl_row = bg * cs0 + fmt.BLOCK_SIZE * fmt.TOK_U8 + pg * fmt.SCL_TRAILER_U8 + # Scales first (vmcnt FIFO; see the dsmla branch). + if fmt.NARROW_SCALE and not fmt.USE_BUFFER_LOAD: + sc = _scale_load( + seg.cache_ptr, + scl_row, + scl_row, + fmt.USE_BUFFER_LOAD, + cfg.gather_l, + fmt.scl_l, + fmt.NG, + cfg.KV_DIM, + False, + 127, + ) + else: + sc = _cache_load( + seg.cache_ptr, + scl_row, + offs_full // 64, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + if fmt.DEQ == "asm": + # 2-byte elements: <2 x i16> = 4 packed fp8 per VGPR out of one + # dword load. Same byte as nope_row (both even), addressed through + # the bf16 view; the layout convert is a rename (shared dim-0 tiling). + row16 = gl.convert_layout(nope_row >> 1, gl.SliceLayout(1, cfg.gather16_l)) + x_u8 = _cache_load( + seg.alt_ptr, + row16, + offs_full16, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + else: + x_u8 = _cache_load( + seg.cache_ptr, + nope_row, + offs_full, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + bgr, pgr, _ = _slots( + cfg, + seg, + k_start + k_rng_rope, + seg_hi, + 0, + False, + cfg.UNI_TILE, + ) + k_rope = _cache_load( + seg.alt_ptr, + bgr * (cs0 // 2) + pgr * fmt.TOK_U16 + fmt.ROPE_U16_OFF, + offs_rope, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + return x_u8, sc, k_rope, valid + + +@gluon.jit +def _stage(cfg, seg, x_u8, sc, k_rope, kv_smem, rope_smem): + """Write one prefetched tile into the LDS buffer(s). The KV buffer is always the + full KV_DIM-wide dequant: "fp8_dsv4_mla" gathers KV_DIM bytes too, the last 64 being + bf16 rope read as garbage fp8 and overwritten by the slice-store below, so + the gather stays pow-2 wide.""" + fmt = seg.fmt + if cfg.FP8_MFMA: + # No dequant: the scale is folded outside the loop (qk_scale on the K + # side, the accumulator on the V side), so what lands in LDS is exactly + # what the gather returned. + kv_smem.store(x_u8.to(gl.float8e4nv, bitcast=True)) + if cfg.ROPE_SEPARATE: + rope_smem.store(k_rope.to(gl.float8e4nv, bitcast=True)) + else: + _deq_store_tile(x_u8, sc, kv_smem, cfg, fmt) + if fmt.KIND == "fp8_dsv4_mla": + kv_smem.slice(fmt.NOPE_DIM, cfg.ROPE_DIM, dim=1).store(k_rope) + elif fmt.KIND == "fp8_dsv32_mla": + rope_smem.store(k_rope) + elif fmt.KIND == "fp8_scalar" and cfg.ROPE_SEPARATE: + rope_smem.store(_fp8_to_bf16(k_rope)) + # "fp8_g64": the whole head is one fp8 tile; nothing else to store. + + +@gluon.jit +def _qkpv( + cfg, + seg, + x_u8, + sc, + k_rope, + valid, + q_dot, + q_rope_dot, + m_i, + l_i, + acc, + head_mask, + qk_scale, + v_scale, + kv_smem, + rope_smem, + k_start=0, + seg_hi=0, +): + """Stage a prefetched fp8 tile into LDS, then QK -> softmax -> PV.""" + _stage(cfg, seg, x_u8, sc, k_rope, kv_smem, rope_smem) + return _qkpv_lds( + cfg, + seg, + valid, + q_dot, + q_rope_dot, + m_i, + l_i, + acc, + head_mask, + qk_scale, + v_scale, + kv_smem, + rope_smem, + k_start, + seg_hi, + ) + + +@gluon.jit +def _qkpv_lds( + cfg, + seg, + valid, + q_dot, + q_rope_dot, + m_i, + l_i, + acc, + head_mask, + qk_scale, + v_scale, + kv_smem, + rope_smem, + k_start=0, + seg_hi=0, +): + """QK -> softmax -> PV over a tile already resident in LDS.""" + neg_inf = float("-inf") + S = _qk_scores(cfg, q_dot, q_rope_dot, kv_smem, rope_smem) + # UNI_TILE folds the tile's range test into the score mask, which is what + # makes the partial last tile correct without a peeled masked body. + COL_MASK: gl.constexpr = cfg.HAS_INVALID or cfg.UNI_TILE + NEED_MASK: gl.constexpr = COL_MASK or (not cfg.HEAD_ALIGNED) + if NEED_MASK: + if COL_MASK: + if cfg.UNI_TILE and not cfg.HAS_INVALID: + # Build the range mask directly in the MFMA layout; converting + # the slot-layout valid vector costs a cross-lane convert per + # tile. + col_mask = ( + k_start + + gl.arange(0, cfg.BLOCK_K, layout=gl.SliceLayout(0, cfg.qk_layout)) + < seg_hi + )[None, :] + else: + col_mask = gl.convert_layout(valid, gl.SliceLayout(0, cfg.qk_layout))[ + None, : + ] + if not cfg.HEAD_ALIGNED: + col_mask = ( + gl.convert_layout(head_mask, gl.SliceLayout(1, cfg.qk_layout))[ + :, None + ] + & col_mask + ) + else: + col_mask = gl.convert_layout(head_mask, gl.SliceLayout(1, cfg.qk_layout))[ + :, None + ] + S = gl.where(col_mask, S, neg_inf) + # Online softmax in the base-2 exponent domain; m_i carries qk_scale + # (= scale * log2e [* k_scale for "fp8_scalar"]). max commutes with a positive + # scale, so scale the row max instead of every element of S; what is left, + # S * qk_scale - m_new, lowers to one FMA, and -inf columns stay -inf. + m_block = _rmax(S, 1) * qk_scale + m_new = _max2(m_i, m_block) + m_new = gl.where(m_new > neg_inf, m_new, 0.0) + p = gl.exp2(S * qk_scale - m_new[:, None]) + alpha = gl.exp2(m_i - m_new) + l_new = l_i * alpha + gl.sum(p, axis=1) + if cfg.ASYNC_LDS and cfg.RELAXED_LOAD: + v = gl.amd.cdna4.async_copy.load_shared_relaxed(kv_smem, cfg.v_layout) + else: + v = kv_smem.load(cfg.v_layout) + if cfg.ASYNC_LDS: + v = v.to(gl.float8e4nv, bitcast=True) + # "fp8_scalar": V was staged as raw fp8 code points; apply the per-tensor scale + # on the small side (p) and leave l scale-free: out = sum(p*s*V)/l exactly. + if seg.fmt.KIND == "fp8_scalar" and not cfg.FP8_MFMA: + p = p * v_scale + if cfg.FP8_MFMA: + p_dot = gl.convert_layout(p.to(gl.float8e4nv), cfg.p_layout) + else: + p_dot = gl.convert_layout(p.to(gl.bfloat16), cfg.p_layout) + alpha_pv = gl.convert_layout(alpha, gl.SliceLayout(1, cfg.pv_layout)) + acc = acc * alpha_pv[:, None] + acc = gl.amd.cdna4.mfma(p_dot, v, acc) + return m_new, l_new, acc + + +# --------------------------------------------------------------------------- +# Non-prefetched path: bf16 segments, and the peeled masked tail when UNI_TILE +# is off (dsv4/uniform only; tensor/dsmla require UNI_TILE). +# --------------------------------------------------------------------------- + + +@gluon.jit +def _decode_tile( + cfg, + seg, + q_dot, + q_rope_dot, + k_start, + hi, + m_i, + l_i, + acc, + head_mask, + qk_scale, + v_scale, + kv_smem, + rope_smem, + offs_full, + offs_full16, + offs_rope, + k_rng_slot, + k_rng_rope, + MASKED: gl.constexpr, +): + """One KV tile -> online-softmax update. MASKED=True is the peeled tail + (fully predicated); full tiles clamp -1 sentinels and mask scores when + HAS_INVALID.""" + neg_inf = float("-inf") + fmt = seg.fmt + cs0 = seg.cs0 + if not fmt.USE_BUFFER_LOAD: + cs0 = cs0.to(gl.int64) # >2 GB cache: 64-bit gather offsets + block_idx, pos, valid1d = _slots( + cfg, + seg, + k_start + k_rng_slot, + hi, + seg.num_rows, + MASKED, + ) + block_idx_g = gl.convert_layout(block_idx, gl.SliceLayout(1, cfg.gather_l)) + pos_g = gl.convert_layout(pos, gl.SliceLayout(1, cfg.gather_l)) + if MASKED: + valid_g = gl.convert_layout(valid1d, gl.SliceLayout(1, cfg.gather_l)) + + if fmt.KIND == "fp8_g64": + NGRP: gl.constexpr = cfg.KV_DIM // 64 + kv_row = block_idx_g * cs0 + pos_g * cfg.KV_DIM + scl_row = block_idx_g * NGRP + scl_col = offs_full // 64 + if MASKED: + x_u8 = _cache_load( + seg.cache_ptr, + kv_row, + offs_full, + fmt.USE_BUFFER_LOAD, + mask=valid_g[:, None], + other=0, + CACHE=cfg.GATHER_CACHE, + ) + sc = _cache_load( + seg.alt_ptr, + scl_row, + scl_col, + fmt.USE_BUFFER_LOAD, + mask=valid_g[:, None], + other=0.0, + CACHE=cfg.GATHER_CACHE, + ) + else: + x_u8 = _cache_load( + seg.cache_ptr, + kv_row, + offs_full, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + sc = _cache_load( + seg.alt_ptr, + scl_row, + scl_col, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + _deq_store_tile(x_u8, sc, kv_smem, cfg, fmt) + elif fmt.KIND == "fp8_dsv4_mla": + nope_row = block_idx_g * cs0 + pos_g * fmt.TOK_U8 + scl_row = ( + block_idx_g * cs0 + fmt.BLOCK_SIZE * fmt.TOK_U8 + pos_g * fmt.SCL_TRAILER_U8 + ) + scl_col = offs_full // 64 + if MASKED: # scales first: see _gather_full + if fmt.NARROW_SCALE and not fmt.USE_BUFFER_LOAD: + exps = _scale_load( + seg.cache_ptr, + scl_row, + valid1d, + fmt.USE_BUFFER_LOAD, + cfg.gather_l, + fmt.scl_l, + fmt.NG, + cfg.KV_DIM, + True, + 127, + ) + else: + exps = _cache_load( + seg.cache_ptr, + scl_row, + scl_col, + fmt.USE_BUFFER_LOAD, + mask=valid_g[:, None], + other=127, + CACHE=cfg.GATHER_CACHE, + ) + if fmt.DEQ == "asm": + row16 = gl.convert_layout( + nope_row >> 1, gl.SliceLayout(1, cfg.gather16_l) + ) + x_u8 = _cache_load( + seg.alt_ptr, + row16, + offs_full16, + fmt.USE_BUFFER_LOAD, + mask=gl.convert_layout(valid1d, gl.SliceLayout(1, cfg.gather16_l))[ + :, None + ], + other=0.0, + CACHE=cfg.GATHER_CACHE, + ) + else: + x_u8 = _cache_load( + seg.cache_ptr, + nope_row, + offs_full, + fmt.USE_BUFFER_LOAD, + mask=valid_g[:, None], + other=0, + CACHE=cfg.GATHER_CACHE, + ) + else: + if fmt.NARROW_SCALE and not fmt.USE_BUFFER_LOAD: + exps = _scale_load( + seg.cache_ptr, + scl_row, + scl_row, + fmt.USE_BUFFER_LOAD, + cfg.gather_l, + fmt.scl_l, + fmt.NG, + cfg.KV_DIM, + False, + 127, + ) + else: + exps = _cache_load( + seg.cache_ptr, + scl_row, + scl_col, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + if fmt.DEQ == "asm": + row16 = gl.convert_layout( + nope_row >> 1, gl.SliceLayout(1, cfg.gather16_l) + ) + x_u8 = _cache_load( + seg.alt_ptr, + row16, + offs_full16, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + else: + x_u8 = _cache_load( + seg.cache_ptr, + nope_row, + offs_full, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + _deq_store_tile(x_u8, exps, kv_smem, cfg, fmt) + block_idx_gr, pos_gr, valid_gr = _slots( + cfg, + seg, + k_start + k_rng_rope, + hi, + seg.num_rows, + MASKED, + ) + rope_row = block_idx_gr * (cs0 // 2) + pos_gr * fmt.TOK_U16 + fmt.ROPE_U16_OFF + if MASKED: + k_rope = _cache_load( + seg.alt_ptr, + rope_row, + offs_rope, + fmt.USE_BUFFER_LOAD, + mask=valid_gr[:, None], + other=0.0, + CACHE=cfg.GATHER_CACHE, + ) + else: + k_rope = _cache_load( + seg.alt_ptr, + rope_row, + offs_rope, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + kv_smem.slice(fmt.NOPE_DIM, cfg.ROPE_DIM, dim=1).store(k_rope) + else: # "bf16" (tensor/dsmla require UNI_TILE, so they never come here) + kv_row2 = block_idx_g * cs0 + pos_g * fmt.TOK_EL + if MASKED: + kv = _cache_load( + seg.alt_ptr, + kv_row2, + offs_full, + fmt.USE_BUFFER_LOAD, + mask=valid_g[:, None], + other=0.0, + CACHE=cfg.GATHER_CACHE, + ) + else: + kv = _cache_load( + seg.alt_ptr, + kv_row2, + offs_full, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + kv_smem.store(kv) + if cfg.ROPE_SEPARATE: + block_idx_gr, pos_gr, valid_gr = _slots( + cfg, + seg, + k_start + k_rng_rope, + hi, + seg.num_rows, + MASKED, + ) + rope_row = block_idx_gr * cs0 + pos_gr * fmt.TOK_EL + cfg.KV_DIM + if MASKED: + k_rope = _cache_load( + seg.alt_ptr, + rope_row, + offs_rope, + fmt.USE_BUFFER_LOAD, + mask=valid_gr[:, None], + other=0.0, + CACHE=cfg.GATHER_CACHE, + ) + else: + k_rope = _cache_load( + seg.alt_ptr, + rope_row, + offs_rope, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + rope_smem.store(k_rope) + + S = _qk_scores(cfg, q_dot, q_rope_dot, kv_smem, rope_smem) + COL_VALID: gl.constexpr = MASKED or cfg.HAS_INVALID + NEED_MASK: gl.constexpr = COL_VALID or (not cfg.HEAD_ALIGNED) + if NEED_MASK: + if COL_VALID: + col_mask = gl.convert_layout(valid1d, gl.SliceLayout(0, cfg.qk_layout))[ + None, : + ] + if not cfg.HEAD_ALIGNED: + col_mask = ( + gl.convert_layout(head_mask, gl.SliceLayout(1, cfg.qk_layout))[ + :, None + ] + & col_mask + ) + else: + col_mask = gl.convert_layout(head_mask, gl.SliceLayout(1, cfg.qk_layout))[ + :, None + ] + S = gl.where(col_mask, S, neg_inf) + + # exp2 softmax with qk_scale folded in; masked cols (-inf) give exp2 = 0. + S = S * qk_scale + m_block = _rmax(S, 1) + m_new = _max2(m_i, m_block) + m_new = gl.where(m_new > neg_inf, m_new, 0.0) # guard all-masked rows + p = gl.exp2(S - m_new[:, None]) + alpha = gl.exp2(m_i - m_new) + l_new = l_i * alpha + gl.sum(p, axis=1) + + v = kv_smem.load(cfg.v_layout) # [BLOCK_K, KV_DIM] + if seg.fmt.KIND == "fp8_scalar": + p = p * v_scale # per-tensor V scale on the small side (see _qkpv) + p_dot = gl.convert_layout(p.to(gl.bfloat16), cfg.p_layout) + alpha_pv = gl.convert_layout(alpha, gl.SliceLayout(1, cfg.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( + cfg, + seg, + q_dot, + q_rope_dot, + lo, + hi, + m_i, + l_i, + acc, + head_mask, + qk_scale, + v_scale, + kv_smem, + rope_smem, +): + offs_full = gl.arange(0, cfg.KV_DIM, layout=gl.SliceLayout(0, cfg.gather_l)) + offs_full16 = gl.arange( + 0, cfg.KV_DIM // 2, layout=gl.SliceLayout(0, cfg.gather16_l) + ) + offs_rope = gl.arange(0, cfg.ROPE_L, layout=gl.SliceLayout(0, cfg.gather_rope_l)) + k_rng_slot = gl.arange(0, cfg.BLOCK_K, layout=cfg.slot_l) + k_rng_rope = gl.arange(0, cfg.BLOCK_K, layout=gl.SliceLayout(1, cfg.gather_rope_l)) + + # [lo, hi_full) are full mask-free tiles; only the peeled tail is masked. + hi_full = lo + ((hi - lo) // cfg.BLOCK_K) * cfg.BLOCK_K + + if seg.fmt.IS_FP8: + # UNI_TILE: the partial tile is just the last iteration (no peeled + # masked copy of the body). + if cfg.UNI_TILE: + n_full = (hi - lo + cfg.BLOCK_K - 1) // cfg.BLOCK_K + else: + n_full = (hi_full - lo) // cfg.BLOCK_K + if n_full > 0: + kn, ks, kr, vld = _gather_full( + cfg, + seg, + lo, + hi, + offs_full, + offs_full16, + offs_rope, + k_rng_slot, + k_rng_rope, + ) + for i in range(1, n_full): + kn2, ks2, kr2, vld2 = _gather_full( + cfg, + seg, + lo + i * cfg.BLOCK_K, + hi, + offs_full, + offs_full16, + offs_rope, + k_rng_slot, + k_rng_rope, + ) + m_i, l_i, acc = _qkpv( + cfg, + seg, + kn, + ks, + kr, + vld, + q_dot, + q_rope_dot, + m_i, + l_i, + acc, + head_mask, + qk_scale, + v_scale, + kv_smem, + rope_smem, + lo + (i - 1) * cfg.BLOCK_K, + hi, + ) + kn, ks, kr, vld = kn2, ks2, kr2, vld2 + m_i, l_i, acc = _qkpv( + cfg, + seg, + kn, + ks, + kr, + vld, + q_dot, + q_rope_dot, + m_i, + l_i, + acc, + head_mask, + qk_scale, + v_scale, + kv_smem, + rope_smem, + lo + (n_full - 1) * cfg.BLOCK_K, + hi, + ) + else: + for k_start in range(lo, hi_full, cfg.BLOCK_K): + m_i, l_i, acc = _decode_tile( + cfg, + seg, + q_dot, + q_rope_dot, + k_start, + hi, + m_i, + l_i, + acc, + head_mask, + qk_scale, + v_scale, + kv_smem, + rope_smem, + offs_full, + offs_full16, + offs_rope, + k_rng_slot, + k_rng_rope, + False, + ) + + if ((not cfg.UNI_TILE) or (not seg.fmt.IS_FP8)) and hi_full < hi: + m_i, l_i, acc = _decode_tile( + cfg, + seg, + q_dot, + q_rope_dot, + hi_full, + hi, + m_i, + l_i, + acc, + head_mask, + qk_scale, + v_scale, + kv_smem, + rope_smem, + offs_full, + offs_full16, + offs_rope, + k_rng_slot, + k_rng_rope, + True, + ) + return m_i, l_i, acc + + +_sparse_mla_repr = make_kernel_repr( + "_sparse_mla", + ["BLOCK_M", "BLOCK_K", "HEAD_SIZE", "NUM_SPLITS", "MAIN_FMT", "ROPE_SEPARATE"], +) + + +@gluon.jit(repr=_sparse_mla_repr) +def _sparse_mla( + # Shapes below: C = queries, H = num_heads, S = HEAD_SIZE (the V width), + # R = ROPE_DIM, nnz = total gathered tokens in a segment's index list. + q_ptr, # [C, H, S (+R when ROPE_SEPARATE)] bf16 + # One segment = a cache plus its index list. The cache is paged + # [num_blocks, BLOCK_SIZE, row] (BLOCK_SIZE = 1 for a flat pool), and + # indices[indptr[t]:indptr[t + 1]] are the rows query t attends to. The two + # cache pointers are the same allocation under different element types; + # which of them is live depends on the format (see Seg). + main_cache_ptr, # main (SWA) cache, u8 view + main_cache_bf16_ptr, # bf16 view of it, or the f32 scale pool ("fp8_g64") + main_indices_ptr, # [nnz_main] int32 row ids + main_indptr_ptr, # [C + 1] int32 + extra_cache_ptr, # top-k segment; aliases main when HAS_EXTRA=False + extra_cache_bf16_ptr, + extra_indices_ptr, # [nnz_extra] int32 + extra_indptr_ptr, # [C + 1] int32 + attn_sink_ptr, # [H] f32, HAS_SINK only + out_ptr, # [C, H, S] bf16, written when NUM_SPLITS == 1 + # Split-K partials, written instead of out_ptr when NUM_SPLITS > 1 (unused + # placeholders otherwise). + part_m_ptr, # [C, NUM_SPLITS, H] f32 row max, base-2 domain + part_l_ptr, # [C, NUM_SPLITS, H] f32 row sum + part_acc_ptr, # [C, NUM_SPLITS, H, S] bf16 or f32, un-normalized + # f32 side-channel per segment: scalar k_scale ("fp8_scalar") or f32 cache view + # ("fp8_dsv32_mla"). None elides the argument, keeping other formats' kernarg + # layouts unchanged. + main_scl_ptr, + extra_scl_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_FMT: gl.constexpr, + EXTRA_FMT: gl.constexpr, + MAIN_BLOCK_SIZE: gl.constexpr, + EXTRA_BLOCK_SIZE: gl.constexpr, + CS0_ALIGN: gl.constexpr, + NOPE_DIM: gl.constexpr, + ROPE_DIM: gl.constexpr, + HEAD_SIZE: gl.constexpr, + ROPE_SEPARATE: gl.constexpr, + BLOCK_M: gl.constexpr, + BLOCK_K: gl.constexpr, + NUM_SPLITS: gl.constexpr, + HEAD_ALIGNED: gl.constexpr, + # NOPE_CHUNK: extent of one dequant piece along CHUNK_AXIS (0 = rows, + # 1 = columns); >= the tile's extent means one shot. + NOPE_CHUNK: gl.constexpr, + CHUNK_AXIS: gl.constexpr, + PART_STORE_CACHE: gl.constexpr, + UNI_TILE: gl.constexpr, + GRID_ORDER: gl.constexpr, + Q_CACHE: gl.constexpr, + # MAIN_SPLITS <= NUM_SPLITS: splitting the SWA window past its tile count + # only manufactures masked partial tiles, so main stops early and extra + # keeps all programs (surplus ones get an empty main range). + MAIN_SPLITS: gl.constexpr, + # ADAPTIVE_SPLITS: re-decide the useful split count per query at runtime. + ADAPTIVE_SPLITS: gl.constexpr, + DEQ: gl.constexpr, # see Fmt.DEQ + # Per-cache buffer/global gate: buffer_load carries a 32-bit offset (2 GB), + # and the two caches are sized independently. + MAIN_USE_BUFFER_LOAD: gl.constexpr, + EXTRA_USE_BUFFER_LOAD: gl.constexpr, + IDX_BUFFER_LOAD: gl.constexpr, + HAS_INVALID: gl.constexpr, + FP8_MFMA: gl.constexpr = False, + # q already quantized to e4m3 by the caller, plus the scalar f32 scale it + # was quantized with. This is the calling convention aiter's asm + # mla_decode_fwd uses, where vLLM passes layer._q_scale. + q_scl_ptr=None, + Q_FP8: gl.constexpr = False, + # [C, H] f32 natural-log LSE, written only when HAS_LSE. + lse_ptr=None, + HAS_LSE: gl.constexpr = False, + # Defaults keep every existing launch byte-identical; the MLA launcher opts in. + GATHER_CACHE: gl.constexpr = _CG, + IDX_CACHE: gl.constexpr = _NO_CACHE, + ASYNC_LDS: gl.constexpr = False, + RELAXED_LOAD: gl.constexpr = True, + PAD_INTERVAL: gl.constexpr = 1024, +): + """One program = (query, split, head-block). Two-loop: main (SWA) then + extra (top-k). NUM_SPLITS==1 writes the output directly; otherwise stores + un-normalized partials for the reduce kernel.""" + NUM_WARPS: gl.constexpr = gl.num_warps() + gl.static_assert( + UNI_TILE or (MAIN_FMT != "fp8_scalar" and MAIN_FMT != "fp8_dsv32_mla"), + "tensor/dsmla formats require UNI_TILE=1", + ) + gl.static_assert( + UNI_TILE or (EXTRA_FMT != "fp8_scalar" and EXTRA_FMT != "fp8_dsv32_mla"), + "tensor/dsmla formats require UNI_TILE=1", + ) + gl.static_assert( + (not ROPE_SEPARATE) or (MAIN_FMT != "fp8_dsv4_mla" and MAIN_FMT != "fp8_g64"), + "dsv4/uniform formats carry rope inside the KV buffer (ROPE_SEPARATE=False)", + ) + gl.static_assert( + MAIN_FMT != "fp8_dsv32_mla" or ROPE_SEPARATE, + "fp8_dsv32_mla is a separated-rope (MLA) format", + ) + # No rope means no rope buffer, and the packed formats all define a rope tail. + gl.static_assert( + ROPE_DIM > 0 or not ROPE_SEPARATE, + "ROPE_DIM=0 is the rope-free geometry; it needs ROPE_SEPARATE=False", + ) + gl.static_assert( + ROPE_DIM > 0 + or ( + MAIN_FMT != "fp8_dsv4_mla" + and MAIN_FMT != "fp8_dsv32_mla" + and ( + not HAS_EXTRA + or (EXTRA_FMT != "fp8_dsv4_mla" and EXTRA_FMT != "fp8_dsv32_mla") + ) + ), + "fp8_dsv4_mla/fp8_dsv32_mla rows carry a rope tail; ROPE_DIM=0 is " + "inconsistent", + ) + # The fp8 path needs one positive scalar scale per cache, since that is what + # folds outside the loop, and OCP e4m3 code points, which is what the matrix + # core reads. + gl.static_assert( + (not FP8_MFMA) + or (MAIN_FMT == "fp8_scalar" and (not HAS_EXTRA or EXTRA_FMT == "fp8_scalar")), + "FP8_MFMA requires the per-tensor fp8 format on every segment", + ) + # The direct-to-LDS path stages raw code points, needs the clamped (branch-free) + # tail, and carries no per-tile validity vector. + # TODO(has_invalid): the sentinel restriction is conservative, not fundamental. + gl.static_assert( + (not ASYNC_LDS) + or (FP8_MFMA and UNI_TILE and not HAS_INVALID and not HAS_EXTRA), + "ASYNC_LDS requires FP8_MFMA + UNI_TILE, no -1 sentinels, one segment", + ) + gl.static_assert( + not (FP8_MFMA and HAS_EXTRA), + "FP8_MFMA defers the V-side scale to the epilogue, so it needs one segment", + ) + # Row bases are block*cs0 + pos*TOK with runtime block/pos, so divisibility + # analysis sees 1-byte alignment unless the driver vouches for cs0. + if CS0_ALIGN > 1: + main_cs0 = gl.multiple_of(main_cs0, CS0_ALIGN) + extra_cs0 = gl.multiple_of(extra_cs0, CS0_ALIGN) + # GRID_ORDER names the launch axes in grid-dim order; dim 0 varies fastest, + # which decides XCD/L2 sharing. + query_idx = gl.program_id(GRID_ORDER.index("q")) + split_id = gl.program_id(GRID_ORDER.index("s")) + pid_h = gl.program_id(GRID_ORDER.index("h")) + + cfg = Cfg( + BLOCK_M, + BLOCK_K, + HEAD_SIZE, + ROPE_DIM, + ROPE_SEPARATE, + NUM_WARPS, + UNI_TILE, + HAS_INVALID, + HEAD_ALIGNED, + IDX_BUFFER_LOAD, + FP8_MFMA, + GATHER_CACHE, + IDX_CACHE, + ASYNC_LDS, + RELAXED_LOAD, + PAD_INTERVAL, + ) + main_fmt = Fmt( + cfg, + MAIN_FMT, + MAIN_BLOCK_SIZE, + MAIN_USE_BUFFER_LOAD, + DEQ, + NOPE_DIM, + NOPE_CHUNK, + CHUNK_AXIS, + ) + extra_fmt = Fmt( + cfg, + EXTRA_FMT, + EXTRA_BLOCK_SIZE, + EXTRA_USE_BUFFER_LOAD, + DEQ, + NOPE_DIM, + NOPE_CHUNK, + CHUNK_AXIS, + ) + + h_off = pid_h * BLOCK_M + + # Segment lengths issue first: they gate the three-deep memory chain + # (indptr -> indices -> cache), and Q is independent of all of it. + main_start = gl.load(main_indptr_ptr + query_idx) + main_end = gl.load(main_indptr_ptr + query_idx + 1) + 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 + else: + extra_start = 0 + extra_len = 0 + main_len = main_end - main_start + + # exp2 softmax + RCP_LN2: gl.constexpr = 1.4426950408889634 + LN2: gl.constexpr = 0.6931471805599453 + qk_scale = scale * RCP_LN2 + main_qk_scale = qk_scale + main_v_scale = 1.0 + if MAIN_FMT == "fp8_scalar": + main_k_scale = gl.load(main_scl_ptr) + main_qk_scale = qk_scale * main_k_scale + main_v_scale = main_k_scale + extra_qk_scale = qk_scale + extra_v_scale = 1.0 + if HAS_EXTRA and EXTRA_FMT == "fp8_scalar": + extra_k_scale = gl.load(extra_scl_ptr) + extra_qk_scale = qk_scale * extra_k_scale + extra_v_scale = extra_k_scale + + # Load Q. ROPE_SEPARATE loads the two pieces separately (the combined + # width is not a pow-2 arange), each converted to its dot layout. + offs_m_q = gl.arange(0, BLOCK_M, layout=gl.SliceLayout(1, cfg.blocked_q)) + offs_d_q = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, cfg.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, cache=Q_CACHE + ) + if FP8_MFMA and not Q_FP8: + # bf16 q: quantize here, one e4m3 scale for this program's whole Q tile + # (nope and rope), so the fold below is one extra factor on qk_scale. + E4M3_MAX: gl.constexpr = 448.0 + q_amax = gl.max(gl.max(gl.abs(q).to(gl.float32), axis=1), axis=0) + q_dot = gl.convert_layout(q, cfg.q_layout) + if ROPE_SEPARATE: + offs_d_qr = gl.arange(0, ROPE_DIM, layout=gl.SliceLayout(0, cfg.blocked_q)) + qr_off = ( + query_idx * q_stride0 + + h_q[:, None] * q_stride1 + + HEAD_SIZE + + offs_d_qr[None, :] + ).to(gl.int32) + q_rope = gl.amd.cdna4.buffer_load( + ptr=q_ptr, offsets=qr_off, mask=h_mask_q[:, None], other=0.0, cache=Q_CACHE + ) + q_rope_dot = gl.convert_layout(q_rope, cfg.q_layout) + else: + q_rope_dot = q_dot # unused (single-buffer QK) -> DCE'd + + if Q_FP8: + # Nothing to quantize + q_scale = gl.load(q_scl_ptr) + if not FP8_MFMA: + # fp8 -> bf16 is exact (3 mantissa bits into 8) + q_dot = gl.convert_layout(q.to(gl.bfloat16), cfg.q_layout) + if ROPE_SEPARATE: + q_rope_dot = gl.convert_layout(q_rope.to(gl.bfloat16), cfg.q_layout) + else: + q_rope_dot = q_dot + main_qk_scale = main_qk_scale * q_scale + extra_qk_scale = extra_qk_scale * q_scale + elif FP8_MFMA: + if ROPE_SEPARATE: + q_amax = gl.maximum( + q_amax, gl.max(gl.max(gl.abs(q_rope).to(gl.float32), axis=1), axis=0) + ) + q_amax = gl.maximum(q_amax, 1e-30) + q_rcp = E4M3_MAX / q_amax + q_dot = gl.convert_layout( + (q.to(gl.float32) * q_rcp).to(gl.float8e4nv), cfg.q_layout + ) + if ROPE_SEPARATE: + q_rope_dot = gl.convert_layout( + (q_rope.to(gl.float32) * q_rcp).to(gl.float8e4nv), cfg.q_layout + ) + else: + q_rope_dot = q_dot + q_scale = q_amax / E4M3_MAX + main_qk_scale = main_qk_scale * q_scale + extra_qk_scale = extra_qk_scale * q_scale + + # head mask in pv-slice layout (for output / partial masking) + offs_m_pv = gl.arange(0, BLOCK_M, layout=gl.SliceLayout(1, cfg.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, cfg.qk_layout) + ) + l_i = gl.zeros([BLOCK_M], gl.float32, layout=gl.SliceLayout(1, cfg.qk_layout)) + acc = gl.zeros([BLOCK_M, HEAD_SIZE], gl.float32, layout=cfg.pv_layout) + + # An fp8 buffer is half the bytes of the bf16 staging it replaces. + SMEM_DT: gl.constexpr = gl.float8e4nv if FP8_MFMA else gl.bfloat16 + # The LDS-DMA converts nothing, so an async buffer's element type has to be the + # cache's own u8; the dot operands bitcast on read (free, layout-preserving). + BUF_DT: gl.constexpr = gl.uint8 if ASYNC_LDS else SMEM_DT + kv_smem = gl.allocate_shared_memory(BUF_DT, [BLOCK_K, HEAD_SIZE], cfg.kv_shared) + if ROPE_SEPARATE: + rope_smem = gl.allocate_shared_memory( + BUF_DT, [BLOCK_K, ROPE_DIM], cfg.rope_shared + ) + else: + rope_smem = kv_smem # never read as the rope buffer in this geometry + + # ADAPTIVE_SPLITS: the host split count is sized for batch averages; in a + # ragged batch the surplus programs would each gather a mostly-masked tile + # and write a full partial. Recompute from this query's own lengths and let + # those programs write a neutral partial (m = -inf) and leave. The reduce + # skips them, so their part_acc never has to be written. + if ADAPTIVE_SPLITS: + m_tiles = (main_len + BLOCK_K - 1) // BLOCK_K + e_tiles = (extra_len + BLOCK_K - 1) // BLOCK_K + work_splits = gl.minimum( + gl.maximum(gl.maximum(m_tiles, e_tiles), 1), NUM_SPLITS + ) + main_splits = gl.minimum(gl.maximum(m_tiles, 1), work_splits) + if split_id >= work_splits: + pm_base = query_idx * pm_stride0 + split_id * pm_stride_s + gl.amd.cdna4.buffer_store( + gl.full( + [BLOCK_M], + float("-inf"), + gl.float32, + layout=gl.SliceLayout(1, cfg.pv_layout), + ), + ptr=part_m_ptr + pm_base, + offsets=h_pv.to(gl.int32), + mask=head_mask_pv, + ) + gl.amd.cdna4.buffer_store( + gl.zeros( + [BLOCK_M], gl.float32, layout=gl.SliceLayout(1, cfg.pv_layout) + ), + ptr=part_l_ptr + pm_base, + offsets=h_pv.to(gl.int32), + mask=head_mask_pv, + ) + return + else: + work_splits = NUM_SPLITS + main_splits = MAIN_SPLITS + + # main (SWA) segment + main_seg = Seg( + main_fmt, + main_cache_ptr, + main_cache_bf16_ptr, + main_scl_ptr if (MAIN_FMT == "fp8_dsv32_mla") else main_cache_ptr, + main_indices_ptr, + main_start, + main_cs0, + main_num_rows, + ) + main_chunk = (main_len + main_splits - 1) // main_splits + main_lo = gl.minimum(split_id * main_chunk, main_len) + main_hi = gl.minimum(main_lo + main_chunk, main_len) + if ASYNC_LDS: + m_i, l_i, acc = _async_segment( + cfg, + main_seg, + q_dot, + q_rope_dot, + main_lo, + main_hi, + m_i, + l_i, + acc, + head_mask_pv, + main_qk_scale, + main_v_scale, + kv_smem, + rope_smem, + ) + else: + m_i, l_i, acc = _process_segment( + cfg, + main_seg, + q_dot, + q_rope_dot, + main_lo, + main_hi, + m_i, + l_i, + acc, + head_mask_pv, + main_qk_scale, + main_v_scale, + kv_smem, + rope_smem, + ) + + if HAS_EXTRA: + extra_seg = Seg( + extra_fmt, + extra_cache_ptr, + extra_cache_bf16_ptr, + extra_scl_ptr if (EXTRA_FMT == "fp8_dsv32_mla") else extra_cache_ptr, + extra_indices_ptr, + extra_start, + extra_cs0, + extra_num_rows, + ) + extra_chunk = (extra_len + work_splits - 1) // work_splits + extra_lo = split_id * extra_chunk + extra_hi = gl.minimum(extra_lo + extra_chunk, extra_len) + m_i, l_i, acc = _process_segment( + cfg, + extra_seg, + q_dot, + q_rope_dot, + extra_lo, + extra_hi, + m_i, + l_i, + acc, + head_mask_pv, + extra_qk_scale, + extra_v_scale, + kv_smem, + rope_smem, + ) + + if FP8_MFMA: + # The fp8 PV dot ran on raw code points, so the V-side scale comes off + # here, once per program instead of once per tile. l is untouched, so + # out = acc*s/l is what the bf16 path computes. + acc = acc * main_v_scale + + # Move the row reductions into pv-slice space for output/partials. + m_pv = gl.convert_layout(m_i, gl.SliceLayout(1, cfg.pv_layout)) + l_pv = gl.convert_layout(l_i, gl.SliceLayout(1, cfg.pv_layout)) + + if NUM_SPLITS == 1: + if HAS_SINK: + # m_pv is in the base-2 exponent domain; lift the sink into it. + sink = ( + gl.amd.cdna4.buffer_load( + ptr=attn_sink_ptr, + offsets=h_pv, + mask=head_mask_pv, + other=float("-inf"), + ).to(gl.float32) + * RCP_LN2 + ) + m_final = _max2(m_pv, sink) + alpha = gl.exp2(m_pv - m_final) + l_final = l_pv * alpha + gl.exp2(sink - m_final) + acc = acc * alpha[:, None] + else: + m_final = m_pv + l_final = l_pv + one_over_l = 1.0 / l_final + out = acc * one_over_l[:, None] + offs_d_o = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, cfg.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], + ) + if HAS_LSE: + # m is base-2, so sum_j exp(s_j) = 2^m * l and ln of it is + # (m + log2 l) * ln2. A fully masked row keeps -inf, not NaN. + gl.amd.cdna4.buffer_store( + (m_final + gl.log2(l_final)) * LN2, + ptr=lse_ptr + query_idx * num_heads, + offsets=h_pv.to(gl.int32), + mask=head_mask_pv, + ) + else: + # Un-normalized partials for the reduce kernel; m stays in the base-2 + # exponent domain (the triton reduce's convention too). + pm_base = query_idx * pm_stride0 + split_id * pm_stride_s + gl.amd.cdna4.buffer_store( + m_pv, + ptr=part_m_ptr + pm_base, + offsets=h_pv.to(gl.int32), + mask=head_mask_pv, + cache=PART_STORE_CACHE, + ) + gl.amd.cdna4.buffer_store( + l_pv, + ptr=part_l_ptr + pm_base, + offsets=h_pv.to(gl.int32), + mask=head_mask_pv, + cache=PART_STORE_CACHE, + ) + offs_d_a = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, cfg.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) + # Follow part_acc's own dtype: bf16 halves the partial HBM traffic. + gl.amd.cdna4.buffer_store( + acc.to(part_acc_ptr.dtype.element_ty), + ptr=part_acc_ptr, + offsets=a_off, + mask=head_mask_pv[:, None], + cache=PART_STORE_CACHE, + ) + + +_sparse_mla_reduce_repr = make_kernel_repr( + "_sparse_mla_reduce", + ["BLOCK_M", "HEAD_SIZE", "NUM_SPLITS"], +) + + +@gluon.jit(repr=_sparse_mla_reduce_repr) +def _sparse_mla_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, + ADAPTIVE_SPLITS: gl.constexpr, + lse_ptr=None, + HAS_LSE: gl.constexpr = False, +): + """Split-KV combine: merge per-split partials, fold the sink, write the + output. Identical for both geometries (partials are HEAD_SIZE = V wide). + Grid: (num_queries, heads_blocks); the combine is pure bandwidth, so + BLOCK_M is sized for workgroup count, not for the attention kernel's tile.""" + NUM_WARPS: gl.constexpr = gl.num_warps() + RCP_LN2: gl.constexpr = 1.4426950408889634 + LN2: gl.constexpr = 0.6931471805599453 + query_idx = gl.program_id(0) + pid_h = gl.program_id(1) + + # Lay the 64 lanes out so a small BLOCK_M spends them on the head dim. + TPW0: gl.constexpr = min(8, BLOCK_M) + TPW1: gl.constexpr = 64 // TPW0 + BLK: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[TPW0, TPW1], + 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") + # Deliberately NOT bounded by the per-query split count: that would make + # these dynamic loops and lose the static unroll + 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 = _max2(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, + cache=".cg", + ).to(gl.float32) + scaled_sink = sink * RCP_LN2 + m_final = _max2(m_final, scaled_sink) # 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, + cache=".cg", + ) + l_s = gl.amd.cdna4.buffer_load( + ptr=part_l_ptr + base, + offsets=h, + mask=head_mask, + other=0.0, + cache=".cg", + ) + 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) + # split's part_acc can be uninitialized, so mask the load + if ADAPTIVE_SPLITS: + acc_mask = head_mask[:, None] & (m_s > neg_inf)[:, None] + else: + acc_mask = head_mask[:, None] + acc_s = gl.amd.cdna4.buffer_load( + ptr=part_acc_ptr, + offsets=a_off, + mask=acc_mask, + other=0.0, + cache=".cg", + ) + acc = acc + w[:, None] * acc_s.to(gl.float32) + + if HAS_SINK: + l_final = l_final + gl.exp2(scaled_sink - m_final) + + if HAS_LSE: + # Same base-2 -> natural conversion as the no-split epilogue. + gl.amd.cdna4.buffer_store( + (m_final + gl.log2(l_final)) * LN2, + ptr=lse_ptr + query_idx * num_heads, + offsets=h.to(gl.int32), + mask=head_mask, + ) + + # One reciprocal per row instead of a per-element f32 divide. + one_over_l = 1.0 / l_final + out = acc * one_over_l[:, 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 c6fafb181e6..423060c670c 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -1,19 +1,10 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Sparse paged-decode attention over a unified KV pool with per-token paged -indices. See ``_triton_kernels/attention/pa_decode_sparse.py`` for the -kernels' caller contract. - -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). +"""Sparse paged-decode attention over a unified or split KV pool with per-token paged +indices. + +TODO: add details once API has settled """ import math @@ -21,11 +12,11 @@ import torch import triton -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.sparse_mla import ( + _sparse_mla as _sparse_mla_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._gluon_kernels.gfx950.attention.sparse_mla import ( + _sparse_mla_reduce as _sparse_mla_reduce_gfx950, ) from aiter.ops.triton._gluon_kernels.gfx1250.attention.pa_decode_sparse import ( _pa_decode_sparse as gluon_pa_decode_sparse, @@ -43,14 +34,34 @@ from aiter.ops.triton.utils.common_utils import max_addressable_bytes from aiter.ops.triton.utils.device_info import get_num_sms from aiter.ops.triton.utils.logger import AiterTritonLogger +from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype DEVICE_ARCH = arch_info.get_arch() +# Fused fp8 x E8M0 -> bf16 upcast check (gluon cdna4.scaled_upcast) +try: + from triton.experimental.gluon.language.amd import cdna4 as _cdna4 + + _HAS_SCALED_UPCAST = hasattr(_cdna4, "scaled_upcast") +except ImportError: + _HAS_SCALED_UPCAST = False + _LOGGER = AiterTritonLogger() _FP8_GROUP_SIZE = 64 -_FP8_DTYPE = torch.float8_e4m3fnuz +_FP8_DTYPE = get_fp8_e4m3_dtype() + + +def _check_out(out, q, dtype): + """Caller-supplied output buffer, or a fresh one. Writing the caller's buffer + directly saves a full [T, H, D] device copy per call.""" + if out is None: + return torch.empty_like(q, dtype=dtype) + assert out.shape == q.shape, f"out shape {tuple(out.shape)} != q {tuple(q.shape)}" + assert out.dtype == dtype, f"out dtype {out.dtype} != {dtype}" + assert out.device == q.device + return out def pa_decode_sparse( @@ -70,6 +81,7 @@ def pa_decode_sparse( extra_cache: torch.Tensor | None = None, extra_indices: torch.Tensor | None = None, extra_indptr: torch.Tensor | None = None, + out: torch.Tensor | None = None, ) -> torch.Tensor: """Sparse paged-decode attention with split-K + widened BLOCK_H. @@ -88,6 +100,8 @@ def pa_decode_sparse( auto-infers to fill ~512 total CTAs while capping below the number of K-blocks, then rounds up to a power of 2. num_stages: software-pipeline depth of the K loop (default 2). + out: optional ``[N, H, D]`` destination. Supplied -> written in place and + returned, which saves the caller a full-size device copy. skip_reduce: when the split-K path is active (``kv_splits > 1``), return the pre-reduce ``(acc_partial, m_partial, l_partial)`` partials instead of launching the reduce kernel. Has no effect when @@ -98,7 +112,8 @@ def pa_decode_sparse( 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 + packed fp8_dsv4_mla (584 B rows) / 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. @@ -120,28 +135,26 @@ def pa_decode_sparse( raise RuntimeError("pa_decode_sparse requires CUDA/HIP tensors") if q.dtype not in (torch.bfloat16, torch.float16): raise RuntimeError(f"pa_decode_sparse expects fp16/bf16 q, got {q.dtype}") + _LOGGER.info( + f"PA_DECODE_SPARSE: q={tuple(q.shape)} unified_kv={tuple(unified_kv.shape)} " + f"{unified_kv.dtype} kv_indices={tuple(kv_indices.shape)}" + ) # 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 + # from the cache: 3D -> packed fp8_dsv4_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: + # gfx950 (CDNA4) reads OCP e4m3 natively. fnuz is the gfx942 encoding, + # so it never appears here and falls through to the triton path below. if unified_kv.ndim == 3: - _ok = kv_scales is None and ( - unified_kv.dtype == torch.uint8 or unified_kv.dtype == q.dtype - ) + # packed / bf16 block cache: it carries its own scales, if any + _ok = kv_scales is None and unified_kv.dtype in (torch.uint8, q.dtype) + elif kv_scales is not None: + _ok = unified_kv.dtype in (torch.float8_e4m3fn, torch.uint8) 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 + _ok = unified_kv.dtype == q.dtype if _ok: cache = ( unified_kv.view(torch.uint8) @@ -162,7 +175,7 @@ def pa_decode_sparse( kv_splits=kv_splits, skip_reduce=skip_reduce, has_invalid=bool(has_invalid), - fp8_fnuz=fp8_fnuz, + out=out, ) assert ( @@ -195,11 +208,8 @@ def pa_decode_sparse( ) T, H, D = q.shape - _LOGGER.info( - f"PA_DECODE_SPARSE T={T} H={H} D={D} " f"total_indices={kv_indices.shape[0]}" - ) - out = torch.empty_like(q) + out = _check_out(out, q, q.dtype) assert kv_indices.dtype == torch.int32 and kv_indices.is_contiguous() assert kv_indptr.dtype == torch.int32 and kv_indptr.is_contiguous() @@ -427,32 +437,20 @@ def _as_int32_contiguous_1d(x: torch.Tensor) -> torch.Tensor: 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. +def _decode_num_splits_occ(num_queries, heads_blocks, avg_main, avg_extra, block_k): + """Split-K count for the gfx950 gluon kernel: fill the machine, but never + split a segment finer than one BLOCK_K tile. """ - 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 + num_sms = get_num_sms() + base_wg = max(1, num_queries * heads_blocks) + cta_cap = max(1, (2 * num_sms) // base_wg) + main_tiles = max(1, math.ceil(avg_main / block_k)) if avg_main > 0 else 0 + extra_tiles = max(1, math.ceil(avg_extra / block_k)) if avg_extra > 0 else 0 + tiles = max(1, main_tiles, extra_tiles) + if base_wg >= num_sms: + # Already at least one workgroup per CU without splitting + return max(1, min(cta_cap, tiles // 4)) + return max(1, min(cta_cap, tiles)) def _pa_decode_sparse_gfx950_gluon( @@ -468,24 +466,25 @@ def _pa_decode_sparse_gfx950_gluon( extra_indptr=None, kv_splits=None, skip_reduce=False, + out=None, 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 + + """Merged gfx950 gluon DSv4 sparse-MLA decode driver. Format from cache.ndim: + 3D [nb, block, 584] -> packed fp8_dsv4_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, + 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). + 2D [pages, D] -> uniform pool: fp8 (uint8) + cache_scales + [pages, D//64] fp32, or bf16 (cache_scales None). """ 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 + # Tuned launch config (gfx950 / MI355). BLOCK_M = heads per MFMA M-tile, and 16 + # is both the MFMA M and the DSv4 head count; BLOCK_K = KV tile; num_warps = + # BLOCK_K // 16, because warps tile the dot-N and MFMA N = 16. + BLOCK_M, BLOCK_K = 16, 64 + num_warps = max(1, BLOCK_K // 16) NOPE_DIM, ROPE_DIM = 448, 64 MAX_BYTES = 2**31 - 1 @@ -502,7 +501,7 @@ def _pa_decode_sparse_gfx950_gluon( 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 + FLAT_POOL = True main_is_fp8 = cache.dtype == torch.uint8 if main_is_fp8: assert cache_scales is not None and cache_scales.dtype == torch.float32 @@ -521,12 +520,13 @@ def _pa_decode_sparse_gfx950_gluon( main_block, extra_block = 1, 1 nope_dim = head_dim main_num_rows = extra_num_rows = cache.shape[0] - cache_bytes = max_addressable_bytes(cache) 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 + # packed fp8_dsv4_mla [nb, block, 584] (UE8M0 block trailer) or bf16 block + # cache. NB: vLLM's fp8_ds_mla also names the 656 B V3.2 layout, which is + # fp8_dsv32_mla and reaches the kernel through sparse_mla.py instead. + FLAT_POOL = False main_is_fp8 = cache.dtype == torch.uint8 main_bf16 = cache.view(torch.bfloat16) if main_is_fp8 else cache has_extra = ( @@ -545,32 +545,60 @@ def _pa_decode_sparse_gfx950_gluon( 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( - max_addressable_bytes(cache), max_addressable_bytes(extra_cache) - ) 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 + # Kernel-side cache-format tags (kernel shared with sparse_mla.py). + if FLAT_POOL: + main_fmt = "fp8_g64" if main_is_fp8 else "bf16" + extra_fmt = main_fmt + else: + main_fmt = "fp8_dsv4_mla" if main_is_fp8 else "bf16" + extra_fmt = "fp8_dsv4_mla" if extra_is_fp8 else "bf16" + + # Alignment hint for the page strides so row gathers can vectorize: the largest + # power of 2 (<= 16) dividing both. + s0, s1 = int(cache.stride(0)), int(extra_cache.stride(0)) + cs0_align = 1 + for a in (16, 8, 4, 2): + if s0 % a == 0 and s1 % a == 0: + cs0_align = a + break + + # Gate each cache on its own span: buffer_load carries a 32-bit offset, and one + # oversized cache must not drop the fast path for the other. The index lists are + # one int32 per gathered token, far under the limit even for a full batch. + main_use_buffer_load = max_addressable_bytes(cache) < MAX_BYTES + extra_use_buffer_load = max_addressable_bytes(extra_cache) < MAX_BYTES + idx_use_buffer_load = ( + max_addressable_bytes(indices) < MAX_BYTES + and max_addressable_bytes(extra_indices) < MAX_BYTES + ) + use_buffer_load = main_use_buffer_load and extra_use_buffer_load HEAD_ALIGNED = num_heads % BLOCK_M == 0 heads_blocks = (num_heads + BLOCK_M - 1) // BLOCK_M - out = torch.empty_like(q, dtype=torch.bfloat16) + out = _check_out(out, q, torch.bfloat16) if kv_splits is not None: num_splits = max(1, int(kv_splits)) else: - num_splits = _decode_num_splits( + num_splits = _decode_num_splits_occ( num_queries, heads_blocks, avg_main, avg_extra, BLOCK_K ) + # Q is read once per query without split-K, and re-read by every split + q_cache = ".cg" if num_splits == 1 else "" + 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) + # bf16 partials halve both the split-K HBM traffic (~31% of the kernel's + # bytes) part_acc = torch.empty( (num_queries, num_splits, num_heads, head_dim), - dtype=torch.float32, + dtype=torch.float32 if skip_reduce else torch.bfloat16, device=q.device, ) pm_stride0, pm_stride_s = part_m.stride(0), part_m.stride(1) @@ -583,8 +611,42 @@ def _pa_decode_sparse_gfx950_gluon( 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 + # Dequant chunking. The gather layout puts 32 of a wave's 64 lanes along a row + # (16 B each, 512 B) and the other 32 on a second row; col_reps is how many + # such spans each lane holds, which is what makes a column split a free + # register rename. + col_reps = head_dim // 512 + chunk_axis = 1 if col_reps >= 4 else 0 + nope_chunk = max(1, BLOCK_K // 4) if chunk_axis == 0 else min(128, head_dim) + + waves_per_eu = 2 + one_wg_per_cu = ( + use_buffer_load and num_queries * heads_blocks * num_splits <= get_num_sms() + ) + if one_wg_per_cu: + waves_per_eu = 1 + + main_splits = num_splits + if has_extra and avg_main > 0: + main_splits = max(1, min(num_splits, math.ceil(avg_main / BLOCK_K))) + + # Per-query split count decided in-kernel; only meaningful when there is more + # than one split to give back. + adaptive_splits = num_splits > 1 + + # Fuse the dsv4 dequant into v_cvt_scalef32_pk_bf16_fp8. The asm fallback + # gathers an extra int16 tile, so it only pays off at one workgroup per CU. + if _HAS_SCALED_UPCAST: + deq = "upcast" + elif one_wg_per_cu: + deq = "asm" + else: + deq = "none" + + # Grid dim 0 varies fastest and XCD assignment is round-robin over the linear + # workgroup id, so the axis order decides what shares an XCD's L2. grid = (num_queries, num_splits, heads_blocks) - _pa_decode_sparse_gfx950[grid]( + _sparse_mla_gfx950[grid]( q, cache, main_bf16, @@ -599,6 +661,10 @@ def _pa_decode_sparse_gfx950_gluon( part_m, part_l, part_acc, + # f32 scale pointers (separated-rope formats only); None is elided, + # keeping the DSv4 kernarg layout unchanged. + None, + None, scale, q.stride(0), q.stride(1), @@ -616,22 +682,35 @@ def _pa_decode_sparse_gfx950_gluon( num_heads, HAS_EXTRA=has_extra, HAS_SINK=has_sink, - MAIN_IS_FP8=main_is_fp8, - EXTRA_IS_FP8=extra_is_fp8, + MAIN_FMT=main_fmt, + EXTRA_FMT=extra_fmt, MAIN_BLOCK_SIZE=main_block, EXTRA_BLOCK_SIZE=extra_block, + CS0_ALIGN=cs0_align, NOPE_DIM=nope_dim, ROPE_DIM=ROPE_DIM, HEAD_SIZE=head_dim, + ROPE_SEPARATE=False, 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, + NOPE_CHUNK=nope_chunk, + CHUNK_AXIS=chunk_axis, + PART_STORE_CACHE="", + Q_CACHE=q_cache, + GRID_ORDER="qsh", + # The partial last tile rides the full-tile body. Gluon inlines, so a peeled + # masked copy would be a second gather+dequant+MFMA body, and its register + # demand spills the tile loop. + UNI_TILE=True, + MAIN_SPLITS=main_splits, + ADAPTIVE_SPLITS=adaptive_splits, + DEQ=deq, + MAIN_USE_BUFFER_LOAD=main_use_buffer_load, + EXTRA_USE_BUFFER_LOAD=extra_use_buffer_load, + IDX_BUFFER_LOAD=idx_use_buffer_load, HAS_INVALID=has_invalid, - FP8_FNUZ=fp8_fnuz, num_warps=num_warps, waves_per_eu=waves_per_eu, ) @@ -641,8 +720,9 @@ def _pa_decode_sparse_gfx950_gluon( if skip_reduce: return part_acc, part_m, part_l - rgrid = (num_queries, heads_blocks) - _pa_decode_sparse_reduce_gfx950[rgrid]( + # One head per reduce workgroup + rgrid = (num_queries, num_heads) + _sparse_mla_reduce_gfx950[rgrid]( part_m, part_l, part_acc, @@ -658,9 +738,10 @@ def _pa_decode_sparse_gfx950_gluon( num_heads, HAS_SINK=has_sink, HEAD_SIZE=head_dim, - BLOCK_M=BLOCK_M, + BLOCK_M=1, NUM_SPLITS=num_splits, - HEAD_ALIGNED=HEAD_ALIGNED, - num_warps=4, + HEAD_ALIGNED=True, + ADAPTIVE_SPLITS=adaptive_splits, + num_warps=1, ) return out diff --git a/aiter/ops/triton/attention/sparse_mla.py b/aiter/ops/triton/attention/sparse_mla.py new file mode 100644 index 00000000000..78ae6294d56 --- /dev/null +++ b/aiter/ops/triton/attention/sparse_mla.py @@ -0,0 +1,705 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import math + +import torch + +from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_mla import ( + _sparse_mla as _sparse_mla_gfx950, +) +from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_mla import ( + _sparse_mla_reduce as _sparse_mla_reduce_gfx950, +) +from aiter.ops.triton.attention.pa_decode_sparse import ( + _as_int32_contiguous_1d, +) +from aiter.ops.triton.utils._triton import arch_info +from aiter.ops.triton.utils.common_utils import max_addressable_bytes +from aiter.ops.triton.utils.device_info import get_num_sms +from aiter.ops.triton.utils.logger import AiterTritonLogger + +_LOGGER = AiterTritonLogger() + + +def _check_out(out, q, kv_lora_rank): + """Caller-supplied output buffer, or a fresh one. A buffer wider than + kv_lora_rank is accepted and written in its leading columns.""" + C, H = q.shape[0], q.shape[1] + if out is None: + return torch.empty((C, H, kv_lora_rank), dtype=torch.bfloat16, device=q.device) + assert ( + out.shape[0] == C and out.shape[1] == H + ), f"out shape {tuple(out.shape)} != [{C}, {H}, >= {kv_lora_rank}]" + assert out.shape[2] >= kv_lora_rank and out.stride(2) == 1 + assert out.dtype == torch.bfloat16 and out.device == q.device + return out + + +def _cache_pointers(fmt, kv, d_qk, kv_scale): + """-> (cache, alt_ptr, scl_ptr, block_size) for an MLA-route format. Pointer + roles follow the kernel's Seg contract (see its docstring).""" + if kv.ndim == 4: # [slots, 1, 1, R], the asm mla_decode_fwd view + kv = kv.reshape(kv.shape[0], kv.shape[3]) + elif kv.ndim == 3 and kv.shape[2] == d_qk: + # vLLM paged cache [nb, block_size, R]: indices are global slot ids and + # rows are contiguous, so the flat view is stride-identical. + if not (kv.stride(2) == 1 and kv.stride(1) == kv.shape[2]): + raise ValueError("a [nb, block, R] pool must have contiguous rows") + kv = kv.reshape(-1, d_qk) + if fmt == "bf16": + return kv, kv, None, 1 + if fmt == "fp8_scalar": + u8 = kv.view(torch.uint8) + return u8, u8, kv_scale.reshape(1), 1 + # fp8_dsv32_mla: 512 fp8 | 4 f32 per-128 scales | 64 bf16 rope = 656 B. + u8 = kv if kv.dtype == torch.uint8 else kv.view(torch.uint8) + if not (u8.stride(2) == 1 and u8.stride(1) == u8.shape[2]): + raise ValueError("fp8_dsv32_mla rows must be contiguous records") + return u8, u8.view(torch.bfloat16), u8.view(torch.float32), u8.shape[1] + + +def _mla_num_splits( + num_queries: int, heads_blocks: int, avg_topk: float, block_k: int = 64 +) -> int: + """Split-K count for the sparse-MLA decode. + + Below one workgroup per CU, split to fill the machine but never past 8. + """ + num_sms = get_num_sms() + base_wg = max(1, num_queries * heads_blocks) + cta_cap = max(1, (2 * num_sms) // base_wg) + tiles = max(1, math.ceil(avg_topk / block_k)) + if base_wg >= num_sms: + return max(1, min(cta_cap, tiles // 16)) + return max(1, min(cta_cap, tiles, 8)) + + +def _async_launch_config( + fp8_dots: bool, + has_invalid: bool, + num_queries: int, + heads_blocks: int, + num_splits: int, + avg_topk: float, + use_buffer_load: bool, + uni_tile: bool = True, + has_extra: bool = False, + block_k: int = 64, +) -> tuple[bool, int, int]: + """-> (ASYNC_LDS, BLOCK_K, waves_per_eu) for this launch. + + BLOCK_K follows what the grid supplies, not the token count: given enough + workgroups to fill four waves/SIMD (prefill) the small tile takes that + occupancy; otherwise (decode, where split-K caps the grid near two + workgroups/CU) the large tile wins instead, by halving the cross-warp softmax + exchange rate per token. + """ + ASYNC_LDS_DEFAULT = True + enabled = ( + ASYNC_LDS_DEFAULT + and fp8_dots + and uni_tile + and not has_invalid + and not has_extra + ) + workgroups = num_queries * heads_blocks * max(1, num_splits) + num_sms = get_num_sms() + if enabled and workgroups >= 4 * num_sms: + return True, 64, 4 + waves_per_eu = 2 + if use_buffer_load and workgroups <= num_sms: + waves_per_eu = 1 + return enabled, (128 if enabled else block_k), waves_per_eu + + +def _resolve_dot_precision(dot_precision: str, fmt: str) -> bool: + if dot_precision not in ("bf16", "fp8"): + raise ValueError( + f"dot_precision must be 'bf16' or 'fp8', got {dot_precision!r}" + ) + if dot_precision == "bf16": + return False + if fmt == "fp8_dsv32_mla": + raise ValueError( + "dot_precision='fp8' does not support the fp8_dsv32_mla cache." + "Use dot_precision='bf16'." + ) + if fmt == "bf16": + raise ValueError("dot_precision='fp8' needs an fp8 cache.") + return True + + +_DSV4_ROW = 448 + 2 * 64 + 8 # 584 B: fp8 nope | bf16 rope | 8 B UE8M0 trailer + + +def _dsv32_row(kv_lora_rank, qk_rope_head_dim): + return kv_lora_rank + 4 * (kv_lora_rank // 128) + 2 * qk_rope_head_dim + + +def _classify_flat(kv, width, slots, kv_scale, what): + """Flat pool rows are one QK row per slot; dtype and kv_scale pick the tag.""" + if kv.dtype == torch.bfloat16: + return "bf16" # kv_scale, if any, is ignored + if kv.dtype == torch.float8_e4m3fnuz: + raise ValueError( + f"{what}: float8_e4m3fnuz is the gfx942 encoding; gfx950 reads OCP e4m3" + ) + if kv.element_size() != 1: + raise ValueError(f"{what}: unsupported cache dtype {kv.dtype}") + if kv_scale is None: + raise ValueError( + f"{what}: a flat fp8 cache needs kv_scale, [1] f32 (fp8_scalar) or " + f"[slots, D // 64] f32 (fp8_g64)" + ) + if kv_scale.dtype != torch.float32: + raise ValueError(f"{what}: kv_scale must be f32, got {kv_scale.dtype}") + if kv_scale.numel() == 1: + return "fp8_scalar" + if width % 64 or tuple(kv_scale.shape) != (slots, width // 64): + raise ValueError( + f"{what}: kv_scale {tuple(kv_scale.shape)} is neither [1] (fp8_scalar) " + f"nor [slots, D // 64] = [{slots}, {width // 64}] (fp8_g64)" + ) + return "fp8_g64" + + +def _classify_cache(q, kv, kv_lora_rank, qk_rope_head_dim, kv_scale, what="kv_buffer"): + """kv -> one of the kernel's format tags, or raise. + + Record width, dtype and the kv_scale shape are the only things that tell the + formats apart, so whatever they do not pin down is rejected here instead of + being decoded as the wrong layout. Host metadata only: safe under graph + capture. + """ + d_qk = q.shape[-1] + if kv.device != q.device: + raise ValueError(f"{what} is on {kv.device}, q is on {q.device}") + if kv.ndim == 4 and (kv.shape[1] != 1 or kv.shape[2] != 1): + raise ValueError( + f"{what}: a 4-D cache must be [slots, 1, 1, R], got {tuple(kv.shape)}" + ) + if kv.ndim in (2, 4): + width = kv.shape[-1] + if width != d_qk: + raise ValueError( + f"{what}: flat cache rows are {width} wide but q is {d_qk}; a flat " + f"pool stores one QK row (kv_lora_rank + qk_rope_head_dim) per slot" + ) + return _classify_flat(kv, width, kv.shape[0], kv_scale, what) + if kv.ndim == 3: + nb, block, width = kv.shape + if width == d_qk: + # [nb, block, R]: a flat pool stored in blocks. Same rules, except the + # per-64 scale vector is only defined for the 2-D pool. + fmt = _classify_flat(kv, width, nb * block, kv_scale, what) + if fmt == "fp8_g64": + raise ValueError(f"{what}: fp8_g64 needs a 2-D [slots, D] pool") + return fmt + if kv_scale is not None: + raise ValueError( + f"{what}: packed caches carry their own scales; kv_scale must be None" + ) + one_byte = kv.element_size() == 1 + if one_byte and width == _DSV4_ROW: + return "fp8_dsv4_mla" + dsv32 = _dsv32_row(kv_lora_rank, qk_rope_head_dim) + if one_byte and qk_rope_head_dim > 0 and width == dsv32: + return "fp8_dsv32_mla" + raise ValueError( + f"{what}: unrecognized 3-D cache, {width}-wide {kv.dtype} records for q " + f"width {d_qk}. Expected {d_qk} (bf16 or fp8_scalar block pool), " + f"{_DSV4_ROW} uint8 (fp8_dsv4_mla) or {dsv32} uint8 (fp8_dsv32_mla with " + f"kv_lora_rank={kv_lora_rank}, qk_rope_head_dim={qk_rope_head_dim})" + ) + raise ValueError(f"{what} must be 2-D, 3-D or 4-D, got {kv.ndim}-D") + + +def _check_geometry(fmt, d_qk, kv_lora_rank, qk_rope_head_dim): + if kv_lora_rank <= 0 or qk_rope_head_dim < 0: + raise ValueError( + f"kv_lora_rank must be > 0 and qk_rope_head_dim >= 0, got " + f"{kv_lora_rank} / {qk_rope_head_dim}" + ) + if fmt == "fp8_dsv4_mla": + if d_qk != 512: + raise ValueError( + f"fp8_dsv4_mla needs q width 512 (448 nope + 64 rope), got {d_qk}" + ) + return + if fmt == "fp8_g64": + return # the whole row is the head; geometry args are not read + if d_qk != kv_lora_rank + qk_rope_head_dim: + hint = ( + " A bf16 or fp8_scalar cache carries no rope information: pass " + "qk_rope_head_dim=0 for rope-inside-the-row or rope-free models." + if fmt in ("bf16", "fp8_scalar") + else "" + ) + raise ValueError( + f"q width {d_qk} != kv_lora_rank {kv_lora_rank} + qk_rope_head_dim " + f"{qk_rope_head_dim}.{hint}" + ) + if fmt == "fp8_dsv32_mla" and kv_lora_rank % 128: + raise ValueError( + f"fp8_dsv32_mla stores one scale per 128 latent columns; kv_lora_rank " + f"{kv_lora_rank} is not a multiple of 128" + ) + + +def _check_index_stream(indptr, indices, num_queries, device, what=""): + """Both are flat integer tensors, so only the length of indptr can catch a + swapped pair. No device reads.""" + for name, t in ((f"{what}kv_indptr", indptr), (f"{what}kv_indices", indices)): + if not torch.is_tensor(t): + raise ValueError(f"{name} must be a tensor, got {type(t).__name__}") + if t.dtype not in (torch.int32, torch.int64): + raise ValueError(f"{name} must be int32 (int64 accepted), got {t.dtype}") + if t.ndim != 1: + raise ValueError(f"{name} must be 1-D, got {tuple(t.shape)}") + if t.device != device: + raise ValueError(f"{name} is on {t.device}, q is on {device}") + if indptr.numel() != num_queries + 1: + raise ValueError( + f"{what}kv_indptr must have C + 1 = {num_queries + 1} entries, got " + f"{indptr.numel()}. kv_indptr and kv_indices are both flat int32; " + "check their order." + ) + + +def _forward_paged( + q, + kv, + kv_indptr, + kv_indices, + softmax_scale, + kv_scale, + kv_splits, + skip_reduce, + has_invalid, + dot_precision, + q_scale, + out, + return_lse, + attn_sink, + extra_kv, + extra_indptr, + extra_indices, +): + """dsv4 and the SWA+top-k two-loop, until the two launchers merge.""" + from aiter.ops.triton.attention.pa_decode_sparse import pa_decode_sparse + + unsupported = [ + name + for name, asked in ( + ("dot_precision='fp8'", dot_precision != "bf16"), + ("return_lse", return_lse), + ) + if asked + ] + if unsupported: + raise ValueError( + f"{', '.join(unsupported)} is not supported for dsv4 / two-loop " + "caches yet." + ) + res = pa_decode_sparse( + q, + kv, + kv_indices, + kv_indptr, + attn_sink, + softmax_scale, + kv_scales=kv_scale, + kv_splits=kv_splits, + has_invalid=has_invalid, + skip_reduce=skip_reduce, + extra_cache=extra_kv, + extra_indices=extra_indices, + extra_indptr=extra_indptr, + out=out, + ) + return res if isinstance(res, tuple) else (res, None) + + +def sparse_mla_fwd( + q: torch.Tensor, + kv_buffer: torch.Tensor, + kv_indptr: torch.Tensor, + kv_indices: torch.Tensor, + softmax_scale: float, + *, + kv_scale: torch.Tensor | None = None, + kv_lora_rank: int = 512, + qk_rope_head_dim: int = 64, + kv_splits: int | None = None, + skip_reduce: bool = False, + has_invalid: bool = False, + dot_precision: str = "bf16", + q_scale: torch.Tensor | None = None, + out: torch.Tensor | None = None, + return_lse: bool = False, + attn_sink: torch.Tensor | None = None, + extra_kv: torch.Tensor | None = None, + extra_indptr: torch.Tensor | None = None, + extra_indices: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Sparse (top-k gathered) MLA attention. + + Everything after softmax_scale is keyword-only: kv_indptr and kv_indices are + both flat int32, so a positional mix-up would type-check and return garbage. + + Supported KV cache formats. The format is inferred from kv_buffer's shape, + dtype and kv_scale; each row gives what the caller has to pass. R is the + QK width, kv_lora_rank + qk_rope_head_dim. + + format kv_buffer kv_scale geometry args + bf16 [slots, R], [nb, block, R], None as the model + [slots, 1, 1, R]; bf16 + fp8_scalar same shapes, fp8 or uint8 [1] f32 as the model + (GLM-5.x; the only format + that runs dot_precision="fp8") + fp8_g64 [slots, D] fp8 or uint8 [slots, D//64] kv_lora_rank=D, + (uniform pool, rope inside) f32 qk_rope_head_dim=0 + fp8_dsv32_mla [nb, block, 656] uint8 None 512 / 64 + (DeepSeek-V3.2, Kimi-K3; vLLM + fp8_ds_mla: 512 fp8 | 4 f32 + per-128 | 64 bf16 rope) + fp8_dsv4_mla [nb, block, 584] uint8 None ignored + (DeepSeek-V4; vLLM fp8_ds_mla: + 448 fp8 | 64 bf16 rope + 8 B + UE8M0 trailer per block) + + Geometry, set by qk_rope_head_dim: separated rope (DeepSeek-V3.2, + GLM-5.1/5.2), where the query is the latent plus an appended rope and V is + the latent; rope-free (GLM-5.3-Flash), qk_rope_head_dim=0, where the query + is the latent alone. DeepSeek-V4 keeps its rope inside the 512-wide row + (448 nope + 64 rope) with V the whole row; that layout is fixed by the cache + format, so the geometry args (kv_lora_rank and qk_rope_head_dim) are not + read for it. A bf16 pool carries no rope information, so a rope-inside or + rope-free model on a bf16 cache must pass qk_rope_head_dim=0. + + Args: + q: [C, H, kv_lora_rank + qk_rope_head_dim] queries, one row per + query token (prefill and decode alike). + qk_rope_head_dim: width of the appended rope. 0 means rope-free, and + the QK contraction becomes a single dot over kv_lora_rank. + kv_buffer: the KV pool, one of the formats in the table above. + kv_indptr: [C + 1] int32 prefix sum of per-query index counts. + kv_indices: flat int32 GLOBAL slot ids into the pool. + softmax_scale: the layer's softmax scale. + kv_scale: f32 cache scale, shape per the table; the packed formats + carry their own scales and take None. + kv_splits: split-K override; default follows the occupancy policy. + skip_reduce: with split-K active, return (part_acc, part_m, part_l) + instead of launching the combine. + has_invalid: index stream carries -1 sentinels (masked out). Default + False. + attn_sink: optional [H] f32 per-head sink: exp(sink) joins the softmax + denominator (and the LSE) as a virtual key. None means no sink, + which is not the same as a zero sink. + extra_kv, extra_indptr, extra_indices: a second segment attended in the + same pass, for the SWA-window + top-k two-loop. All three together. + dot_precision: what the QK and PV matrix-core ops run in. + + "bf16" (default): the KV tile is dequantized to bf16 on its way + into LDS and both dots are bf16. Works with every cache format. + "fp8": the cache's own code points go to the fp8 matrix core with no + dequant, and the per-tensor scale folds outside the tile loop. + tensor scale fp8 kv cache only. + + q is adapted to the choice. bf16 q is quantized in the kernel + prologue, one scale per (query, head-block) tile; fp8 q is passed + straight through, and fp8 q under "bf16" dots is widened in-kernel. + + q_scale: scalar f32, required when q arrives already fp8 (the scale it + was quantized with; the aiter asm convention). Ignored for bf16 q, + which the kernel quantizes itself when dot_precision="fp8". + out: optional [C, H, >= kv_lora_rank] bf16 destination. + return_lse: also return the natural-log log-sum-exp, [C, H] f32, for + merging partials across context-parallel ranks. A fully masked row + reports -inf. + + Returns: + (out, lse), out is + [C, H, kv_lora_rank] bf16 (the latent V), lse is None unless return_lse. + """ + if q.ndim != 3: + raise ValueError(f"expected q=[C, H, d_qk], got {tuple(q.shape)}") + num_queries, num_heads, d_qk = q.shape + fmt = _classify_cache(q, kv_buffer, kv_lora_rank, qk_rope_head_dim, kv_scale) + _LOGGER.info( + f"SPARSE_MLA_FWD: q={tuple(q.shape)} kv_buffer={tuple(kv_buffer.shape)} " + f"{kv_buffer.dtype} kv_indices={tuple(kv_indices.shape)} fmt={fmt}" + ) + _check_geometry(fmt, d_qk, kv_lora_rank, qk_rope_head_dim) + _check_index_stream(kv_indptr, kv_indices, num_queries, q.device) + if attn_sink is not None and attn_sink.numel() != num_heads: + raise ValueError( + f"attn_sink must have one entry per head ({num_heads}), got " + f"{tuple(attn_sink.shape)}" + ) + extra = (extra_kv, extra_indptr, extra_indices) + has_extra = all(a is not None for a in extra) + if any(a is not None for a in extra) and not has_extra: + raise ValueError("extra_kv, extra_indptr and extra_indices go together") + if has_extra: + # The SWA-window + top-k two-loop: both segments are block caches the + # dsv4 decoder reads. A flat main pool would silently drop the extra one. + if kv_buffer.ndim != 3 or fmt not in ("fp8_dsv4_mla", "bf16"): + raise ValueError( + "the two-loop needs a 3-D block cache (fp8_dsv4_mla or bf16) as " + f"kv_buffer, got {fmt}" + ) + extra_fmt = _classify_cache( + q, extra_kv, kv_lora_rank, qk_rope_head_dim, None, what="extra_kv" + ) + if extra_kv.ndim != 3 or extra_fmt not in ("fp8_dsv4_mla", "bf16"): + raise ValueError( + f"extra_kv must be a 3-D fp8_dsv4_mla or bf16 block cache, got " + f"{extra_fmt}" + ) + _check_index_stream( + extra_indptr, extra_indices, num_queries, q.device, what="extra_" + ) + if fmt in ("fp8_dsv4_mla", "fp8_g64") or has_extra: + return _forward_paged( + q, + kv_buffer, + kv_indptr, + kv_indices, + softmax_scale, + kv_scale, + kv_splits, + skip_reduce, + has_invalid, + dot_precision, + q_scale, + out, + return_lse, + attn_sink, + extra_kv, + extra_indptr, + extra_indices, + ) + assert arch_info.get_arch() == "gfx950", "sparse_mla_fwd is gfx950-only" + q_is_fp8 = q.dtype == torch.float8_e4m3fn + if q.dtype not in (torch.bfloat16, torch.float8_e4m3fn): + raise ValueError( + f"q must be bf16 or float8_e4m3fn, got {q.dtype}" + + ( + " (fnuz is a different encoding from what the matrix core reads)" + if "fnuz" in str(q.dtype) + else "" + ) + ) + if q_is_fp8: + # Caller-quantized q, tensor scale + if q_scale is None: + raise ValueError("fp8 q needs q_scale (the scale it was quantized with)") + if q_scale.numel() != 1: + raise ValueError( + f"q_scale must be a single per-tensor value, got " + f"{tuple(q_scale.shape)}; per-head q scales are not supported" + ) + q_scale = q_scale.reshape(1).to(torch.float32).contiguous() + cache, alt, scl, block_size = _cache_pointers(fmt, kv_buffer, d_qk, kv_scale) + fp8_dots = _resolve_dot_precision(dot_precision, fmt) + if not q_is_fp8: + # q_scale describes an fp8 q's encoding. With bf16 q the kernel quantizes + # per (query, head-block) tile when the dots are fp8, so a caller-supplied + # scale has nothing to apply to; callers pass layer._q_scale regardless. + q_scale = None + kv_indices = _as_int32_contiguous_1d(kv_indices) + kv_indptr = _as_int32_contiguous_1d(kv_indptr) + has_sink = attn_sink is not None + if has_sink: + attn_sink = attn_sink.reshape(-1).to(torch.float32).contiguous() + else: + # The kernel still wants a live pointer for the compile-time-elided slot. + attn_sink = torch.empty(1, device=q.device, dtype=torch.float32) + + # Tuned launch config (gfx950 / MI355). H < 16 runs natively at + # BLOCK_M = next_pow2(H) instead of padding heads + block_m = 16 if num_heads >= 16 else max(8, 1 << (num_heads - 1).bit_length()) + block_k = 64 + num_warps = block_k // 16 + + num_rows = cache.shape[0] * block_size if cache.ndim >= 2 else cache.shape[0] + avg_topk = kv_indices.numel() / max(1, num_queries) + + # Alignment hint for cs0 so row gathers can vectorize. + cs0_align = 1 + for a in (16, 8, 4, 2): + if int(cache.stride(0)) % a == 0: + cs0_align = a + break + # buffer_load carries a signed 32-bit offset + MAX_BYTES = 2**31 - 1 + use_buffer_load = max_addressable_bytes(cache) < MAX_BYTES + idx_use_buffer_load = max_addressable_bytes(kv_indices) < MAX_BYTES + + head_aligned = num_heads % block_m == 0 + heads_blocks = (num_heads + block_m - 1) // block_m + out = _check_out(out, q, kv_lora_rank) + if return_lse and skip_reduce: + raise ValueError("return_lse needs the combine, so skip_reduce cannot be set") + # A live pointer either way, like attn_sink below. + lse = ( + torch.empty((num_queries, num_heads), dtype=torch.float32, device=q.device) + if return_lse + else torch.empty(1, dtype=torch.float32, device=q.device) + ) + + if kv_splits is not None: + num_splits = max(1, int(kv_splits)) + else: + num_splits = _mla_num_splits(num_queries, heads_blocks, avg_topk, 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) + # bf16 partials halve the split-K HBM traffic; skip_reduce hands the + # partials back to the caller and keeps f32. + part_acc = torch.empty( + (num_queries, num_splits, num_heads, kv_lora_rank), + dtype=torch.float32 if skip_reduce else torch.bfloat16, + 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 + + # Dequant chunking. The gather layout puts 32 of a wave's 64 lanes along a row + # (16 B each, 512 B) and the other 32 on a second row; col_reps is how many + # such spans each lane holds, which is what makes a column split a free + # register rename. + col_reps = kv_lora_rank // 512 + chunk_axis = 1 if col_reps >= 4 else 0 + nope_chunk = max(1, block_k // 4) if chunk_axis == 0 else min(128, kv_lora_rank) + async_lds_on, block_k, waves_per_eu = _async_launch_config( + fp8_dots, + has_invalid, + num_queries, + heads_blocks, + num_splits, + avg_topk, + use_buffer_load, + uni_tile=True, + has_extra=False, + block_k=block_k, + ) + + # Q is read once per query without split-K, and re-read by every split + q_cache = ".cg" if num_splits == 1 else "" + grid = (num_queries, num_splits, heads_blocks) + _sparse_mla_gfx950[grid]( + q, + cache, + alt, + kv_indices, + kv_indptr, + cache, # extra_* segment: unread placeholders (HAS_EXTRA=False) + alt, + kv_indices, + kv_indptr, + attn_sink, + out, + part_m, + part_l, + part_acc, + scl, # f32 side-channel: k_scale ("fp8_scalar") / f32 view ("fp8_dsv32_mla") + scl, # extra segment's slot, unread (HAS_EXTRA=False) + float(softmax_scale), + q.stride(0), + q.stride(1), + out.stride(0), + out.stride(1), + cache.stride(0), + cache.stride(0), + num_rows, + num_rows, + pm_stride0, + pm_stride_s, + pa_stride0, + pa_stride_s, + pa_stride_h, + num_heads, + HAS_EXTRA=False, + HAS_SINK=has_sink, + MAIN_FMT=fmt, + EXTRA_FMT=fmt, + MAIN_BLOCK_SIZE=block_size, + EXTRA_BLOCK_SIZE=block_size, + CS0_ALIGN=cs0_align, + NOPE_DIM=kv_lora_rank, + ROPE_DIM=qk_rope_head_dim, + HEAD_SIZE=kv_lora_rank, + ROPE_SEPARATE=qk_rope_head_dim > 0, + BLOCK_M=block_m, + BLOCK_K=block_k, + NUM_SPLITS=num_splits, + HEAD_ALIGNED=head_aligned, + NOPE_CHUNK=nope_chunk, + CHUNK_AXIS=chunk_axis, + PART_STORE_CACHE="", + UNI_TILE=True, + GRID_ORDER="qsh", + Q_CACHE=q_cache, + MAIN_SPLITS=num_splits, + ADAPTIVE_SPLITS=num_splits > 1, + DEQ="none", + MAIN_USE_BUFFER_LOAD=use_buffer_load, + EXTRA_USE_BUFFER_LOAD=use_buffer_load, + IDX_BUFFER_LOAD=idx_use_buffer_load, + HAS_INVALID=has_invalid, + FP8_MFMA=fp8_dots, + ASYNC_LDS=async_lds_on, + GATHER_CACHE="", + q_scl_ptr=q_scale, + Q_FP8=q_is_fp8, + lse_ptr=lse, + HAS_LSE=return_lse, + num_warps=num_warps, + waves_per_eu=waves_per_eu, + ) + + if num_splits == 1: + return out, (lse if return_lse else None) + if skip_reduce: + return part_acc, part_m, part_l + + # One head per reduce workgroup + rgrid = (num_queries, num_heads) + _sparse_mla_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=kv_lora_rank, + BLOCK_M=1, + NUM_SPLITS=num_splits, + HEAD_ALIGNED=True, + ADAPTIVE_SPLITS=num_splits > 1, + lse_ptr=lse, + HAS_LSE=return_lse, + num_warps=1, + ) + return out, (lse if return_lse else None) diff --git a/op_tests/op_benchmarks/triton/bench_sparse_mla.py b/op_tests/op_benchmarks/triton/bench_sparse_mla.py new file mode 100644 index 00000000000..a21ec66546e --- /dev/null +++ b/op_tests/op_benchmarks/triton/bench_sparse_mla.py @@ -0,0 +1,272 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Benchmark for sparse_mla_fwd (gfx950 gluon, separated-rope MLA). + +The cache is flushed between iterations by default. Leaving it warm lets the loop +re-read its KV and flatters decode shapes badly. + +Usage: + python op_tests/op_benchmarks/triton/bench_sparse_mla.py + python op_tests/op_benchmarks/triton/bench_sparse_mla.py --num_seqs 1 --num_tokens 8192 + python op_tests/op_benchmarks/triton/bench_sparse_mla.py --num_seqs 64 --metric bandwidth +""" + +import argparse + +import torch +import triton +from torch.autograd import DeviceType +from torch.profiler import ProfilerActivity, profile + +from aiter.ops.triton.attention.sparse_mla import sparse_mla_fwd +from aiter.ops.triton.utils._triton import arch_info +from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype +from op_tests.op_benchmarks.triton.utils.benchmark_utils import ( + get_caller_name_no_ext, +) + +KV_LORA_RANK = 512 +QK_ROPE_HEAD_DIM = 64 +D_QK = KV_LORA_RANK + QK_ROPE_HEAD_DIM +E4M3_DTYPE = get_fp8_e4m3_dtype() +E4M3_MAX = torch.finfo(E4M3_DTYPE).max + +# Picks up the main kernel and _sparse_mla_reduce, nothing else. +KERNEL_MATCH = "_sparse_mla" +FLUSH_BYTES = 512 << 20 + +_flush_buf = None + + +def flush_l2(): + global _flush_buf + if _flush_buf is None: + _flush_buf = torch.empty(FLUSH_BYTES, dtype=torch.uint8, device="cuda") + _flush_buf.fill_(1) + + +def device_time_ms(func, warmup=25, rep=100, flush=True): + for _ in range(warmup): + if flush: + flush_l2() + func() + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for _ in range(rep): + if flush: + flush_l2() + func() + torch.cuda.synchronize() + total = sum( + e.self_device_time_total + for e in prof.key_averages() + if e.device_type == DeviceType.CUDA + and KERNEL_MATCH in e.key + and e.self_device_time_total > 0 + ) + if total == 0: + raise RuntimeError(f"no kernel matching {KERNEL_MATCH!r} was profiled") + return total / rep / 1e3 + + +def bytes_moved(num_tokens, num_heads, nnz, kv_elem_bytes): + """Bytes the launch actually moves, counting rows as gathered. + + Split-K partials are left out: the split count is the kernel's own decision, + so when it does split, that traffic lands in the time but not here. + """ + kv = nnz * D_QK * kv_elem_bytes # the gather, and the bulk of it + idx = nnz * 4 # int32 index stream, read once + q = num_tokens * num_heads * D_QK * kv_elem_bytes + out = num_tokens * num_heads * KV_LORA_RANK * 2 + return kv + idx + q + out + + +def build_case(num_seqs, num_tokens, num_heads, context, topk, device="cuda"): + """Build one launch the way vLLM's sparse-MLA path lays it out. + + Every sequence owns its own slice of the pool, so there is no cross-request + KV sharing to inflate the cache hit rate. Per-token KV length is + min(pos + 1, topk), which is what the index converter emits. + """ + torch.manual_seed(0) + gen = torch.Generator().manual_seed(1) + pool = num_seqs * context + + kv = torch.randn(pool, D_QK, dtype=torch.bfloat16, device=device) * 0.125 + q = ( + torch.randn( + num_seqs * num_tokens, + num_heads, + D_QK, + dtype=torch.bfloat16, + device=device, + ) + * 0.125 + ) + + rows, lens = [], [] + for s in range(num_seqs): + base = s * context + for t in range(num_tokens): + pos = context - 1 if num_tokens == 1 else t + if pos + 1 <= topk: + sel = torch.arange(pos + 1, dtype=torch.int64) # everything fits + else: + sel = torch.randperm(pos + 1, generator=gen)[:topk] + rows.append((sel + base).to(torch.int32)) + lens.append(sel.numel()) + + indices = torch.cat(rows).to(device) + indptr = torch.zeros(len(lens) + 1, dtype=torch.int32) + indptr[1:] = torch.tensor(lens, dtype=torch.int32).cumsum(0) + + # fp8 dots want the flat per-tensor format; bf16 dots read kv as it is. + scale = (kv.float().abs().amax() / E4M3_MAX).clamp_min(1e-30).reshape(1) + kv_fp8 = ((kv.float() / scale).clamp(-E4M3_MAX, E4M3_MAX).to(E4M3_DTYPE)).view( + torch.uint8 + ) + + return q, kv, kv_fp8, scale.to(device), indices, indptr.to(device) + + +def run_benchmark(args): + if args.metric == "time": + ylabel = "Time (ms)" + elif args.metric == "throughput": + ylabel = "TFLOPs" + else: + ylabel = "GB/s" + + x_vals_list = [ + [ + "decode" if tokens == 1 else "prefill", + seqs, + tokens, + args.num_heads, + args.context, + args.topk, + ] + for seqs in args.num_seqs + for tokens in args.num_tokens + ] + + benchmark = triton.testing.Benchmark( + x_names=["phase", "num_seqs", "num_tokens", "num_heads", "context", "topk"], + x_vals=x_vals_list, + line_arg="dots", + line_vals=["bf16", "fp8"], + line_names=["bf16 dots", "fp8 dots"], + styles=[("green", "-"), ("blue", "-")], + ylabel=ylabel, + plot_name=get_caller_name_no_ext(), + args={"metric": args.metric}, + ) + + @triton.testing.perf_report([benchmark]) + def bench_sparse_mla( + phase, num_seqs, num_tokens, num_heads, context, topk, dots, metric, **kwargs + ): + q, kv, kv_fp8, kv_scale, indices, indptr = build_case( + num_seqs, num_tokens, num_heads, context, topk + ) + sm_scale = D_QK**-0.5 + out = torch.empty( + q.shape[0], + num_heads, + KV_LORA_RANK, + dtype=torch.bfloat16, + device=q.device, + ) + if dots == "fp8": + # Hand q over already quantized, the way production does. + cache, scale = kv_fp8, kv_scale + q_scale = (q.float().abs().amax() / E4M3_MAX).clamp_min(1e-30).reshape(1) + q = (q.float() / q_scale).clamp(-E4M3_MAX, E4M3_MAX).to(E4M3_DTYPE) + else: + cache, scale, q_scale = kv, None, None + + def func(): + # out= is supplied, so the returned pair is not needed + sparse_mla_fwd( + q, + cache, + indptr, + indices, + sm_scale, + kv_scale=scale, + q_scale=q_scale, + dot_precision=dots, + out=out, + ) + + time_ms = device_time_ms(func, flush=not args.no_flush_l2) + + nnz = int(indptr[-1].item()) + num_tokens = q.shape[0] + # QK reads the whole row, PV only the latent half + flops = 2.0 * num_heads * nnz * (D_QK + KV_LORA_RANK) + moved = bytes_moved(num_tokens, num_heads, nnz, 1 if dots == "fp8" else 2) + + if metric == "time": + return time_ms + if metric == "throughput": + return flops / (time_ms * 1e-3) / 1e12 + return moved / (time_ms * 1e-3) / 1e9 + + bench_sparse_mla.run(save_path="." if args.o else None, print_data=True) + + +def main(): + parser = argparse.ArgumentParser( + description="Sparse MLA Benchmark", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--num_seqs", + type=int, + nargs="+", + default=[1, 8, 64, 256], + help="sequences in the launch; every value is benchmarked", + ) + parser.add_argument( + "--num_tokens", + type=int, + nargs="+", + default=[1], + help="query tokens per sequence. 1 is a decode step, more is prefill; " + "every value is benchmarked against every --num_seqs", + ) + parser.add_argument("--num_heads", type=int, default=16, help="q heads (TP4=16)") + parser.add_argument("--context", type=int, default=8192, help="context length") + parser.add_argument("--topk", type=int, default=2048, help="indexer top-k") + parser.add_argument( + "--no_flush_l2", + action="store_true", + help="skip the cache flush between iterations, so the gather runs warm", + ) + parser.add_argument( + "-o", action="store_true", help="Write performance results to CSV file" + ) + parser.add_argument( + "--metric", + type=str, + choices=["time", "throughput", "bandwidth"], + default="time", + help="metric to plot", + ) + args = parser.parse_args() + over = [t for t in args.num_tokens if t > args.context] + if over: + raise SystemExit( + f"--num_tokens {over} exceeds --context {args.context}: a sequence " + "cannot prefill more tokens than its context holds" + ) + if arch_info.get_arch() != "gfx950": + raise SystemExit(f"sparse_mla_fwd is gfx950-only (got {arch_info.get_arch()})") + run_benchmark(args) + + +if __name__ == "__main__": + main() diff --git a/op_tests/triton_tests/attention/test_mla_v4_triton.py b/op_tests/triton_tests/attention/test_mla_v4_triton.py index be966fa696e..96a91d7f30b 100644 --- a/op_tests/triton_tests/attention/test_mla_v4_triton.py +++ b/op_tests/triton_tests/attention/test_mla_v4_triton.py @@ -10,6 +10,7 @@ from aiter.benchmark_data_init import fill, make_generator from aiter.ops.triton.attention.pa_decode_sparse import pa_decode_sparse from aiter.ops.triton.utils._triton import arch_info +from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype from aiter.test_common import ( benchmark, checkAllclose, @@ -455,13 +456,13 @@ def test_pa_decode_sparse_vs_reference( # --------------------------------------------------------------------------- _FP8_GROUP_SIZE = 64 -_FP8_DTYPE = torch.float8_e4m3fnuz +_FP8_DTYPE = get_fp8_e4m3_dtype() def _quantize_kv_fp8(unified_kv, group_size=_FP8_GROUP_SIZE): """Quantize bf16/fp16 unified_kv to (fp8, scales) with 1xGROUP_SIZE block scaling. - Returns (kv_fp8, kv_scales) where kv_fp8 is float8_e4m3fnuz and + Returns (kv_fp8, kv_scales) where kv_fp8 is _FP8_DTYPE and kv_scales is [total_pages, D // group_size] fp32. """ total_pages, D = unified_kv.shape 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 72bdd8e7191..b4f89a6836b 100644 --- a/op_tests/triton_tests/attention/test_pa_decode_sparse.py +++ b/op_tests/triton_tests/attention/test_pa_decode_sparse.py @@ -9,6 +9,7 @@ from aiter.ops.triton.attention.pa_decode_sparse import pa_decode_sparse from aiter.ops.triton.utils._triton import arch_info +from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype from aiter.test_common import checkAllclose @@ -311,18 +312,14 @@ def test_pa_decode_sparse_vs_reference( ) -# --------------------------------------------------------------------------- -# FP8 KV cache quantization helpers -# --------------------------------------------------------------------------- - _FP8_GROUP_SIZE = 64 -_FP8_DTYPE = torch.float8_e4m3fnuz +_FP8_DTYPE = get_fp8_e4m3_dtype() def _quantize_kv_fp8(unified_kv, group_size=_FP8_GROUP_SIZE): """Quantize bf16/fp16 unified_kv to (fp8, scales) with 1xGROUP_SIZE block scaling. - Returns (kv_fp8, kv_scales) where kv_fp8 is float8_e4m3fnuz and + Returns (kv_fp8, kv_scales) where kv_fp8 is _FP8_DTYPE and kv_scales is [total_pages, D // group_size] fp32. """ total_pages, D = unified_kv.shape diff --git a/op_tests/triton_tests/attention/test_sparse_mla.py b/op_tests/triton_tests/attention/test_sparse_mla.py new file mode 100644 index 00000000000..7ac5026a566 --- /dev/null +++ b/op_tests/triton_tests/attention/test_sparse_mla.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Tests for sparse_mla_fwd (gfx950 gluon MLA). + +Covers both geometries: separated rope (DSV3.2, GLM-5.1, GLM-5.2) and +rope-free (GLM-5.3-Flash), where the query is the latent alone. +""" + +import pytest +import torch + +from aiter.ops.triton.utils._triton import arch_info +from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype + +if arch_info.get_arch() == "gfx950": + import aiter.ops.triton.attention.sparse_mla as smd + from aiter.ops.triton.attention.sparse_mla import sparse_mla_fwd + +# The packed fp8_ds_mla record is OCP e4m3 by definition of the format. +FP8_DTYPE = get_fp8_e4m3_dtype() +FP8_MAX = torch.finfo(FP8_DTYPE).max +KV_LORA, ROPE = 512, 64 +D_QK = KV_LORA + ROPE + + +def _skip_unless_gfx950(): + if arch_info.get_arch() != "gfx950": + pytest.skip("sparse_mla_fwd is gfx950-only") + + +def quantize_flat_fp8(kv): + """vLLM's production layout: whole row fp8 with one per-tensor scale.""" + scale = (kv.float().abs().amax() / FP8_MAX).clamp_min(1e-30).reshape(1) + q8 = (kv.float() / scale).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE) + return q8.view(torch.uint8), scale.to(torch.float32) + + +def pack_ds_mla(kv, block_size=64): + """vLLM fp8_ds_mla: 656 B rows [512 fp8 | 4 x f32 per-128 | 64 bf16 rope].""" + T = kv.shape[0] + nb = (T + block_size - 1) // block_size + cache = torch.zeros(nb * block_size, 656, dtype=torch.uint8, device=kv.device) + x = kv[:, :KV_LORA].float().reshape(T, 4, 128) + scale = (x.abs().amax(-1) / FP8_MAX).clamp_min(1e-30) + q8 = (x / scale[..., None]).to(FP8_DTYPE).view(torch.uint8) + cache[:T, :KV_LORA] = q8.reshape(T, KV_LORA) + cache[:T, 512:528].view(torch.float32).copy_(scale) + cache[:T, 528:].view(torch.bfloat16).copy_(kv[:, KV_LORA:]) + return cache.view(nb, block_size, 656) + + +def dequant_flat_fp8(u8, scale): + return u8.view(FP8_DTYPE).float() * scale + + +def dequant_ds_mla(cache): + flat = cache.reshape(-1, 656) + x = flat[:, :KV_LORA].view(FP8_DTYPE).float().reshape(-1, 4, 128) + sc = flat[:, 512:528].view(torch.float32) + nope = (x * sc[..., None]).reshape(-1, KV_LORA) + rope = flat[:, 528:].view(torch.bfloat16).float() + return torch.cat([nope, rope], dim=-1) + + +def make_indices(C, pool, topk, device, seed, ragged): + # Built on CPU + g = torch.Generator(device="cpu").manual_seed(seed) + lens = ( + torch.randint(1, topk + 1, (C,), generator=g, device="cpu") + if ragged + else torch.full((C,), topk, device="cpu") + ) + rows = [ + torch.randperm(pool, generator=g, device="cpu")[: int(lens[c])] + for c in range(C) + ] + indptr = torch.zeros(C + 1, dtype=torch.int64, device="cpu") + indptr[1:] = lens.cumsum(0) + return ( + torch.cat(rows).to(torch.int32).to(device), + indptr.to(torch.int32).to(device), + ) + + +def reference(q, kv_truth, indices, indptr, sm_scale): + """f32 attention over the dequantized cache.""" + outs = [] + for c in range(q.shape[0]): + kvs = kv_truth[indices[indptr[c] : indptr[c + 1]].long()].float() + s = torch.einsum("hd,td->ht", q[c].float(), kvs) * sm_scale + p = torch.softmax(s, dim=-1) + outs.append(torch.einsum("ht,td->hd", p, kvs[:, :KV_LORA])) + return torch.stack(outs) + + +def rel_err(out, ref): + return ((out.float() - ref).abs().max() / ref.abs().max()).item() + + +def _build(fmt, C, H, topk, pool, ragged, seed=0, rope=ROPE): + torch.manual_seed(seed) + d_qk = KV_LORA + rope + q = torch.randn(C, H, d_qk, dtype=torch.bfloat16, device="cuda") * 0.125 + kv = torch.randn(pool, d_qk, dtype=torch.bfloat16, device="cuda") * 0.125 + idx, ptr = make_indices(C, pool, topk, "cuda", seed + 7, ragged) + ks = None + if fmt == "bf16": + cache, truth = kv, kv.float() + elif fmt == "tensor": + cache, ks = quantize_flat_fp8(kv) + truth = dequant_flat_fp8(cache, ks) + else: + cache = pack_ds_mla(kv) + truth = dequant_ds_mla(cache)[:pool] + return q, cache, ks, idx, ptr, truth + + +def _run_and_check( + fmt, C, H, topk, ragged, tol=2e-2, pool=1 << 16, rope=ROPE, **kwargs +): + sm = (KV_LORA + rope) ** -0.5 + q, cache, ks, idx, ptr, truth = _build(fmt, C, H, topk, pool, ragged, rope=rope) + ref = reference(q, truth.to(torch.bfloat16), idx, ptr, sm) + out, _ = sparse_mla_fwd( + q, cache, ptr, idx, sm, kv_scale=ks, qk_rope_head_dim=rope, **kwargs + ) + e = rel_err(out, ref) + assert e < tol, f"{fmt} C={C} H={H} topk={topk} ragged={ragged}: rel-err {e:.3e}" + + +@pytest.mark.parametrize( + "fmt,dots,tol", + [("bf16", "bf16", 2e-2), ("tensor", "bf16", 2e-2), ("tensor", "fp8", 7e-2)], + ids=["bf16", "tensor", "tensor_fp8"], +) +@pytest.mark.parametrize("H", [8, 16]) +@pytest.mark.parametrize( + "C,topk,ragged,pool", + [(8, 2048, False, 1 << 16), (8, 500, True, 1 << 16), (2048, 256, True, 1 << 13)], + ids=["topk2048", "ragged500", "prefill"], +) +def test_sparse_mla(fmt, dots, tol, H, C, topk, ragged, pool): + _skip_unless_gfx950() + _run_and_check( + fmt, C=C, H=H, topk=topk, ragged=ragged, pool=pool, tol=tol, dot_precision=dots + ) + + +def test_ds_mla_format(): + _skip_unless_gfx950() + _run_and_check("dsmla", C=8, H=16, topk=2048, ragged=True) + + +def reference_lse(q, kv_truth, indices, indptr, sm_scale): + """Natural-log LSE per (query, head), matching mla_decode_fwd's convention.""" + rows = [] + for c in range(q.shape[0]): + kvs = kv_truth[indices[indptr[c] : indptr[c + 1]].long()].float() + s = torch.einsum("hd,td->ht", q[c].float(), kvs) * sm_scale + rows.append(torch.logsumexp(s, dim=-1)) + return torch.stack(rows) + + +@pytest.mark.parametrize("fmt", ["bf16", "tensor"]) +@pytest.mark.parametrize( + "C,topk,ragged,splits", + [(8, 2048, False, None), (8, 500, True, None), (2048, 256, True, 1)], + ids=["split", "split_ragged", "nosplit"], +) +def test_return_lse(fmt, C, topk, ragged, splits): + _skip_unless_gfx950() + H, pool = 16, 1 << 16 + sm = D_QK**-0.5 + q, cache, ks, idx, ptr, truth = _build(fmt, C, H, topk, pool, ragged) + kwargs = {} if splits is None else {"kv_splits": splits} + out, lse = sparse_mla_fwd( + q, cache, ptr, idx, sm, kv_scale=ks, return_lse=True, **kwargs + ) + assert lse is not None + assert lse.shape == (C, H) and lse.dtype == torch.float32 + ref = reference_lse(q, truth.to(torch.bfloat16), idx, ptr, sm) + e = (lse - ref).abs().max().item() / ref.abs().max().item() + assert e < 2e-2, f"{fmt} C={C} splits={splits}: lse rel-err {e:.3e}" + # the output must not change just because the LSE was asked for + out_no, lse_no = sparse_mla_fwd(q, cache, ptr, idx, sm, kv_scale=ks, **kwargs) + assert lse_no is None + assert torch.equal(out.view(torch.int16), out_no.view(torch.int16)) + + +def test_global_load_path(): + """A pool whose addressable span passes buffer_load's 2 GB offset limit.""" + _skip_unless_gfx950() + live = 1 << 16 + sm = D_QK**-0.5 + q, cache, ks, idx, ptr, truth = _build("tensor", 8, 16, 2048, live, ragged=False) + big = torch.empty(3_800_000, D_QK, dtype=cache.dtype, device=cache.device) + big[:live] = cache + assert smd.max_addressable_bytes(big) >= 2**31 - 1 # past buffer_load's offset + out, _ = sparse_mla_fwd(q, big, ptr, idx, sm, kv_scale=ks) + e = rel_err(out, reference(q, truth.to(torch.bfloat16), idx, ptr, sm)) + assert e < 2e-2, f"global load path: rel-err {e:.3e}" + + +@pytest.mark.parametrize( + "fmt,dots,tol", + [("bf16", "bf16", 2e-2), ("tensor", "bf16", 2e-2), ("tensor", "fp8", 7e-2)], + ids=["bf16", "tensor", "tensor_fp8"], +) +@pytest.mark.parametrize("H", [8, 16]) +@pytest.mark.parametrize( + "C,topk,ragged,pool", + [(8, 2048, False, 1 << 16), (8, 500, True, 1 << 16), (2048, 256, True, 1 << 13)], + ids=["topk2048", "ragged500", "prefill"], +) +def test_sparse_mla_rope_free(fmt, dots, tol, H, C, topk, ragged, pool): + _skip_unless_gfx950() + _run_and_check( + fmt, + C=C, + H=H, + topk=topk, + ragged=ragged, + pool=pool, + tol=tol, + rope=0, + dot_precision=dots, + )