From ea21fa0a6c5d40916d2a7112812727f54217ebd7 Mon Sep 17 00:00:00 2001 From: Sumin Hong Date: Wed, 30 Sep 2026 10:35:27 +0000 Subject: [PATCH 1/3] [Kernel] Add HIP BF16 sparse MLA for gfx950 Expose Full140 v1 and hybrid v2 through one CSR attention API for BF16 D512 latent K/V and H16/H64. Reuse HipKittens, support caller-owned output and split workspace, and preserve current-stream and graph semantics. Include FP64 and graph tests, dispatch-boundary coverage, a pinned-VGPR assembly check, and a reproducible same-checkout Gluon benchmark. Signed-off-by: Sumin Hong Co-authored-by: OpenAI Codex --- aiter/__init__.py | 1 + aiter/jit/optCompilerConfig.json | 25 + aiter/ops/sparse_mla.py | 163 ++++ .../kernels/mla/hk/sparse_mla_bf16_common.cuh | 878 ++++++++++++++++++ .../mla/hk/sparse_mla_bf16_full_lds.cuh | 374 ++++++++ .../mla/hk/sparse_mla_bf16_full_lds_pv.cuh | 155 ++++ .../mla/hk/sparse_mla_bf16_headtiles.cuh | 502 ++++++++++ .../kernels/mla/hk/sparse_mla_bf16_reduce.cuh | 107 +++ csrc/kernels/mla/hk/sparse_mla_bf16_wave4.cuh | 493 ++++++++++ csrc/kernels/mla/sparse_mla_bf16.cu | 218 +++++ docs/sparse_mla_bf16.md | 125 +++ op_tests/check_sparse_mla_bf16_isa.py | 112 +++ .../op_benchmarks/bench_sparse_mla_bf16.py | 406 ++++++++ op_tests/test_sparse_mla_bf16.py | 260 ++++++ 14 files changed, 3819 insertions(+) create mode 100644 aiter/ops/sparse_mla.py create mode 100644 csrc/kernels/mla/hk/sparse_mla_bf16_common.cuh create mode 100644 csrc/kernels/mla/hk/sparse_mla_bf16_full_lds.cuh create mode 100644 csrc/kernels/mla/hk/sparse_mla_bf16_full_lds_pv.cuh create mode 100644 csrc/kernels/mla/hk/sparse_mla_bf16_headtiles.cuh create mode 100644 csrc/kernels/mla/hk/sparse_mla_bf16_reduce.cuh create mode 100644 csrc/kernels/mla/hk/sparse_mla_bf16_wave4.cuh create mode 100644 csrc/kernels/mla/sparse_mla_bf16.cu create mode 100644 docs/sparse_mla_bf16.md create mode 100644 op_tests/check_sparse_mla_bf16_isa.py create mode 100644 op_tests/op_benchmarks/bench_sparse_mla_bf16.py create mode 100644 op_tests/test_sparse_mla_bf16.py diff --git a/aiter/__init__.py b/aiter/__init__.py index c5bce4c858e..838d362d67d 100644 --- a/aiter/__init__.py +++ b/aiter/__init__.py @@ -108,6 +108,7 @@ def getLogger(): from .ops.aiter_operator import * from .ops.activation import * from .ops.attention import * + from .ops.sparse_mla import sparse_mla_bf16_fwd # noqa: F401 from .ops.custom import * from .ops.custom_all_reduce import * from .ops.quick_all_reduce import * diff --git a/aiter/jit/optCompilerConfig.json b/aiter/jit/optCompilerConfig.json index c087cd5a3d0..9d06815c128 100644 --- a/aiter/jit/optCompilerConfig.json +++ b/aiter/jit/optCompilerConfig.json @@ -1,4 +1,29 @@ { + "module_sparse_mla_bf16": { + "srcs": [ + "f'{AITER_CSRC_DIR}/kernels/mla/sparse_mla_bf16.cu'", + "f'{AITER_CSRC_DIR}/kernels/mla/hk/sparse_mla_bf16_common.cuh'", + "f'{AITER_CSRC_DIR}/kernels/mla/hk/sparse_mla_bf16_full_lds.cuh'", + "f'{AITER_CSRC_DIR}/kernels/mla/hk/sparse_mla_bf16_full_lds_pv.cuh'", + "f'{AITER_CSRC_DIR}/kernels/mla/hk/sparse_mla_bf16_headtiles.cuh'", + "f'{AITER_CSRC_DIR}/kernels/mla/hk/sparse_mla_bf16_reduce.cuh'", + "f'{AITER_CSRC_DIR}/kernels/mla/hk/sparse_mla_bf16_wave4.cuh'" + ], + "flags_extra_cc": ["'-DENABLE_CK=0'"], + "flags_extra_hip": [ + "'-DKITTENS_CDNA4'", + "'-fno-gpu-flush-denormals-to-zero'" + ], + "extra_ldflags": "None", + "extra_include": [], + "verbose": "False", + "torch_exclude": "True", + "blob_gen_cmd": "''", + "third_party": [ + "'HipKittens'" + ], + "is_experimental": "True" + }, "module_aiter_core": { "srcs": [ "f'{AITER_CSRC_DIR}/pybind/aiter_core_pybind.cu'" diff --git a/aiter/ops/sparse_mla.py b/aiter/ops/sparse_mla.py new file mode 100644 index 00000000000..7c42a88fa7e --- /dev/null +++ b/aiter/ops/sparse_mla.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: MIT +"""HIP sparse MLA for BF16 D512 latent K/V on MI350 (gfx950).""" + +import torch + +from aiter.jit.core import compile_ops +from aiter.jit.utils.chip_info import get_gfx_runtime +from csrc.cpp_itfs.torch_utils import direct_register_custom_op + + +@compile_ops( + "module_sparse_mla_bf16", fc_name="sparse_mla_bf16_fwd_out", ffi_type="ctypes" +) +def _jit_sparse_mla_bf16_fwd_out( + q: torch.Tensor, + kv_buffer: torch.Tensor, + kv_indptr: torch.Tensor, + kv_indices: torch.Tensor, + out: torch.Tensor, + lse: torch.Tensor | None, + partial_o: torch.Tensor | None, + partial_lse: torch.Tensor | None, + softmax_scale: float, + kv_splits: int, + version: int, +) -> None: ... + + +def _sparse_mla_bf16_fwd_out( + q: torch.Tensor, + kv_buffer: torch.Tensor, + kv_indptr: torch.Tensor, + kv_indices: torch.Tensor, + out: torch.Tensor, + lse: torch.Tensor | None, + partial_o: torch.Tensor | None, + partial_lse: torch.Tensor | None, + softmax_scale: float, + kv_splits: int, + version: int, +) -> None: + # Runtime discovery stays inside the opaque op, outside Dynamo tracing. + if get_gfx_runtime() != "gfx950": + raise NotImplementedError("sparse_mla_bf16_fwd requires gfx950") + _jit_sparse_mla_bf16_fwd_out( + q, + kv_buffer, + kv_indptr, + kv_indices, + out, + lse, + partial_o, + partial_lse, + softmax_scale, + kv_splits, + version, + ) + + +def _fake_out( + q, + kv_buffer, + kv_indptr, + kv_indices, + out, + lse, + partial_o, + partial_lse, + softmax_scale, + kv_splits, + version, +): + return None + + +direct_register_custom_op( + "sparse_mla_bf16_fwd_out", + _sparse_mla_bf16_fwd_out, + ["out", "lse", "partial_o", "partial_lse"], + fake_impl=_fake_out, +) + + +def sparse_mla_bf16_fwd( + q: torch.Tensor, + kv_buffer: torch.Tensor, + kv_indptr: torch.Tensor, + kv_indices: torch.Tensor, + softmax_scale: float, + *, + kv_splits: int = 1, + version: int = 2, + out: torch.Tensor | None = None, + return_lse: bool = False, + lse: torch.Tensor | None = None, + workspace: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Sparse attention with a shared BF16 K/V latent and no appended RoPE. + + ``q`` is [Q, H, 512], H in {16, 64}. ``kv_buffer`` is [slots, 512], + [pages, page_size, 512], or [slots, 1, 1, 512]. Paged pools must be viewable + as flat rows without copying. Q/KV permit padded outer strides and BF16 + storage offsets; their last dimension must have stride one. + + ``kv_indptr`` is contiguous int32 [Q+1], and ``kv_indices`` is a contiguous + int32 vector of GLOBAL slot ids. Offsets must start at zero, be monotone, + and end at indices.numel(). These device values are caller invariants; + the hot path does not synchronize to validate them. Negative and + out-of-range slot ids are masked. Duplicates retain their multiplicity. + Empty rows produce O=0 and natural-log LSE=-inf. + + ``version=1`` selects Full140 for H64 (139,520 B LDS) and the one-wave + H16 kernel. ``version=2`` selects the hybrid: H16 uses four waves, and + H64 uses four head tiles when Q*kv_splits <= 128, otherwise four waves. + Both H64 versions and v2 H16 use the fully unrolled split-K reducer. + ``kv_splits`` is an explicit power of two from 1 through 32. + + O is contiguous BF16 [Q,H,512]; optional LSE is contiguous FP32 [Q,H]. + For split-K, workspace is contiguous FP32 ([Q,S,H,512], [Q,S,H]). + Supply out/lse/workspace to reuse all storage without allocation. Writable + storage must not overlap other arguments. Warm the JIT before graph capture. + Only gfx950 is supported; unsupported configurations raise an error. + """ + if q.ndim != 3: + raise ValueError("q must have shape [Q,H,512]") + if version not in (1, 2) or kv_splits not in (1, 2, 4, 8, 16, 32): + raise ValueError("version must be 1/2 and kv_splits must be 1/2/4/8/16/32") + queries, heads, width = q.shape + if kv_buffer.ndim == 4 and kv_buffer.shape[1:3] == (1, 1): + kv_buffer = kv_buffer.view(kv_buffer.shape[0], kv_buffer.shape[-1]) + elif kv_buffer.ndim == 3: + kv_buffer = kv_buffer.view(-1, kv_buffer.shape[-1]) + if out is None: + out = torch.empty((queries, heads, width), dtype=q.dtype, device=q.device) + if lse is not None and not return_lse: + raise ValueError("supplying lse requires return_lse=True") + if return_lse and lse is None: + lse = torch.empty((queries, heads), dtype=torch.float32, device=q.device) + partial_o = partial_lse = None + if kv_splits > 1: + if workspace is None: + partial_o = torch.empty( + (queries, kv_splits, heads, 512), dtype=torch.float32, device=q.device + ) + partial_lse = torch.empty( + (queries, kv_splits, heads), dtype=torch.float32, device=q.device + ) + else: + partial_o, partial_lse = workspace + torch.ops.aiter.sparse_mla_bf16_fwd_out( + q, + kv_buffer, + kv_indptr, + kv_indices, + out, + lse, + partial_o, + partial_lse, + float(softmax_scale), + kv_splits, + version, + ) + return out, lse diff --git a/csrc/kernels/mla/hk/sparse_mla_bf16_common.cuh b/csrc/kernels/mla/hk/sparse_mla_bf16_common.cuh new file mode 100644 index 00000000000..07a41673df9 --- /dev/null +++ b/csrc/kernels/mla/hk/sparse_mla_bf16_common.cuh @@ -0,0 +1,878 @@ +// SPDX-License-Identifier: MIT +// Derived from ROCm/aiter PR #3459 (e961cb0c), HipKittens d3cd9b31. +// Sparse MLA with a shared BF16 K/V latent: D512, no appended RoPE, H16/H64. +#pragma once + +#include "hk_mla_v40_fwd_decode_gen1_common.cuh" +#include +#include + +namespace sparse_mla_bf16 { + +template +struct Traits +{ + static_assert(Heads == 16 || Heads == 64); + static constexpr uint32_t kBlockM = Heads, kNumWarps = Heads / 16; + static constexpr uint32_t kBlockN = 32, kBlockK = 32, kTileM = 16; + static constexpr uint32_t kQkHeadDim = 512, kVoHeadDim = 512; + static constexpr uint32_t kQkNopeHeadDim = 512, kQkRopeHeadDim = 0; + static constexpr uint32_t kNumThreads = 64 * kNumWarps; + static constexpr uint32_t kRoundMode = 0; // nearest even, including BF16 ties + // H16 keeps one KV tile for four LDS-limited workgroups/CU; H64 uses two + // tiles to overlap next-KV DMA with current MFMA. Q staging overlays this + // allocation only in the prologue; terminal O bounce overlays dead KV0. + static constexpr uint32_t kKvBytes = 32 * 512 * 2; + static constexpr uint32_t kKvBuffers = Heads == 64 ? 2u : 1u; + static constexpr uint32_t kIndexOffset = kKvBuffers * kKvBytes; + static constexpr uint32_t kLdsBytes = kIndexOffset + 256; + static_assert(kLdsBytes <= 80 * 1024); + static_assert(kNumWarps * 4352 <= kKvBytes); +}; + +struct Params +{ + const uint16_t* q; + const uint16_t* kv; + const int32_t* indptr; + const int32_t* indices; + uint16_t* out; + float* lse; + float* partial_o; + float* partial_lse; + int32_t queries; + int32_t splits; + int64_t slots; + int64_t q_stride0; + int64_t q_stride1; + int64_t kv_stride0; + float scale; +}; + +// These BF16 LDS readers are copied from PR #3459 unchanged. They intentionally +// have no FP8/scale/RoPE loaders or traits: the new writer supplies their exact +// sb8-permuted, row-bank-swizzled 32x512 BF16 layout. +template +struct Bf16KvManager +{ + static constexpr uint32_t kLoadBlockN = 32, kSubBlockRows = 16; + static constexpr uint32_t kSubBlockCols = 32, kSubBlockBytes = 1024; + static constexpr uint32_t kNumRowTiles = 2; + __device__ __forceinline__ static constexpr uint32_t sub_block_byte_offset(uint32_t row_tile, + uint32_t col_tile) + { + return (col_tile * kNumRowTiles + row_tile) * kSubBlockBytes; + } + // ---- LDS -> VGPR readout for QK / PV mfmas ----------------------------- + // QK A-tile load: ds_read_b128 of one 16 x 32 bf16 sub-block into 4 vgprs. + // (kRowOffset, kColOffset) selects which (row_tile, col_tile) of the pong; + // the per-lane offset within the sub-block follows the mfma_f32_16x16x32_bf16 + // A-operand layout (lane = (row_in_tile, group_in_row)). + // kRowOffset spans the FULL logical tile (0/16 for kBlockN=32; 0/16/32/48 for + // kBlockN=64). The sub-tile (kRowOffset/32) is folded into the ds_read imm as + // kSubIdx*kSubPongBytes so A and B share the same base ptr -- avoids a per-call + // base+kSubPong address VGPR that CSE holds across QK and overlaps a pinned reg. + template + __device__ __forceinline__ static void load_k_to_gpr(RT& dst, const uintptr_t p_lds_kv) + { + static_assert(kRowOffset % kSubBlockRows == 0, + "load_k_to_gpr: kRowOffset must be a multiple of 16."); + static_assert(kColOffset % kSubBlockCols == 0, + "load_k_to_gpr: kColOffset must be a multiple of 32."); + static_assert(kRowOffset < T::kBlockN, "load_k_to_gpr: kRowOffset out of range."); + static_assert(kColOffset < T::kQkHeadDim, "load_k_to_gpr: kColOffset out of range."); + // Sub-tile B (kRowOffset >= 32) -> +kSubPongBytes in the ds_read imm. + constexpr uint32_t kSubIdx = kRowOffset / kLoadBlockN; + constexpr uint32_t kRowInSub = kRowOffset % kLoadBlockN; + constexpr uint32_t kSubPongBytes = kLoadBlockN * T::kQkHeadDim * sizeof(hk::bf16); + + // mfma_f32_16x16x32_bf16 A-operand layout: lane t holds 8 bf16 from + // row r = lane%16, cols [c, c+8) where c = (lane/16) * 8. + // 8 bf16 = 4 dwords -> ds_read_b128. + constexpr uint32_t kMfmaRows = 16; + constexpr uint32_t kMfmaElemPerThr = 8; + + const uint32_t lane_idx = opus::lane_id(); + const uint32_t row = lane_idx % kMfmaRows; + const uint32_t col = (lane_idx / kMfmaRows) * kMfmaElemPerThr; + + // Un-swizzle: writer XORs intra-sub-block byte position by 32 on + // sub-tile-rows 1 & 3 (rows 4..7 and 12..15) to break the 2-way bank + // conflict; the reader applies the same XOR on the col-byte component. + const uint32_t row_bank_swap = ((row >> 2) & 1u) << 5; + const uint32_t in_sb_byte = + row * (kSubBlockCols * sizeof(hk::bf16)) + ((col * sizeof(hk::bf16)) ^ row_bank_swap); + + // Constexpr sub-block selector (compiles to immediate offset); the sub-tile + // (A/B) contributes kSubIdx*kSubPongBytes so the base ptr stays p_lds_kv. + constexpr uint32_t kFixedOffset = + sub_block_byte_offset(kRowInSub / kSubBlockRows, kColOffset / kSubBlockCols) + + kSubIdx * kSubPongBytes; + + // RT must hold a single 4-vgpr range (16 bf16 mfma A-tile = 4 vgprs/lane). + using range_type = hkdart::get_nth_range_t; + static_assert(range_type::lo + 3 == range_type::hi, + "ds_read_b128 requires 4 consecutive registers"); + + const uintptr_t p_lds_kv_lane = p_lds_kv + in_sb_byte; + hkm::ds_read_b128(static_cast(p_lds_kv_lane), kFixedOffset); + } + + // PV A-tile load: ds_read_b64_tr_b16 (bf16 transpose-read) of one 16-row x 16-col + // bf16 patch from a sub-block, results land in 2 dwords/lane in (GPR, GPR+1). + // + // PV math is V^T @ P^T = O^T computed via mma_ABt(oaccu, kv, p_mfma) (= kv @ p_mfma^T, + // matching the QK convention of K^T @ Q^T = P^T). So `kv` is the A operand of + // v_mfma_f32_16x16x32_bf16, holding V^T values reorganized into the mfma A layout. + // + // Within each 16-lane group, lane t's 4 bf16 (4 bf16 = 1 b64 = 2 dwords/lane) are + // (after HW transpose): + // output_lane[g*16 + l] holds V[g*4+0..g*4+3, kColOffset + l] + // for g = lane_group_idx (0..3), l = lane_in_group (0..15). I.e. each lane gets + // 4 K-rows of one V-col. Caller stitches two row halves (kRowOffset = 0, then 16) + // into a single mfma A operand spanning 8 K-rows (= mfma K = 0..7). + // + // Per-lane source address (within the selected 16x32 sub-block): + // in_sb_byte = (lane >> 2) * (kSubBlockCols * sizeof(bf16)) + (lane & 3) * 8 + // = lane_row * 64 + lane_col_quad * 8 + // (row stride = 32 bf16 cols * 2 B = 64 B; each "col_quad" = 4 bf16 = 8 B.) + // + // Compile-time fixed_offset selects: + // * the sub-block (row_tile = kRowOffset/16, col_tile = kColOffset/32), and + // * the 16-col half within that 32-col sub-block: kColOffset%32 -> +0 or +32 B. + // + // Un-swizzle: writer XORs intra-sub-block byte position by 32 on rows + // whose sub-tile-row index is odd (rows 4..7 and 12..15 within the 16-row + // sub-block). The reader applies the same XOR. With this swizzle both + // cycles of ds_read_b64_tr_b16 (lanes 0..31 covering rows 0..7, lanes + // 32..63 covering rows 8..15) hit 32 distinct conflict slots -- fully + // conflict-free for this per-wave instruction mapping. + // kRowOffset spans the FULL logical tile (0/16 = sub-tile A, 32/48 = sub-tile B). + // The sub-tile (kRowOffset/32) is folded into the ds_read imm as + // kSubIdx*kSubPongBytes so the base ptr stays p_lds_v (no +kSubPong VGPR). + template + __device__ __forceinline__ static void load_transposed_v_to_gpr(const uintptr_t p_lds_v) + { + static_assert((kRowOffset % kSubBlockRows == 0u) && (kRowOffset < T::kBlockN), + "load_transposed_v_to_gpr: kRowOffset must be 0/16 (A) or 32/48 (B)."); + static_assert( + (kColOffset % 16u == 0u) && (kColOffset < T::kVoHeadDim), + "load_transposed_v_to_gpr: kColOffset must be a multiple of 16, < kVoHeadDim."); + + constexpr uint32_t kSubIdx = kRowOffset / kLoadBlockN; // 0 (A) or 1 (B) + constexpr uint32_t kRowInSub = kRowOffset % kLoadBlockN; // 0 or 16 + constexpr uint32_t kSubPongBytes = kLoadBlockN * T::kQkHeadDim * sizeof(hk::bf16); + constexpr uint32_t kRowTile = kRowInSub / kSubBlockRows; // 0 or 1 + constexpr uint32_t kColTile = kColOffset / kSubBlockCols; // 0..15 + constexpr uint32_t kColInSbBytes = (kColOffset % kSubBlockCols) * sizeof(hk::bf16); + // Bank-swizzle re-expressed as a conditional ±32 delta so that + // (kFixedOffset + kColInSbBytes) stays fully constexpr in the + // ds_read_b64_tr_b16 immediate offset: XOR-by-32 against a constexpr + // value flips bit 5, equivalent to "+32 if bit was 0, else -32". + // The sign is compile-time (from kColInSbBytes's bit 5); only the + // boolean `is_swz` (1 bit per lane) is runtime. Avoids materialising + // kColInSbBytes as a runtime VGPR (vs. the plain XOR formulation), + // which freed 2 unpinned VGPRs in the audit. + constexpr int32_t kSwzDelta = (kColInSbBytes & 32u) ? -32 : +32; + constexpr uint32_t kFixedOffset = + sub_block_byte_offset(kRowTile, kColTile) + kColInSbBytes + kSubIdx * kSubPongBytes; + + const uint32_t lane_idx = opus::lane_id(); + const uint32_t row_in_sb = lane_idx >> 2; + const uint32_t is_swz = (row_in_sb >> 2) & 1u; + const uint32_t in_sb = row_in_sb * (kSubBlockCols * sizeof(hk::bf16)) + + (lane_idx & 3u) * 8u + is_swz * static_cast(kSwzDelta); + const uint32_t addr = static_cast(p_lds_v) + in_sb; + + hkm::ds_read_b64_tr_b16(addr, kFixedOffset); + } + + // bf16 ds_read_b64_tr_b16 already lands in the mfma A-operand layout -- no + // intra-lane v_swap_b32 fixup needed (V32's fp8 path interleaved cols c and + // c+16 into the same 2 GPRs and required a swap; the b16 transpose does not). + // Kept as a no-op for caller parity with KvManager8bitsV3. + template + __device__ __forceinline__ static void finalize_load_transposed_v_to_gpr() + { + } +}; + +// PR PV schedule, changed only to accept the BF16-only LDS manager. +template +__device__ __forceinline__ void +pv_gemm(Bf16KvManager& kv_manager, const uintptr_t p_lds_v, const float rescale) +{ + using R = HkMlaV40Regs; + using comp_t = typename R::comp_t; + constexpr uint32_t k_o_begin = R::k_o_begin; + constexpr uint32_t k_v0_begin = R::k_v0_begin; + constexpr uint32_t k_k0_begin = R::k_k0_begin; + constexpr uint32_t k_k1_begin = R::k_k1_begin; + constexpr uint32_t k_k2_begin = R::k_k2_begin; + // kBlockN=64 (m16x8 double-tile) contracts BOTH 32-row sub-tiles in one call: + // sub-tile A (rows 0:32) uses p_mfma_a, B (rows 32:64) uses p_mfma_b. kBlockN=32 + // has only A. Picked per-iter by row_base (see below), no caller arg. + typename R::p_mfma_a_t p_mfma_a; + typename R::p_mfma_b_t p_mfma_b; + typename R::pv_v_0_t pv_v_0; + typename R::pv_v_1_t pv_v_1; + typename R::pv_v_2_t pv_v_2; + typename R::pv_v_3_t pv_v_3; + + // D-tiling: each D-iter emits 2 oaccu 16x16 sub-tiles (32 D-cols); 512/32 = 16 + // D-iters per 32-row KV sub-tile. kBlockN=64 runs BOTH sub-tiles in one call, so + // num_pv_iter = kNumKvSub * 16. iter [0,16) = sub-tile A, [16,32) = sub-tile B. + constexpr uint32_t kDIters = T::kVoHeadDim / (2u * T::kTileM); // 16 + constexpr uint32_t num_pv_iter = R::kNumKvSub * kDIters; // 16 or 32 + // Flat V base tiles across all sub-tiles: 32 per sub-tile (S_0..S_31 each). + constexpr uint32_t kVTilesPerSub = 2u * kDIters; // 32 + constexpr uint32_t kNumVTiles = R::kNumKvSub * kVTilesPerSub; // 32 or 64 + + auto pk_mul_pair = [&](float r, auto base_c) { + constexpr uint32_t base = decltype(base_c)::value; + const float2 r2 = {r, r}; + asm volatile("v_pk_mul_f32 v[%0:%1], %2, v[%0:%1]" : : "n"(base), "n"(base + 1), "v"(r2)); + }; + auto mul_pair = [&](float r, auto base_c) { + constexpr uint32_t base = decltype(base_c)::value; + asm volatile("v_mul_f32_e32 v[%0], %1, v[%0]" : : "n"(base), "v"(r)); + asm volatile("v_mul_f32_e32 v[%0], %1, v[%0]" : : "n"(base + 1), "v"(r)); + }; + + // Issue both ds_read_b64 for base tile S_jj (within its sub-tile) into a + // round-robin slot (4 slots). kRowBase selects the KV sub-tile (0=A, 32=B), + // folded into the V ds_read imm by load_transposed_v_to_gpr. + auto load_S = [&]() { + constexpr uint32_t base = (slot == 0u) ? k_v0_begin + : (slot == 1u) ? k_k0_begin + : (slot == 2u) ? k_k1_begin + : k_k2_begin; + kv_manager.template load_transposed_v_to_gpr(p_lds_v); + kv_manager.template load_transposed_v_to_gpr(p_lds_v); + }; + // mfma: oaccu_dst (+)= pv_v_{slot}^T @ p_mfma[sub-tile]. 3-arg init only on the + // FIRST sub-tile of a first-iter call (kRowBase==0); sub-tile B always accumulates + // onto A's oaccu (same oaccu tile, second N-half of P). + auto do_mma = [&](OA& oaccu_dst) { + auto run = [&](auto& v) { + if constexpr(kIsFirstIter && kRowBase == 0u) + hk::mma_ABt(oaccu_dst, v, p_mfma_a); + else if constexpr(kRowBase == 0u) + hk::mma_ABt(oaccu_dst, v, p_mfma_a, oaccu_dst); + else + hk::mma_ABt(oaccu_dst, v, p_mfma_b, oaccu_dst); + }; + if constexpr(slot == 0u) + run(pv_v_0); + else if constexpr(slot == 1u) + run(pv_v_1); + else if constexpr(slot == 2u) + run(pv_v_2); + else + run(pv_v_3); + }; + + if constexpr(kDoRescale) + { + opus::static_for<2>([&](auto s) { + pk_mul_pair(rescale, opus::number{}); + pk_mul_pair(rescale, opus::number{}); + }); + } + + // Prologue: preload the first 3 flat V tiles (S_0,S_1,S_2 of sub-tile A) into + // slots 0,1,2. ONE prologue for the whole call -- sub-tile B's S_0..S_2 are + // prefetched by sub-tile A's tail iters via the flat prefetch index below. + load_S.template operator()<0u, 0u, 0u>(); + load_S.template operator()<0u, 1u, 1u>(); + load_S.template operator()<0u, 2u, 2u>(); + + opus::static_for([&](auto i) { + constexpr uint32_t iter = i.value; + constexpr uint32_t col_idx = iter % kDIters; // 0..15 D-tile within sub-tile + constexpr uint32_t row_base = (iter / kDIters) * 32u; // 0 (A) or 32 (B) + constexpr bool has_next = (iter + 1u) < num_pv_iter; + // Rescale only folds on the FIRST sub-tile pass (row_base==0) and only when a + // NEXT D-tile in that pass exists (col_idx+1 < kDIters) -- else oaccu overflows. + constexpr bool resc_next = kDoRescale && (row_base == 0u) && (col_idx + 1u < kDIters); + constexpr uint32_t next_oaccu_base = k_o_begin + (col_idx + 1u) * 8u; + + // Flat V-tile index across all sub-tiles: prefetch runs on this so sub-tile A's + // tail prefetches sub-tile B's S_0.. (ONE prologue). Consume uses col_idx. + constexpr uint32_t jf_lo = 2u * iter; // flat mfma idx (a) + constexpr uint32_t jf_hi = 2u * iter + 1u; // flat mfma idx (b) + // 4-slot round-robin (flat tile T -> slot T%4). Reading slot j%4 while refilling + // slot (j+3)%4 keeps the refill ds_read's dst != the mfma's src reg. + constexpr uint32_t slot_lo = jf_lo % 4u; + constexpr uint32_t slot_hi = jf_hi % 4u; + + constexpr uint32_t oaccu_base = k_o_begin + col_idx * 8u; + using oaccu_a_r = + hkdart::split_many_t>, + 4>; + using oaccu_b_r = + hkdart::split_many_t>, + 4>; + hk::art oaccu_a; + hk::art oaccu_b; + + // Prefetch S at flat index jf+3 -> its own sub-tile row_base + local jj%32. + auto prefetch = [&]() { + if constexpr(jf < kNumVTiles) + { + constexpr uint32_t pf_row = (jf / kVTilesPerSub) * 32u; + constexpr uint32_t pf_jj = jf % kVTilesPerSub; + load_S.template operator()(); + } + }; + + // ---- mfma_a (flat jf_lo, reads slot_lo) ---- + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(has_next ? 4 : 2, -1)); + prefetch.template operator()(); + do_mma.template operator()(oaccu_a); + if constexpr(resc_next) + { + mul_pair(rescale, opus::number{}); + mul_pair(rescale, opus::number{}); + } + + // ---- mfma_b (flat jf_hi, reads slot_hi) ---- + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(has_next ? 4 : 0, -1)); + prefetch.template operator()(); + do_mma.template operator()(oaccu_b); + if constexpr(resc_next) + { + mul_pair(rescale, opus::number{}); + mul_pair(rescale, opus::number{}); + } + }); +} + +// Pinned-register access is deliberately opaque, as in the upstream hk::art +// helpers. The entry's compiler scratch cap and clobber map are part of its ABI. +template +__device__ __forceinline__ float read_f32() +{ + float x; + asm volatile("v_mov_b32 %0, v[%1]" : "=v"(x) : "n"(Reg)); + return x; +} +template +__device__ __forceinline__ int32_t read_i32() +{ + int32_t x; + asm volatile("v_mov_b32 %0, v[%1]" : "=v"(x) : "n"(Reg)); + return x; +} +template +__device__ __forceinline__ void write_f32(float x) +{ + asm volatile("v_mov_b32 v[%0], %1" : : "n"(Reg), "v"(x)); +} +template +__device__ __forceinline__ void write_u32(uint32_t x) +{ + asm volatile("v_mov_b32 v[%0], %1" : : "n"(Reg), "v"(x)); +} +__device__ __forceinline__ void drain() +{ + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, 0)); + __builtin_amdgcn_sched_barrier(0); +} + +// The common aligned layout moves an entire eight-BF16 MFMA packet in one +// 128-bit load. The ABI still admits only two-byte alignment: padded/misaligned +// tensors use the exact scalar-pair path below. All values are copied bitwise. +template +__device__ __forceinline__ void load_q(const Params& p) +{ + using R = HkMlaV40Regs; + const uint32_t lane = threadIdx.x & 63u; + const uint32_t head = (threadIdx.x >> 6) * 16u + (lane & 15u); + const uint64_t offset = + uint64_t(blockIdx.x) * uint64_t(p.q_stride0) + uint64_t(head) * uint64_t(p.q_stride1); + const uint16_t* row = p.q + offset; + const bool aligned = ((reinterpret_cast(p.q) | (uint64_t(p.q_stride0) * 2u) | + (uint64_t(p.q_stride1) * 2u)) & + 15u) == 0; + opus::static_for<16>([&](auto c_) { + constexpr uint32_t c = c_.value; + const uint32_t col = sb8_inv_perm_col_elems(c * 32u + (lane >> 4) * 8u); + if(aligned) + { + // sb8 changes the packet order, never the eight values inside it. + const v4ui packet = *reinterpret_cast(row + col); + opus::static_for<4>([&](auto j_) { + constexpr uint32_t j = j_.value; + write_u32(packet[j]); + }); + } + else + { + opus::static_for<4>([&](auto j_) { + constexpr uint32_t j = j_.value; + const uint32_t pair = + uint32_t(row[col + j * 2u]) | (uint32_t(row[col + j * 2u + 1u]) << 16); + write_u32(pair); + }); + } + // Do not hoist an entire 512-column Q row into compiler-managed VGPRs. + __builtin_amdgcn_sched_barrier(0); + }); + drain(); +} + +// Prologue-only Q staging reuses the not-yet-live KV ping/pong allocation. +// H64 occupies exactly64KiB (four private16KiB wave regions); H16 occupies16KiB. +// After every wave reads its pinned Q and crosses the terminal barrier, these +// bytes are dead and KV may overwrite them. No extra steady-state LDS is used. +template +__device__ __forceinline__ void load_q_dma(const Params& p, uint16_t* staging) +{ + using R = HkMlaV40Regs; + static_assert(T::kNumWarps * 16u * 512u * 2u <= T::kIndexOffset); + const bool aligned = ((reinterpret_cast(p.q) | (uint64_t(p.q_stride0) * 2u) | + (uint64_t(p.q_stride1) * 2u)) & + 15u) == 0; + if(!aligned) + { + load_q(p); // preserve the original two-byte-alignment ABI + return; + } + const uint32_t lane = threadIdx.x & 63u; + const uint32_t wave = __builtin_amdgcn_readfirstlane(threadIdx.x >> 6); + const uint32_t r = lane >> 2, quad = lane & 3u; + const uint32_t head = wave * 16u + r; + const uint64_t row_offset = + uint64_t(blockIdx.x) * uint64_t(p.q_stride0) + uint64_t(head) * uint64_t(p.q_stride1); + const uint32_t base = __builtin_amdgcn_readfirstlane( + static_cast(reinterpret_cast(staging))) + + wave * 16384u; + // Queue all16 packets without a VMEM wait between chunks. Every source and + // destination is aligned; Q values remain their supplied BF16 bit patterns. + for(uint32_t c = 0; c < 16u; ++c) + { + const uint32_t L = c * 32u + ((quad * 8u) ^ (((r >> 2) & 1u) * 16u)); + const uint32_t col = sb8_inv_perm_col_elems(L); + const uint64_t address = reinterpret_cast(p.q + row_offset + uint64_t(col)); + const uint32_t destination = base + c * 1024u; + asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tglobal_load_lds_dwordx4 %1, off" + : + : "s"(destination), "v"(address) + : "m0", "memory"); + } + drain(); + __syncthreads(); // all Q packets have reached their wave-private LDS region + const uint32_t row = lane & 15u, col = (lane >> 4) * 8u; + const uint32_t read_address = base + row * 64u + ((col * 2u) ^ (((row >> 2) & 1u) * 32u)); + opus::static_for<16>([&](auto c_) { + constexpr uint32_t c = c_.value; + hkm::ds_read_b128(read_address, c * 1024u); + }); + drain(); + __syncthreads(); // Q LDS lifetime ends before any wave starts writing KV +} + +template +__device__ __forceinline__ void +stage_kv(const Params& p, uint16_t* tile, int32_t* selected, int64_t start, int64_t end) +{ + const uint32_t tid = threadIdx.x; + if(tid < 32) + { + // CSR entry and physical slot are both guarded before their load. + int32_t slot = -1; + if(int64_t(start) + tid < int64_t(end)) + slot = p.indices[int64_t(start) + tid]; + selected[tid] = slot; + } + __syncthreads(); + // Keep each load address 64 bit through multiplication. An invalid slot + // does not produce a pointer, including when the entire cache is empty. + const bool aligned = + ((reinterpret_cast(p.kv) | (uint64_t(p.kv_stride0) * 2u)) & 15u) == 0; + if(aligned) + { + // Eight contiguous BF16 values map to one aligned 16-byte LDS packet: + // sb8 permutation and bank XOR preserve all intra-packet bits. This + // reduces address/index work and global/LDS instructions by eight. + // The PR's K ring and P registers are dead during this stage. Reuse + // v44..v59 as four in-flight vector packets while Q64..127 and O128..255 + // remain resident. Compiler scratch stays strictly below v44. + using R = HkMlaV40Regs; + constexpr uint32_t first_reg = R::k_scratch_budget; + constexpr uint32_t packets = 4; + static_assert(first_reg == 44 && first_reg + packets * 4u <= R::k_q_vgpr_begin); + static_assert((32u * 64u) % (packets * T::kNumThreads) == 0); + for(uint32_t first = tid; first < 32u * 64u; first += packets * T::kNumThreads) + { + // Full EXEC initializes every lane before any predicated global + // load. Invalid lanes thus publish zero without forming a pointer. + opus::static_for([&](auto reg) { + asm volatile("v_mov_b32 v[%0], 0" : : "n"(first_reg + reg.value)); + }); + opus::static_for([&](auto i_) { + constexpr uint32_t i = i_.value, reg = first_reg + i * 4u; + uint32_t packet = first + i * T::kNumThreads; + // Prevent four independently strength-reduced address streams + // from staying live across the pinned prefetch and escaping + // compiler scratch v0..43 (the same break-CSE idiom as PR O). + asm volatile("" : "+v"(packet)); + const uint32_t row = packet / 64u, col = (packet % 64u) * 8u; + const int32_t slot = selected[row]; + if(slot >= 0 && int64_t(slot) < p.slots) + { + const uint64_t offset = uint64_t(slot) * uint64_t(p.kv_stride0) + uint64_t(col); + const uint64_t address = reinterpret_cast(p.kv + offset); + // Flat 64-bit address; never truncate to a cache-wide + // buffer descriptor or a 32-bit byte offset. + asm volatile("global_load_dwordx4 v[%0:%1], %2, off" + : + : "n"(reg), "n"(reg + 3u), "v"(address) + : "memory"); + } + }); + // Some packets can issue no VMEM instruction (all lanes invalid), + // so a fixed partial vmcnt threshold would be incorrect here. + drain(); + opus::static_for([&](auto i_) { + constexpr uint32_t i = i_.value, reg = first_reg + i * 4u; + uint32_t packet = first + i * T::kNumThreads; + // Prevent four independently strength-reduced address streams + // from staying live across the pinned prefetch and escaping + // compiler scratch v0..43 (the same break-CSE idiom as PR O). + asm volatile("" : "+v"(packet)); + const uint32_t row = packet / 64u, col = (packet % 64u) * 8u; + const uint32_t L = sb8_perm_col_elems(col); + const uint32_t r = row & 15u; + const uint32_t byte = (L / 32u * 2u + row / 16u) * 1024u + r * 64u + + ((2u * (L & 31u)) ^ (((r >> 2) & 1u) * 32u)); + hkm::ds_write_b128(reinterpret_cast(tile) + byte, 0); + }); + drain(); // retire every LDS writer before its pinned source is reused + } + } + else + { + for(uint32_t linear = tid; linear < 32u * 512u; linear += T::kNumThreads) + { + const uint32_t row = linear / 512u, col = linear % 512u; + const int32_t slot = selected[row]; + uint16_t value = 0; + if(slot >= 0 && int64_t(slot) < p.slots) + value = p.kv[uint64_t(slot) * uint64_t(p.kv_stride0) + uint64_t(col)]; + const uint32_t L = sb8_perm_col_elems(col); + const uint32_t r = row & 15u; + const uint32_t byte = (L / 32u * 2u + row / 16u) * 1024u + r * 64u + + ((2u * (L & 31u)) ^ (((r >> 2) & 1u) * 32u)); + tile[byte / 2u] = value; + } + } + drain(); + __syncthreads(); +} + +// Gfx950 full-64-bit global-to-LDS DMA. M0 supplies a wave-uniform destination +// base; lane t lands sixteen bytes at M0+t*16. The PR subblock mapping undoes +// bank XOR and sb8 on the global source side, making the destination contiguous. +// Counter semantics: DMA completion is VMEM (vmcnt), not LGKM. This helper only +// retires ordinary LDS operations, deliberately leaving valid DMA in flight. +template +__device__ __forceinline__ void +stage_kv_async(const Params& p, uint16_t* tile, int32_t* selected, int64_t start, int64_t end) +{ + const uint32_t tid = threadIdx.x; + if(tid < 32u) + { + int32_t slot = -1; + if(start + tid < end) + slot = p.indices[start + tid]; + selected[tid] = slot; + } + __syncthreads(); // index publication precedes every DMA; none is issued yet + const uint32_t lane = tid & 63u; + const uint32_t wave = __builtin_amdgcn_readfirstlane(tid >> 6); + const uint32_t base = + __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(tile))); + for(uint32_t block = wave; block < 32u; block += T::kNumWarps) + { + const uint32_t r = lane >> 2, quad = lane & 3u; + const uint32_t row = (block & 1u) * 16u + r; + const uint32_t L = (block >> 1) * 32u + ((quad * 8u) ^ (((r >> 2) & 1u) * 16u)); + const uint32_t col = sb8_inv_perm_col_elems(L); + const int32_t slot = selected[row]; + const uint32_t destination = base + block * 1024u; + if(slot >= 0 && int64_t(slot) < p.slots) + { + const uint64_t offset = uint64_t(slot) * uint64_t(p.kv_stride0) + uint64_t(col); + const uint64_t address = reinterpret_cast(p.kv + offset); + // M0->VMEM hazard gap follows LLVM's gfx950 DMA lowering. Source + // addresses are only formed inside the validity predicate. + asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tglobal_load_lds_dwordx4 %1, off" + : + : "s"(destination), "v"(address) + : "m0", "memory"); + } + else + { + // Inactive DMA lanes do not zero memory. Clear exactly their + // private sixteen-byte landing so stale previous-pong data cannot + // enter QK/PV, including all-invalid first/middle/final tiles. + const v4ui zero = {0, 0, 0, 0}; + *reinterpret_cast(reinterpret_cast(tile) + block * 1024u + + lane * 16u) = zero; + } + } + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, -1)); + __builtin_amdgcn_sched_barrier(0); +} + +// The QK ring and issue/wait/MFMA order are the PR m16x4 schedule. K comes from +// the same sb8 LDS layout; Q uses exactly the corresponding D-axis permutation. +template +__device__ __forceinline__ void qk_gemm(uintptr_t tile) +{ + using R = HkMlaV40Regs; + using mfma_ab_t = hk::bf16; + typename R::p_comp_lo_t lo; + typename R::p_comp_hi_t hi; + auto ring = [](uint32_t s) constexpr { + return s % 3u == 0u ? R::k_k0_begin : s % 3u == 1u ? R::k_k1_begin : R::k_k2_begin; + }; + auto issue = [&]() { + constexpr uint32_t base = ring(S); + using kr = hkdart::split_many_t>, 4>; + hk::art k; + Bf16KvManager::template load_k_to_gpr<(S % 2u) * 16u, (S / 2u) * 32u>(k, tile); + }; + auto mma = [&]() { + constexpr uint32_t base = ring(S), qb = R::k_q_vgpr_begin + (S / 2u) * 4u; + using kr = hkdart::split_many_t>, 4>; + using qr = hkdart::split_many_t>, 4>; + hk::art k; + hk::art q; + if constexpr(S == 0) + hk::mma_ABt(lo, k, q); + else if constexpr(S == 1) + hk::mma_ABt(hi, k, q); + else if constexpr((S % 2u) == 0) + hk::mma_ABt(lo, k, q, lo); + else + hk::mma_ABt(hi, k, q, hi); + }; + opus::static_for<3>([&](auto s) { issue.template operator()(); }); + opus::static_for<32>([&](auto s_) { + constexpr uint32_t s = s_.value, remain = 31u - s; + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(remain < 2u ? remain : 2u, -1)); + mma.template operator()(); + if constexpr(s + 3u < 32u) + issue.template operator()(); + }); + // The PR hides the final MFMA-to-VALU RAW latency behind its next-KV + // conversion/staging. Our synchronous BF16 path has no such intervening + // work. LLVM cannot see the numeric pinned-register dependence, and an + // lgkmcnt/vmcnt drain does not wait for MFMA arithmetic. Reserve a + // conservative 16 issue cycles before softmax first reads p56..p63. + // Without this gap the final D-axis slice is missing in one-hot Q tests. + asm volatile("s_nop 7\n\ts_nop 7"); + __builtin_amdgcn_sched_barrier(0); +} + +template +__device__ __forceinline__ void +softmax(const Params& p, const int32_t* selected, float& maximum, float& denominator) +{ + using R = HkMlaV40Regs; + constexpr float log2e = 1.4426950408889634074f; + const uint32_t group = (threadIdx.x & 63u) >> 4; + // QK's K ring is dead after its explicit MFMA latency gap. Prefetch the + // eight mask indices into v44..51 while keeping live scores in v56..63. + // One LGKM drain replaces eight serialized scalar reads; pending next-KV + // DMA (VMEM counter) remains in flight. PV may reuse the K ring afterward. + constexpr uint32_t mask_reg = R::k_scratch_budget; + static_assert(mask_reg + 8u <= R::k_p_comp_begin); + const uint32_t mask_address = + static_cast(reinterpret_cast(selected)) + group * 4u * sizeof(int32_t); + hkm::ds_read_b128(mask_address, 0); + hkm::ds_read_b128(mask_address, 16u * sizeof(int32_t)); + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, -1)); + __builtin_amdgcn_sched_barrier(0); + float local_max = -INFINITY; + opus::static_for<8>([&](auto i_) { + constexpr uint32_t i = i_.value; + const int32_t slot = read_i32(); + float score = read_f32(); + // Scaling is applied after the supplied BF16 QK FP32 accumulation. + score = (score * p.scale) * log2e; + if(slot < 0 || int64_t(slot) >= p.slots) + score = -INFINITY; + write_f32(score); + local_max = fmaxf(local_max, score); + }); + local_max = fmaxf(local_max, __shfl_xor(local_max, 16, 64)); + local_max = fmaxf(local_max, __shfl_xor(local_max, 32, 64)); + const float new_max = fmaxf(maximum, local_max); + const float alpha = denominator > 0.0f ? __builtin_amdgcn_exp2f(maximum - new_max) : 0.0f; + maximum = new_max; + typename R::oaccu_t o; + hk::mul_vgpr(o, o, alpha); + float sum = 0.0f; + opus::static_for<8>([&](auto i_) { + constexpr uint32_t i = i_.value; + const float score = read_f32(); + // All-empty tiles never evaluate exp2(-inf - -inf). + const float weight = score == -INFINITY ? 0.0f : __builtin_amdgcn_exp2f(score - maximum); + write_f32(weight); + sum += weight; + }); + sum += __shfl_xor(sum, 16, 64); + sum += __shfl_xor(sum, 32, 64); + denominator = denominator * alpha + sum; + // Exactly the upstream low-to-high P packing order and BF16 MFMA operands. + pack_2f32_to_bf16_pair_pinned(); + pack_2f32_to_bf16_pair_pinned(); + pack_2f32_to_bf16_pair_pinned(); + pack_2f32_to_bf16_pair_pinned(); +} + +template +__device__ __forceinline__ void sparse_mla_legacy(Params p) +{ + using T = Traits; + using R = HkMlaV40Regs; + static_assert(R::k_scratch_budget == 44); + // Q/O live throughout; K/P ranges intentionally overlay PV operands. + hkdart::clobber(); + hkdart::clobber(); + hkdart::clobber(); + hkdart::clobber(); + hkdart::clobber(); + hkdart::clobber(); + extern __shared__ __align__(16) unsigned char storage[]; + auto* tile = reinterpret_cast(storage); + auto* next_tile = reinterpret_cast(storage + T::kKvBytes); + auto* selected = reinterpret_cast(storage + T::kIndexOffset); + auto* next_selected = selected + 32; + const uintptr_t output_bounce = reinterpret_cast(storage); + const uint32_t query = blockIdx.x, partition = blockIdx.y; + const int32_t row_start = p.indptr[query], row_end = p.indptr[query + 1u]; + const int32_t length = row_end - row_start; + const int32_t begin = row_start + int64_t(length) * partition / p.splits; + const int32_t end = row_start + int64_t(length) * (partition + 1u) / p.splits; + load_q_dma(p, tile); + typename R::oaccu_t o; + hk::zero(o); + float maximum = -INFINITY, denominator = 0.0f; + Bf16KvManager manager; + const bool async_aligned = + T::kKvBuffers == 2u && + ((reinterpret_cast(p.kv) | (uint64_t(p.kv_stride0) * 2u)) & 15u) == 0; + if(async_aligned && begin < end) + { + stage_kv_async(p, tile, selected, begin, end); + drain(); + __syncthreads(); // tile0 prologue: every DMA/zero has completed + } + for(int64_t pos = begin; pos < end; pos += 32) + { + if(async_aligned) + { + if(pos + 32 < end) + stage_kv_async(p, next_tile, next_selected, pos + 32, end); + // No VMEM wait here: next-pong DMA overlaps current-pong compute. + } + else + { + // Preserve the proven BF16-alignment-only path synchronously. + stage_kv(p, tile, selected, pos, end); + } + const uintptr_t tile_addr = reinterpret_cast(tile); + qk_gemm(tile_addr); + softmax(p, selected, maximum, denominator); + // Reuse PR's exact PV four-slot LDS/MFMA schedule. Rescale occurred once + // above, so this call only accumulates into initialized FP32 O registers. + pv_gemm(manager, tile_addr, 1.0f); + drain(); // next DMA complete AND every current LDS reader retired + __syncthreads(); // publish next tile; old current may now become next + if(async_aligned) + { + auto* old_tile = tile; + tile = next_tile; + next_tile = old_tile; + auto* old_selected = selected; + selected = next_selected; + next_selected = old_selected; + } + } + const float inv_sum = denominator > 0.0f ? 1.0f / denominator : 0.0f; + hk::mul_vgpr(o, o, inv_sum); + constexpr float ln2 = 0.6931471805599453094f; + const float lse = + denominator > 0.0f ? (maximum + __builtin_amdgcn_logf(denominator)) * ln2 : -INFINITY; + const uint32_t lane = threadIdx.x & 63u, wave = threadIdx.x >> 6; + const uint64_t slot = Split ? uint64_t(query) * uint64_t(p.splits) + partition : query; + if(lane < 16u) + { + const uint64_t h = uint64_t(wave) * 16u + lane; + if constexpr(Split) + p.partial_lse[slot * Heads + h] = lse; + else if(p.lse) + p.lse[slot * Heads + h] = lse; + } + // Rebase the descriptor to one query/partition (64-bit) before the upstream + // manager's local 32-bit output offsets. This also works above 4 GiB. + if constexpr(Split) + { + OManager32bitsV4Gen1Swizzle output; + float* dst = p.partial_o + slot * Heads * 512u; + opus::static_for<8>([&](auto i) { + output.template output_to_vram_pair( + dst, wave, 0, 0, output_bounce, Heads); + drain(); + }); + } + else + { + OManager16bitsV4Gen1Swizzle output; + auto* dst = reinterpret_cast(p.out + slot * Heads * 512u); + opus::static_for<8>([&](auto i) { + output.template output_to_vram_pair( + dst, wave, 0, 0, output_bounce, Heads); + drain(); + }); + } +} + +// Guarded normalized-FP32 partial reducer. Each thread owns one D element; +// all-empty partitions contribute exactly zero without exp(-inf - -inf). +__global__ __launch_bounds__(256) void reduce(Params p, int heads) +{ + const uint32_t query = blockIdx.x, head = blockIdx.y; + const uint32_t d0 = threadIdx.x; + const uint64_t first = uint64_t(query) * uint64_t(p.splits) * uint64_t(heads) + head; + float maximum = -INFINITY; + for(int s = 0; s < p.splits; ++s) + maximum = fmaxf(maximum, p.partial_lse[first + uint64_t(s) * heads]); + float sum = 0.0f, a0 = 0.0f, a1 = 0.0f; + for(int s = 0; s < p.splits; ++s) + { + const uint64_t row = first + uint64_t(s) * heads; + const float value = p.partial_lse[row]; + const float weight = value == -INFINITY ? 0.0f : expf(value - maximum); + sum += weight; + a0 += weight * p.partial_o[row * 512u + d0]; + a1 += weight * p.partial_o[row * 512u + d0 + 256u]; + } + const float reciprocal = sum > 0.0f ? 1.0f / sum : 0.0f; + const uint64_t out_row = (uint64_t(query) * uint64_t(heads) + head) * 512u; + // Explicit nearest-even conversion, including subnormal and tie semantics. + uint32_t pair; + const float v0 = a0 * reciprocal, v1 = a1 * reciprocal; + asm volatile("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(pair) : "v"(v0), "v"(v1)); + p.out[out_row + d0] = uint16_t(pair); + p.out[out_row + d0 + 256u] = uint16_t(pair >> 16); + if(threadIdx.x == 0 && p.lse) + p.lse[uint64_t(query) * heads + head] = sum > 0.0f ? maximum + logf(sum) : -INFINITY; +} + +} // namespace sparse_mla_bf16 diff --git a/csrc/kernels/mla/hk/sparse_mla_bf16_full_lds.cuh b/csrc/kernels/mla/hk/sparse_mla_bf16_full_lds.cuh new file mode 100644 index 00000000000..e578e2debe0 --- /dev/null +++ b/csrc/kernels/mla/hk/sparse_mla_bf16_full_lds.cuh @@ -0,0 +1,374 @@ +// SPDX-License-Identifier: MIT +// Full-LDS H64 implementation: Q, two KV tiles and P stay in 139,520 B LDS. +// Included inside namespace sparse_mla_bf16. +struct FullLdsTraits : Traits<64> +{ + static constexpr uint32_t kNumWarps = 8, kNumThreads = 512; + static constexpr uint32_t kIndexOffset = 139264, kLdsBytes = 139520; +}; + +struct FullLdsRegs +{ + using comp_t = float; + template + using ranges = hkdart::split_many_t>, 4>; + template + using bf16_fragment = hk::art>; + // All compiler-managed instructions must remain below v40. QK operands + // occupy40:55; scores60:63. Softmax reuses40:47 for eight scores and48:55 + // for masks. PV uses40:55 for four V slots and60:63 for packed full-N32 P. + static constexpr uint32_t k_o_begin = 64, kNumKvSub = 1; + static constexpr uint32_t k_v0_begin = 40, k_k0_begin = 44; + static constexpr uint32_t k_k1_begin = 48, k_k2_begin = 52; + using pinned_ranges = ranges<40, 88>; + using oaccu_t = hk::art>; + using score_t = hk::art>; + using p_mfma_a_t = bf16_fragment<60>; + using p_mfma_b_t = bf16_fragment<60>; // unused by the N32 PV schedule + using pv_v_0_t = bf16_fragment<40>; + using pv_v_1_t = bf16_fragment<44>; + using pv_v_2_t = bf16_fragment<48>; + using pv_v_3_t = bf16_fragment<52>; +}; + +#include "sparse_mla_bf16_full_lds_pv.cuh" + +// LDS communications use explicit DS operations and a barrier that does not +// implicitly drain VMEM. Next-KV DMA must be allowed to overlap these phases. +__device__ __forceinline__ void full_lds_barrier() +{ + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, -1)); + __builtin_amdgcn_sched_barrier(0); + asm volatile("s_barrier" ::: "memory"); +} + +// Same packet mapping/validity/zeroing as stage_kv_async; H64 row reuse only. +__device__ __forceinline__ void full_lds_stage_kv_async( + const Params& p, uint16_t* tile, int32_t* selected, int64_t start, int64_t end) +{ + using T = FullLdsTraits; + static_assert(T::kNumWarps == 8u); + const uint32_t tid = threadIdx.x; + if(tid < 32u) + { + int32_t slot = -1; + if(start + tid < end) + slot = p.indices[start + tid]; + selected[tid] = slot; + } + __syncthreads(); // index publication precedes every DMA; none is issued yet + const uint32_t lane = tid & 63u; + const uint32_t wave = __builtin_amdgcn_readfirstlane(tid >> 6); + const uint32_t base = + __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(tile))); + // Four D32 packets share one selected row in this eight-wave layout. + const uint32_t r = lane >> 2, quad = lane & 3u; + const uint32_t row = (wave & 1u) * 16u + r; + const int32_t slot = selected[row]; + const bool valid = slot >= 0 && int64_t(slot) < p.slots; + // Each lane's predicate is invariant across all four packets. Restore + // EXEC for every lane before the common LDS drain below; no CTA barrier + // is entered from either divergent branch. + if(valid) + { + const uint64_t offset = uint64_t(slot) * uint64_t(p.kv_stride0); + const uint64_t row_base = reinterpret_cast(p.kv + offset); + opus::static_for<4>([&](auto packet_) { + const uint32_t block = wave + packet_.value * T::kNumWarps; + const uint32_t L = (block >> 1) * 32u + ((quad * 8u) ^ (((r >> 2) & 1u) * 16u)); + const uint32_t col = sb8_inv_perm_col_elems(L); + const uint32_t destination = base + block * 1024u; + const uint64_t address = row_base + uint64_t(col) * sizeof(uint16_t); + asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tglobal_load_lds_dwordx4 %1, off" + : + : "s"(destination), "v"(address) + : "m0", "memory"); + }); + } + else + { + opus::static_for<4>([&](auto packet_) { + const uint32_t block = wave + packet_.value * T::kNumWarps; + // Invalid lanes own the same private 16 B landing as valid lanes. + // Zero through the dead QK operands to avoid hoisted compiler + // constants crossing the pinned register boundary. + asm volatile("v_mov_b32 v40, 0\n\tv_mov_b32 v41, 0\n\t" + "v_mov_b32 v42, 0\n\tv_mov_b32 v43, 0" :: + : "memory"); + hkm::ds_write_b128<40>(static_cast(reinterpret_cast(tile)) + + block * 1024u + lane * 16u, + 0); + }); + } + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, -1)); + __builtin_amdgcn_sched_barrier(0); +} + +__device__ __forceinline__ void full_lds_load_q(const Params& p, uint16_t* q_lds) +{ + const uint32_t tid = threadIdx.x, lane = tid & 63u; + const uint32_t wave = __builtin_amdgcn_readfirstlane(tid >> 6); + const uint32_t base = + __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(q_lds))); + const bool aligned = ((reinterpret_cast(p.q) | (uint64_t(p.q_stride0) * 2u) | + (uint64_t(p.q_stride1) * 2u)) & + 15u) == 0; + if(aligned) + { + // 4 head groups * 16 D32 blocks. Eight waves jointly load Q exactly + // once; paired compute waves later read the same16KiB head-group area. + for(uint32_t block = wave; block < 64u; block += 8u) + { + const uint32_t group = block / 16u, c = block % 16u; + const uint32_t row = lane >> 2, quad = lane & 3u; + const uint32_t col = + sb8_inv_perm_col_elems(c * 32u + ((quad * 8u) ^ (((row >> 2) & 1u) * 16u))); + const uint64_t offset = uint64_t(blockIdx.x) * uint64_t(p.q_stride0) + + uint64_t(group * 16u + row) * uint64_t(p.q_stride1) + col; + const uint64_t address = reinterpret_cast(p.q + offset); + const uint32_t destination = base + block * 1024u; + asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tglobal_load_lds_dwordx4 %1, off" + : + : "s"(destination), "v"(address) + : "m0", "memory"); + } + } + else + { + // Minimum contract alignment is2B, including odd BF16 row/head strides. + for(uint32_t linear = tid; linear < 64u * 512u; linear += 512u) + { + uint32_t item = linear; + asm volatile("" : "+v"(item)); + const uint32_t head = item / 512u, col = item % 512u; + const uint64_t offset = uint64_t(blockIdx.x) * uint64_t(p.q_stride0) + + uint64_t(head) * uint64_t(p.q_stride1) + col; + const uint16_t value = p.q[offset]; + const uint32_t L = sb8_perm_col_elems(col), row = head & 15u; + const uint32_t byte = (head / 16u) * 16384u + (L / 32u) * 1024u + row * 64u + + ((2u * (L & 31u)) ^ (((row >> 2) & 1u) * 32u)); + q_lds[byte / 2u] = value; + } + } + drain(); + __syncthreads(); +} + +template +__device__ __forceinline__ void full_lds_qk(uintptr_t q_base, uintptr_t kv_base) +{ + using T = FullLdsTraits; + const uint32_t lane = threadIdx.x & 63u, row = lane & 15u; + const uint32_t col = (lane >> 4) * 8u; + const uint32_t q_address = static_cast(q_base) + (threadIdx.x >> 7) * 16384u + + row * 64u + ((col * 2u) ^ (((row >> 2) & 1u) * 32u)); + FullLdsRegs::score_t score; + auto issue = [&]() { + constexpr uint32_t qr = 40u + (C % 2u) * 4u, kr = 48u + (C % 2u) * 4u; + typename FullLdsRegs::template bf16_fragment k; + hkm::ds_read_b128(q_address, C * 1024u); + Bf16KvManager::template load_k_to_gpr(k, kv_base); + }; + auto mma = [&]() { + typename FullLdsRegs::template bf16_fragment<40u + (C % 2u) * 4u> q; + typename FullLdsRegs::template bf16_fragment<48u + (C % 2u) * 4u> k; + if constexpr(C == 0u) + hk::mma_ABt(score, k, q); + else + hk::mma_ABt(score, k, q, score); + }; + issue.template operator()<0u>(); + issue.template operator()<1u>(); + opus::static_for<16>([&](auto c_) { + constexpr uint32_t c = c_.value; + // Each future fragment consists of exactly Q+K LDS reads. Wait until + // the oldest pair is ready, leaving the next pair in flight. + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(c < 15u ? 2 : 0, -1)); + mma.template operator()(); + if constexpr(c + 2u < 16u) + issue.template operator()(); + }); + // Memory waits do not scoreboard the opaque numeric MFMA accumulator. + asm volatile("s_nop 7\n\ts_nop 7"); + __builtin_amdgcn_sched_barrier(0); +} + +__device__ __forceinline__ void full_lds_softmax( + const Params& p, const int32_t* selected, uintptr_t scores, float& maximum, float& denominator) +{ + const uint32_t lane = threadIdx.x & 63u, group = lane >> 4; + const uint32_t head = (threadIdx.x >> 7) * 16u + (lane & 15u); + const uint32_t score_address = + static_cast(scores) + (head / 16u) * 2048u + (group * 16u + (head & 15u)) * 16u; + const uint32_t mask_address = + static_cast(reinterpret_cast(selected)) + group * 16u; + // Both paired waves read the same N32 scores and perform the same softmax. + // This deliberately duplicates its arithmetic in the first candidate but + // needs no extra state slab, no BF16 score approximation and no broadcast. + hkm::ds_read_b128<40>(score_address, 0); + hkm::ds_read_b128<44>(score_address, 1024); + hkm::ds_read_b128<48>(mask_address, 0); + hkm::ds_read_b128<52>(mask_address, 64); + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, -1)); + __builtin_amdgcn_sched_barrier(0); + constexpr float log2e = 1.4426950408889634074f; + float local_max = -INFINITY; + opus::static_for<8>([&](auto i_) { + constexpr uint32_t i = i_.value; + const int32_t slot = read_i32<48u + i>(); + float score = (read_f32<40u + i>() * p.scale) * log2e; + if(slot < 0 || int64_t(slot) >= p.slots) + score = -INFINITY; + write_f32<40u + i>(score); + local_max = fmaxf(local_max, score); + }); + local_max = fmaxf(local_max, __shfl_xor(local_max, 16, 64)); + local_max = fmaxf(local_max, __shfl_xor(local_max, 32, 64)); + const float next_max = fmaxf(maximum, local_max); + const float alpha = denominator > 0.0f ? __builtin_amdgcn_exp2f(maximum - next_max) : 0.0f; + maximum = next_max; + FullLdsRegs::oaccu_t o; + hk::mul_vgpr(o, o, alpha); + float sum = 0.0f; + opus::static_for<8>([&](auto i_) { + constexpr uint32_t i = i_.value; + const float score = read_f32<40u + i>(); + const float weight = score == -INFINITY ? 0.0f : __builtin_amdgcn_exp2f(score - maximum); + write_f32<40u + i>(weight); + sum += weight; + }); + sum += __shfl_xor(sum, 16, 64); + sum += __shfl_xor(sum, 32, 64); + denominator = denominator * alpha + sum; + pack_2f32_to_bf16_pair_pinned<60, 40>(); + pack_2f32_to_bf16_pair_pinned<61, 42>(); + pack_2f32_to_bf16_pair_pinned<62, 44>(); + pack_2f32_to_bf16_pair_pinned<63, 46>(); +} + +template +__device__ __forceinline__ void full_lds_output(const Params& p, uint64_t slot, uintptr_t bounce) +{ + using T = FullLdsTraits; // physical output pitch remains D512 + const uint32_t head_group = threadIdx.x >> 7; + // Preserve independent LDS bounce regions for both waves of each pair. + // Float output needs8*4352=34816B, which fits dead Q64KiB but not one KV32KiB. + const uintptr_t half_bounce = bounce + Half * 4u * 4352u; + if constexpr(Split) + { + OManager32bitsV4Gen1Swizzle output; + float* dst = p.partial_o + slot * 64u * 512u; + opus::static_for<4>([&](auto i) { + output.template output_to_vram_pair<64u + i.value * 16u, + Half * 256u + i.value * 64u, + false>(dst, head_group, 0, 0, half_bounce, 64); + drain(); + }); + } + else + { + OManager16bitsV4Gen1Swizzle output; + auto* dst = reinterpret_cast(p.out + slot * 64u * 512u); + opus::static_for<4>([&](auto i) { + output.template output_to_vram_pair<64u + i.value * 16u, + Half * 256u + i.value * 64u, + false>(dst, head_group, 0, 0, half_bounce, 64); + drain(); + }); + } +} + +template +__device__ __forceinline__ void full_lds_h64(Params p) +{ + using T = FullLdsTraits; + static_assert(T::kNumWarps == 8u && T::kLdsBytes == 139520u); + hkdart::clobber(); + extern __shared__ __align__(16) unsigned char storage[]; + auto* q_lds = reinterpret_cast(storage); + auto* tile = reinterpret_cast(storage + 65536u); + auto* next_tile = reinterpret_cast(storage + 98304u); + const uintptr_t scores = reinterpret_cast(storage + 131072u); + auto* selected = reinterpret_cast(storage + 139264u); + auto* next_selected = selected + 32u; + const uint32_t query = blockIdx.x, partition = blockIdx.y; + const int32_t row_start = p.indptr[query], row_end = p.indptr[query + 1u]; + const int32_t length = row_end - row_start; + const int64_t begin = int64_t(row_start) + int64_t(length) * partition / p.splits; + const int64_t end = int64_t(row_start) + int64_t(length) * (partition + 1u) / p.splits; + full_lds_load_q(p, q_lds); + FullLdsRegs::oaccu_t o; + hk::zero(o); + float maximum = -INFINITY, denominator = 0.0f; + const uint32_t lane = threadIdx.x & 63u, wave = threadIdx.x >> 6; + const bool aligned = + ((reinterpret_cast(p.kv) | (uint64_t(p.kv_stride0) * 2u)) & 15u) == 0; + Bf16KvManager manager; + if(aligned && begin < end) + { + full_lds_stage_kv_async(p, tile, selected, begin, end); + drain(); + __syncthreads(); + } + for(int64_t pos = begin; pos < end; pos += 32) + { + if(aligned) + { + if(pos + 32 < end) + full_lds_stage_kv_async(p, next_tile, next_selected, pos + 32, end); + } + else + { + // v44:59 staging payloads are dead Q/K slots in this map; live + // O64:127 remains untouched. Retains minimum2B input alignment. + stage_kv(p, tile, selected, pos, end); + } + const uintptr_t kv_address = reinterpret_cast(tile); + if(wave & 1u) + full_lds_qk<1u>(reinterpret_cast(q_lds), kv_address); + else + full_lds_qk<0u>(reinterpret_cast(q_lds), kv_address); + const uint32_t head = (wave / 2u) * 16u + (lane & 15u); + const uint32_t key = (wave & 1u) * 16u + (lane >> 4) * 4u; + hkm::ds_write_b128<60>(static_cast(scores) + (head / 16u) * 2048u + + ((key / 4u) * 16u + (head & 15u)) * 16u, + 0); + full_lds_barrier(); // both N16 halves published; next KV stays pending + full_lds_softmax(p, selected, scores, maximum, denominator); + if(wave & 1u) + full_lds_pv_gemm<256u, false, false, T>(manager, kv_address, 1.0f); + else + full_lds_pv_gemm<0u, false, false, T>(manager, kv_address, 1.0f); + drain(); + __syncthreads(); // all score/current reads retired and next DMA complete + if(aligned) + { + auto* old_tile = tile; + tile = next_tile; + next_tile = old_tile; + auto* old_selected = selected; + selected = next_selected; + next_selected = old_selected; + } + } + const float inv_sum = denominator > 0.0f ? 1.0f / denominator : 0.0f; + hk::mul_vgpr(o, o, inv_sum); + constexpr float ln2 = 0.6931471805599453094f; + const float lse = + denominator > 0.0f ? (maximum + __builtin_amdgcn_logf(denominator)) * ln2 : -INFINITY; + const uint64_t slot = Split ? uint64_t(query) * uint64_t(p.splits) + partition : query; + if((wave & 1u) == 0 && lane < 16u) + { + const uint64_t head = uint64_t(wave / 2u) * 16u + lane; + if constexpr(Split) + p.partial_lse[slot * 64u + head] = lse; + else if(p.lse) + p.lse[slot * 64u + head] = lse; + } + // No Q read survives the terminal loop barrier (or prologue barrier when + // empty), so its64KiB region is now safe for the independent O bounce slots. + if(wave & 1u) + full_lds_output(p, slot, reinterpret_cast(storage)); + else + full_lds_output(p, slot, reinterpret_cast(storage)); +} diff --git a/csrc/kernels/mla/hk/sparse_mla_bf16_full_lds_pv.cuh b/csrc/kernels/mla/hk/sparse_mla_bf16_full_lds_pv.cuh new file mode 100644 index 00000000000..e64bf00ade5 --- /dev/null +++ b/csrc/kernels/mla/hk/sparse_mla_bf16_full_lds_pv.cuh @@ -0,0 +1,155 @@ +// SPDX-License-Identifier: MIT +// Included inside namespace sparse_mla_bf16, after FullLdsRegs. +// PR PV four-slot schedule, compact register map and D256 per wave. +// ColBase chooses the sb8-preserving physical output half. +template +__device__ __forceinline__ void +full_lds_pv_gemm(Bf16KvManager& kv_manager, const uintptr_t p_lds_v, const float rescale) +{ + using R = FullLdsRegs; + using comp_t = typename R::comp_t; + constexpr uint32_t k_o_begin = R::k_o_begin; + constexpr uint32_t k_v0_begin = R::k_v0_begin; + constexpr uint32_t k_k0_begin = R::k_k0_begin; + constexpr uint32_t k_k1_begin = R::k_k1_begin; + constexpr uint32_t k_k2_begin = R::k_k2_begin; + // kBlockN=64 (m16x8 double-tile) contracts BOTH 32-row sub-tiles in one call: + // sub-tile A (rows 0:32) uses p_mfma_a, B (rows 32:64) uses p_mfma_b. kBlockN=32 + // has only A. Picked per-iter by row_base (see below), no caller arg. + typename R::p_mfma_a_t p_mfma_a; + typename R::p_mfma_b_t p_mfma_b; + typename R::pv_v_0_t pv_v_0; + typename R::pv_v_1_t pv_v_1; + typename R::pv_v_2_t pv_v_2; + typename R::pv_v_3_t pv_v_3; + + // D-tiling: each D-iter emits 2 oaccu 16x16 sub-tiles (32 D-cols); 512/32 = 16 + // D-iters per 32-row KV sub-tile. kBlockN=64 runs BOTH sub-tiles in one call, so + // num_pv_iter = kNumKvSub * 16. iter [0,16) = sub-tile A, [16,32) = sub-tile B. + constexpr uint32_t kDIters = 256u / (2u * T::kTileM); // 8 per D256 wave + constexpr uint32_t num_pv_iter = R::kNumKvSub * kDIters; // 16 or 32 + // Flat V base tiles across all sub-tiles: 32 per sub-tile (S_0..S_31 each). + constexpr uint32_t kVTilesPerSub = 2u * kDIters; // 32 + constexpr uint32_t kNumVTiles = R::kNumKvSub * kVTilesPerSub; // 32 or 64 + + auto pk_mul_pair = [&](float r, auto base_c) { + constexpr uint32_t base = decltype(base_c)::value; + const float2 r2 = {r, r}; + asm volatile("v_pk_mul_f32 v[%0:%1], %2, v[%0:%1]" : : "n"(base), "n"(base + 1), "v"(r2)); + }; + auto mul_pair = [&](float r, auto base_c) { + constexpr uint32_t base = decltype(base_c)::value; + asm volatile("v_mul_f32_e32 v[%0], %1, v[%0]" : : "n"(base), "v"(r)); + asm volatile("v_mul_f32_e32 v[%0], %1, v[%0]" : : "n"(base + 1), "v"(r)); + }; + + // Issue both ds_read_b64 for base tile S_jj (within its sub-tile) into a + // round-robin slot (4 slots). kRowBase selects the KV sub-tile (0=A, 32=B), + // folded into the V ds_read imm by load_transposed_v_to_gpr. + auto load_S = [&]() { + constexpr uint32_t base = (slot == 0u) ? k_v0_begin + : (slot == 1u) ? k_k0_begin + : (slot == 2u) ? k_k1_begin + : k_k2_begin; + kv_manager.template load_transposed_v_to_gpr( + p_lds_v); + kv_manager.template load_transposed_v_to_gpr( + p_lds_v); + }; + // mfma: oaccu_dst (+)= pv_v_{slot}^T @ p_mfma[sub-tile]. 3-arg init only on the + // FIRST sub-tile of a first-iter call (kRowBase==0); sub-tile B always accumulates + // onto A's oaccu (same oaccu tile, second N-half of P). + auto do_mma = [&](OA& oaccu_dst) { + auto run = [&](auto& v) { + if constexpr(kIsFirstIter && kRowBase == 0u) + hk::mma_ABt(oaccu_dst, v, p_mfma_a); + else if constexpr(kRowBase == 0u) + hk::mma_ABt(oaccu_dst, v, p_mfma_a, oaccu_dst); + else + hk::mma_ABt(oaccu_dst, v, p_mfma_b, oaccu_dst); + }; + if constexpr(slot == 0u) + run(pv_v_0); + else if constexpr(slot == 1u) + run(pv_v_1); + else if constexpr(slot == 2u) + run(pv_v_2); + else + run(pv_v_3); + }; + + if constexpr(kDoRescale) + { + opus::static_for<2>([&](auto s) { + pk_mul_pair(rescale, opus::number{}); + pk_mul_pair(rescale, opus::number{}); + }); + } + + // Prologue: preload the first 3 flat V tiles (S_0,S_1,S_2 of sub-tile A) into + // slots 0,1,2. ONE prologue for the whole call -- sub-tile B's S_0..S_2 are + // prefetched by sub-tile A's tail iters via the flat prefetch index below. + load_S.template operator()<0u, 0u, 0u>(); + load_S.template operator()<0u, 1u, 1u>(); + load_S.template operator()<0u, 2u, 2u>(); + + opus::static_for([&](auto i) { + constexpr uint32_t iter = i.value; + constexpr uint32_t col_idx = iter % kDIters; // 0..15 D-tile within sub-tile + constexpr uint32_t row_base = (iter / kDIters) * 32u; // 0 (A) or 32 (B) + constexpr bool has_next = (iter + 1u) < num_pv_iter; + // Rescale only folds on the FIRST sub-tile pass (row_base==0) and only when a + // NEXT D-tile in that pass exists (col_idx+1 < kDIters) -- else oaccu overflows. + constexpr bool resc_next = kDoRescale && (row_base == 0u) && (col_idx + 1u < kDIters); + constexpr uint32_t next_oaccu_base = k_o_begin + (col_idx + 1u) * 8u; + + // Flat V-tile index across all sub-tiles: prefetch runs on this so sub-tile A's + // tail prefetches sub-tile B's S_0.. (ONE prologue). Consume uses col_idx. + constexpr uint32_t jf_lo = 2u * iter; // flat mfma idx (a) + constexpr uint32_t jf_hi = 2u * iter + 1u; // flat mfma idx (b) + // 4-slot round-robin (flat tile T -> slot T%4). Reading slot j%4 while refilling + // slot (j+3)%4 keeps the refill ds_read's dst != the mfma's src reg. + constexpr uint32_t slot_lo = jf_lo % 4u; + constexpr uint32_t slot_hi = jf_hi % 4u; + + constexpr uint32_t oaccu_base = k_o_begin + col_idx * 8u; + using oaccu_a_r = + hkdart::split_many_t>, + 4>; + using oaccu_b_r = + hkdart::split_many_t>, + 4>; + hk::art oaccu_a; + hk::art oaccu_b; + + // Prefetch S at flat index jf+3 -> its own sub-tile row_base + local jj%32. + auto prefetch = [&]() { + if constexpr(jf < kNumVTiles) + { + constexpr uint32_t pf_row = (jf / kVTilesPerSub) * 32u; + constexpr uint32_t pf_jj = jf % kVTilesPerSub; + load_S.template operator()(); + } + }; + + // ---- mfma_a (flat jf_lo, reads slot_lo) ---- + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(has_next ? 4 : 2, -1)); + prefetch.template operator()(); + do_mma.template operator()(oaccu_a); + if constexpr(resc_next) + { + mul_pair(rescale, opus::number{}); + mul_pair(rescale, opus::number{}); + } + + // ---- mfma_b (flat jf_hi, reads slot_hi) ---- + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(has_next ? 4 : 0, -1)); + prefetch.template operator()(); + do_mma.template operator()(oaccu_b); + if constexpr(resc_next) + { + mul_pair(rescale, opus::number{}); + mul_pair(rescale, opus::number{}); + } + }); +} diff --git a/csrc/kernels/mla/hk/sparse_mla_bf16_headtiles.cuh b/csrc/kernels/mla/hk/sparse_mla_bf16_headtiles.cuh new file mode 100644 index 00000000000..f2b0da8ce3d --- /dev/null +++ b/csrc/kernels/mla/hk/sparse_mla_bf16_headtiles.cuh @@ -0,0 +1,502 @@ +// SPDX-License-Identifier: MIT +// Experimental H16/H64 head tiling: each head16 CTA uses QK N64 / PV D128 per wave. +// Reuses PR #3459 HipKittens BF16 LDS readers, four-slot PV schedule, O managers. +#pragma once +#include "sparse_mla_bf16_common.cuh" + +namespace sparse_mla_bf16::headtiles { +struct Traits +{ + static constexpr uint32_t kBlockN = 64, kBlockK = 32, kTileM = 16; + static constexpr uint32_t kQkHeadDim = 512, kVoHeadDim = 512; + static constexpr uint32_t kNumWarps = 4, kNumThreads = 256, kRoundMode = 0; + static constexpr uint32_t kKvBytes = 65536, kPBytes = 2048; + static constexpr uint32_t kPOffset = kKvBytes; + static constexpr uint32_t kReduceOffset = kPOffset + kPBytes; + static constexpr uint32_t kIndexOffset = kReduceOffset + 512; + static constexpr uint32_t kLdsBytes = kIndexOffset + 256; + static_assert(kLdsBytes == 68352); +}; +template +using ranges = hkdart::split_many_t>, 4>; +struct Regs +{ + using comp_t = float; + static constexpr uint32_t kNumKvSub = 2; + static constexpr uint32_t k_o_begin = 128, k_v0_begin = 52; + static constexpr uint32_t k_k0_begin = 40, k_k1_begin = 44, k_k2_begin = 48; + using all_ranges = ranges<40, 159>; + using o_ranges = ranges<128, 159>; + using oaccu_t = hk::art; + using score_t = hk::art>; + using p_mfma_a_t = hk::art>; + using p_mfma_b_t = hk::art>; + using pv_v_0_t = hk::art>; + using pv_v_1_t = hk::art>; + using pv_v_2_t = hk::art>; + using pv_v_3_t = hk::art>; +}; + +template +__device__ __forceinline__ void +pv_gemm(Bf16KvManager& kv_manager, const uintptr_t p_lds_v, const float rescale) +{ + using R = Regs; + using comp_t = typename R::comp_t; + constexpr uint32_t k_o_begin = R::k_o_begin; + constexpr uint32_t k_v0_begin = R::k_v0_begin; + constexpr uint32_t k_k0_begin = R::k_k0_begin; + constexpr uint32_t k_k1_begin = R::k_k1_begin; + constexpr uint32_t k_k2_begin = R::k_k2_begin; + // kBlockN=64 (m16x8 double-tile) contracts BOTH 32-row sub-tiles in one call: + // sub-tile A (rows 0:32) uses p_mfma_a, B (rows 32:64) uses p_mfma_b. kBlockN=32 + // has only A. Picked per-iter by row_base (see below), no caller arg. + typename R::p_mfma_a_t p_mfma_a; + typename R::p_mfma_b_t p_mfma_b; + typename R::pv_v_0_t pv_v_0; + typename R::pv_v_1_t pv_v_1; + typename R::pv_v_2_t pv_v_2; + typename R::pv_v_3_t pv_v_3; + + // Each wave owns D128: four D32 iterations per N32 half, eight total. + // Iterations0:4 use keys0:32; iterations4:8 use keys32:64. + constexpr uint32_t kDIters = 128u / (2u * T::kTileM); // 4 + constexpr uint32_t num_pv_iter = R::kNumKvSub * kDIters; // 8 + // Eight V base tiles per N32 half, sixteen total. + constexpr uint32_t kVTilesPerSub = 2u * kDIters; // 8 + constexpr uint32_t kNumVTiles = R::kNumKvSub * kVTilesPerSub; // 16 + + auto pk_mul_pair = [&](float r, auto base_c) { + constexpr uint32_t base = decltype(base_c)::value; + const float2 r2 = {r, r}; + asm volatile("v_pk_mul_f32 v[%0:%1], %2, v[%0:%1]" : : "n"(base), "n"(base + 1), "v"(r2)); + }; + auto mul_pair = [&](float r, auto base_c) { + constexpr uint32_t base = decltype(base_c)::value; + asm volatile("v_mul_f32_e32 v[%0], %1, v[%0]" : : "n"(base), "v"(r)); + asm volatile("v_mul_f32_e32 v[%0], %1, v[%0]" : : "n"(base + 1), "v"(r)); + }; + + // Issue both ds_read_b64 for base tile S_jj (within its sub-tile) into a + // round-robin slot (4 slots). kRowBase selects the KV sub-tile (0=A, 32=B), + // folded into the V ds_read imm by load_transposed_v_to_gpr. + auto load_S = [&]() { + constexpr uint32_t base = (slot == 0u) ? k_v0_begin + : (slot == 1u) ? k_k0_begin + : (slot == 2u) ? k_k1_begin + : k_k2_begin; + kv_manager.template load_transposed_v_to_gpr(p_lds_v); + kv_manager.template load_transposed_v_to_gpr(p_lds_v); + }; + // mfma: oaccu_dst (+)= pv_v_{slot}^T @ p_mfma[sub-tile]. 3-arg init only on the + // FIRST sub-tile of a first-iter call (kRowBase==0); sub-tile B always accumulates + // onto A's oaccu (same oaccu tile, second N-half of P). + auto do_mma = [&](OA& oaccu_dst) { + auto run = [&](auto& v) { + if constexpr(kIsFirstIter && kRowBase == 0u) + hk::mma_ABt(oaccu_dst, v, p_mfma_a); + else if constexpr(kRowBase == 0u) + hk::mma_ABt(oaccu_dst, v, p_mfma_a, oaccu_dst); + else + hk::mma_ABt(oaccu_dst, v, p_mfma_b, oaccu_dst); + }; + if constexpr(slot == 0u) + run(pv_v_0); + else if constexpr(slot == 1u) + run(pv_v_1); + else if constexpr(slot == 2u) + run(pv_v_2); + else + run(pv_v_3); + }; + + if constexpr(kDoRescale) + { + opus::static_for<2>([&](auto s) { + pk_mul_pair(rescale, opus::number{}); + pk_mul_pair(rescale, opus::number{}); + }); + } + + // Prologue: preload the first 3 flat V tiles (S_0,S_1,S_2 of sub-tile A) into + // slots 0,1,2. ONE prologue for the whole call -- sub-tile B's S_0..S_2 are + // prefetched by sub-tile A's tail iters via the flat prefetch index below. + load_S.template operator()<0u, 0u, 0u>(); + load_S.template operator()<0u, 1u, 1u>(); + load_S.template operator()<0u, 2u, 2u>(); + + opus::static_for([&](auto i) { + constexpr uint32_t iter = i.value; + constexpr uint32_t col_idx = iter % kDIters; // 0..3 D-tile within sub-tile + constexpr uint32_t row_base = (iter / kDIters) * 32u; // 0 (A) or 32 (B) + constexpr bool has_next = (iter + 1u) < num_pv_iter; + // Rescale only folds on the FIRST sub-tile pass (row_base==0) and only when a + // NEXT D-tile in that pass exists (col_idx+1 < kDIters) -- else oaccu overflows. + constexpr bool resc_next = kDoRescale && (row_base == 0u) && (col_idx + 1u < kDIters); + constexpr uint32_t next_oaccu_base = k_o_begin + (col_idx + 1u) * 8u; + + // Flat V-tile index across all sub-tiles: prefetch runs on this so sub-tile A's + // tail prefetches sub-tile B's S_0.. (ONE prologue). Consume uses col_idx. + constexpr uint32_t jf_lo = 2u * iter; // flat mfma idx (a) + constexpr uint32_t jf_hi = 2u * iter + 1u; // flat mfma idx (b) + // 4-slot round-robin (flat tile T -> slot T%4). Reading slot j%4 while refilling + // slot (j+3)%4 keeps the refill ds_read's dst != the mfma's src reg. + constexpr uint32_t slot_lo = jf_lo % 4u; + constexpr uint32_t slot_hi = jf_hi % 4u; + + constexpr uint32_t oaccu_base = k_o_begin + col_idx * 8u; + using oaccu_a_r = + hkdart::split_many_t>, + 4>; + using oaccu_b_r = + hkdart::split_many_t>, + 4>; + hk::art oaccu_a; + hk::art oaccu_b; + + // Prefetch S at flat index jf+3 -> its own sub-tile row_base + local jj%8. + auto prefetch = [&]() { + if constexpr(jf < kNumVTiles) + { + constexpr uint32_t pf_row = (jf / kVTilesPerSub) * 32u; + constexpr uint32_t pf_jj = jf % kVTilesPerSub; + load_S.template operator()(); + } + }; + + // ---- mfma_a (flat jf_lo, reads slot_lo) ---- + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(has_next ? 4 : 2, -1)); + prefetch.template operator()(); + do_mma.template operator()(oaccu_a); + if constexpr(resc_next) + { + mul_pair(rescale, opus::number{}); + mul_pair(rescale, opus::number{}); + } + + // ---- mfma_b (flat jf_hi, reads slot_hi) ---- + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(has_next ? 4 : 0, -1)); + prefetch.template operator()(); + do_mma.template operator()(oaccu_b); + if constexpr(resc_next) + { + mul_pair(rescale, opus::number{}); + mul_pair(rescale, opus::number{}); + } + }); +} + +// All waves need the same head16 Q block selected by grid.z. Only one cooperative 16KiB global/LDS +// copy occurs; the four wave-private register copies are populated after the barrier. +__device__ __forceinline__ void load_q(const Params& p, uint16_t* tile) +{ + const uint32_t lane = threadIdx.x & 63u; + const uint32_t wave = __builtin_amdgcn_readfirstlane(threadIdx.x >> 6); + const bool aligned = ((reinterpret_cast(p.q) | uint64_t(p.q_stride0) * 2u | + uint64_t(p.q_stride1) * 2u) & + 15u) == 0; + if(aligned) + { + const uint32_t row = lane >> 2, quad = lane & 3u; + const uint64_t offset = uint64_t(blockIdx.x) * uint64_t(p.q_stride0) + + uint64_t(blockIdx.z * 16u + row) * uint64_t(p.q_stride1); + const uint32_t base = __builtin_amdgcn_readfirstlane( + static_cast(reinterpret_cast(tile))); + for(uint32_t c = wave; c < 16u; c += 4u) + { + const uint32_t L = c * 32u + ((quad * 8u) ^ (((row >> 2) & 1u) * 16u)); + const uint32_t col = sb8_inv_perm_col_elems(L); + const uint64_t address = reinterpret_cast(p.q + offset + col); + const uint32_t destination = base + c * 1024u; + asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tglobal_load_lds_dwordx4 %1, off" + : + : "s"(destination), "v"(address) + : "m0", "memory"); + } + drain(); + __syncthreads(); + const uint32_t r = lane & 15u, col = (lane >> 4) * 8u; + const uint32_t address = base + r * 64u + ((col * 2u) ^ (((r >> 2) & 1u) * 32u)); + opus::static_for<16>( + [&](auto c) { hkm::ds_read_b128<64u + c.value * 4u>(address, c.value * 1024u); }); + } + else + { + const uint32_t head = lane & 15u; + const uint64_t offset = uint64_t(blockIdx.x) * uint64_t(p.q_stride0) + + uint64_t(blockIdx.z * 16u + head) * uint64_t(p.q_stride1); + opus::static_for<16>([&](auto c) { + const uint32_t col = sb8_inv_perm_col_elems(c.value * 32u + (lane >> 4) * 8u); + opus::static_for<4>([&](auto j) { + const uint32_t pair = uint32_t(p.q[offset + col + j.value * 2u]) | + (uint32_t(p.q[offset + col + j.value * 2u + 1u]) << 16); + write_u32<64u + c.value * 4u + j.value>(pair); + }); + __builtin_amdgcn_sched_barrier(0); + }); + } + drain(); + __syncthreads(); // every replicated Q register read retires before KV overwrite +} + +// Two PR 32-row subtiles, each preserving the exact sb8 + row-bank swizzle. +// This initial N64 geometry intentionally uses a single synchronous KV buffer. +__device__ __forceinline__ void +stage_kv(const Params& p, uint16_t* tile, int32_t* selected, int64_t start, int64_t end) +{ + const uint32_t tid = threadIdx.x, lane = tid & 63u; + if(tid < 64u) + { + int32_t slot = -1; + if(start + tid < end) + slot = p.indices[start + tid]; + selected[tid] = slot; + } + __syncthreads(); + const bool aligned = + ((reinterpret_cast(p.kv) | uint64_t(p.kv_stride0) * 2u) & 15u) == 0; + if(aligned) + { + const uint32_t wave = __builtin_amdgcn_readfirstlane(tid >> 6); + const uint32_t base = __builtin_amdgcn_readfirstlane( + static_cast(reinterpret_cast(tile))); + for(uint32_t block = wave; block < 64u; block += 4u) + { + const uint32_t local = block & 31u, r = lane >> 2, quad = lane & 3u; + const uint32_t row = (block >> 5) * 32u + (local & 1u) * 16u + r; + const uint32_t L = (local >> 1) * 32u + ((quad * 8u) ^ (((r >> 2) & 1u) * 16u)); + const uint32_t col = sb8_inv_perm_col_elems(L); + const int32_t slot = selected[row]; + const uint32_t destination = base + block * 1024u; + if(slot >= 0 && int64_t(slot) < p.slots) + { + const uint64_t offset = uint64_t(slot) * uint64_t(p.kv_stride0) + uint64_t(col); + const uint64_t address = reinterpret_cast(p.kv + offset); + asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tglobal_load_lds_dwordx4 %1, off" + : + : "s"(destination), "v"(address) + : "m0", "memory"); + } + else + { + const v4ui zero = {0, 0, 0, 0}; + *reinterpret_cast(reinterpret_cast(tile) + block * 1024u + + lane * 16u) = zero; + } + } + } + else + { + for(uint32_t linear = tid; linear < 64u * 512u; linear += 256u) + { + const uint32_t row = linear / 512u, col = linear % 512u; + const int32_t slot = selected[row]; + uint16_t value = 0; + if(slot >= 0 && int64_t(slot) < p.slots) + value = p.kv[uint64_t(slot) * uint64_t(p.kv_stride0) + uint64_t(col)]; + const uint32_t L = sb8_perm_col_elems(col), r = row & 15u; + const uint32_t byte = (row / 32u) * 32768u + + (L / 32u * 2u + (row % 32u) / 16u) * 1024u + r * 64u + + ((2u * (L & 31u)) ^ (((r >> 2) & 1u) * 32u)); + tile[byte / 2u] = value; + } + } + drain(); + __syncthreads(); +} + +__device__ __forceinline__ void qk(uintptr_t tile) +{ + // Wave w computes its 16 selected tokens against all 16 query heads. + const uint32_t wave = __builtin_amdgcn_readfirstlane(threadIdx.x >> 6); + const uintptr_t wave_tile = tile + (wave / 2u) * 32768u + (wave % 2u) * 1024u; + Regs::score_t score; + auto issue = [&]() { + constexpr uint32_t base = 40u + (S % 3u) * 4u; + hk::art> k; + Bf16KvManager::template load_k_to_gpr<0u, S * 32u>(k, wave_tile); + }; + auto mma = [&]() { + constexpr uint32_t base = 40u + (S % 3u) * 4u, qb = 64u + S * 4u; + hk::art> k; + hk::art> q; + if constexpr(S == 0) + hk::mma_ABt(score, k, q); + else + hk::mma_ABt(score, k, q, score); + }; + opus::static_for<3>([&](auto s) { issue.template operator()(); }); + opus::static_for<16>([&](auto s) { + constexpr uint32_t remain = 15u - s.value; + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(remain < 2u ? remain : 2u, -1)); + mma.template operator()(); + if constexpr(s.value + 3u < 16u) + issue.template operator()(); + }); + asm volatile("s_nop 7\n\ts_nop 7"); // opaque MFMA score dependency is not compiler-visible + __builtin_amdgcn_sched_barrier(0); +} + +__device__ __forceinline__ void softmax(const Params& p, + const int32_t* selected, + float* exchange, + uint16_t* p_lds, + float& maximum, + float& denominator) +{ + const uint32_t lane = threadIdx.x & 63u, wave = threadIdx.x >> 6; + const uint32_t head = lane & 15u, group = lane >> 4; + const uint32_t first_key = wave * 16u + group * 4u; + const uint32_t index_address = + static_cast(reinterpret_cast(selected)) + first_key * sizeof(int32_t); + hkm::ds_read_b128<40>(index_address, 0); + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, -1)); + float local_max = -INFINITY; + opus::static_for<4>([&](auto i) { + const int32_t slot = read_i32<40u + i.value>(); + float x = (read_f32<56u + i.value>() * p.scale) * 1.4426950408889634074f; + if(slot < 0 || int64_t(slot) >= p.slots) + x = -INFINITY; + write_f32<56u + i.value>(x); + local_max = fmaxf(local_max, x); + }); + local_max = fmaxf(local_max, __shfl_xor(local_max, 16, 64)); + local_max = fmaxf(local_max, __shfl_xor(local_max, 32, 64)); + if(lane < 16u) + exchange[wave * 16u + head] = local_max; + __syncthreads(); + float all_max = -INFINITY; + opus::static_for<4>([&](auto w) { all_max = fmaxf(all_max, exchange[w.value * 16u + head]); }); + const float new_max = fmaxf(maximum, all_max); + const float alpha = denominator > 0.0f ? __builtin_amdgcn_exp2f(maximum - new_max) : 0.0f; + maximum = new_max; + Regs::oaccu_t o; + hk::mul_vgpr(o, o, alpha); + float local_sum = 0.0f; + opus::static_for<4>([&](auto i) { + const float x = read_f32<56u + i.value>(); + const float weight = x == -INFINITY ? 0.0f : __builtin_amdgcn_exp2f(x - maximum); + write_f32<56u + i.value>(weight); + local_sum += weight; + }); + local_sum += __shfl_xor(local_sum, 16, 64); + local_sum += __shfl_xor(local_sum, 32, 64); + if(lane < 16u) + exchange[64u + wave * 16u + head] = local_sum; + // Each wave owns a disjoint N16 slab of BF16 P. Normal [head,key] LDS layout + // makes the subsequent full-N64 operand copies explicit and bit-preserving. + v2ui packed; + const float p0 = read_f32<56>(), p1 = read_f32<57>(); + const float p2 = read_f32<58>(), p3 = read_f32<59>(); + asm volatile("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(packed[0]) : "v"(p0), "v"(p1)); + asm volatile("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(packed[1]) : "v"(p2), "v"(p3)); + const uintptr_t address = reinterpret_cast(p_lds) + (head * 64u + first_key) * 2u; + hkm::ds_write_b64(packed, static_cast(address), 0); + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, -1)); + __syncthreads(); // all P and four partial denominators are visible + float all_sum = 0.0f; + opus::static_for<4>([&](auto w) { all_sum += exchange[64u + w.value * 16u + head]; }); + denominator = denominator * alpha + all_sum; + const uint32_t p_address = + static_cast(reinterpret_cast(p_lds)) + (head * 64u + group * 4u) * 2u; + // PR's transposed-V operand orders the two N16 halves as group*4 and +16. + hkm::ds_read_b64<56>(p_address, 0); + hkm::ds_read_b64<58>(p_address, 32); + hkm::ds_read_b64<60>(p_address, 64); + hkm::ds_read_b64<62>(p_address, 96); + drain(); + __syncthreads(); // denominator/max consumers retire before next iteration reuses exchange +} + +template +__device__ __forceinline__ void run(Params p) +{ + hkdart::clobber(); + extern __shared__ __align__(16) unsigned char storage[]; + auto* tile = reinterpret_cast(storage); + auto* p_lds = reinterpret_cast(storage + Traits::kPOffset); + auto* exchange = reinterpret_cast(storage + Traits::kReduceOffset); + auto* selected = reinterpret_cast(storage + Traits::kIndexOffset); + const uint32_t query = blockIdx.x, partition = blockIdx.y; + const uint32_t head_base = blockIdx.z * 16u; + static_assert(TotalHeads == 16 || TotalHeads == 64); + const int32_t first = p.indptr[query], last = p.indptr[query + 1u], length = last - first; + const int32_t begin = first + int64_t(length) * partition / p.splits; + const int32_t end = first + int64_t(length) * (partition + 1u) / p.splits; + load_q(p, tile); + Regs::oaccu_t o; + hk::zero(o); + float maximum = -INFINITY, denominator = 0.0f; + Bf16KvManager manager; + const uint32_t wave = __builtin_amdgcn_readfirstlane(threadIdx.x >> 6); + for(int64_t pos = begin; pos < end; pos += 64) + { + stage_kv(p, tile, selected, pos, end); + qk(reinterpret_cast(tile)); + softmax(p, selected, exchange, p_lds, maximum, denominator); + // D128/wave is a multiple of the sb8 permutation's 64-element period. + const uintptr_t v_tile = reinterpret_cast(tile) + wave * 8192u; + sparse_mla_bf16::headtiles::pv_gemm(manager, v_tile, 1.0f); + asm volatile("s_nop 7\n\ts_nop 7"); + drain(); + __syncthreads(); + } + const float inverse = denominator > 0.0f ? 1.0f / denominator : 0.0f; + hk::mul_vgpr(o, o, inverse); + const float lse = denominator > 0.0f + ? (maximum + __builtin_amdgcn_logf(denominator)) * 0.6931471805599453094f + : -INFINITY; + const uint32_t lane = threadIdx.x & 63u; + const uint64_t output_slot = Split ? uint64_t(query) * uint64_t(p.splits) + partition : query; + if(wave == 0u && lane < 16u) + { + if constexpr(Split) + p.partial_lse[output_slot * TotalHeads + head_base + lane] = lse; + else if(p.lse) + p.lse[output_slot * TotalHeads + head_base + lane] = lse; + } + // PR O manager preserves global head stride512. Offset its descriptor to + // this wave's D128, set head-wave id to zero, and give each wave a private + // dead-KV LDS bounce so its data cannot race another output wave. + const uintptr_t bounce = reinterpret_cast(tile) + wave * 4352u; + // grid.z introduces another head-block base. The PR buffer descriptor must + // be an SGPR tuple, so state its per-wave uniformity explicitly in64bits. + auto uniform_pointer = [](uint64_t address) { + const uint32_t lo = __builtin_amdgcn_readfirstlane(uint32_t(address)); + const uint32_t hi = __builtin_amdgcn_readfirstlane(uint32_t(address >> 32)); + return uint64_t(lo) | (uint64_t(hi) << 32); + }; + if constexpr(Split) + { + OManager32bitsV4Gen1Swizzle output; + float* dst = reinterpret_cast(uniform_pointer(reinterpret_cast( + p.partial_o + (output_slot * TotalHeads + head_base) * 512u + wave * 128u))); + opus::static_for<2>([&](auto i) { + output.template output_to_vram_pair<128u + i.value * 16u, i.value * 64u, false>( + dst, 0, 0, 0, bounce, TotalHeads); + drain(); + }); + } + else + { + OManager16bitsV4Gen1Swizzle output; + auto* dst = reinterpret_cast(uniform_pointer(reinterpret_cast( + p.out + (output_slot * TotalHeads + head_base) * 512u + wave * 128u))); + opus::static_for<2>([&](auto i) { + output.template output_to_vram_pair<128u + i.value * 16u, i.value * 64u, false>( + dst, 0, 0, 0, bounce, TotalHeads); + drain(); + }); + } +} +} // namespace sparse_mla_bf16::headtiles + +namespace sparse_mla_bf16 { +template +__global__ __launch_bounds__(256) + __attribute__((amdgpu_num_vgpr(40))) void sparse_mla_headtiles(Params p) +{ + static_assert(Heads == 64); + headtiles::run(p); +} +} // namespace sparse_mla_bf16 diff --git a/csrc/kernels/mla/hk/sparse_mla_bf16_reduce.cuh b/csrc/kernels/mla/hk/sparse_mla_bf16_reduce.cuh new file mode 100644 index 00000000000..ad1b3395c80 --- /dev/null +++ b/csrc/kernels/mla/hk/sparse_mla_bf16_reduce.cuh @@ -0,0 +1,107 @@ +// SPDX-License-Identifier: MIT +// Fully unrolled, one-wave split-K reducer shared by the v1/v2 implementations. +#pragma once +#include + +namespace sparse_mla_bf16 { + +template +__global__ __launch_bounds__(64) void reduce(Params p) +{ + static_assert(Heads == 16 || Heads == 64); + static_assert(Splits == 2 || Splits == 4 || Splits == 8 || Splits == 16 || Splits == 32); + const uint32_t query = blockIdx.x, head = blockIdx.y; + const uint32_t d = threadIdx.x * 8u; + const uint64_t first = uint64_t(query) * Splits * Heads + head; + float maximum = -INFINITY; +#pragma unroll + for(int s = 0; s < Splits; ++s) + maximum = fmaxf(maximum, p.partial_lse[first + uint64_t(s) * Heads]); + float sum = 0.0f; + float accumulator[8] = {}; + // Every partial row is 2048 B; caller-provided workspace may be unaligned. + // Retain a scalar-load fallback for callers of the raw C ABI. + const bool aligned_partial = (reinterpret_cast(p.partial_o) & 15u) == 0; +#pragma unroll + for(int s = 0; s < Splits; ++s) + { + const uint64_t row = first + uint64_t(s) * Heads; + const float value = p.partial_lse[row]; + // Preserve the baseline summation order and empty-partition semantics. + const float weight = value == -INFINITY ? 0.0f : expf(value - maximum); + sum += weight; + const float* source = p.partial_o + row * 512u + d; + float values[8]; + if(aligned_partial) + { + const float4 lo = *reinterpret_cast(source); + const float4 hi = *reinterpret_cast(source + 4); + values[0] = lo.x; + values[1] = lo.y; + values[2] = lo.z; + values[3] = lo.w; + values[4] = hi.x; + values[5] = hi.y; + values[6] = hi.z; + values[7] = hi.w; + } + else + { +#pragma unroll + for(int i = 0; i < 8; ++i) + values[i] = source[i]; + } +#pragma unroll + for(int i = 0; i < 8; ++i) + accumulator[i] += weight * values[i]; + } + const float reciprocal = sum > 0.0f ? 1.0f / sum : 0.0f; + uint32_t pairs[4]; +#pragma unroll + for(int i = 0; i < 4; ++i) + { + const float lo = accumulator[2 * i] * reciprocal; + const float hi = accumulator[2 * i + 1] * reciprocal; + asm volatile("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(pairs[i]) : "v"(lo), "v"(hi)); + } + uint16_t* destination = p.out + (uint64_t(query) * Heads + head) * 512u + d; + if((reinterpret_cast(destination) & 15u) == 0) + { + const uint4 packed{pairs[0], pairs[1], pairs[2], pairs[3]}; + *reinterpret_cast(destination) = packed; + } + else + { + // The API permits an output storage offset of one BF16 (2 B alignment). +#pragma unroll + for(int i = 0; i < 4; ++i) + { + destination[2 * i] = uint16_t(pairs[i]); + destination[2 * i + 1] = uint16_t(pairs[i] >> 16); + } + } + if(threadIdx.x == 0 && p.lse) + p.lse[uint64_t(query) * Heads + head] = sum > 0.0f ? maximum + logf(sum) : -INFINITY; +} + +template +inline hipError_t launch_reduce64(Params p, hipStream_t stream) +{ +#define SPARSE_MLA_REDUCE_CASE(S) \ + case S: \ + hipLaunchKernelGGL((reduce), dim3(p.queries, Heads), dim3(64), 0, stream, p); \ + break + switch(p.splits) + { + SPARSE_MLA_REDUCE_CASE(2); + SPARSE_MLA_REDUCE_CASE(4); + SPARSE_MLA_REDUCE_CASE(8); + SPARSE_MLA_REDUCE_CASE(16); + SPARSE_MLA_REDUCE_CASE(32); + default: return hipErrorInvalidValue; + } +#undef SPARSE_MLA_REDUCE_CASE + return hipGetLastError(); +} + +} // namespace sparse_mla_bf16 diff --git a/csrc/kernels/mla/hk/sparse_mla_bf16_wave4.cuh b/csrc/kernels/mla/hk/sparse_mla_bf16_wave4.cuh new file mode 100644 index 00000000000..fae437157e7 --- /dev/null +++ b/csrc/kernels/mla/hk/sparse_mla_bf16_wave4.cuh @@ -0,0 +1,493 @@ +// SPDX-License-Identifier: MIT +// Experimental H16 four-wave geometry: QK partitions N64, PV partitions D512. +// Reuses PR #3459 HipKittens BF16 LDS readers, four-slot PV schedule, O managers. +#pragma once +#include "sparse_mla_bf16_common.cuh" + +namespace sparse_mla_bf16::wave4 { +struct Traits +{ + static constexpr uint32_t kBlockN = 64, kBlockK = 32, kTileM = 16; + static constexpr uint32_t kQkHeadDim = 512, kVoHeadDim = 512; + static constexpr uint32_t kNumWarps = 4, kNumThreads = 256, kRoundMode = 0; + static constexpr uint32_t kKvBytes = 65536, kPBytes = 2048; + static constexpr uint32_t kPOffset = kKvBytes; + static constexpr uint32_t kReduceOffset = kPOffset + kPBytes; + static constexpr uint32_t kIndexOffset = kReduceOffset + 512; + static constexpr uint32_t kLdsBytes = kIndexOffset + 256; + static_assert(kLdsBytes == 68352); +}; +template +using ranges = hkdart::split_many_t>, 4>; +struct Regs +{ + using comp_t = float; + static constexpr uint32_t kNumKvSub = 2; + static constexpr uint32_t k_o_begin = 128, k_v0_begin = 52; + static constexpr uint32_t k_k0_begin = 40, k_k1_begin = 44, k_k2_begin = 48; + using all_ranges = ranges<40, 159>; + using o_ranges = ranges<128, 159>; + using oaccu_t = hk::art; + using score_t = hk::art>; + using p_mfma_a_t = hk::art>; + using p_mfma_b_t = hk::art>; + using pv_v_0_t = hk::art>; + using pv_v_1_t = hk::art>; + using pv_v_2_t = hk::art>; + using pv_v_3_t = hk::art>; +}; + +template +__device__ __forceinline__ void +pv_gemm(Bf16KvManager& kv_manager, const uintptr_t p_lds_v, const float rescale) +{ + using R = Regs; + using comp_t = typename R::comp_t; + constexpr uint32_t k_o_begin = R::k_o_begin; + constexpr uint32_t k_v0_begin = R::k_v0_begin; + constexpr uint32_t k_k0_begin = R::k_k0_begin; + constexpr uint32_t k_k1_begin = R::k_k1_begin; + constexpr uint32_t k_k2_begin = R::k_k2_begin; + // kBlockN=64 (m16x8 double-tile) contracts BOTH 32-row sub-tiles in one call: + // sub-tile A (rows 0:32) uses p_mfma_a, B (rows 32:64) uses p_mfma_b. kBlockN=32 + // has only A. Picked per-iter by row_base (see below), no caller arg. + typename R::p_mfma_a_t p_mfma_a; + typename R::p_mfma_b_t p_mfma_b; + typename R::pv_v_0_t pv_v_0; + typename R::pv_v_1_t pv_v_1; + typename R::pv_v_2_t pv_v_2; + typename R::pv_v_3_t pv_v_3; + + // D-tiling: each D-iter emits 2 oaccu 16x16 sub-tiles (32 D-cols); 512/32 = 16 + // D-iters per 32-row KV sub-tile. kBlockN=64 runs BOTH sub-tiles in one call, so + // num_pv_iter = kNumKvSub * 16. iter [0,16) = sub-tile A, [16,32) = sub-tile B. + constexpr uint32_t kDIters = 128u / (2u * T::kTileM); // 16 + constexpr uint32_t num_pv_iter = R::kNumKvSub * kDIters; // 16 or 32 + // Flat V base tiles across all sub-tiles: 32 per sub-tile (S_0..S_31 each). + constexpr uint32_t kVTilesPerSub = 2u * kDIters; // 32 + constexpr uint32_t kNumVTiles = R::kNumKvSub * kVTilesPerSub; // 32 or 64 + + auto pk_mul_pair = [&](float r, auto base_c) { + constexpr uint32_t base = decltype(base_c)::value; + const float2 r2 = {r, r}; + asm volatile("v_pk_mul_f32 v[%0:%1], %2, v[%0:%1]" : : "n"(base), "n"(base + 1), "v"(r2)); + }; + auto mul_pair = [&](float r, auto base_c) { + constexpr uint32_t base = decltype(base_c)::value; + asm volatile("v_mul_f32_e32 v[%0], %1, v[%0]" : : "n"(base), "v"(r)); + asm volatile("v_mul_f32_e32 v[%0], %1, v[%0]" : : "n"(base + 1), "v"(r)); + }; + + // Issue both ds_read_b64 for base tile S_jj (within its sub-tile) into a + // round-robin slot (4 slots). kRowBase selects the KV sub-tile (0=A, 32=B), + // folded into the V ds_read imm by load_transposed_v_to_gpr. + auto load_S = [&]() { + constexpr uint32_t base = (slot == 0u) ? k_v0_begin + : (slot == 1u) ? k_k0_begin + : (slot == 2u) ? k_k1_begin + : k_k2_begin; + kv_manager.template load_transposed_v_to_gpr(p_lds_v); + kv_manager.template load_transposed_v_to_gpr(p_lds_v); + }; + // mfma: oaccu_dst (+)= pv_v_{slot}^T @ p_mfma[sub-tile]. 3-arg init only on the + // FIRST sub-tile of a first-iter call (kRowBase==0); sub-tile B always accumulates + // onto A's oaccu (same oaccu tile, second N-half of P). + auto do_mma = [&](OA& oaccu_dst) { + auto run = [&](auto& v) { + if constexpr(kIsFirstIter && kRowBase == 0u) + hk::mma_ABt(oaccu_dst, v, p_mfma_a); + else if constexpr(kRowBase == 0u) + hk::mma_ABt(oaccu_dst, v, p_mfma_a, oaccu_dst); + else + hk::mma_ABt(oaccu_dst, v, p_mfma_b, oaccu_dst); + }; + if constexpr(slot == 0u) + run(pv_v_0); + else if constexpr(slot == 1u) + run(pv_v_1); + else if constexpr(slot == 2u) + run(pv_v_2); + else + run(pv_v_3); + }; + + if constexpr(kDoRescale) + { + opus::static_for<2>([&](auto s) { + pk_mul_pair(rescale, opus::number{}); + pk_mul_pair(rescale, opus::number{}); + }); + } + + // Prologue: preload the first 3 flat V tiles (S_0,S_1,S_2 of sub-tile A) into + // slots 0,1,2. ONE prologue for the whole call -- sub-tile B's S_0..S_2 are + // prefetched by sub-tile A's tail iters via the flat prefetch index below. + load_S.template operator()<0u, 0u, 0u>(); + load_S.template operator()<0u, 1u, 1u>(); + load_S.template operator()<0u, 2u, 2u>(); + + opus::static_for([&](auto i) { + constexpr uint32_t iter = i.value; + constexpr uint32_t col_idx = iter % kDIters; // 0..15 D-tile within sub-tile + constexpr uint32_t row_base = (iter / kDIters) * 32u; // 0 (A) or 32 (B) + constexpr bool has_next = (iter + 1u) < num_pv_iter; + // Rescale only folds on the FIRST sub-tile pass (row_base==0) and only when a + // NEXT D-tile in that pass exists (col_idx+1 < kDIters) -- else oaccu overflows. + constexpr bool resc_next = kDoRescale && (row_base == 0u) && (col_idx + 1u < kDIters); + constexpr uint32_t next_oaccu_base = k_o_begin + (col_idx + 1u) * 8u; + + // Flat V-tile index across all sub-tiles: prefetch runs on this so sub-tile A's + // tail prefetches sub-tile B's S_0.. (ONE prologue). Consume uses col_idx. + constexpr uint32_t jf_lo = 2u * iter; // flat mfma idx (a) + constexpr uint32_t jf_hi = 2u * iter + 1u; // flat mfma idx (b) + // 4-slot round-robin (flat tile T -> slot T%4). Reading slot j%4 while refilling + // slot (j+3)%4 keeps the refill ds_read's dst != the mfma's src reg. + constexpr uint32_t slot_lo = jf_lo % 4u; + constexpr uint32_t slot_hi = jf_hi % 4u; + + constexpr uint32_t oaccu_base = k_o_begin + col_idx * 8u; + using oaccu_a_r = + hkdart::split_many_t>, + 4>; + using oaccu_b_r = + hkdart::split_many_t>, + 4>; + hk::art oaccu_a; + hk::art oaccu_b; + + // Prefetch S at flat index jf+3 -> its own sub-tile row_base + local jj%32. + auto prefetch = [&]() { + if constexpr(jf < kNumVTiles) + { + constexpr uint32_t pf_row = (jf / kVTilesPerSub) * 32u; + constexpr uint32_t pf_jj = jf % kVTilesPerSub; + load_S.template operator()(); + } + }; + + // ---- mfma_a (flat jf_lo, reads slot_lo) ---- + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(has_next ? 4 : 2, -1)); + prefetch.template operator()(); + do_mma.template operator()(oaccu_a); + if constexpr(resc_next) + { + mul_pair(rescale, opus::number{}); + mul_pair(rescale, opus::number{}); + } + + // ---- mfma_b (flat jf_hi, reads slot_hi) ---- + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(has_next ? 4 : 0, -1)); + prefetch.template operator()(); + do_mma.template operator()(oaccu_b); + if constexpr(resc_next) + { + mul_pair(rescale, opus::number{}); + mul_pair(rescale, opus::number{}); + } + }); +} + +// All waves need the same H16 Q. Only one cooperative 16KiB global/LDS copy +// occurs; the four wave-private register copies are populated after the barrier. +__device__ __forceinline__ void load_q(const Params& p, uint16_t* tile) +{ + const uint32_t lane = threadIdx.x & 63u; + const uint32_t wave = __builtin_amdgcn_readfirstlane(threadIdx.x >> 6); + const bool aligned = ((reinterpret_cast(p.q) | uint64_t(p.q_stride0) * 2u | + uint64_t(p.q_stride1) * 2u) & + 15u) == 0; + if(aligned) + { + const uint32_t row = lane >> 2, quad = lane & 3u; + const uint64_t offset = + uint64_t(blockIdx.x) * uint64_t(p.q_stride0) + uint64_t(row) * uint64_t(p.q_stride1); + const uint32_t base = __builtin_amdgcn_readfirstlane( + static_cast(reinterpret_cast(tile))); + for(uint32_t c = wave; c < 16u; c += 4u) + { + const uint32_t L = c * 32u + ((quad * 8u) ^ (((row >> 2) & 1u) * 16u)); + const uint32_t col = sb8_inv_perm_col_elems(L); + const uint64_t address = reinterpret_cast(p.q + offset + col); + const uint32_t destination = base + c * 1024u; + asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tglobal_load_lds_dwordx4 %1, off" + : + : "s"(destination), "v"(address) + : "m0", "memory"); + } + drain(); + __syncthreads(); + const uint32_t r = lane & 15u, col = (lane >> 4) * 8u; + const uint32_t address = base + r * 64u + ((col * 2u) ^ (((r >> 2) & 1u) * 32u)); + opus::static_for<16>( + [&](auto c) { hkm::ds_read_b128<64u + c.value * 4u>(address, c.value * 1024u); }); + } + else + { + const uint32_t head = lane & 15u; + const uint64_t offset = + uint64_t(blockIdx.x) * uint64_t(p.q_stride0) + uint64_t(head) * uint64_t(p.q_stride1); + opus::static_for<16>([&](auto c) { + const uint32_t col = sb8_inv_perm_col_elems(c.value * 32u + (lane >> 4) * 8u); + opus::static_for<4>([&](auto j) { + const uint32_t pair = uint32_t(p.q[offset + col + j.value * 2u]) | + (uint32_t(p.q[offset + col + j.value * 2u + 1u]) << 16); + write_u32<64u + c.value * 4u + j.value>(pair); + }); + __builtin_amdgcn_sched_barrier(0); + }); + } + drain(); + __syncthreads(); // every replicated Q register read retires before KV overwrite +} + +// Two PR 32-row subtiles, each preserving the exact sb8 + row-bank swizzle. +// This initial N64 geometry intentionally uses a single synchronous KV buffer. +__device__ __forceinline__ void +stage_kv(const Params& p, uint16_t* tile, int32_t* selected, int64_t start, int64_t end) +{ + const uint32_t tid = threadIdx.x, lane = tid & 63u; + if(tid < 64u) + { + int32_t slot = -1; + if(start + tid < end) + slot = p.indices[start + tid]; + selected[tid] = slot; + } + __syncthreads(); + const bool aligned = + ((reinterpret_cast(p.kv) | uint64_t(p.kv_stride0) * 2u) & 15u) == 0; + if(aligned) + { + const uint32_t wave = __builtin_amdgcn_readfirstlane(tid >> 6); + const uint32_t base = __builtin_amdgcn_readfirstlane( + static_cast(reinterpret_cast(tile))); + for(uint32_t block = wave; block < 64u; block += 4u) + { + const uint32_t local = block & 31u, r = lane >> 2, quad = lane & 3u; + const uint32_t row = (block >> 5) * 32u + (local & 1u) * 16u + r; + const uint32_t L = (local >> 1) * 32u + ((quad * 8u) ^ (((r >> 2) & 1u) * 16u)); + const uint32_t col = sb8_inv_perm_col_elems(L); + const int32_t slot = selected[row]; + const uint32_t destination = base + block * 1024u; + if(slot >= 0 && int64_t(slot) < p.slots) + { + const uint64_t offset = uint64_t(slot) * uint64_t(p.kv_stride0) + uint64_t(col); + const uint64_t address = reinterpret_cast(p.kv + offset); + asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tglobal_load_lds_dwordx4 %1, off" + : + : "s"(destination), "v"(address) + : "m0", "memory"); + } + else + { + const v4ui zero = {0, 0, 0, 0}; + *reinterpret_cast(reinterpret_cast(tile) + block * 1024u + + lane * 16u) = zero; + } + } + } + else + { + for(uint32_t linear = tid; linear < 64u * 512u; linear += 256u) + { + const uint32_t row = linear / 512u, col = linear % 512u; + const int32_t slot = selected[row]; + uint16_t value = 0; + if(slot >= 0 && int64_t(slot) < p.slots) + value = p.kv[uint64_t(slot) * uint64_t(p.kv_stride0) + uint64_t(col)]; + const uint32_t L = sb8_perm_col_elems(col), r = row & 15u; + const uint32_t byte = (row / 32u) * 32768u + + (L / 32u * 2u + (row % 32u) / 16u) * 1024u + r * 64u + + ((2u * (L & 31u)) ^ (((r >> 2) & 1u) * 32u)); + tile[byte / 2u] = value; + } + } + drain(); + __syncthreads(); +} + +__device__ __forceinline__ void qk(uintptr_t tile) +{ + // Wave w computes its 16 selected tokens against all 16 query heads. + const uint32_t wave = __builtin_amdgcn_readfirstlane(threadIdx.x >> 6); + const uintptr_t wave_tile = tile + (wave / 2u) * 32768u + (wave % 2u) * 1024u; + Regs::score_t score; + auto issue = [&]() { + constexpr uint32_t base = 40u + (S % 3u) * 4u; + hk::art> k; + Bf16KvManager::template load_k_to_gpr<0u, S * 32u>(k, wave_tile); + }; + auto mma = [&]() { + constexpr uint32_t base = 40u + (S % 3u) * 4u, qb = 64u + S * 4u; + hk::art> k; + hk::art> q; + if constexpr(S == 0) + hk::mma_ABt(score, k, q); + else + hk::mma_ABt(score, k, q, score); + }; + opus::static_for<3>([&](auto s) { issue.template operator()(); }); + opus::static_for<16>([&](auto s) { + constexpr uint32_t remain = 15u - s.value; + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(remain < 2u ? remain : 2u, -1)); + mma.template operator()(); + if constexpr(s.value + 3u < 16u) + issue.template operator()(); + }); + asm volatile("s_nop 7\n\ts_nop 7"); // opaque MFMA score dependency is not compiler-visible + __builtin_amdgcn_sched_barrier(0); +} + +__device__ __forceinline__ void softmax(const Params& p, + const int32_t* selected, + float* exchange, + uint16_t* p_lds, + float& maximum, + float& denominator) +{ + const uint32_t lane = threadIdx.x & 63u, wave = threadIdx.x >> 6; + const uint32_t head = lane & 15u, group = lane >> 4; + const uint32_t first_key = wave * 16u + group * 4u; + const uint32_t index_address = + static_cast(reinterpret_cast(selected)) + first_key * sizeof(int32_t); + hkm::ds_read_b128<40>(index_address, 0); + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, -1)); + float local_max = -INFINITY; + opus::static_for<4>([&](auto i) { + const int32_t slot = read_i32<40u + i.value>(); + float x = (read_f32<56u + i.value>() * p.scale) * 1.4426950408889634074f; + if(slot < 0 || int64_t(slot) >= p.slots) + x = -INFINITY; + write_f32<56u + i.value>(x); + local_max = fmaxf(local_max, x); + }); + local_max = fmaxf(local_max, __shfl_xor(local_max, 16, 64)); + local_max = fmaxf(local_max, __shfl_xor(local_max, 32, 64)); + if(lane < 16u) + exchange[wave * 16u + head] = local_max; + __syncthreads(); + float all_max = -INFINITY; + opus::static_for<4>([&](auto w) { all_max = fmaxf(all_max, exchange[w.value * 16u + head]); }); + const float new_max = fmaxf(maximum, all_max); + const float alpha = denominator > 0.0f ? __builtin_amdgcn_exp2f(maximum - new_max) : 0.0f; + maximum = new_max; + Regs::oaccu_t o; + hk::mul_vgpr(o, o, alpha); + float local_sum = 0.0f; + opus::static_for<4>([&](auto i) { + const float x = read_f32<56u + i.value>(); + const float weight = x == -INFINITY ? 0.0f : __builtin_amdgcn_exp2f(x - maximum); + write_f32<56u + i.value>(weight); + local_sum += weight; + }); + local_sum += __shfl_xor(local_sum, 16, 64); + local_sum += __shfl_xor(local_sum, 32, 64); + if(lane < 16u) + exchange[64u + wave * 16u + head] = local_sum; + // Each wave owns a disjoint N16 slab of BF16 P. Normal [head,key] LDS layout + // makes the subsequent full-N64 operand copies explicit and bit-preserving. + v2ui packed; + const float p0 = read_f32<56>(), p1 = read_f32<57>(); + const float p2 = read_f32<58>(), p3 = read_f32<59>(); + asm volatile("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(packed[0]) : "v"(p0), "v"(p1)); + asm volatile("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(packed[1]) : "v"(p2), "v"(p3)); + const uintptr_t address = reinterpret_cast(p_lds) + (head * 64u + first_key) * 2u; + hkm::ds_write_b64(packed, static_cast(address), 0); + __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, -1)); + __syncthreads(); // all P and four partial denominators are visible + float all_sum = 0.0f; + opus::static_for<4>([&](auto w) { all_sum += exchange[64u + w.value * 16u + head]; }); + denominator = denominator * alpha + all_sum; + const uint32_t p_address = + static_cast(reinterpret_cast(p_lds)) + (head * 64u + group * 4u) * 2u; + // PR's transposed-V operand orders the two N16 halves as group*4 and +16. + hkm::ds_read_b64<56>(p_address, 0); + hkm::ds_read_b64<58>(p_address, 32); + hkm::ds_read_b64<60>(p_address, 64); + hkm::ds_read_b64<62>(p_address, 96); + drain(); + __syncthreads(); // denominator/max consumers retire before next iteration reuses exchange +} + +template +__device__ __forceinline__ void run(Params p) +{ + hkdart::clobber(); + extern __shared__ __align__(16) unsigned char storage[]; + auto* tile = reinterpret_cast(storage); + auto* p_lds = reinterpret_cast(storage + Traits::kPOffset); + auto* exchange = reinterpret_cast(storage + Traits::kReduceOffset); + auto* selected = reinterpret_cast(storage + Traits::kIndexOffset); + const uint32_t query = blockIdx.x, partition = blockIdx.y; + const int32_t first = p.indptr[query], last = p.indptr[query + 1u], length = last - first; + const int32_t begin = first + int64_t(length) * partition / p.splits; + const int32_t end = first + int64_t(length) * (partition + 1u) / p.splits; + load_q(p, tile); + Regs::oaccu_t o; + hk::zero(o); + float maximum = -INFINITY, denominator = 0.0f; + Bf16KvManager manager; + const uint32_t wave = __builtin_amdgcn_readfirstlane(threadIdx.x >> 6); + for(int64_t pos = begin; pos < end; pos += 64) + { + stage_kv(p, tile, selected, pos, end); + qk(reinterpret_cast(tile)); + softmax(p, selected, exchange, p_lds, maximum, denominator); + // D128/wave is a multiple of the sb8 permutation's 64-element period. + const uintptr_t v_tile = reinterpret_cast(tile) + wave * 8192u; + sparse_mla_bf16::wave4::pv_gemm(manager, v_tile, 1.0f); + asm volatile("s_nop 7\n\ts_nop 7"); + drain(); + __syncthreads(); + } + const float inverse = denominator > 0.0f ? 1.0f / denominator : 0.0f; + hk::mul_vgpr(o, o, inverse); + const float lse = denominator > 0.0f + ? (maximum + __builtin_amdgcn_logf(denominator)) * 0.6931471805599453094f + : -INFINITY; + const uint32_t lane = threadIdx.x & 63u; + const uint64_t output_slot = Split ? uint64_t(query) * uint64_t(p.splits) + partition : query; + if(wave == 0u && lane < 16u) + { + if constexpr(Split) + p.partial_lse[output_slot * 16u + lane] = lse; + else if(p.lse) + p.lse[output_slot * 16u + lane] = lse; + } + // PR O manager preserves global head stride512. Offset its descriptor to + // this wave's D128, set head-wave id to zero, and give each wave a private + // dead-KV LDS bounce so its data cannot race another output wave. + const uintptr_t bounce = reinterpret_cast(tile) + wave * 4352u; + if constexpr(Split) + { + OManager32bitsV4Gen1Swizzle output; + float* dst = p.partial_o + output_slot * 16u * 512u + wave * 128u; + opus::static_for<2>([&](auto i) { + output.template output_to_vram_pair<128u + i.value * 16u, i.value * 64u, false>( + dst, 0, 0, 0, bounce, 16); + drain(); + }); + } + else + { + OManager16bitsV4Gen1Swizzle output; + auto* dst = reinterpret_cast(p.out + output_slot * 16u * 512u + wave * 128u); + opus::static_for<2>([&](auto i) { + output.template output_to_vram_pair<128u + i.value * 16u, i.value * 64u, false>( + dst, 0, 0, 0, bounce, 16); + drain(); + }); + } +} +} // namespace sparse_mla_bf16::wave4 + +namespace sparse_mla_bf16 { +template +__global__ __launch_bounds__(256) __attribute__((amdgpu_num_vgpr(40))) void sparse_mla_v2(Params p) +{ + if constexpr(Heads == 16) + wave4::run(p); + else + sparse_mla_legacy(p); +} +} // namespace sparse_mla_bf16 diff --git a/csrc/kernels/mla/sparse_mla_bf16.cu b/csrc/kernels/mla/sparse_mla_bf16.cu new file mode 100644 index 00000000000..ac2050105ea --- /dev/null +++ b/csrc/kernels/mla/sparse_mla_bf16.cu @@ -0,0 +1,218 @@ +// SPDX-License-Identifier: MIT +// BF16, D512 sparse MLA for gfx950. v1: full LDS; v2: hybrid + unrolled reducer. +// aiter_ctypes_error.h requires the common declarations supplied by aiter_tensor.h. +// clang-format off +#include "aiter_tensor.h" +#include "aiter_ctypes_error.h" +#include "hk/sparse_mla_bf16_wave4.cuh" +#include "hk/sparse_mla_bf16_headtiles.cuh" +#include "hk/sparse_mla_bf16_reduce.cuh" +// clang-format on +#include +#include + +AITER_CTYPES_ERROR_DEF + +namespace sparse_mla_bf16 { +#include "hk/sparse_mla_bf16_full_lds.cuh" + +template +__global__ __launch_bounds__(Heads == 64 ? 512 : 64) + __attribute__((amdgpu_num_vgpr(40))) void sparse_mla_v1(Params p) +{ + if constexpr(Heads == 64) + full_lds_h64(p); + else + sparse_mla_legacy(p); +} + +template +hipError_t launch_main(Params p, int version, hipStream_t stream) +{ + const dim3 grid(p.queries, p.splits); + if(version == 1) + { + constexpr int threads = Heads == 64 ? 512 : 64; + constexpr int lds = Heads == 64 ? FullLdsTraits::kLdsBytes : Traits<16>::kLdsBytes; + hipLaunchKernelGGL((sparse_mla_v1), grid, dim3(threads), lds, stream, p); + } + else if(Heads == 64 && uint64_t(p.queries) * p.splits <= 128u) + { + hipLaunchKernelGGL((sparse_mla_headtiles<64, Split>), + dim3(p.queries, p.splits, 4), + dim3(256), + headtiles::Traits::kLdsBytes, + stream, + p); + } + else + { + constexpr int lds = Heads == 16 ? wave4::Traits::kLdsBytes : Traits<64>::kLdsBytes; + hipLaunchKernelGGL((sparse_mla_v2), grid, dim3(256), lds, stream, p); + } + return hipGetLastError(); +} + +template +hipError_t launch(Params p, int version, hipStream_t stream) +{ + if(p.splits == 1) + return launch_main(p, version, stream); + hipError_t status = launch_main(p, version, stream); + if(status != hipSuccess) + return status; + if constexpr(Heads == 16) + { + if(version == 1) + { + hipLaunchKernelGGL(reduce, dim3(p.queries, Heads), dim3(256), 0, stream, p, Heads); + return hipGetLastError(); + } + } + return launch_reduce64(p, stream); +} +} // namespace sparse_mla_bf16 + +namespace { +void check_tensor(const aiter_tensor_t* t, const char* name, AiterDtype dtype, int device) +{ + AITER_CHECK(t && t->is_gpu() && t->device_id == device, name, ": expected the query device"); + AITER_CHECK(t->dtype() == dtype, name, ": unexpected dtype"); + AITER_CHECK(t->numel() == 0 || t->data_ptr(), name, ": null storage"); +} + +void check_shape(const aiter_tensor_t* t, const char* name, std::initializer_list shape) +{ + AITER_CHECK(t->dim() == int(shape.size()), name, ": unexpected rank"); + int d = 0; + for(int64_t extent : shape) + AITER_CHECK(t->size(d++) == extent, name, ": unexpected shape"); +} + +// Conservatively reject overlapping storage spans. Inputs can alias each other; +// writable tensors must not alias inputs or another writable tensor. +uintptr_t end_address(const aiter_tensor_t* t) +{ + if(!t || !t->numel()) + return 0; + size_t span = 1; + for(int d = 0; d < t->dim(); ++d) + span += (t->size(d) - 1) * t->stride(d); + return reinterpret_cast(t->data_ptr()) + span * t->element_size(); +} +void check_no_overlap(const aiter_tensor_t* a, const aiter_tensor_t* b) +{ + if(!a || !b || !a->numel() || !b->numel()) + return; + AITER_CHECK(reinterpret_cast(a->data_ptr()) >= end_address(b) || + reinterpret_cast(b->data_ptr()) >= end_address(a), + "sparse_mla_bf16: output/workspace storage must not overlap"); +} +} // namespace + +AITER_CTYPES_DEFINE_ENTRYPOINT_VOID(sparse_mla_bf16_fwd_out, + (aiter_tensor_t * q, + aiter_tensor_t* kv_buffer, + aiter_tensor_t* kv_indptr, + aiter_tensor_t* kv_indices, + aiter_tensor_t* out, + aiter_tensor_t* lse, + aiter_tensor_t* partial_o, + aiter_tensor_t* partial_lse, + float softmax_scale, + int64_t kv_splits, + int64_t version, + hipStream_t stream), + (q, + kv_buffer, + kv_indptr, + kv_indices, + out, + lse, + partial_o, + partial_lse, + softmax_scale, + kv_splits, + version, + stream)) +{ + AITER_CHECK(q && q->is_gpu(), "q must be on a HIP device"); + const int device = q->device_id; + HipDeviceGuard guard(device); + AITER_CHECK(get_gpu_arch() == "gfx950", "sparse_mla_bf16 requires gfx950"); + AITER_CHECK(version == 1 || version == 2, "version must be 1 or 2"); + AITER_CHECK(kv_splits >= 1 && kv_splits <= 32 && (kv_splits & (kv_splits - 1)) == 0, + "kv_splits must be one of 1, 2, 4, 8, 16, 32"); + AITER_CHECK(std::isfinite(softmax_scale) && softmax_scale > 0, "invalid softmax_scale"); + check_tensor(q, "q", AITER_DTYPE_bf16, device); + AITER_CHECK(q->dim() == 3 && q->size(2) == 512 && (q->size(1) == 16 || q->size(1) == 64), + "expected q=[Q,16|64,512]"); + const int64_t queries = q->size(0), heads = q->size(1); + AITER_CHECK(queries <= INT_MAX, "too many query rows"); + AITER_CHECK(q->stride(2) == 1 && q->stride(1) >= 512 && + q->stride(0) >= (heads - 1) * q->stride(1) + 512, + "q must have non-overlapping, unit-inner-stride rows"); + check_tensor(kv_buffer, "kv_buffer", AITER_DTYPE_bf16, device); + AITER_CHECK(kv_buffer->dim() == 2 && kv_buffer->size(1) == 512 && kv_buffer->stride(1) == 1 && + kv_buffer->stride(0) >= 512, + "expected kv_buffer=[slots,512] with unit inner stride"); + check_tensor(kv_indptr, "kv_indptr", AITER_DTYPE_i32, device); + check_tensor(kv_indices, "kv_indices", AITER_DTYPE_i32, device); + check_shape(kv_indptr, "kv_indptr", {queries + 1}); + AITER_CHECK(kv_indices->dim() == 1 && kv_indptr->is_contiguous() && kv_indices->is_contiguous(), + "CSR tensors must be contiguous vectors"); + AITER_CHECK(kv_indices->numel() <= INT_MAX, "CSR offsets must fit int32"); + check_tensor(out, "out", AITER_DTYPE_bf16, device); + check_shape(out, "out", {queries, heads, 512}); + AITER_CHECK(out->is_contiguous(), "out must be contiguous"); + if(lse) + { + check_tensor(lse, "lse", AITER_DTYPE_fp32, device); + check_shape(lse, "lse", {queries, heads}); + AITER_CHECK(lse->is_contiguous(), "lse must be contiguous"); + } + if(kv_splits > 1) + { + check_tensor(partial_o, "partial_o", AITER_DTYPE_fp32, device); + check_tensor(partial_lse, "partial_lse", AITER_DTYPE_fp32, device); + check_shape(partial_o, "partial_o", {queries, kv_splits, heads, 512}); + check_shape(partial_lse, "partial_lse", {queries, kv_splits, heads}); + AITER_CHECK(partial_o->is_contiguous() && partial_lse->is_contiguous(), + "workspace must be contiguous"); + } + else + { + partial_o = partial_lse = nullptr; + } + const aiter_tensor_t* writes[] = {out, lse, partial_o, partial_lse}; + for(int i = 0; i < 4; ++i) + { + for(auto input : {q, kv_buffer, kv_indptr, kv_indices}) + check_no_overlap(writes[i], input); + for(int j = 0; j < i; ++j) + check_no_overlap(writes[i], writes[j]); + } + // CSR offsets are caller-owned device data: monotone, starting at zero and + // ending at indices.numel(). No D2H reads or synchronization during capture. + // Out-of-range slot ids (including -1) are masked in the kernel. + if(!queries) + return; + sparse_mla_bf16::Params p{static_cast(q->data_ptr()), + static_cast(kv_buffer->data_ptr()), + static_cast(kv_indptr->data_ptr()), + static_cast(kv_indices->data_ptr()), + static_cast(out->data_ptr()), + lse ? static_cast(lse->data_ptr()) : nullptr, + partial_o ? static_cast(partial_o->data_ptr()) : nullptr, + partial_lse ? static_cast(partial_lse->data_ptr()) : nullptr, + int32_t(queries), + int32_t(kv_splits), + kv_buffer->size(0), + q->stride(0), + q->stride(1), + kv_buffer->stride(0), + softmax_scale}; + const auto status = heads == 16 ? sparse_mla_bf16::launch<16>(p, version, stream) + : sparse_mla_bf16::launch<64>(p, version, stream); + HIP_CALL_LAUNCH(status); +} diff --git a/docs/sparse_mla_bf16.md b/docs/sparse_mla_bf16.md new file mode 100644 index 00000000000..4492eccb084 --- /dev/null +++ b/docs/sparse_mla_bf16.md @@ -0,0 +1,125 @@ +# BF16 sparse MLA on gfx950 + +`aiter.sparse_mla_bf16_fwd` computes attention over a per-query CSR list of +global cache slots. K and V are the same BF16 latent, D=512, with no appended +RoPE dimensions. Both prefill and decode use this contract. H16 and H64 are +supported; GLM-5.3-Flash is one caller of this geometry. + +This is an explicit HIP entry point. The existing Gluon `sparse_mla_fwd` +continues to handle its broader set of formats and geometries. An engine can +select this entry point after checking its device, dtype, head count and +geometry, and retain its existing implementation for other configurations. + +## API and storage + +```python +import aiter +import torch + +# q: [Q,H,512] BF16; kv: [slots,512] BF16 +# ptr: [Q+1] int32; idx: [nnz] int32 GLOBAL physical slot ids +splits = 8 # choose outside capture; supported: 1,2,4,8,16,32 +out = torch.empty_like(q, memory_format=torch.contiguous_format) +lse = torch.empty(q.shape[:2], device=q.device, dtype=torch.float32) +workspace = ( + torch.empty((q.shape[0], splits, q.shape[1], 512), device=q.device), + torch.empty((q.shape[0], splits, q.shape[1]), device=q.device), +) +result, result_lse = aiter.sparse_mla_bf16_fwd( + q, kv, ptr, idx, softmax_scale, + version=2, kv_splits=splits, + out=out, return_lse=True, lse=lse, workspace=workspace, +) +``` + +O is contiguous BF16 `[Q,H,512]`. LSE is optional, contiguous FP32 `[Q,H]` +and uses the natural logarithm. Split-K workspace contains normalized FP32 +partial O and natural-log partial LSE, contiguous `[Q,S,H,512]` and `[Q,S,H]`. +It is unused for S=1. The convenience wrapper allocates omitted buffers; +supplying all buffers makes the call allocation-free after JIT warmup. +Warm the extension before capturing a graph. Calls use the query device's +current stream and do not synchronize to inspect CSR values. + +Q/KV allow padded outer strides and BF16 storage offsets, with stride one +in the last dimension. KV also accepts flat-view-compatible +`[pages,page_size,512]` and `[slots,1,1,512]` layouts. Paged inputs are viewed, +not copied. Out/workspace storage must not overlap inputs or one another. +An output with a larger physical head stride requires the engine to use a +contiguous latent output and then its existing projection/copy path. + +The caller supplies monotone CSR offsets starting at zero and ending at +`idx.numel()`. Slot ids outside `[0,slots)` are masked; duplicates retain +their multiplicity. Empty or fully masked rows yield O=0, LSE=-inf. +Causal/window selection is expressed through CSR. There is no implicit +top-k truncation, attention sink, cache quantization, or separate RoPE input. +In particular, a 2051-entry row preserves all 2051 entries. + +## Implementations + +| Version | H16 main | H64 main | Split-K reduction | +|---|---|---|---| +| 1 | One wave, Q in VGPR, 33,024 B LDS | Full140: eight waves, Q in LDS, 139,520 B LDS | Fully unrolled one-wave H64; original H16 reducer | +| 2 | Four waves, 68,352 B LDS | Four head tiles when Q×S≤128 (68,352 B/CTA), otherwise four waves (65,792 B/CTA) | Fully unrolled one-wave H16/H64 | + +The version selector keeps both implementations reproducible under one API. +The default version is 2; the default split count is the explicit S=1, +not an autotuned choice. Use the supplied sweep to choose a split count +for the deployment workload. No automatic routing changes to the existing +Gluon API are part of this operator. + +The extension reuses AITER's existing HipKittens dependency and pinned register +helpers from [PR #3459](https://github.com/ROCm/aiter/pull/3459). No new external +dependency is introduced. This module uses AITER's CK-free headers. Like the +other HK modules, it is excluded from the +ordinary all-op prebuild unless `AITER_ENABLE_EXPERIMENTAL=1`; explicit JIT +calls can build it directly. The module preserves FP32 denormals, overriding +AITER's usual flush-to-zero flag. Compiler-managed VGPRs must stay below v40. + +## Reproduce correctness and timing + +```bash +python -m pytest -q op_tests/test_sparse_mla_bf16.py +python -m op_tests.op_benchmarks.bench_sparse_mla_bf16 \ + --verify --output sparse_mla_o_only.jsonl +python -m op_tests.op_benchmarks.bench_sparse_mla_bf16 \ + --verify --return-lse --output sparse_mla_o_lse.jsonl +# Independent repeat with reversed case/variant order and fresh allocations: +python -m op_tests.op_benchmarks.bench_sparse_mla_bf16 \ + --verify --reverse --output sparse_mla_o_only_repeat.jsonl +``` + +The benchmark defaults to H16/H64, Q=1/32/128/256/512/1024/2048, K=2048/2051, +S=1/2/4/8/16/32. Each query owns a disjoint 32K-slot pool. Q=2048 uses a +64 GiB KV allocation, so run one benchmark on an otherwise idle MI355X. +Native v1/v2 share input, final output and workspace addresses. The Gluon +comparison comes from the same checkout and includes explicit splits and +its shipped split policy. Internal Gluon workspaces differ from the native +workspace; their traffic is part of the timed operation. + +Timing uses alternating variant order and 60 graph/event samples per point. +Warm and cold results are separate. Cold flushes 512 MiB before the graph's +start event. JIT and Python allocation/dispatch time are excluded. Every +case checks sampled query rows against independent FP64 before and after +timing, and hashes all inputs before/after. Correctness failures are reported +and excluded from ratios; a native failure aborts the run. + +Report the complete per-shape table, matched splits, best tested splits and +Gluon's shipped policy separately. JSONL includes raw samples, median/p90, +source/binary hashes, logical GB/s and QK+PV TFLOP/s. Logical traffic counts +Q, one logical gather of KV, indices, O and optional LSE; it excludes repeated +reads and partial buffers and is not measured HBM traffic. A roofline derived +from it is a lower-bound model, not proof of a memory bottleneck. + +For execution-path validation, run a small sweep under `rocprofv3 +--kernel-trace --hip-trace` with `--purpose trace --samples 1 --warmup 0`. +Trace durations are diagnostic and must not be used as performance results. +Build `sparse_mla_bf16.cu` with `--save-temps` and check the retained assembly: + +```bash +python op_tests/check_sparse_mla_bf16_isa.py path/to/sparse_mla_bf16-hip-amdgcn-amd-amdhsa-gfx950.s +``` + +Operator correctness and latency are distinct from model quality and serving +performance. An engine integration should validate its logical-to-physical +index mapping, selected-set policy (including any tail fitting), graph buffer +reuse and fallback, then measure the model with a stable baseline. diff --git a/op_tests/check_sparse_mla_bf16_isa.py b/op_tests/check_sparse_mla_bf16_isa.py new file mode 100644 index 00000000000..2844873f20a --- /dev/null +++ b/op_tests/check_sparse_mla_bf16_isa.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: MIT +"""Gate retained gfx950 compiler assembly before running pinned-VGPR kernels. + +Compile sparse_mla_bf16.cu with --save-temps, then pass its device .s here. +Checks compiler-managed VGPR exclusion, FP32 denormals, spills and private memory. +This is a conservative build check, not a proof of synchronization or liveness. +""" + +import argparse +import hashlib +import json +import re +from pathlib import Path + + +def audit(path): + text = path.read_text() + metadata = {} + for block in re.split(r"(?m)^ - (?=\.)", text.split("\t.amdgpu_metadata", 1)[-1])[ + 1: + ]: + fields = dict(re.findall(r"(?m)^\s*\.(\w+):\s*([^\n]+)", block)) + if "name" in fields: + metadata[fields["name"].strip()] = fields + errors, kernels = [], {} + current, inline = None, False + for number, line in enumerate(text.splitlines(), 1): + function = re.match(r"\s*\.type\s+(\S+),@function", line) + if function: + if inline: + errors.append(f"unclosed inline asm before {number}") + current, inline = function[1], False + if "sparse_mla_bf16" in current and current in metadata: + kernels[current] = {"compiler_vgpr_max": -1, "denorm32": None} + if re.search(r"#(?:ASMSTART|APP)\b", line): + inline = True + continue + if re.search(r"#(?:ASMEND|NO_APP)\b", line): + inline = False + continue + if current not in kernels: + continue + info = kernels[current] + mode = re.search(r"\.amdhsa_float_denorm_mode_32\s+(\d+)", line) + if mode: + info["denorm32"] = int(mode[1]) + instruction = re.match(r"^\s+([a-z][a-z0-9_]*)\s*(.*)$", line) + if not instruction: + continue + if instruction[1].startswith("scratch_"): + errors.append(f"{current}:{number}: scratch instruction") + if not inline: + registers = [ + int(b or a or c, 0) + for a, b, c in re.findall( + r"\bv(?:\[(0x[\da-fA-F]+|\d+)(?::(0x[\da-fA-F]+|\d+))?\]|(\d+))", + line, + ) + ] + info["compiler_vgpr_max"] = max(info["compiler_vgpr_max"], *registers, -1) + main = [name for name in kernels if "sparse_mla_v" in name] + tiles = [name for name in kernels if "sparse_mla_headtiles" in name] + if len(main) != 8 or len(tiles) != 2: + errors.append( + f"expected 8 v1/v2 and 2 head-tiled mains, got {len(main)} and {len(tiles)}" + ) + if len(kernels) != 21: + errors.append(f"expected 21 kernels, got {len(kernels)}") + for name, info in kernels.items(): + if name in main + tiles and info["compiler_vgpr_max"] >= 40: + errors.append(f"{name}: compiler crosses pinned v40 boundary") + if info["denorm32"] != 3: + errors.append(f"{name}: FP32 denormals are not preserved") + for field in ( + "vgpr_count", + "agpr_count", + "sgpr_count", + "vgpr_spill_count", + "sgpr_spill_count", + "private_segment_fixed_size", + "group_segment_fixed_size", + ): + value = metadata[name].get(field, "") + info[field] = int(value) if value.isdecimal() else None + if info[field] is None: + errors.append(f"{name}: missing {field}") + for field in ( + "vgpr_spill_count", + "sgpr_spill_count", + "private_segment_fixed_size", + ): + if info[field] != 0: + errors.append(f"{name}: {field} must be zero") + return { + "passed": not errors, + "errors": errors, + "kernels": kernels, + "assembly_sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + } + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("assembly", type=Path) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + result = audit(args.assembly) + report = json.dumps(result, indent=2) + "\n" + if args.output: + args.output.write_text(report) + print(report) + raise SystemExit(not result["passed"]) diff --git a/op_tests/op_benchmarks/bench_sparse_mla_bf16.py b/op_tests/op_benchmarks/bench_sparse_mla_bf16.py new file mode 100644 index 00000000000..8e395ab5771 --- /dev/null +++ b/op_tests/op_benchmarks/bench_sparse_mla_bf16.py @@ -0,0 +1,406 @@ +# SPDX-License-Identifier: MIT +"""Compare HIP v1/v2 and same-checkout Gluon on disjoint per-query KV pools. + +python op_tests/op_benchmarks/bench_sparse_mla_bf16.py --verify --output result.jsonl +Events inside captured graphs time main + reducer, excluding Python/JIT overhead. +The cold run flushes 512 MiB before the start event. Logical GB/s is not measured +HBM traffic. Sampled FP64 accuracy is checked before and after every timed case. +""" + +import argparse +import gc +import hashlib +import json +import math +from pathlib import Path +import statistics +import sys + +import torch +import triton + +import aiter +from aiter.jit.core import get_user_jit_dir +from aiter.ops.triton.attention import sparse_mla as gluon +from op_tests.test_sparse_mla_bf16 import check, reference + + +def ints(value): + return [int(x) for x in value.split(",")] + + +def tensor_hash(tensor): + data = tensor.detach().contiguous().view(torch.uint8).flatten() + digest = hashlib.sha256() + for start in range(0, data.numel(), 8 << 20): + digest.update(data[start : start + (8 << 20)].cpu().numpy().tobytes()) + return digest.hexdigest() + + +def file_hash(path): + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--heads", type=ints, default=[16, 64]) + parser.add_argument( + "--queries", type=ints, default=[1, 32, 128, 256, 512, 1024, 2048] + ) + parser.add_argument("--topk", type=ints, default=[2048, 2051]) + parser.add_argument("--splits", type=ints, default=[1, 2, 4, 8, 16, 32]) + parser.add_argument("--context", type=int, default=32768) + parser.add_argument("--samples", type=int, default=60) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--check-rows", type=int, default=4) + parser.add_argument("--cache-modes", default="warm,cold") + parser.add_argument("--return-lse", action="store_true") + parser.add_argument("--purpose", choices=["timing", "trace"], default="timing") + parser.add_argument("--reverse", action="store_true") + parser.add_argument("--verify", action="store_true", required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.output.exists(): + raise ValueError("output already exists; use a fresh run") + assert args.samples >= 1 and args.check_rows >= 1 + assert args.purpose == "trace" or (args.samples >= 30 and args.warmup >= 10) + assert all(q > 0 for q in args.queries) + assert set(args.heads) <= {16, 64} and set(args.splits) <= {1, 2, 4, 8, 16, 32} + assert set(args.cache_modes.split(",")) <= {"warm", "cold"} + assert all(0 < k <= args.context for k in args.topk) + assert torch.version.hip and torch.cuda.device_count() == 1 + torch.cuda.set_device(0) + torch.set_num_threads(4) + props = torch.cuda.get_device_properties(0) + assert props.gcnArchName.split(":")[0] == "gfx950" + + def emit(row): + with args.output.open("a") as output: + output.write(json.dumps(row, allow_nan=False) + "\n") + keys = ( + "type", + "heads", + "queries", + "topk", + "variant", + "splits", + "cache", + "median_us", + "status", + ) + print(json.dumps({k: row[k] for k in keys if k in row}), flush=True) + + source = Path(__file__).resolve() + jit = Path(get_user_jit_dir()) / "module_sparse_mla_bf16.so" + emit( + { + "type": "config", + "command": sys.argv, + "device": str(props), + "torch": str(torch.__version__), + "triton": triton.__version__, + "hip": torch.version.hip, + "benchmark_sha256": file_hash(source), + "aiter": aiter.__file__, + "gluon_source_sha256": file_hash(Path(gluon.__file__)), + "native_library_sha256": file_hash(jit) if jit.exists() else None, + "scope": "captured attention main+reduce; not model/engine latency", + "pool_layout": "disjoint", + "output_contract": "O+LSE" if args.return_lse else "O-only", + "args": vars(args) | {"output": str(args.output)}, + } + ) + stream = torch.cuda.Stream() + flush = torch.empty(512 << 20, device="cuda", dtype=torch.uint8) + cases = [(h, q, k) for h in args.heads for q in args.queries for k in args.topk] + if args.reverse: + cases.reverse() + total = failed = 0 + for heads, queries, topk in cases: + with torch.cuda.stream(stream): + torch.manual_seed(1701 + heads + queries + topk) + q = ( + torch.randn((queries, heads, 512), device="cuda", dtype=torch.bfloat16) + * 0.5 + ) + kv = torch.empty( + (queries * args.context, 512), device="cuda", dtype=torch.bfloat16 + ) + kv.normal_().mul_(0.5) + idx = torch.cat( + [ + torch.randperm(args.context, device="cuda")[:topk].to(torch.int32) + + row * args.context + for row in range(queries) + ] + ) + ptr = torch.arange(queries + 1, device="cuda", dtype=torch.int32) * topk + out = torch.empty_like(q) + lse = torch.empty(q.shape[:2], device="cuda") if args.return_lse else None + max_s = max(args.splits) + po = torch.empty((queries * max_s * heads * 512,), device="cuda") + pl = torch.empty((queries * max_s * heads,), device="cuda") + rows = sorted( + { + round(i * (queries - 1) / max(1, min(queries, args.check_rows) - 1)) + for i in range(min(queries, args.check_rows)) + } + ) + selected = torch.tensor(rows, device="cuda", dtype=torch.long) + oracle_ptr = ( + torch.arange(len(rows) + 1, device="cuda", dtype=torch.int32) * topk + ) + expected = reference( + q[selected], + kv, + oracle_ptr, + idx.view(queries, topk)[selected].flatten(), + 1 / 16, + ) + tensors = {"q": q, "kv": kv, "indptr": ptr, "indices": idx} + hashes = {name: tensor_hash(value) for name, value in tensors.items()} + case = { + "heads": heads, + "queries": queries, + "topk": topk, + "pool_bytes": kv.numel() * 2, + "input_sha256": hashes, + "input_addresses": { + name: value.data_ptr() for name, value in tensors.items() + }, + "out_address": out.data_ptr(), + "checked_rows": rows, + "output_contract": "O+LSE" if args.return_lse else "O-only", + } + variants = [(f"v{v}", s) for s in args.splits for v in (1, 2)] + variants += [("gluon", s) for s in args.splits] + [("gluon_default", None)] + if args.reverse: + variants.reverse() + runners = [] + for variant, splits in variants: + state = [None] + if variant.startswith("v"): + workspace = ( + ( + po[: queries * splits * heads * 512].view( + queries, splits, heads, 512 + ), + pl[: queries * splits * heads].view(queries, splits, heads), + ) + if splits > 1 + else None + ) + + def invoke( + v=int(variant[1]), + s=splits, + workspace=workspace, + state=state, + inputs=(q, kv, ptr, idx, out, lse), + ): + q, kv, ptr, idx, out, lse = inputs + state[0] = aiter.sparse_mla_bf16_fwd( + q, + kv, + ptr, + idx, + 1 / 16, + version=v, + kv_splits=s, + out=out, + return_lse=args.return_lse, + lse=lse, + workspace=workspace, + ) + + actual_splits = splits + else: + + def invoke(s=splits, state=state, inputs=(q, kv, ptr, idx, out)): + q, kv, ptr, idx, out = inputs + state[0] = gluon.sparse_mla_fwd( + q, + kv, + ptr, + idx, + 1 / 16, + kv_lora_rank=512, + qk_rope_head_dim=0, + kv_splits=s, + has_invalid=False, + out=out, + return_lse=args.return_lse, + ) + + actual_splits = splits or gluon._mla_num_splits( + queries, heads // 16, topk + ) + key = {"variant": variant, "splits": actual_splits} + invoke() + torch.cuda.synchronize() + try: + before = check( + ( + state[0][0][selected], + state[0][1][selected] if state[0][1] is not None else None, + ), + expected, + ) + except AssertionError as error: + failed += 1 + emit( + { + "type": "accuracy_failure", + **case, + **key, + "status": "excluded", + "error": str(error), + } + ) + if variant.startswith("v"): + raise + continue + start = torch.cuda.Event(enable_timing=True, external=True) + end = torch.cuda.Event(enable_timing=True, external=True) + start.record() + end.record() + stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + start.record() + invoke() + end.record() + graph.replay() + stream.synchronize() + check( + ( + state[0][0][selected], + state[0][1][selected] if state[0][1] is not None else None, + ), + expected, + ) + runners.append( + { + **key, + "graph": graph, + "start": start, + "end": end, + "state": state, + "before_rms": before, + } + ) + for mode in args.cache_modes.split(","): + for i in range(args.warmup): + for r in (runners if i % 2 == 0 else runners[::-1]): + if mode == "cold": + flush.fill_(1) + r["graph"].replay() + stream.synchronize() + measured = [[] for _ in runners] + for i in range(args.samples): + order = ( + range(len(runners)) + if i % 2 == 0 + else reversed(range(len(runners))) + ) + for j in order: + r = runners[j] + if mode == "cold": + flush.fill_(1) + r["graph"].replay() + r["end"].synchronize() + value = r["start"].elapsed_time(r["end"]) * 1000 + assert math.isfinite(value) and value > 0 + measured[j].append(value) + for r, samples in zip(runners, measured): + r["graph"].replay() + stream.synchronize() + output, output_lse = r["state"][0] + after = check( + ( + output[selected], + output_lse[selected] if output_lse is not None else None, + ), + expected, + ) + median = statistics.median(samples) + logical_bytes = ( + queries * (heads * 512 * 4 + topk * (512 * 2 + 4)) + + (queries + 1) * 4 + ) + if args.return_lse: + logical_bytes += queries * heads * 4 + emit( + { + "type": ( + "timing" + if args.purpose == "timing" + else "diagnostic_timing" + ), + **case, + "variant": r["variant"], + "splits": r["splits"], + "cache": mode, + "samples_us": samples, + "median_us": median, + "p90_us": sorted(samples)[ + math.ceil(len(samples) * 0.9) - 1 + ], + "logical_bytes": logical_bytes, + "logical_gbps": logical_bytes / median / 1000, + "qk_pv_tflops": 4 + * queries + * heads + * topk + * 512 + / median + / 1e6, + "before_rms": r["before_rms"], + "after_rms": after, + "status": "passed", + } + ) + total += 1 + assert hashes == { + name: tensor_hash(value) for name, value in tensors.items() + } + emit( + { + "type": "case_complete", + **case, + "status": "passed", + "inputs_unchanged": True, + } + ) + runners.clear() + del ( + q, + kv, + ptr, + idx, + out, + lse, + po, + pl, + tensors, + invoke, + state, + graph, + r, + output, + output_lse, + ) + gc.collect() + torch.cuda.empty_cache() + emit( + { + "type": "complete", + "status": "passed", + "timing_records": total, + "accuracy_exclusions": failed, + } + ) + + +if __name__ == "__main__": + main() diff --git a/op_tests/test_sparse_mla_bf16.py b/op_tests/test_sparse_mla_bf16.py new file mode 100644 index 00000000000..405515d4067 --- /dev/null +++ b/op_tests/test_sparse_mla_bf16.py @@ -0,0 +1,260 @@ +# SPDX-License-Identifier: MIT +"""BF16 D512 sparse MLA: independent FP64, ragged CSR, strides and graph replay. + +Run: python -m pytest -q op_tests/test_sparse_mla_bf16.py +""" + +import pytest +import torch + +from aiter import sparse_mla_bf16_fwd +from aiter.jit.utils.chip_info import get_gfx_runtime + +pytestmark = pytest.mark.skipif(get_gfx_runtime() != "gfx950", reason="requires gfx950") +SPLITS = [1, 2, 4, 8, 16, 32] + + +def reference(q, kv, indptr, indices, scale): + """Reference the mathematical operation, independently of tile/split layout.""" + output = torch.zeros_like(q, dtype=torch.float64) + lse = torch.full(q.shape[:2], -torch.inf, device=q.device, dtype=torch.float64) + offsets = indptr.cpu().tolist() + for row, (start, end) in enumerate(zip(offsets, offsets[1:])): + selected = indices[start:end].long() + selected = selected[(selected >= 0) & (selected < kv.shape[0])] + if not selected.numel(): + continue + values = kv[selected].double() + scores = q[row].double() @ values.T * scale + output[row] = scores.softmax(-1) @ values + lse[row] = scores.logsumexp(-1) + return output, lse + + +def check(actual, expected): + out, lse = actual + ref_o, ref_lse = expected + torch.testing.assert_close(out.double(), ref_o, atol=1e-2, rtol=1e-2) + rms = (out.double() - ref_o).square().mean(-1).sqrt() + relative = rms / ref_o.square().mean(-1).sqrt().clamp_min(1e-3) + assert (relative <= 1e-2).all(), relative.max().item() + if lse is not None: + torch.testing.assert_close(lse.double(), ref_lse, atol=1e-4, rtol=1e-5) + return relative.max().item() if relative.numel() else 0.0 + + +def fixture(heads, topk, padded=False): + torch.manual_seed(321 + heads + topk) + queries, slots = 4, max(65, topk + 7) + if padded: + # Offset by one BF16, with padded query/head/cache rows. + q_store = torch.randn(queries, heads, 520, device="cuda", dtype=torch.bfloat16) + kv_store = torch.randn(slots, 520, device="cuda", dtype=torch.bfloat16) + q, kv = q_store[..., 1:513], kv_store[:, 1:513] + else: + q = torch.randn(queries, heads, 512, device="cuda", dtype=torch.bfloat16) + kv = torch.randn(slots, 512, device="cuda", dtype=torch.bfloat16) + lengths = [topk, max(0, topk - 3), 0, int(topk > 0)] + ptr = torch.tensor( + [0, *torch.tensor(lengths).cumsum(0).tolist()], device="cuda", dtype=torch.int32 + ) + idx = torch.randint(slots, (sum(lengths),), device="cuda", dtype=torch.int32) + if topk > 1: + idx[:topk:11] = -1 + idx[1:topk:17] = slots + 1 + # Duplicate rows intentionally retain multiplicity. + if topk > 4: + idx[2:4] = 3 + return q, kv, ptr, idx + + +@pytest.mark.parametrize("version", [1, 2]) +@pytest.mark.parametrize("heads", [16, 64]) +@pytest.mark.parametrize("splits", SPLITS) +@pytest.mark.parametrize("topk", [0, 1, 31, 32, 33, 63, 65, 2048, 2051]) +@pytest.mark.parametrize("padded", [False, True]) +def test_fp64(version, heads, splits, topk, padded): + q, kv, ptr, idx = fixture(heads, topk, padded) + q_count = q.shape[0] + out_storage = torch.full((q.numel() + 2,), 123, device=q.device, dtype=q.dtype) + out = out_storage[1:-1].view(q.shape) + lse = torch.empty(q.shape[:2], device=q.device, dtype=torch.float32) + workspace = None + if splits > 1: + count = q_count * splits * heads + # Four-byte-aligned partial O exercises the scalar fallback. + partial = torch.empty(count * 512 + 1, device=q.device, dtype=torch.float32)[1:] + workspace = ( + partial.view(q_count, splits, heads, 512), + torch.empty((q_count, splits, heads), device=q.device), + ) + result = sparse_mla_bf16_fwd( + q, + kv, + ptr, + idx, + 1 / 16, + version=version, + kv_splits=splits, + out=out, + lse=lse, + return_lse=True, + workspace=workspace, + ) + assert result[0].data_ptr() == out.data_ptr() + assert result[1].data_ptr() == lse.data_ptr() + check(result, reference(q, kv, ptr, idx, 1 / 16)) + assert out_storage[0] == 123 and out_storage[-1] == 123 + + +@pytest.mark.parametrize("version", [1, 2]) +@pytest.mark.parametrize("heads", [16, 64]) +@pytest.mark.parametrize("splits", SPLITS) +@pytest.mark.parametrize("return_lse", [False, True]) +def test_graph_current_stream(version, heads, splits, return_lse): + q, kv, ptr, idx = fixture(heads, 65) + out = torch.empty_like(q) + lse = torch.empty(q.shape[:2], device=q.device) if return_lse else None + workspace = ( + ( + torch.empty((4, splits, heads, 512), device=q.device), + torch.empty((4, splits, heads), device=q.device), + ) + if splits > 1 + else None + ) + + def run(): + return sparse_mla_bf16_fwd( + q, + kv, + ptr, + idx, + 1 / 16, + version=version, + kv_splits=splits, + out=out, + lse=lse, + return_lse=return_lse, + workspace=workspace, + ) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + run() + # Replay must consume new device values, not captured query/index data. + q.mul_(0.5) + idx.copy_(idx.roll(1)) + for _ in range(30): + graph.replay() + expected = reference(q, kv, ptr, idx, 1 / 16) + torch.cuda.current_stream().wait_stream(stream) + check((out, lse), expected) + + +@pytest.mark.parametrize("version", [1, 2]) +@pytest.mark.parametrize("heads", [16, 64]) +@pytest.mark.parametrize("splits", [1, 32]) +def test_subnormal_single_key(version, heads, splits): + q = torch.zeros((1, heads, 512), device="cuda", dtype=torch.bfloat16) + bits = torch.tensor( + [1, 2, 127, -32767, -32766, -32641, 0, 0], device="cuda", dtype=torch.int16 + ) + kv = bits.repeat(64).view(torch.bfloat16).view(1, 512) + ptr = torch.tensor([0, 1], device="cuda", dtype=torch.int32) + idx = torch.zeros(1, device="cuda", dtype=torch.int32) + out, lse = sparse_mla_bf16_fwd( + q, kv, ptr, idx, 1 / 16, version=version, kv_splits=splits, return_lse=True + ) + assert torch.equal( + out.view(torch.int16), kv.expand(heads, 512)[None].view(torch.int16) + ) + assert (lse == 0).all() + + +@pytest.mark.parametrize("splits", SPLITS) +@pytest.mark.parametrize("above_boundary", [False, True]) +def test_h64_hybrid_dispatch_boundary(splits, above_boundary): + # Q*S=128 uses head tiles; the next query switches to the four-wave main. + # The small ragged fixture alone stays on the head-tiled branch for all S. + queries = 128 // splits + int(above_boundary) + torch.manual_seed(812) + q = torch.randn((queries, 64, 512), dtype=torch.bfloat16, device="cuda") + kv = torch.randn((97, 512), dtype=torch.bfloat16, device="cuda") + ptr = torch.arange(queries + 1, dtype=torch.int32, device="cuda") * 65 + idx = torch.randint(97, (queries, 65), dtype=torch.int32, device="cuda") + idx[::2, 16:32] = -1 + idx[-1] = -1 + result = sparse_mla_bf16_fwd( + q, + kv, + ptr, + idx.flatten(), + 1 / 16, + version=2, + kv_splits=splits, + return_lse=True, + ) + check(result, reference(q, kv, ptr, idx.flatten(), 1 / 16)) + + +@pytest.mark.parametrize("version", [1, 2]) +@pytest.mark.parametrize("pool_shape", ["flat", "paged", "mla"]) +def test_pool_views_empty_q_and_compile(version, pool_shape): + q, kv, ptr, idx = fixture(16, 33) + # Slice to a multiple of the page size; invalid ids are masked by contract. + kv = kv[:64] + pool = {"flat": kv, "paged": kv.view(8, 8, 512), "mla": kv[:, None, None]}[ + pool_shape + ] + run = torch.compile(sparse_mla_bf16_fwd, backend="eager", fullgraph=True) + actual = run( + q, pool, ptr, idx, 1 / 16, version=version, kv_splits=2, return_lse=True + ) + check(actual, reference(q, kv, ptr, idx, 1 / 16)) + empty = sparse_mla_bf16_fwd( + q[:0], + pool, + ptr[:1], + idx[:0], + 1 / 16, + version=version, + kv_splits=2, + return_lse=True, + ) + assert empty[0].shape == (0, 16, 512) and empty[1].shape == (0, 16) + + +@pytest.mark.parametrize( + "change", + ["dtype", "width", "heads", "splits", "version", "scale", "alias", "workspace"], +) +def test_reject_invalid_metadata(change): + q, kv, ptr, idx = fixture(16, 33) + args = dict(kv_splits=2) + if change == "dtype": + q = q.float() + elif change == "width": + q = q[..., :256] + elif change == "heads": + q = q[:, :8] + elif change == "splits": + args["kv_splits"] = 3 + elif change == "version": + args["version"] = 3 + elif change == "alias": + args["out"] = q + elif change == "workspace": + args["workspace"] = ( + torch.empty(1, device=q.device), + torch.empty(1, device=q.device), + ) + with pytest.raises((RuntimeError, ValueError)): + sparse_mla_bf16_fwd( + q, kv, ptr, idx, float("nan") if change == "scale" else 1 / 16, **args + ) From 59f833e2051c185365d4928154489bf1f765106e Mon Sep 17 00:00:00 2001 From: Sumin Hong Date: Fri, 16 Oct 2026 00:14:03 +0000 Subject: [PATCH 2/3] Add FP8 sparse MLA with fused BF16 query conversion on gfx950 Signed-off-by: Sumin Hong --- aiter/__init__.py | 5 +- aiter/jit/optCompilerConfig.json | 37 +- aiter/ops/sparse_mla.py | 174 ++-- csrc/kernels/mla/fp8/kv_layout.cuh | 24 + csrc/kernels/mla/fp8/layout.cuh | 15 + csrc/kernels/mla/fp8/online_common.cuh | 645 +++++++++++++ csrc/kernels/mla/fp8/order.cuh | 38 + csrc/kernels/mla/fp8/pad16_factored.cuh | 109 +++ csrc/kernels/mla/fp8/params.cuh | 43 + csrc/kernels/mla/fp8/partial_io.cuh | 34 + csrc/kernels/mla/fp8/partial_layout.cuh | 34 + csrc/kernels/mla/fp8/prefetch.cu | 629 +++++++++++++ csrc/kernels/mla/fp8/pv_pipe.cuh | 46 + csrc/kernels/mla/fp8/q_encode.cuh | 109 +++ csrc/kernels/mla/fp8/reduce.cuh | 110 +++ csrc/kernels/mla/fp8/rotating.cu | 365 ++++++++ csrc/kernels/mla/fp8/small.cu | 847 +++++++++++++++++ csrc/kernels/mla/fp8/softmax_lane.cuh | 21 + csrc/kernels/mla/fp8/streaming.cu | 303 ++++++ .../kernels/mla/hk/sparse_mla_bf16_common.cuh | 878 ------------------ .../mla/hk/sparse_mla_bf16_full_lds.cuh | 374 -------- .../mla/hk/sparse_mla_bf16_full_lds_pv.cuh | 155 ---- .../mla/hk/sparse_mla_bf16_headtiles.cuh | 502 ---------- .../kernels/mla/hk/sparse_mla_bf16_reduce.cuh | 107 --- csrc/kernels/mla/hk/sparse_mla_bf16_wave4.cuh | 493 ---------- csrc/kernels/mla/sparse_mla_bf16.cu | 218 ----- csrc/kernels/mla/sparse_mla_fp8.cu | 162 ++++ docs/sparse_mla_bf16.md | 125 --- docs/sparse_mla_fp8.md | 118 +++ op_tests/check_sparse_mla_bf16_isa.py | 112 --- op_tests/check_sparse_mla_fp8_isa.py | 95 ++ .../op_benchmarks/bench_sparse_mla_bf16.py | 406 -------- .../op_benchmarks/bench_sparse_mla_fp8.py | 210 +++++ op_tests/test_sparse_mla_bf16.py | 260 ------ op_tests/test_sparse_mla_fp8.py | 191 ++++ 35 files changed, 4273 insertions(+), 3721 deletions(-) create mode 100644 csrc/kernels/mla/fp8/kv_layout.cuh create mode 100644 csrc/kernels/mla/fp8/layout.cuh create mode 100644 csrc/kernels/mla/fp8/online_common.cuh create mode 100644 csrc/kernels/mla/fp8/order.cuh create mode 100644 csrc/kernels/mla/fp8/pad16_factored.cuh create mode 100644 csrc/kernels/mla/fp8/params.cuh create mode 100644 csrc/kernels/mla/fp8/partial_io.cuh create mode 100644 csrc/kernels/mla/fp8/partial_layout.cuh create mode 100644 csrc/kernels/mla/fp8/prefetch.cu create mode 100644 csrc/kernels/mla/fp8/pv_pipe.cuh create mode 100644 csrc/kernels/mla/fp8/q_encode.cuh create mode 100644 csrc/kernels/mla/fp8/reduce.cuh create mode 100644 csrc/kernels/mla/fp8/rotating.cu create mode 100644 csrc/kernels/mla/fp8/small.cu create mode 100644 csrc/kernels/mla/fp8/softmax_lane.cuh create mode 100644 csrc/kernels/mla/fp8/streaming.cu delete mode 100644 csrc/kernels/mla/hk/sparse_mla_bf16_common.cuh delete mode 100644 csrc/kernels/mla/hk/sparse_mla_bf16_full_lds.cuh delete mode 100644 csrc/kernels/mla/hk/sparse_mla_bf16_full_lds_pv.cuh delete mode 100644 csrc/kernels/mla/hk/sparse_mla_bf16_headtiles.cuh delete mode 100644 csrc/kernels/mla/hk/sparse_mla_bf16_reduce.cuh delete mode 100644 csrc/kernels/mla/hk/sparse_mla_bf16_wave4.cuh delete mode 100644 csrc/kernels/mla/sparse_mla_bf16.cu create mode 100644 csrc/kernels/mla/sparse_mla_fp8.cu delete mode 100644 docs/sparse_mla_bf16.md create mode 100644 docs/sparse_mla_fp8.md delete mode 100644 op_tests/check_sparse_mla_bf16_isa.py create mode 100644 op_tests/check_sparse_mla_fp8_isa.py delete mode 100644 op_tests/op_benchmarks/bench_sparse_mla_bf16.py create mode 100644 op_tests/op_benchmarks/bench_sparse_mla_fp8.py delete mode 100644 op_tests/test_sparse_mla_bf16.py create mode 100644 op_tests/test_sparse_mla_fp8.py diff --git a/aiter/__init__.py b/aiter/__init__.py index 838d362d67d..1749444af0c 100644 --- a/aiter/__init__.py +++ b/aiter/__init__.py @@ -108,7 +108,10 @@ def getLogger(): from .ops.aiter_operator import * from .ops.activation import * from .ops.attention import * - from .ops.sparse_mla import sparse_mla_bf16_fwd # noqa: F401 + from .ops.sparse_mla import ( # noqa: F401 + sparse_mla_fp8_fwd, + allocate_sparse_mla_fp8_workspace, + ) from .ops.custom import * from .ops.custom_all_reduce import * from .ops.quick_all_reduce import * diff --git a/aiter/jit/optCompilerConfig.json b/aiter/jit/optCompilerConfig.json index 9d06815c128..6dea140984b 100644 --- a/aiter/jit/optCompilerConfig.json +++ b/aiter/jit/optCompilerConfig.json @@ -1,17 +1,28 @@ { - "module_sparse_mla_bf16": { - "srcs": [ - "f'{AITER_CSRC_DIR}/kernels/mla/sparse_mla_bf16.cu'", - "f'{AITER_CSRC_DIR}/kernels/mla/hk/sparse_mla_bf16_common.cuh'", - "f'{AITER_CSRC_DIR}/kernels/mla/hk/sparse_mla_bf16_full_lds.cuh'", - "f'{AITER_CSRC_DIR}/kernels/mla/hk/sparse_mla_bf16_full_lds_pv.cuh'", - "f'{AITER_CSRC_DIR}/kernels/mla/hk/sparse_mla_bf16_headtiles.cuh'", - "f'{AITER_CSRC_DIR}/kernels/mla/hk/sparse_mla_bf16_reduce.cuh'", - "f'{AITER_CSRC_DIR}/kernels/mla/hk/sparse_mla_bf16_wave4.cuh'" + "module_sparse_mla_fp8": { + "srcs": [ + "f'{AITER_CSRC_DIR}/kernels/mla/sparse_mla_fp8.cu'", + "f'{AITER_CSRC_DIR}/kernels/mla/fp8/kv_layout.cuh'", + "f'{AITER_CSRC_DIR}/kernels/mla/fp8/layout.cuh'", + "f'{AITER_CSRC_DIR}/kernels/mla/fp8/online_common.cuh'", + "f'{AITER_CSRC_DIR}/kernels/mla/fp8/order.cuh'", + "f'{AITER_CSRC_DIR}/kernels/mla/fp8/pad16_factored.cuh'", + "f'{AITER_CSRC_DIR}/kernels/mla/fp8/params.cuh'", + "f'{AITER_CSRC_DIR}/kernels/mla/fp8/partial_io.cuh'", + "f'{AITER_CSRC_DIR}/kernels/mla/fp8/partial_layout.cuh'", + "f'{AITER_CSRC_DIR}/kernels/mla/fp8/prefetch.cu'", + "f'{AITER_CSRC_DIR}/kernels/mla/fp8/pv_pipe.cuh'", + "f'{AITER_CSRC_DIR}/kernels/mla/fp8/q_encode.cuh'", + "f'{AITER_CSRC_DIR}/kernels/mla/fp8/reduce.cuh'", + "f'{AITER_CSRC_DIR}/kernels/mla/fp8/rotating.cu'", + "f'{AITER_CSRC_DIR}/kernels/mla/fp8/small.cu'", + "f'{AITER_CSRC_DIR}/kernels/mla/fp8/softmax_lane.cuh'", + "f'{AITER_CSRC_DIR}/kernels/mla/fp8/streaming.cu'" + ], + "flags_extra_cc": [ + "'-DENABLE_CK=0'" ], - "flags_extra_cc": ["'-DENABLE_CK=0'"], "flags_extra_hip": [ - "'-DKITTENS_CDNA4'", "'-fno-gpu-flush-denormals-to-zero'" ], "extra_ldflags": "None", @@ -19,9 +30,7 @@ "verbose": "False", "torch_exclude": "True", "blob_gen_cmd": "''", - "third_party": [ - "'HipKittens'" - ], + "third_party": [], "is_experimental": "True" }, "module_aiter_core": { diff --git a/aiter/ops/sparse_mla.py b/aiter/ops/sparse_mla.py index 7c42a88fa7e..f4e9a0d90aa 100644 --- a/aiter/ops/sparse_mla.py +++ b/aiter/ops/sparse_mla.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: MIT -"""HIP sparse MLA for BF16 D512 latent K/V on MI350 (gfx950).""" +"""HIP sparse MLA with FP8 latent K/V and BF16 output on gfx950.""" import torch @@ -8,156 +8,178 @@ from csrc.cpp_itfs.torch_utils import direct_register_custom_op +def _split_count(queries: int) -> int: + if queries <= 6: + return 17 + if queries <= 12: + return 9 + if queries <= 21: + return 6 + if queries <= 23: + return 4 + if queries <= 42: + return 6 + if queries <= 64: + return 4 + if queries <= 128: + return 2 + return 1 + + +def allocate_sparse_mla_fp8_workspace( + q: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Allocate opaque scratch storage for ``sparse_mla_fp8_fwd``. + + Reuse this pair for the same Q shape and device. No initialization or GPU + kernel is required. The layout is private and may vary with query count. + """ + if q.ndim != 3 or q.shape[1:] != (64, 512) or not 0 <= q.shape[0] <= 2048: + raise ValueError("expected q=[Q,64,512], 0 <= Q <= 2048") + queries = q.shape[0] + splits = _split_count(queries) + rows = queries * splits * 64 if splits > 1 else 0 + return ( + torch.empty((rows * 512,), dtype=torch.bfloat16, device=q.device), + torch.empty((rows,), dtype=torch.float32, device=q.device), + ) + + @compile_ops( - "module_sparse_mla_bf16", fc_name="sparse_mla_bf16_fwd_out", ffi_type="ctypes" + "module_sparse_mla_fp8", fc_name="sparse_mla_fp8_fwd_out", ffi_type="ctypes" ) -def _jit_sparse_mla_bf16_fwd_out( +def _jit_sparse_mla_fp8_fwd_out( q: torch.Tensor, kv_buffer: torch.Tensor, - kv_indptr: torch.Tensor, kv_indices: torch.Tensor, out: torch.Tensor, lse: torch.Tensor | None, - partial_o: torch.Tensor | None, - partial_lse: torch.Tensor | None, + partial_o: torch.Tensor, + partial_lse: torch.Tensor, softmax_scale: float, - kv_splits: int, - version: int, + q_scale: float, + kv_scale: float, ) -> None: ... -def _sparse_mla_bf16_fwd_out( +def _sparse_mla_fp8_fwd_out( q: torch.Tensor, kv_buffer: torch.Tensor, - kv_indptr: torch.Tensor, kv_indices: torch.Tensor, out: torch.Tensor, lse: torch.Tensor | None, - partial_o: torch.Tensor | None, - partial_lse: torch.Tensor | None, + partial_o: torch.Tensor, + partial_lse: torch.Tensor, softmax_scale: float, - kv_splits: int, - version: int, + q_scale: float, + kv_scale: float, ) -> None: - # Runtime discovery stays inside the opaque op, outside Dynamo tracing. + # Runtime discovery stays inside the opaque operation, outside Dynamo. if get_gfx_runtime() != "gfx950": - raise NotImplementedError("sparse_mla_bf16_fwd requires gfx950") - _jit_sparse_mla_bf16_fwd_out( + raise NotImplementedError("sparse_mla_fp8_fwd requires gfx950") + if q.dtype not in (torch.bfloat16, torch.float8_e4m3fn): + raise ValueError("q must be BF16 or float8_e4m3fn") + if kv_buffer.dtype != torch.float8_e4m3fn: + raise ValueError("kv_buffer must be float8_e4m3fn") + _jit_sparse_mla_fp8_fwd_out( q, kv_buffer, - kv_indptr, kv_indices, out, lse, partial_o, partial_lse, softmax_scale, - kv_splits, - version, + q_scale, + kv_scale, ) def _fake_out( q, kv_buffer, - kv_indptr, kv_indices, out, lse, partial_o, partial_lse, softmax_scale, - kv_splits, - version, + q_scale, + kv_scale, ): return None direct_register_custom_op( - "sparse_mla_bf16_fwd_out", - _sparse_mla_bf16_fwd_out, + "sparse_mla_fp8_fwd_out", + _sparse_mla_fp8_fwd_out, ["out", "lse", "partial_o", "partial_lse"], fake_impl=_fake_out, ) -def sparse_mla_bf16_fwd( +def sparse_mla_fp8_fwd( q: torch.Tensor, kv_buffer: torch.Tensor, - kv_indptr: torch.Tensor, kv_indices: torch.Tensor, softmax_scale: float, *, - kv_splits: int = 1, - version: int = 2, + q_scale: float = 1.0, + kv_scale: float = 1.0, out: torch.Tensor | None = None, return_lse: bool = False, lse: torch.Tensor | None = None, workspace: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Sparse attention with a shared BF16 K/V latent and no appended RoPE. - - ``q`` is [Q, H, 512], H in {16, 64}. ``kv_buffer`` is [slots, 512], - [pages, page_size, 512], or [slots, 1, 1, 512]. Paged pools must be viewable - as flat rows without copying. Q/KV permit padded outer strides and BF16 - storage offsets; their last dimension must have stride one. - - ``kv_indptr`` is contiguous int32 [Q+1], and ``kv_indices`` is a contiguous - int32 vector of GLOBAL slot ids. Offsets must start at zero, be monotone, - and end at indices.numel(). These device values are caller invariants; - the hot path does not synchronize to validate them. Negative and - out-of-range slot ids are masked. Duplicates retain their multiplicity. - Empty rows produce O=0 and natural-log LSE=-inf. - - ``version=1`` selects Full140 for H64 (139,520 B LDS) and the one-wave - H16 kernel. ``version=2`` selects the hybrid: H16 uses four waves, and - H64 uses four head tiles when Q*kv_splits <= 128, otherwise four waves. - Both H64 versions and v2 H16 use the fully unrolled split-K reducer. - ``kv_splits`` is an explicit power of two from 1 through 32. - - O is contiguous BF16 [Q,H,512]; optional LSE is contiguous FP32 [Q,H]. - For split-K, workspace is contiguous FP32 ([Q,S,H,512], [Q,S,H]). - Supply out/lse/workspace to reuse all storage without allocation. Writable - storage must not overlap other arguments. Warm the JIT before graph capture. - Only gfx950 is supported; unsupported configurations raise an error. + """Sparse attention with a shared FP8 K/V latent and no appended RoPE. + + Q is BF16 or E4M3FN [Q,64,512], with Q <= 2048. Q rows must be aligned + to 16 bytes, non-overlapping, and have unit inner stride. KV is contiguous + E4M3FN [slots,512], [pages,page_size,512], or [slots,1,1,512]. + + ``kv_indices`` is contiguous int32 [Q,K] or [Q,1,K], K <= 2176. Its + entries are physical slot IDs, not page IDs. Negative/out-of-range IDs + are masked; duplicates count separately. Empty rows return zero output + and negative-infinity natural-log LSE. + + ``q_scale`` and ``kv_scale`` are finite positive scalar decoding scales. + BF16 Q is rounded to E4M3FN as Q/q_scale inside attention. FP8 Q is used + directly. One FP8 QK and FP8 PV are computed, with FP32 accumulation, + softmax and LSE. Split partial output and final output are BF16. + + Supply ``out``, optional ``lse``, and a pair from + ``allocate_sparse_mla_fp8_workspace(q)`` to reuse storage. Writable tensors + must not overlap any other argument. Warm the JIT before HIP Graph capture; + dispatch uses host shape metadata and replay reads updated input values. + This opt-in operator supports gfx950 only and has no engine fallback. """ - if q.ndim != 3: - raise ValueError("q must have shape [Q,H,512]") - if version not in (1, 2) or kv_splits not in (1, 2, 4, 8, 16, 32): - raise ValueError("version must be 1/2 and kv_splits must be 1/2/4/8/16/32") - queries, heads, width = q.shape + if q.ndim != 3 or q.shape[1:] != (64, 512) or not 0 <= q.shape[0] <= 2048: + raise ValueError("expected q=[Q,64,512], 0 <= Q <= 2048") if kv_buffer.ndim == 4 and kv_buffer.shape[1:3] == (1, 1): kv_buffer = kv_buffer.view(kv_buffer.shape[0], kv_buffer.shape[-1]) elif kv_buffer.ndim == 3: kv_buffer = kv_buffer.view(-1, kv_buffer.shape[-1]) + if kv_indices.ndim == 3 and kv_indices.shape[1] == 1: + kv_indices = kv_indices.view(kv_indices.shape[0], kv_indices.shape[2]) if out is None: - out = torch.empty((queries, heads, width), dtype=q.dtype, device=q.device) + out = torch.empty(q.shape, dtype=torch.bfloat16, device=q.device) if lse is not None and not return_lse: raise ValueError("supplying lse requires return_lse=True") if return_lse and lse is None: - lse = torch.empty((queries, heads), dtype=torch.float32, device=q.device) - partial_o = partial_lse = None - if kv_splits > 1: - if workspace is None: - partial_o = torch.empty( - (queries, kv_splits, heads, 512), dtype=torch.float32, device=q.device - ) - partial_lse = torch.empty( - (queries, kv_splits, heads), dtype=torch.float32, device=q.device - ) - else: - partial_o, partial_lse = workspace - torch.ops.aiter.sparse_mla_bf16_fwd_out( + lse = torch.empty(q.shape[:2], dtype=torch.float32, device=q.device) + if workspace is None: + workspace = allocate_sparse_mla_fp8_workspace(q) + partial_o, partial_lse = workspace + torch.ops.aiter.sparse_mla_fp8_fwd_out( q, kv_buffer, - kv_indptr, kv_indices, out, lse, partial_o, partial_lse, float(softmax_scale), - kv_splits, - version, + float(q_scale), + float(kv_scale), ) return out, lse diff --git a/csrc/kernels/mla/fp8/kv_layout.cuh b/csrc/kernels/mla/fp8/kv_layout.cuh new file mode 100644 index 00000000000..65201be503f --- /dev/null +++ b/csrc/kernels/mla/fp8/kv_layout.cuh @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: MIT +#ifndef KV_XOR +#define KV_XOR 0 +#endif +#ifndef KV_PAD +#define KV_PAD 0 +#endif +static_assert(KV_PAD >= 0 && KV_PAD % 16 == 0); +constexpr int KV_STRIDE = 1024 + KV_PAD; +__device__ __forceinline__ int kv_xor(int row) +{ + if constexpr(KV_XOR == 1) + return ((row >> 2) & 3) * 16; + if constexpr(KV_XOR == 2) + return (((row >> 2) ^ (row >> 5)) & 3) * 16; + return (row & 3) * 16; +} +template +__device__ __forceinline__ int address(int row, int col) +{ + constexpr int stride = Rows == 128 ? KV_STRIDE : 1024; + int mask = Rows == 128 ? kv_xor(row) : (row & 3) * 16; + return ((col / 64) * (Rows / 16) + row / 16) * stride + (row % 16) * 64 + ((col % 64) ^ mask); +} diff --git a/csrc/kernels/mla/fp8/layout.cuh b/csrc/kernels/mla/fp8/layout.cuh new file mode 100644 index 00000000000..dc8c2bcf90b --- /dev/null +++ b/csrc/kernels/mla/fp8/layout.cuh @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: MIT +#pragma once +#ifdef __HIPCC__ +#define LAYOUT_INLINE __host__ __device__ __forceinline__ +#else +#define LAYOUT_INLINE inline __attribute__((always_inline)) +#endif +template +LAYOUT_INLINE int kv_address(int row, int col) +{ + static_assert(Pad >= 0 && Pad % 16 == 0); + return ((col / 64) * (Rows / 16) + row / 16) * (1024 + Pad) + (row % 16) * 64 + + ((col % 64) ^ ((row & 3) * 16)); +} +#undef LAYOUT_INLINE diff --git a/csrc/kernels/mla/fp8/online_common.cuh b/csrc/kernels/mla/fp8/online_common.cuh new file mode 100644 index 00000000000..7587acdbf58 --- /dev/null +++ b/csrc/kernels/mla/fp8/online_common.cuh @@ -0,0 +1,645 @@ +// SPDX-License-Identifier: MIT +#pragma once +#include "params.cuh" +#include "partial_io.cuh" +#include "q_encode.cuh" +#include "softmax_lane.cuh" +#define STREAM_ONLINE 1 +#define STREAM_QK_VGPR 1 +#define STREAM_PV_VGPR 1 +#define STREAM_PV_IMMEDIATE 1 +#define HEAD32_ORDER 2 +#define KV_CHUNK_PAD 16 +#define BANK_QK_ADDRESS 1 +#define BANK_PV_ADDRESS 0 +namespace sparse_mla_fp8::online { + +// SPDX-License-Identifier: MIT +// H64 model, CTA32 N128 FP8 MLA. Two waves per 16 heads share QK N and PV D. +#include "layout.cuh" +#include "order.cuh" +#ifndef KV_CHUNK_PAD +#define KV_CHUNK_PAD 0 +#endif +constexpr int KV_CHUNK_STRIDE = 1024 + KV_CHUNK_PAD; +#ifndef BANK_QK_ADDRESS +#define BANK_QK_ADDRESS 0 +#endif +#ifndef BANK_PV_ADDRESS +#define BANK_PV_ADDRESS 0 +#endif +#ifndef HEAD32_ORDER +#define HEAD32_ORDER 0 +#endif +using I8 = int __attribute__((ext_vector_type(8))); +using U4 = uint32_t __attribute__((ext_vector_type(4))); +using F4 = float __attribute__((ext_vector_type(4))); +using Params = sparse_mla_fp8::Params; +#ifndef STREAM_QK_VGPR +#define STREAM_QK_VGPR 0 +#endif +#ifndef STREAM_PV_VGPR +#define STREAM_PV_VGPR 0 +#endif +constexpr float LOG2E = 1.4426950408889634f, LN2 = .6931471805599453f; +__device__ __forceinline__ F4 mma(I8 a, I8 b, F4 c) +{ + return __builtin_amdgcn_mfma_scale_f32_16x16x128_f8f6f4(a, b, c, 0, 0, 0, 127, 0, 127); +} +// Same scaled MFMA and unit E8M0 exponent as the intrinsic. Explicit +// read/write VGPR accumulator prevents unnecessary AGPR/VGPR duplication. +__device__ __forceinline__ F4 mma_vgpr(I8 a, I8 b, F4 c) +{ + int scale = 127; + // Table 38 input RAW wait remains inside asm after compiler operand moves. + asm volatile( + "s_nop 1\n\tv_mfma_scale_f32_16x16x128_f8f6f4 %0, %1, %2, %0, %3, %3 op_sel_hi:[0,0,0]" + : "+v"(c) + : "v"(a), "v"(b), "v"(scale)); + return c; +} +__device__ __forceinline__ F4 mma_qk(I8 a, I8 b, F4 c) +{ + if constexpr(STREAM_QK_VGPR) + return mma_vgpr(a, b, c); + return mma(a, b, c); +} +__device__ __forceinline__ F4 mma_pv(I8 a, I8 b, F4 c) +{ + if constexpr(STREAM_PV_VGPR) + return mma_vgpr(a, b, c); + return mma(a, b, c); +} +__device__ __forceinline__ uint32_t pack4(F4 x) +{ +#pragma unroll + for(int j = 0; j < 4; ++j) + x[j] = fminf(448.f, fmaxf(-448.f, x[j])); + int w = __builtin_amdgcn_cvt_pk_fp8_f32(x[0], x[1], 0, 0); + return __builtin_amdgcn_cvt_pk_fp8_f32(x[2], x[3], w, 1); +} +// Each 1 KiB chunk is 16 rows x 64 bytes. XOR preserves 16-byte vectors. +template +__device__ __forceinline__ int address(int row, int col) +{ + return kv_address(row, col); +} +template +__device__ __forceinline__ I8 fragment(const uint8_t* base, int row, int col) +{ + I8 x; + const U4 a = *reinterpret_cast(base + address(row, col)); + const U4 b = *reinterpret_cast(base + address(row, col + 16)); +#pragma unroll + for(int j = 0; j < 4; ++j) + { + x[j] = a[j]; + x[j + 4] = b[j]; + } + return x; +} +// Constant tile offsets factored from the common per-lane base. +// This changes only integer address expressions, not the QK arithmetic tree. +template +__device__ __forceinline__ I8 bank_qfragment(const uint8_t* base, int row, int col) +{ + constexpr int off = (16 * K + N) * KV_CHUNK_STRIDE; + const U4 a = *reinterpret_cast(base + address<128>(row, col) + off); + const U4 b = *reinterpret_cast(base + address<128>(row, col + 16) + off); + I8 x; +#pragma unroll + for(int j = 0; j < 4; ++j) + { + x[j] = a[j]; + x[j + 4] = b[j]; + } + return x; +} +template +__device__ __forceinline__ void +bank_qk_step(const uint8_t* base, int row, int col, I8 qr, F4 (&score)[Fragments]) +{ + score[0] = mma_qk(bank_qfragment(base, row, col), qr, score[0]); + score[1] = mma_qk(bank_qfragment(base, row, col), qr, score[1]); + if constexpr(Fragments == 4) + { + score[2] = mma_qk(bank_qfragment(base, row, col), qr, score[2]); + score[3] = mma_qk(bank_qfragment(base, row, col), qr, score[3]); + } +} +__device__ __forceinline__ int p_address(int head, int key) +{ + return head * 128 + (key ^ ((head & 7) * 16)); +} +__device__ __forceinline__ I8 p_fragment(const uint8_t* base, int head, int col) +{ + I8 x; + U4 a = *reinterpret_cast(base + p_address(head, col)); + U4 b = *reinterpret_cast(base + p_address(head, col + 16)); +#pragma unroll + for(int j = 0; j < 4; ++j) + { + x[j] = a[j]; + x[j + 4] = b[j]; + } + return x; +} +__device__ __forceinline__ I8 v_fragment(const uint8_t* base, int col, int key) +{ + I8 x{}; +#pragma unroll + for(int j = 0; j < 8; ++j) + { + uint32_t w = 0; +#pragma unroll + for(int b = 0; b < 4; ++b) + w |= uint32_t(base[address<128>(key + j * 4 + b, col)]) << (8 * b); + x[j] = w; + } + return x; +} +using I2 = int __attribute__((ext_vector_type(2))); +__device__ __forceinline__ I8 v_fragment_tr(const uint8_t* base, int col_base) +{ + const int lane = threadIdx.x & 63; + I8 x; +#pragma unroll + for(int i = 0; i < 4; ++i) + { + int row = (lane / 16) * 32 + i * 8 + (lane % 16) / 2; + int col = col_base + (lane % 2) * 8; + uint32_t addr = uint32_t(reinterpret_cast(base + address<128>(row, col))); + I2 r; + asm volatile("ds_read_b64_tr_b8 %0, %1" : "=v"(r) : "v"(addr) : "memory"); + x[i * 2] = r[0]; + x[i * 2 + 1] = r[1]; + } + // Tie the loaded operand to the wait. A memory-only barrier can move + // after the last register-only MFMA when its results stay in registers. + asm volatile("s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)" : "+v"(x) : : "memory"); + return x; +} +__device__ __forceinline__ void direct_load(uint8_t* dest, const uint8_t* src) +{ + uint32_t lds = __builtin_amdgcn_readfirstlane(uint32_t(reinterpret_cast(dest))); + uint64_t addr = reinterpret_cast(src); + asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tglobal_load_lds_dwordx4 %1, off" + : + : "s"(lds), "v"(addr) + : "m0", "memory"); +} +#ifndef HEAD32_FUSE_SUM_P +#define HEAD32_FUSE_SUM_P 0 +#endif +#ifndef HEAD32_SUM_LOCAL +#define HEAD32_SUM_LOCAL 0 +#endif +#ifndef STREAM_ONLINE +#define STREAM_ONLINE 0 +#endif +#ifndef STREAM_PV_IMMEDIATE +#define STREAM_PV_IMMEDIATE 0 +#endif +#ifndef STREAM_PV_FENCE +#define STREAM_PV_FENCE 0 +#endif +#ifndef STREAM_BRANCHLESS +#define STREAM_BRANCHLESS 0 +#endif + +// Within the 8-row transpose piece, the XOR swizzle is unchanged. +// Fold the piece's row/chunk offset into the instruction immediate. +template +__device__ __forceinline__ I2 bank_pv_read(const uint8_t* base) +{ + const int lane = threadIdx.x & 63; + int row = (lane / 16) * 32 + (lane % 16) / 2; + int col = (Tile % 4) * 16 + (lane % 2) * 8; + uint32_t addr = uint32_t(reinterpret_cast(base + address<128>(row, col))); + constexpr int off = ((Tile / 4) * 8 + Piece / 2) * KV_CHUNK_STRIDE + (Piece % 2) * 512; + static_assert(off <= 65535); + I2 x; + asm volatile("ds_read_b64_tr_b8 %0, %1 offset:%2" : "=v"(x) : "v"(addr), "n"(off) : "memory"); + return x; +} +template +__device__ __forceinline__ I8 stream_vfragment(const uint8_t* base) +{ + if constexpr(!STREAM_PV_IMMEDIATE) + return v_fragment_tr(base, Tile * 16); + else + { + if constexpr(BankPV) + { + I8 x; + I2 a = bank_pv_read(base), b = bank_pv_read(base); + I2 c = bank_pv_read(base), d = bank_pv_read(base); + x[0] = a[0]; + x[1] = a[1]; + x[2] = b[0]; + x[3] = b[1]; + x[4] = c[0]; + x[5] = c[1]; + x[6] = d[0]; + x[7] = d[1]; + asm volatile("s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)" : "+v"(x) : : "memory"); + return x; + } + const int lane = threadIdx.x & 63; + I8 x; +#pragma unroll + for(int i = 0; i < 4; ++i) + { + int row = (lane / 16) * 32 + i * 8 + (lane % 16) / 2; + int col = (Tile % 4) * 16 + (lane % 2) * 8; + uint32_t addr = uint32_t(reinterpret_cast(base + address<128>(row, col))); + I2 r; + asm volatile("ds_read_b64_tr_b8 %0, %1 offset:%2" + : "=v"(r) + : "v"(addr), "n"((Tile / 4) * 8 * KV_CHUNK_STRIDE) + : "memory"); + x[i * 2] = r[0]; + x[i * 2 + 1] = r[1]; + } + asm volatile("s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)" : "+v"(x) : : "memory"); + return x; + } +} + +template +__device__ __forceinline__ void stream_pv_tile(const uint8_t* kvs, + I8 prob, + F4 (&merged)[32 / WavesPerHead], + int iteration, + float alpha, + float beta, + float scale, + int half) +{ + I8 v = stream_vfragment(kvs + half * (64 / WavesPerHead) * KV_CHUNK_STRIDE); + if constexpr(STREAM_ONLINE) + { + merged[Tile] = mma_pv(v, prob, merged[Tile] * alpha); + if constexpr(STREAM_PV_FENCE) + asm volatile("" : "+v"(merged[Tile]) : : "memory"); + return; + } + F4 out = mma_pv(v, prob, F4{}); + out *= scale; + if constexpr(STREAM_BRANCHLESS) + { +#pragma unroll + for(int j = 0; j < 4; ++j) + merged[Tile][j] = fmaf(out[j], iteration == 0 ? 1.f : beta, merged[Tile][j] * alpha); + } + else + { + if(iteration == 0) + merged[Tile] = out; + else + { +#pragma unroll + for(int j = 0; j < 4; ++j) + merged[Tile][j] = fmaf(out[j], beta, merged[Tile][j] * alpha); + } + } + if constexpr(STREAM_PV_FENCE) + asm volatile("" : "+v"(merged[Tile]) : : "memory"); +} + +// Compile-time body choices. The head/K geometries change data reuse, not the +template +__device__ __forceinline__ int p_address(int head, int key) +{ + return head * BlockK + (key ^ ((head & (BlockK / 16 - 1)) * 16)); +} +template +__device__ __forceinline__ I8 p_fragment(const uint8_t* ps, int head, int col) +{ + I8 x{}; + if(col < BlockK) + { + U4 a = *reinterpret_cast(ps + p_address(head, col)); + U4 b = *reinterpret_cast(ps + p_address(head, col + 16)); +#pragma unroll + for(int j = 0; j < 4; ++j) + { + x[j] = a[j]; + x[j + 4] = b[j]; + } + } + return x; +} +template +__device__ __forceinline__ I8 k_fragment(const uint8_t* base, int row, int col) +{ + constexpr int offset = (K * (BlockK / 8) + N) * KV_CHUNK_STRIDE; + const U4 a = *reinterpret_cast(base + address(row, col) + offset); + const U4 b = *reinterpret_cast(base + address(row, col + 16) + offset); + I8 x; +#pragma unroll + for(int j = 0; j < 4; ++j) + { + x[j] = a[j]; + x[j + 4] = b[j]; + } + return x; +} +template +__device__ __forceinline__ void +qk_step(const uint8_t* kvs, int row, int col, I8 q, F4 (&score)[Fragments]) +{ + score[N] = mma_qk(k_fragment(kvs, row, col), q, score[N]); + if constexpr(N + 1 < Fragments) + qk_step(kvs, row, col, q, score); +} +template +__device__ __forceinline__ I8 q_load(Params p, int64_t off) +{ + I8 x; + U4 a, b; + if constexpr(BF16) + { + const auto* src = reinterpret_cast(p.q) + off; + a = q_encode::pack16(src, p.q_scale); + b = q_encode::pack16(src + 16, p.q_scale); + } + else + { + a = *reinterpret_cast(p.q + off); + b = *reinterpret_cast(p.q + off + 16); + } +#pragma unroll + for(int j = 0; j < 4; ++j) + { + x[j] = a[j]; + x[j + 4] = b[j]; + } + return x; +} +template +__device__ __forceinline__ int q_address(int row, int col) +{ + constexpr int rank = __builtin_popcount(Mask & ((1 << K) - 1)); + return rank * HeadRows * 128 + kv_address(row, col); +} +template +__device__ __forceinline__ I8 q_fragment(Params p, + const I8 (&retained)[4], + uint8_t* qs, + int64_t off, + int iteration, + int row, + int col, + int half) +{ + if constexpr(!(Mask & (1 << K))) + return retained[K]; + else + { + const int low = q_address(row, col), + high = q_address(row, col + 16); + if(iteration == 0) + { + I8 x = q_load(p, off + K * 128); + if(__builtin_amdgcn_readfirstlane(half) == 0) + { + U4 a, b; +#pragma unroll + for(int j = 0; j < 4; ++j) + { + a[j] = x[j]; + b[j] = x[j + 4]; + } + *reinterpret_cast(qs + low) = a; + *reinterpret_cast(qs + high) = b; + } + return x; + } + U4 a, b; + uint32_t addr0 = uint32_t(reinterpret_cast(qs + low)), + addr1 = uint32_t(reinterpret_cast(qs + high)); + asm volatile("ds_read_b128 %0, %2\n\tds_read_b128 %1, %3" + : "=&v"(a), "=&v"(b) + : "v"(addr0), "v"(addr1) + : "memory"); + // Tie the loaded operands to the wait; the old qreload failure showed + // that a memory-only wait is insufficient for register-only MFMA use. + asm volatile("s_waitcnt lgkmcnt(0)" : "+v"(a), "+v"(b) : : "memory"); + I8 x; +#pragma unroll + for(int j = 0; j < 4; ++j) + { + x[j] = a[j]; + x[j + 4] = b[j]; + } + return x; + } +} +// No implicit wait: callers can batch or pipeline several four-read operands. +template +__device__ __forceinline__ I8 v_read(const uint8_t* base) +{ + const int lane = threadIdx.x & 63; + I8 x{}; + if constexpr(BlockK == 128) + { +#pragma unroll + for(int i = 0; i < 4; ++i) + { + int row = (lane / 16) * 32 + i * 8 + (lane % 16) / 2; + int col = (Tile % 4) * 16 + (lane % 2) * 8; + uint32_t addr = uint32_t(reinterpret_cast(base + address(row, col))); + I2 pair; + asm volatile("ds_read_b64_tr_b8 %0, %1 offset:%2" + : "=v"(pair) + : "v"(addr), "n"((Tile / 4) * (BlockK / 16) * KV_CHUNK_STRIDE) + : "memory"); + x[i * 2] = pair[0]; + x[i * 2 + 1] = pair[1]; + } + } + else if(lane < 32) + { +#pragma unroll + for(int i = 0; i < 4; ++i) + { + int row = (lane / 16) * 32 + i * 8 + (lane % 16) / 2; + int col = (Tile % 4) * 16 + (lane % 2) * 8; + uint32_t addr = uint32_t(reinterpret_cast(base + address(row, col))); + I2 pair; + asm volatile("ds_read_b64_tr_b8 %0, %1 offset:%2" + : "=v"(pair) + : "v"(addr), "n"((Tile / 4) * (BlockK / 16) * KV_CHUNK_STRIDE) + : "memory"); + x[i * 2] = pair[0]; + x[i * 2 + 1] = pair[1]; + } + } + return x; +} +template +__device__ __forceinline__ void pv(const uint8_t* kvs, I8 prob, F4 (&merged)[Tiles], float alpha) +{ + I8 a = v_read(kvs); + if constexpr(Mode == 1) + { + static_assert(Tile + 1 < Tiles); + I8 b = v_read(kvs); + asm volatile("s_waitcnt lgkmcnt(0)" : "+v"(a), "+v"(b) : : "memory"); + merged[Tile] = mma_pv(a, prob, merged[Tile] * alpha); + merged[Tile + 1] = mma_pv(b, prob, merged[Tile + 1] * alpha); + if constexpr(Tile + 2 < Tiles) + pv(kvs, prob, merged, alpha); + } + else + { + asm volatile("s_waitcnt lgkmcnt(0)" : "+v"(a) : : "memory"); + merged[Tile] = mma_pv(a, prob, merged[Tile] * alpha); + if constexpr(Tile + 1 < Tiles) + pv(kvs, prob, merged, alpha); + } +} +template +__device__ __forceinline__ void +pv_pipeline(const uint8_t* kvs, I8 prob, F4 (&merged)[Tiles], float alpha, I8 current) +{ + if constexpr(Tile + 1 < Tiles) + { + I8 next = v_read(kvs); + // Only the next operand's four ordered LDS reads may remain pending. + asm volatile("s_waitcnt lgkmcnt(4)" : "+v"(current) : : "memory"); + merged[Tile] = mma_pv(current, prob, merged[Tile] * alpha); + pv_pipeline(kvs, prob, merged, alpha, next); + } + else + { + asm volatile("s_waitcnt lgkmcnt(0)" : "+v"(current) : : "memory"); + merged[Tile] = mma_pv(current, prob, merged[Tile] * alpha); + } +} + +// K64 PV consumes two real K32 operands; no zero-padded K128 work. +// Clang official signature is (int64_t,int64_t,F4,const int,const int,const int). +// Keep an actual builtin expression for compiler signature/type verification. +// build01 emitted this path; H32 selected costly AGPR transfer pairs, preserved. +__device__ __forceinline__ F4 mma32_signature(I2 a, I2 b, F4 accum) +{ + return __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8( + __builtin_bit_cast(int64_t, a), __builtin_bit_cast(int64_t, b), accum, 0, 0, 0); +} +__device__ __forceinline__ F4 mma32(I2 a, I2 b, F4 accum) +{ + asm volatile("v_mfma_f32_16x16x32_fp8_fp8 %0, %1, %2, %0" : "+v"(accum) : "v"(a), "v"(b)); + return accum; +} +template +__device__ __forceinline__ I2 p32(const uint8_t* ps, int head, int group) +{ + return *reinterpret_cast(ps + p_address<64>(head, Sub * 32 + group * 8)); +} +template +__device__ __forceinline__ I2 v32(const uint8_t* base) +{ + const int lane = threadIdx.x & 63; + const int row = Sub * 32 + (lane / 16) * 8 + (lane % 16) / 2; + const int col = (Tile % 4) * 16 + (lane % 2) * 8; + uint32_t addr = uint32_t(reinterpret_cast(base + address<64>(row, col))); + I2 x; + asm volatile("ds_read_b64_tr_b8 %0, %1 offset:%2" + : "=v"(x) + : "v"(addr), "n"((Tile / 4) * 4 * KV_CHUNK_STRIDE) + : "memory"); + return x; +} +template +__device__ __forceinline__ void +pv32(const uint8_t* kvs, I2 p0, I2 p1, F4 (&merged)[Tiles], float alpha) +{ + I2 a = v32(kvs), b = v32(kvs); + if constexpr(Pair) + { + static_assert(Tile + 1 < Tiles); + I2 c = v32(kvs), d = v32(kvs); + asm volatile("s_waitcnt lgkmcnt(0)" : "+v"(a), "+v"(b), "+v"(c), "+v"(d) : : "memory"); + merged[Tile] = mma32(b, p1, mma32(a, p0, merged[Tile] * alpha)); + merged[Tile + 1] = mma32(d, p1, mma32(c, p0, merged[Tile + 1] * alpha)); + if constexpr(Tile + 2 < Tiles) + pv32(kvs, p0, p1, merged, alpha); + } + else + { + asm volatile("s_waitcnt lgkmcnt(0)" : "+v"(a), "+v"(b) : : "memory"); + merged[Tile] = mma32(b, p1, mma32(a, p0, merged[Tile] * alpha)); + if constexpr(Tile + 1 < Tiles) + pv32(kvs, p0, p1, merged, alpha); + } +} + +// Stage one future K128 tile into its own KV and selected-ID buffers. +// Separate current/next masks survive until their matching tile is consumed. +// The caller waits for VMEM and all current readers before reusing a buffer. +template +__device__ __forceinline__ void +stage_async(Params p, uint8_t* kv, int* selected, int start, int end, int tile) +{ + static_assert((Waves == 8 || Waves == 4) && BlockK == 128); + const int tid = threadIdx.x, lane = tid & 63, wave = tid / 64; + if(tid < BlockK) + { + const int at = start + tile * BlockK + tid; + const int id = at < end ? p.indices[at] : -1; + selected[tid] = (id >= 0 && id < p.slots) ? id : -1; + } + __syncthreads(); + const int id0 = selected[(wave % (BlockK / 16)) * 16 + lane / 4]; + const int id1 = Waves == 4 ? selected[(wave + 4) * 16 + lane / 4] : -1; +#pragma unroll + for(int i = 0; i < (BlockK / 2) / Waves; ++i) + { + const int chunk = wave + i * Waves, r = (chunk % (BlockK / 16)) * 16 + lane / 4; + const int col = (chunk / (BlockK / 16)) * 64 + ((lane % 4) * 16 ^ ((r & 3) * 16)); + const int id = (Waves == 4 && (i & 1)) ? id1 : id0; + if(id >= 0) + direct_load(kv + chunk * KV_CHUNK_STRIDE, p.kv + int64_t(id) * 512 + col); + else + *reinterpret_cast(kv + chunk * KV_CHUNK_STRIDE + lane * 16) = U4{}; + } +} + +template +__device__ __forceinline__ float domain_exp2(float x) +{ + if constexpr(Gradual) + { + const bool tiny = x < -126.f; + const float shifted = x + (tiny ? 64.f : 0.f); + return __builtin_amdgcn_exp2f(shifted) * (tiny ? 0x1p-64f : 1.f); + } + else + return __builtin_amdgcn_exp2f(x); +} +template +__device__ __forceinline__ void +stage_iddirect(Params p, uint8_t* kv, int* selected, int start, int end, int tile) +{ + static_assert(Waves == 8 && BlockK == 128); + const int lane = threadIdx.x & 63, wave = threadIdx.x / 64; + const int row = wave * 16 + lane / 4; + const int at = start + tile * BlockK + row; + const int raw = at < end ? p.indices[at] : -1; + const int id = (raw >= 0 && raw < p.slots) ? raw : -1; + if((lane & 3) == 0) + selected[row] = id; +#pragma unroll + for(int i = 0; i < (BlockK / 2) / Waves; ++i) + { + const int chunk = wave + i * Waves; + const int col = i * 64 + ((lane % 4) * 16 ^ ((row & 3) * 16)); + if(id >= 0) + direct_load(kv + chunk * KV_CHUNK_STRIDE, p.kv + int64_t(id) * 512 + col); + else + *reinterpret_cast(kv + chunk * KV_CHUNK_STRIDE + lane * 16) = U4{}; + } +} +} // namespace sparse_mla_fp8::online diff --git a/csrc/kernels/mla/fp8/order.cuh b/csrc/kernels/mla/fp8/order.cuh new file mode 100644 index 00000000000..ee97fc855ac --- /dev/null +++ b/csrc/kernels/mla/fp8/order.cuh @@ -0,0 +1,38 @@ +// SPDX-License-Identifier: MIT +#pragma once +#include +#ifdef __HIPCC__ +#define ORDER_INLINE __host__ __device__ __forceinline__ +#else +#define ORDER_INLINE inline __attribute__((always_inline)) +#endif +struct OrderIds +{ + uint32_t query, part, head; +}; +template +ORDER_INLINE OrderIds order_ids(uint32_t x, uint32_t y, uint32_t z, uint32_t q) +{ + static_assert(Order <= 3); + if constexpr(Order == 0) + return {x, y, z}; + uint32_t work = x + q * (y + S * z); + if constexpr(Order >= 2) + { + const uint32_t total = q * S * Heads, per = (total + 7) / 8; + const uint32_t tall = (total % 8) ? total % 8 : 8; + const uint32_t xcd = work % 8, local = work / 8; + work = xcd < tall ? xcd * per + local : tall * per + (xcd - tall) * (per - 1) + local; + } + if constexpr(Order == 3 && S > 1) + { + const uint32_t chunk = (q + 7) / 8, per = chunk * S * Heads; + const uint32_t block = work / per, i = work % per; + const uint32_t remaining = q - block * chunk; + const uint32_t rows = remaining < chunk ? remaining : chunk; + const uint32_t r = i % (rows * Heads); + return {block * chunk + r / Heads, i / (rows * Heads), r % Heads}; + } + return {work / (S * Heads), (work / Heads) % S, work % Heads}; +} +#undef ORDER_INLINE diff --git a/csrc/kernels/mla/fp8/pad16_factored.cuh b/csrc/kernels/mla/fp8/pad16_factored.cuh new file mode 100644 index 00000000000..03da63d60ea --- /dev/null +++ b/csrc/kernels/mla/fp8/pad16_factored.cuh @@ -0,0 +1,109 @@ +// SPDX-License-Identifier: MIT +#pragma once +__host__ __device__ constexpr int padded_address(int row, int col) +{ + return ((col / 64) * 8 + row / 16) * 1040 + (row % 16) * 64 + ((col % 64) ^ ((row & 3) * 16)); +} +__host__ __device__ constexpr int padded_pv_base(int lane, int wave, int tile) +{ + return (16 * wave + 2 * (lane / 16)) * 1040 + ((lane % 16) / 2) * 64 + + ((8 * (lane & 1)) ^ (16 * (((lane % 16) / 2) & 3)) ^ (16 * (tile & 3))); +} +__host__ __device__ constexpr int padded_pv_offset(int tile, int piece) +{ + return (8 * (tile / 4) + piece / 2) * 1040 + 512 * (piece & 1); +} +__host__ __device__ constexpr int padded_qk_base(int lane, int wave, int vec) +{ + return (8 * ((lane / 16) / 2) + 2 * wave) * 1040 + (lane % 16) * 64 + + ((((lane / 16) % 2) * 32 + vec) ^ (((lane % 16) & 3) * 16)); +} +template +__device__ __forceinline__ I8 padded_k_fragment(const uint8_t* kvs, int lane, int wave) +{ + const uint8_t* a = kvs + padded_qk_base(lane, wave, 0); + const uint8_t* b = kvs + padded_qk_base(lane, wave, 16); + constexpr int Offset = (16 * K + Half) * 1040; + U4 lo = *reinterpret_cast(a + Offset), hi = *reinterpret_cast(b + Offset); + I8 out; +#pragma unroll + for(int j = 0; j < 4; ++j) + { + out[j] = lo[j]; + out[j + 4] = hi[j]; + } + return out; +} +template +__device__ __forceinline__ void +padded_qk_step(const uint8_t* qs, const uint8_t* kvs, int head, int group, int wave, F4& a, F4& b) +{ + const I8 q = fragment<16>(qs, head, K * 128 + group * 32); + const int lane = threadIdx.x & 63; + a = mma(padded_k_fragment(kvs, lane, wave), q, a); + b = mma(padded_k_fragment(kvs, lane, wave), q, b); + if constexpr(K < 3) + padded_qk_step(qs, kvs, head, group, wave, a, b); +} +template +__device__ __forceinline__ I8 padded_v_fragment(const uint8_t* kvs, int wave) +{ + const int lane = threadIdx.x & 63; + const uint32_t base = + uint32_t(reinterpret_cast(kvs + padded_pv_base(lane, wave, Tile))); + I2 r0, r1, r2, r3; + asm volatile("ds_read_b64_tr_b8 %0, %1 offset:%2" + : "=v"(r0) + : "v"(base), "n"(padded_pv_offset(Tile, 0)) + : "memory"); + asm volatile("ds_read_b64_tr_b8 %0, %1 offset:%2" + : "=v"(r1) + : "v"(base), "n"(padded_pv_offset(Tile, 1)) + : "memory"); + asm volatile("ds_read_b64_tr_b8 %0, %1 offset:%2" + : "=v"(r2) + : "v"(base), "n"(padded_pv_offset(Tile, 2)) + : "memory"); + asm volatile("ds_read_b64_tr_b8 %0, %1 offset:%2" + : "=v"(r3) + : "v"(base), "n"(padded_pv_offset(Tile, 3)) + : "memory"); + I8 x = {r0[0], r0[1], r1[0], r1[1], r2[0], r2[1], r3[0], r3[1]}; + asm volatile("s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)" : "+v"(x) : : "memory"); + return x; +} +template +__device__ __forceinline__ void padded_pv_step(const uint8_t* kvs, + I8 prob, + int wave, + int iteration, + float alpha, + float beta, + float scale, + F4 (&merged)[8]) +{ + I8 v = padded_v_fragment(kvs, wave); + F4 out = mma(v, prob, F4{}); + out *= scale; + if(iteration == 0) + merged[Tile] = out; + else + { +#pragma unroll + for(int j = 0; j < 4; ++j) + merged[Tile][j] = fmaf(out[j], beta, merged[Tile][j] * alpha); + } + if constexpr(Tile < 7) + padded_pv_step(kvs, prob, wave, iteration, alpha, beta, scale, merged); +} + +// FP32 partial normalization. This deliberately changes FP8 P rounding. +template +__device__ __forceinline__ void +padded_online_pv_step(const uint8_t* kvs, I8 prob, int wave, float alpha, F4 (&merged)[8]) +{ + I8 v = padded_v_fragment(kvs, wave); + merged[Tile] = online_pv_mma(v, prob, merged[Tile] * alpha); + if constexpr(Tile < 7) + padded_online_pv_step(kvs, prob, wave, alpha, merged); +} diff --git a/csrc/kernels/mla/fp8/params.cuh b/csrc/kernels/mla/fp8/params.cuh new file mode 100644 index 00000000000..90125378740 --- /dev/null +++ b/csrc/kernels/mla/fp8/params.cuh @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: MIT +#pragma once +#include +#include +#include +namespace sparse_mla_fp8 { +struct Params +{ + const uint8_t *q, *kv; + const int32_t* indices; + uint16_t* partial_o; + float* partial_lse; + uint8_t* out; + float* lse; + int queries, slots, capacity; + float q_scale, kv_scale, attention_scale; + static constexpr float p_scale = 1.f / 256.f; + int64_t q_stride_row, q_stride_head; + bool bf16_q; +}; +inline int split_count(int queries) +{ + if(queries <= 6) + return 17; + if(queries <= 12) + return 9; + if(queries <= 21) + return 6; + if(queries <= 23) + return 4; + if(queries <= 42) + return 6; + if(queries <= 64) + return 4; + if(queries <= 128) + return 2; + return 1; +} +hipError_t launch_small(Params p, hipStream_t stream); +hipError_t launch_prefetch(Params p, hipStream_t stream); +hipError_t launch_rotating(Params p, hipStream_t stream); +hipError_t launch_streaming(Params p, hipStream_t stream); +} // namespace sparse_mla_fp8 diff --git a/csrc/kernels/mla/fp8/partial_io.cuh b/csrc/kernels/mla/fp8/partial_io.cuh new file mode 100644 index 00000000000..ca45a9076ef --- /dev/null +++ b/csrc/kernels/mla/fp8/partial_io.cuh @@ -0,0 +1,34 @@ +// SPDX-License-Identifier: MIT +// Only the split partial storage changes; all arithmetic remains FP32. +using PF4 = float __attribute__((ext_vector_type(4))); +using PI2 = int __attribute__((ext_vector_type(2))); +template +__device__ __forceinline__ void partial_store(uint16_t* ptr, PF4 x) +{ + uint32_t lo, hi; + asm("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(lo) : "v"(x[0]), "v"(x[1])); + asm("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(hi) : "v"(x[2]), "v"(x[3])); + PI2 v{int(lo), int(hi)}; + if constexpr(NT) + __builtin_nontemporal_store(v, reinterpret_cast(ptr)); + else + *reinterpret_cast(ptr) = v; +} +template +__device__ __forceinline__ PF4 partial_load(const uint16_t* ptr) +{ + const PI2* p = reinterpret_cast(ptr); + PI2 v; + if constexpr(NT) + v = __builtin_nontemporal_load(p); + else + v = *p; + PF4 x; +#pragma unroll + for(int i = 0; i < 4; ++i) + { + uint32_t bits = (uint32_t(v[i / 2]) >> ((i % 2) * 16)) << 16; + x[i] = __builtin_bit_cast(float, bits); + } + return x; +} diff --git a/csrc/kernels/mla/fp8/partial_layout.cuh b/csrc/kernels/mla/fp8/partial_layout.cuh new file mode 100644 index 00000000000..5be972686bf --- /dev/null +++ b/csrc/kernels/mla/fp8/partial_layout.cuh @@ -0,0 +1,34 @@ +// SPDX-License-Identifier: MIT +// Physical layout, preserving contiguous four-float vectors. Shared by S1/S2. +template +__host__ __device__ __forceinline__ int64_t partial_address(int64_t row, int col) +{ + static_assert(Tile == 0 || Tile == 16 || Tile == 32); + if constexpr(Tile == 0) + return row * 512 + col; + else + return (row & ~int64_t(15)) * 512 + int64_t(col / Tile) * (16 * Tile) + (row & 15) * Tile + + col % Tile; +} +struct CtaTile +{ + int q, b, c; +}; +template +__host__ __device__ __forceinline__ uint32_t mapped_cta(uint32_t pid, uint32_t total) +{ + if constexpr(!Remap) + return pid; + // Balanced eight-XCD permutation also covers odd Q * S17, with no padding. + const uint32_t quotient = total / 8, remainder = total % 8, xcd = pid % 8, local = pid / 8; + return xcd * quotient + (xcd < remainder ? xcd : remainder) + local; +} +template +__device__ __forceinline__ CtaTile cta_tile() +{ + if constexpr(!Remap) + return {int(blockIdx.x), int(blockIdx.y), int(blockIdx.z)}; + const uint32_t pid = blockIdx.x + gridDim.x * (blockIdx.y + gridDim.y * blockIdx.z); + const uint32_t j = mapped_cta(pid, gridDim.x * B * C); + return {int(j / (B * C)), int((j / C) % B), int(j % C)}; +} diff --git a/csrc/kernels/mla/fp8/prefetch.cu b/csrc/kernels/mla/fp8/prefetch.cu new file mode 100644 index 00000000000..4cd5ec67500 --- /dev/null +++ b/csrc/kernels/mla/fp8/prefetch.cu @@ -0,0 +1,629 @@ +// SPDX-License-Identifier: MIT +#include "online_common.cuh" +#include "reduce.cuh" +#define PF_BUFFERS 2 +#define QRAW_MODE 1 +#define PACKED_MUL 1 +#define REDUCER_FAST_EXP 1 +#define LSE_STRIDE 9 +#define LSE_VECTOR 0 +namespace sparse_mla_fp8::online { +template +__global__ __launch_bounds__(64) void reduce_compact(Params p) +{ + constexpr int V = Cols / 16; + static_assert(!SingleHead || (Cols == 16 || Cols == 32)); + int lane = threadIdx.x, head = SingleHead ? int(blockIdx.y) : int(blockIdx.y) * 16 + lane / 4; + int col = SingleHead ? int(blockIdx.z) * (256 * V) + lane * 4 + : int(blockIdx.z) * Cols + (lane % 4) * 4; + int64_t row = int64_t(blockIdx.x) * 64 + head; + float accum[V][4]{}; + float z = 0.f, m = -INFINITY; + if constexpr(Mode == 0) + { +#pragma unroll + for(int s = 0; s < Splits; ++s) + { + float l = p.partial_lse[(int64_t(blockIdx.x) * 64 + head) * LSE_STRIDE + s]; + float d = exp2f(-fabsf((m == -INFINITY && l == -INFINITY) ? 0.f : m - l) * LOG2E); + float alpha = m == -INFINITY ? 0.f : (m >= l ? 1.f : d), + beta = l == -INFINITY ? 0.f : (l >= m ? 1.f : d); + z = z * alpha + beta; + m = fmaxf(m, l); +#pragma unroll + for(int v = 0; v < V; ++v) + { + F4 x = partial_load(p.partial_o + + ((int64_t(blockIdx.x) * Splits + s) * 64 + head) * 512 + col + + v * (SingleHead ? 256 : 16)); +#pragma unroll + for(int j = 0; j < 4; ++j) + accum[v][j] = fmaf(x[j], beta, accum[v][j] * alpha); + } + } + } + else if constexpr(Mode == 1) + { + float ls[Splits]; +#if LSE_VECTOR + static_assert(Splits == 9 && LSE_STRIDE == 16); + const float* lp = p.partial_lse + (int64_t(blockIdx.x) * 64 + head) * LSE_STRIDE; + const F4 a = *reinterpret_cast(lp), b = *reinterpret_cast(lp + 4); +#pragma unroll + for(int j = 0; j < 4; ++j) + { + ls[j] = a[j]; + ls[4 + j] = b[j]; + } + ls[8] = lp[8]; +#endif +#pragma unroll + for(int s = 0; s < Splits; ++s) + { +#if !LSE_VECTOR + ls[s] = p.partial_lse[(int64_t(blockIdx.x) * 64 + head) * LSE_STRIDE + s]; +#endif + m = fmaxf(m, ls[s]); + } +#if REDUCER_FAST_EXP + // A single head makes every LSE and this decision wave-uniform. + // For a finite spread <=80, every exp2 argument lies in [-116,0]. + // All other inputs retain the original gradual-underflow implementation. + float min_l = INFINITY; +#pragma unroll + for(int s = 0; s < Splits; ++s) + min_l = fminf(min_l, ls[s]); + if(__builtin_amdgcn_readfirstlane(int(SingleHead && m - min_l <= 80.f))) + { +#pragma unroll + for(int s = 0; s < Splits; ++s) + { + float w = __builtin_amdgcn_exp2f((ls[s] - m) * LOG2E); + z += w; +#pragma unroll + for(int v = 0; v < V; ++v) + { + F4 x = partial_load(p.partial_o + + ((int64_t(blockIdx.x) * Splits + s) * 64 + head) * 512 + + col + v * (SingleHead ? 256 : 16)); +#pragma unroll + for(int j = 0; j < 4; ++j) + x[j] = w > 0 ? x[j] : 0.f; + using Pair = float __attribute__((ext_vector_type(2))); + Pair factor{w, w}, lo, hi; + Pair xlo{x[0], x[1]}, xhi{x[2], x[3]}, alo{accum[v][0], accum[v][1]}, + ahi{accum[v][2], accum[v][3]}; + asm("v_pk_fma_f32 %0, %1, %2, %3" : "=v"(lo) : "v"(xlo), "v"(factor), "v"(alo)); + asm("v_pk_fma_f32 %0, %1, %2, %3" : "=v"(hi) : "v"(xhi), "v"(factor), "v"(ahi)); + accum[v][0] = lo[0]; + accum[v][1] = lo[1]; + accum[v][2] = hi[0]; + accum[v][3] = hi[1]; + } + } + } + else +#endif + { +#pragma unroll + for(int s = 0; s < Splits; ++s) + { + float w = m == -INFINITY ? 0.f : exp2f((ls[s] - m) * LOG2E); + z += w; +#pragma unroll + for(int v = 0; v < V; ++v) + { + F4 x = partial_load(p.partial_o + + ((int64_t(blockIdx.x) * Splits + s) * 64 + head) * 512 + + col + v * (SingleHead ? 256 : 16)); +#pragma unroll + for(int j = 0; j < 4; ++j) + x[j] = w > 0 ? x[j] : 0.f; + using Pair = float __attribute__((ext_vector_type(2))); + Pair factor{w, w}, lo, hi; + Pair xlo{x[0], x[1]}, xhi{x[2], x[3]}, alo{accum[v][0], accum[v][1]}, + ahi{accum[v][2], accum[v][3]}; + asm("v_pk_fma_f32 %0, %1, %2, %3" : "=v"(lo) : "v"(xlo), "v"(factor), "v"(alo)); + asm("v_pk_fma_f32 %0, %1, %2, %3" : "=v"(hi) : "v"(xhi), "v"(factor), "v"(ahi)); + accum[v][0] = lo[0]; + accum[v][1] = lo[1]; + accum[v][2] = hi[0]; + accum[v][3] = hi[1]; + } + } + } + } + if constexpr(Mode >= 2) + { +#pragma unroll + for(int s = 0; s < Splits; ++s) + m = fmaxf(m, p.partial_lse[(int64_t(blockIdx.x) * 64 + head) * LSE_STRIDE + s]); +#pragma unroll(Mode == 3 ? 4 : 1) + for(int s = 0; s < Splits; ++s) + { + float l = p.partial_lse[(int64_t(blockIdx.x) * 64 + head) * LSE_STRIDE + s]; + float w = m == -INFINITY ? 0.f : exp2f((l - m) * LOG2E); + z += w; +#pragma unroll + for(int v = 0; v < V; ++v) + { + F4 x = partial_load(p.partial_o + + ((int64_t(blockIdx.x) * Splits + s) * 64 + head) * 512 + col + + v * (SingleHead ? 256 : 16)); +#pragma unroll + for(int j = 0; j < 4; ++j) + accum[v][j] = fmaf(w > 0 ? x[j] : 0.f, w, accum[v][j]); + } + } + } + float inv = z > 0 ? 1.f / (z * (OutputFP8 ? 1.f : 1.f)) : 0.f; +#pragma unroll + for(int v = 0; v < V; ++v) + { + F4 x; +#if PACKED_MUL + using Pair = float __attribute__((ext_vector_type(2))); + Pair lo, hi, factor{inv, inv}, alo{accum[v][0], accum[v][1]}, ahi{accum[v][2], accum[v][3]}; + asm("v_pk_mul_f32 %0, %1, %2" : "=v"(lo) : "v"(alo), "v"(factor)); + asm("v_pk_mul_f32 %0, %1, %2" : "=v"(hi) : "v"(ahi), "v"(factor)); + x = F4{lo[0], lo[1], hi[0], hi[1]}; +#else + for(int j = 0; j < 4; ++j) + x[j] = accum[v][j] * inv; +#endif + for(int j = 0; j < 4; ++j) {} + if constexpr(OutputFP8) + *reinterpret_cast(p.out + row * 512 + col + v * (SingleHead ? 256 : 16)) = + pack4(x); + else + { + uint32_t lo, hi; + asm("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(lo) : "v"(x[0]), "v"(x[1])); + asm("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(hi) : "v"(x[2]), "v"(x[3])); + *reinterpret_cast(reinterpret_cast(p.out) + row * 512 + col + + v * (SingleHead ? 256 : 16)) = I2{int(lo), int(hi)}; + } + } + if(p.lse && blockIdx.z == 0 && (SingleHead ? lane == 0 : lane % 4 == 0)) + p.lse[row] = z > 0 ? m + log2f(z) * LN2 : -INFINITY; +} +template +__global__ __launch_bounds__(Waves * 64) void prefetch_attention(Params p) +{ + + constexpr int WPH = Waves * 16 / HeadRows, Fragments = BlockK / (WPH * 16), PVTiles = 32 / WPH; + static_assert((HeadRows == 32 || HeadRows == 64) && (WPH == 1 || WPH == 2 || WPH == 4)); + static_assert(BlockK == 128 || (BlockK == 64 && (Waves == 4 || Waves == 8))); + __shared__ __align__(16) uint8_t kv_storage[PF_BUFFERS][(BlockK / 2) * KV_CHUNK_STRIDE]; + __shared__ __align__(16) uint8_t ps[HeadRows * BlockK]; + __shared__ __align__(16) uint8_t qs[HeadRows * 512]; + __shared__ float maxs[WPH >= 2 ? Waves * 16 : 1], sums[WPH >= 2 ? Waves * 16 : 1]; + __shared__ int selected_storage[PF_BUFFERS][BlockK]; + const int tid = threadIdx.x, lane = tid & 63, wave = tid / 64, head = lane & 15, + group = lane / 16; + const auto work = order_ids( + blockIdx.x, blockIdx.y, blockIdx.z, Compact ? 8 : p.queries); + const int query = work.query, part = work.part, half = wave % WPH, hgroup = wave / WPH; + const int hbase = work.head * HeadRows + hgroup * 16; + const int start = FixedCapacity ? query * FixedCapacity : (query * p.capacity), + end = FixedCapacity ? start + FixedCapacity : ((query + 1) * p.capacity); + const int total_tiles = (end - start + BlockK - 1) / BlockK, quotient = total_tiles / Splits, + remainder = total_tiles % Splits; + const int tile_begin = part * quotient + (part < remainder ? part : remainder), + tile_count = quotient + (part < remainder ? 1 : 0); + const int64_t qoff = + int64_t(query) * p.q_stride_row + int64_t(hbase + head) * p.q_stride_head + group * 32; + // Fixed S9/H32/W8 has at most two K128 tiles for capacity <=2176. + // Publish selected IDs once, issue the initial KV loads, then encode Q + // into an independent LDS allocation while the KV requests are in flight. + static_assert(Splits >= 9 && Splits <= 17 && HeadRows == 32 && Waves == 8 && BlockK == 128 && + QMask == 0 && !DoubleBuffer); + static_assert(PF_BUFFERS == 1 || PF_BUFFERS == 2); + int next_id = -1; + constexpr int QChunks = HeadRows / 16 * 8; + static_assert(HeadRows == 32 && Waves == 8 && QChunks / Waves == 2); + q_encode::Packed4 qraw_lo[2], qraw_hi[2], qpacked[2]; + if constexpr(QBF16 && QUnitScale) + { +#pragma unroll + for(int chunk = wave; chunk < QChunks; chunk += Waves) + { + const int row = (chunk % (HeadRows / 16)) * 16 + lane / 4; + const int col = (chunk / (HeadRows / 16)) * 64 + ((lane % 4) * 16 ^ ((row & 3) * 16)); + const int64_t off = int64_t(query) * p.q_stride_row + + int64_t(work.head * HeadRows + row) * p.q_stride_head + col; + const uint16_t* source = reinterpret_cast(p.q) + off; + qraw_lo[chunk / Waves] = *reinterpret_cast(source); + qraw_hi[chunk / Waves] = *reinterpret_cast(source + 8); + } + } + if(tid < BlockK) + { +#pragma unroll + for(int buffer = 0; buffer < PF_BUFFERS; ++buffer) + { + const int at = start + (tile_begin + buffer) * BlockK + tid; + const int id = buffer < tile_count && at < end ? p.indices[at] : -1; + selected_storage[buffer][tid] = (id >= 0 && id < p.slots) ? id : -1; + } + if constexpr(PF_BUFFERS == 1) + if(tile_count > 1) + { + const int at = start + (tile_begin + 1) * BlockK + tid; + next_id = at < end ? p.indices[at] : -1; + } + } + if constexpr(QBF16 && QUnitScale) + { + // Consume raw operands before any KV request. Compiler dependency and + // memory fences require Q loads to finish here without a vmcnt(N) guess. + asm volatile("" + : "+v"(qraw_lo[0]), "+v"(qraw_hi[0]), "+v"(qraw_lo[1]), "+v"(qraw_hi[1]) + : + : "memory"); + if constexpr(QRAW_MODE == 2) + { + qpacked[0] = q_encode::pack_loaded16(qraw_lo[0], qraw_hi[0]); + qpacked[1] = q_encode::pack_loaded16(qraw_lo[1], qraw_hi[1]); + asm volatile("" : "+v"(qpacked[0]), "+v"(qpacked[1]) : : "memory"); + } + } + __syncthreads(); +#pragma unroll + for(int buffer = 0; buffer < PF_BUFFERS; ++buffer) + if(buffer < tile_count) + { + const int id = selected_storage[buffer][wave * 16 + lane / 4]; +#pragma unroll + for(int i = 0; i < (BlockK / 2) / Waves; ++i) + { + const int chunk = wave + i * Waves, r = (chunk % (BlockK / 16)) * 16 + lane / 4; + const int col = (chunk / (BlockK / 16)) * 64 + ((lane % 4) * 16 ^ ((r & 3) * 16)); + if(id >= 0) + direct_load(kv_storage[buffer] + chunk * KV_CHUNK_STRIDE, + p.kv + int64_t(id) * 512 + col); + else + *reinterpret_cast(kv_storage[buffer] + chunk * KV_CHUNK_STRIDE + + lane * 16) = U4{}; + } + } + uint8_t* qstage = qs; + if constexpr(QBF16 && QUnitScale) + { + // Define the conversion inputs after KV issue so conversion cannot be + // hoisted across the volatile direct-load operations. + if constexpr(QRAW_MODE == 1) + asm volatile("" + : "+v"(qraw_lo[0]), "+v"(qraw_hi[0]), "+v"(qraw_lo[1]), "+v"(qraw_hi[1]) + : + : "memory"); + else + asm volatile("" : "+v"(qpacked[0]), "+v"(qpacked[1]) : : "memory"); + } +#pragma unroll + for(int chunk = wave; chunk < QChunks; chunk += Waves) + { + const int row = (chunk % (HeadRows / 16)) * 16 + lane / 4; + const int col = (chunk / (HeadRows / 16)) * 64 + ((lane % 4) * 16 ^ ((row & 3) * 16)); + const int64_t off = int64_t(query) * p.q_stride_row + + int64_t(work.head * HeadRows + row) * p.q_stride_head + col; + U4 value; + if constexpr(QBF16 && QUnitScale) + { + if constexpr(QRAW_MODE == 1) + value = q_encode::pack_loaded16(qraw_lo[chunk / Waves], qraw_hi[chunk / Waves]); + else + value = qpacked[chunk / Waves]; + } + else if constexpr(QBF16) + value = q_encode::pack16(reinterpret_cast(p.q) + off, + p.q_scale); + else + value = *reinterpret_cast(p.q + off); + *reinterpret_cast(qstage + chunk * 1024 + lane * 16) = value; + } + asm volatile("s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)" : : : "memory"); + __syncthreads(); + I8 retained[4]; +#pragma unroll + for(int k = 0; k < 4; ++k) + { + const int col = k * 128 + group * 32, row = hgroup * 16 + head; + U4 a, b; + const uint32_t lo = + uint32_t(reinterpret_cast(qstage + kv_address(row, col))); + const uint32_t hi = + uint32_t(reinterpret_cast(qstage + kv_address(row, col + 16))); + asm volatile("ds_read_b128 %0, %2\n\tds_read_b128 %1, %3" + : "=&v"(a), "=&v"(b) + : "v"(lo), "v"(hi) + : "memory"); +#pragma unroll + for(int j = 0; j < 4; ++j) + { + retained[k][j] = a[j]; + retained[k][j + 4] = b[j]; + } + } + asm volatile("s_waitcnt lgkmcnt(0)" + : "+v"(retained[0]), "+v"(retained[1]), "+v"(retained[2]), "+v"(retained[3]) + : + : "memory"); + F4 merged[PVTiles]{}; + float merge_max = -INFINITY, merge_sum = 0.f; +#pragma unroll 1 + for(int iteration = 0; iteration < tile_count; ++iteration) + { + const int split = tile_begin + iteration; + uint8_t* kvs = kv_storage[PF_BUFFERS == 2 ? iteration : 0]; + int* selected = selected_storage[PF_BUFFERS == 2 ? iteration : 0]; + if(iteration >= PF_BUFFERS) + { + if(tid < BlockK) + { + const int at = start + split * BlockK + tid; + const int id = iteration == 0 ? (at < end ? p.indices[at] : -1) : next_id; + selected[tid] = (id >= 0 && id < p.slots) ? id : -1; + } + __syncthreads(); + if(tid < BlockK && iteration + 1 < tile_count) + { + const int at = start + (split + 1) * BlockK + tid; + next_id = at < end ? p.indices[at] : -1; + } + const int id0 = selected[(wave % (BlockK / 16)) * 16 + lane / 4]; + const int id1 = + (Waves == 4 && BlockK == 128) ? selected[(wave + 4) * 16 + lane / 4] : -1; +#pragma unroll + for(int i = 0; i < (BlockK / 2) / Waves; ++i) + { + const int chunk = wave + i * Waves, r = (chunk % (BlockK / 16)) * 16 + lane / 4; + const int c = (chunk / (BlockK / 16)) * 64 + ((lane % 4) * 16 ^ ((r & 3) * 16)); + const int id = (Waves == 4 && BlockK == 128 && (i & 1)) ? id1 : id0; + if(id >= 0) + direct_load(kvs + chunk * KV_CHUNK_STRIDE, p.kv + int64_t(id) * 512 + c); + else + *reinterpret_cast(kvs + chunk * KV_CHUNK_STRIDE + lane * 16) = U4{}; + } + asm volatile("s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)" : : : "memory"); + __syncthreads(); + } + F4 score[Fragments]{}; +#define TUNE_QK(K) \ + qk_step( \ + kvs, \ + half * (BlockK / WPH) + head, \ + group * 32, \ + q_fragment( \ + p, retained, qs, qoff, iteration, hgroup * 16 + head, group * 32, half), \ + score) + TUNE_QK(0); + TUNE_QK(1); + TUNE_QK(2); + TUNE_QK(3); +#undef TUNE_QK + // Inline MFMA does not expose its result latency to the scheduler. + // Fence every H32 score operand before non-MFMA consumers; otherwise + // a different compiler schedule may read it before the 12-cycle delay. + if constexpr(Fragments == 2) + asm volatile("s_nop 7\n\ts_nop 3" : "+v"(score[0]), "+v"(score[1]) : : "memory"); + float m = -INFINITY; + const float scale = (p.q_scale * p.kv_scale * p.attention_scale) * LOG2E; +#pragma unroll + for(int n = 0; n < Fragments; ++n) + { +#pragma unroll + for(int j = 0; j < 4; ++j) + { + const int key = half * (BlockK / WPH) + n * 16 + group * 4 + j; + score[n][j] = selected[key] >= 0 ? score[n][j] * scale : -INFINITY; + m = fmaxf(m, score[n][j]); + const int slot = split * BlockK + key; + } + } + m = softmax_four_groups(m); + if constexpr(WPH >= 2) + { + if(group == 0) + maxs[wave * 16 + head] = m; + __syncthreads(); + if constexpr(WPH == 4) + { + const int first = (wave & ~3) * 16 + head; + m = fmaxf(fmaxf(maxs[first], maxs[first + 16]), + fmaxf(maxs[first + 32], maxs[first + 48])); + } + else + m = fmaxf(m, maxs[(wave ^ 1) * 16 + head]); + } + const float newmax = fmaxf(merge_max, m); + const float alpha = merge_max == -INFINITY ? 0.f : domain_exp2(merge_max - newmax); + merge_max = newmax; + m = newmax; + float zz[Fragments / 2]{}; +#pragma unroll + for(int pair = 0; pair < Fragments / 2; ++pair) + { +#pragma unroll + for(int j = 0; j < 4; ++j) + { + score[2 * pair][j] = + m == -INFINITY ? 0.f : domain_exp2(score[2 * pair][j] - m); + score[2 * pair + 1][j] = + m == -INFINITY ? 0.f : domain_exp2(score[2 * pair + 1][j] - m); + zz[pair] += score[2 * pair][j] + score[2 * pair + 1][j]; + } + zz[pair] = softmax_four_groups(zz[pair]); + } + float local_z = zz[0]; + if constexpr(Fragments >= 4) + local_z += zz[1]; + if constexpr(Fragments == 8) + local_z += zz[2] + zz[3]; + if constexpr(WPH >= 2) + if(group == 0) + sums[wave * 16 + head] = local_z; +// Publish P together with the sum exchange, saving the separate barrier. +#pragma unroll + for(int n = 0; n < Fragments; ++n) + { + const int key = half * (BlockK / WPH) + n * 16 + group * 4; + uint32_t prob = pack4(score[n] * (1.f / p.p_scale)); + *reinterpret_cast(ps + hgroup * 16 * BlockK + p_address(head, key)) = + prob; + const int slot = split * BlockK + key; + } + float z = local_z; + if constexpr(WPH >= 2) + { + __syncthreads(); + if constexpr(WPH == 4) + { + // Same tree as two pairs inside each original WPH2 wave, + // followed by the original cross-wave sum: (z0+z1)+(z2+z3). + const int first = (wave & ~3) * 16 + head; + z = (sums[first] + sums[first + 16]) + (sums[first + 32] + sums[first + 48]); + } + else + { + const float other = sums[(wave ^ 1) * 16 + head]; + z = half == 0 ? local_z + other : other + local_z; + } + } + else + { + // Each wave owns sixteen heads; no other wave reads its P region. + asm volatile("s_waitcnt lgkmcnt(0)" : : : "memory"); + } + merge_sum = merge_sum * alpha + z; + if(group == 0 && half == 0 && (BlockK == 128 || (split % 2) == 0)) + { + const int tile128 = (split * BlockK) / 128; + const int64_t off = (int64_t(query) * 17 + tile128) * 64 + hbase + head; + } + const uint8_t* pvbase = kvs + half * (8 / WPH) * (BlockK / 16) * KV_CHUNK_STRIDE; + if constexpr(PVMode >= 3) + { + static_assert(BlockK == 64 && (PVMode == 3 || PVMode == 4)); + I2 p0 = p32<0>(ps + hgroup * 16 * BlockK, head, group); + I2 p1 = p32<1>(ps + hgroup * 16 * BlockK, head, group); + asm volatile("s_waitcnt lgkmcnt(0)" : "+v"(p0), "+v"(p1) : : "memory"); + pv32<0, PVTiles, PVMode == 4>(pvbase, p0, p1, merged, alpha); + } + else + { + I8 prob = p_fragment(ps + hgroup * 16 * BlockK, head, group * 32); + asm volatile("s_waitcnt lgkmcnt(0)" : "+v"(prob) : : "memory"); + if constexpr(PVMode == 2) + { + static_assert(BlockK == 128); + I8 first = v_read<0, 128>(pvbase); + pv_pipeline<0, PVTiles>(pvbase, prob, merged, alpha, first); + } + else + pv<0, PVTiles, BlockK, PVMode>(pvbase, prob, merged, alpha); + } + if constexpr(PF_BUFFERS == 2) + if(iteration + 1 < tile_count) + { + // Every wave must finish reading P/max/sum before their reuse. + asm volatile("s_waitcnt lgkmcnt(0)" : : : "memory"); + __syncthreads(); + } + } + const float inv = merge_sum > 0 ? p.p_scale * p.kv_scale / merge_sum : 0.f; +#pragma unroll + for(int tile = 0; tile < PVTiles; ++tile) + { + F4 x = merged[tile] * inv; + const int col = half * (512 / WPH) + tile * 16 + group * 4; + if constexpr(Splits == 1) + { + const int64_t off = (int64_t(query) * 64 + hbase + head) * 512 + col; + + uint32_t lo, hi; + asm("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(lo) : "v"(x[0]), "v"(x[1])); + asm("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(hi) : "v"(x[2]), "v"(x[3])); + *reinterpret_cast(reinterpret_cast(p.out) + off) = I2{int(lo), int(hi)}; + } + else + { + const int64_t off = ((int64_t(query) * Splits + part) * 64 + hbase + head) * 512 + col; + partial_store(p.partial_o + off, x); + } + } + if(group == 0 && half == 0) + { + const float l = merge_sum > 0 ? (merge_max + log2f(merge_sum)) * LN2 : -INFINITY; + if constexpr(Splits == 1) + { + if(p.lse) + p.lse[int64_t(query) * 64 + hbase + head] = l; + } + else if constexpr(Compact) + p.partial_lse[(int64_t(query) * 64 + hbase + head) * 9 + part] = l; + else + p.partial_lse[(int64_t(query) * Splits + part) * 64 + hbase + head] = l; + } +} + +template +hipError_t launch_prefetch_typed(Params p, hipStream_t stream) +{ + hipLaunchKernelGGL( + (prefetch_attention<9, 32, 8, 128, 0, 2, false, BF16, Unit, false, Capacity, Compact>), + dim3(p.queries, 9, 2), + dim3(512), + 0, + stream, + p); + auto status = hipGetLastError(); + if(status != hipSuccess) + return status; + if constexpr(Compact) + hipLaunchKernelGGL((reduce_compact<9, 1, 16, true, false>), + dim3(p.queries, 64, 2), + dim3(64), + 0, + stream, + p); + else + hipLaunchKernelGGL( + (reduce_splits<9, 1, 16, true, false>), dim3(p.queries, 64, 2), dim3(64), 0, stream, p); + return hipGetLastError(); +} +template +hipError_t prefetch(Params p, hipStream_t stream) +{ + if(p.queries == 8) + { + if(p.capacity == 2112) + return launch_prefetch_typed(p, stream); + return launch_prefetch_typed(p, stream); + } + if(p.capacity == 2112) + return launch_prefetch_typed(p, stream); + return launch_prefetch_typed(p, stream); +} +} // namespace sparse_mla_fp8::online +namespace sparse_mla_fp8 { +hipError_t launch_prefetch(Params p, hipStream_t stream) +{ + if(!p.bf16_q) + return online::prefetch(p, stream); + if(p.q_scale == 1.f) + return online::prefetch(p, stream); + return online::prefetch(p, stream); +} +} // namespace sparse_mla_fp8 diff --git a/csrc/kernels/mla/fp8/pv_pipe.cuh b/csrc/kernels/mla/fp8/pv_pipe.cuh new file mode 100644 index 00000000000..c1cab1152d1 --- /dev/null +++ b/csrc/kernels/mla/fp8/pv_pipe.cuh @@ -0,0 +1,46 @@ +// SPDX-License-Identifier: MIT +// PV fragment batching only for the selected stage1 template. +#ifndef PV_PIPE_READ +#define PV_PIPE_READ 0 +#endif +#ifndef PV_PIPE_WAIT +#define PV_PIPE_WAIT 2 +#endif +static_assert(PV_PIPE_READ == 0 || PV_PIPE_READ == 1 || PV_PIPE_READ == 2 || PV_PIPE_READ == 4); +static_assert(PV_PIPE_WAIT == 1 || PV_PIPE_WAIT == 2); +__device__ __forceinline__ I8 pv_pipe_fragment(const uint8_t* base, int col_base) +{ + const int lane = threadIdx.x & 63; + I8 x; +#pragma unroll + for(int i = 0; i < 4; ++i) + { + int row = (lane / 16) * 32 + i * 8 + (lane % 16) / 2; + int col = col_base + (lane % 2) * 8; + uint32_t addr = uint32_t(reinterpret_cast(base + address<128>(row, col))); + I2 r; + asm volatile("ds_read_b64_tr_b8 %0, %1" : "=v"(r) : "v"(addr) : "memory"); + x[i * 2] = r[0]; + x[i * 2 + 1] = r[1]; + } + return x; +} +template +__device__ __forceinline__ void pv_pipe_wait(I8 (&v)[N]) +{ +#if PV_PIPE_WAIT == 2 +#define PV_PIPE_WAIT_TEXT "s_waitcnt lgkmcnt(0)" +#else +#define PV_PIPE_WAIT_TEXT "s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)" +#endif + if constexpr(N == 1) + asm volatile(PV_PIPE_WAIT_TEXT : "+v"(v[0]) : : "memory"); + if constexpr(N == 2) + asm volatile(PV_PIPE_WAIT_TEXT : "+v"(v[0]), "+v"(v[1]) : : "memory"); + if constexpr(N == 4) + asm volatile(PV_PIPE_WAIT_TEXT + : "+v"(v[0]), "+v"(v[1]), "+v"(v[2]), "+v"(v[3]) + : + : "memory"); +#undef PV_PIPE_WAIT_TEXT +} diff --git a/csrc/kernels/mla/fp8/q_encode.cuh b/csrc/kernels/mla/fp8/q_encode.cuh new file mode 100644 index 00000000000..042fed6b2e2 --- /dev/null +++ b/csrc/kernels/mla/fp8/q_encode.cuh @@ -0,0 +1,109 @@ +// SPDX-License-Identifier: MIT +#pragma once +#include +// E4M3FN RNE, with signed zero, gradual underflow, and NaN on overflow. +// Same format semantics as torch.Tensor.to(torch.float8_e4m3fn); no saturation. +// Bit conversion follows the public c10 Float8_e4m3fn scalar algorithm. +namespace q_encode { +union FloatBits +{ + float value; + uint32_t bits; +}; +__host__ __device__ __forceinline__ float bf16_float(uint16_t x) +{ + FloatBits v; + v.bits = uint32_t(x) << 16; + return v.value; +} +__host__ __device__ __forceinline__ uint8_t from_fp32(float x) +{ + FloatBits v; + v.value = x; + const uint32_t sign = v.bits & 0x80000000u; + uint32_t bits = v.bits ^ sign, result; + if(bits >= (1087u << 20)) + result = 0x7f; + else if(bits < (121u << 23)) + { + v.bits = bits; + FloatBits offset; + offset.bits = 141u << 23; + v.value = v.value + offset.value; + result = v.bits - offset.bits; + } + else + { + const uint32_t odd = (bits >> 20) & 1; + bits += uint32_t(7 - 127) * (1u << 23) + 0x7ffffu + odd; + result = bits >> 20; + } + return uint8_t(result | (sign >> 24)); +} +__device__ __forceinline__ uint32_t pack4(const uint16_t* src, float scale) +{ + uint32_t result = 0; +#pragma unroll + for(int j = 0; j < 4; ++j) + { + result |= uint32_t(from_fp32(bf16_float(src[j]) / scale)) << (8 * j); + // Bound scalar-conversion temporaries in the non-unit-scale path. + // The empty asm preserves encoding while keeping later loads ordered. + asm volatile("" : "+v"(result) : : "memory"); + } + return result; +} +// A scale-one fast path: packed BF16 loads, direct packed BF16 -> FP8 RNE, then +// explicit signed NaN encoding for the non-saturating E4M3FN contract. Values +// with abs>464 overflow; the exact midpoint464 rounds to even448. This fixup +// also covers every Inf/NaN bit pattern and preserves its sign. +using Packed4 = uint32_t __attribute__((ext_vector_type(4))); +__host__ __device__ __forceinline__ uint32_t fast_pair(uint32_t bits) +{ + const uint16_t lo = uint16_t(bits), hi = uint16_t(bits >> 16); + uint32_t packed; +#ifdef __HIP_DEVICE_COMPILE__ + // gfx950 packed conversion uses only the scale exponent; unity is exact. + // Do not use this as a replacement for arbitrary FP32 division by q_scale. + using BF16x2 = __bf16 __attribute__((ext_vector_type(2))); + using I16x2 = short __attribute__((ext_vector_type(2))); + const BF16x2 source = __builtin_bit_cast(BF16x2, bits); + const I16x2 result = __builtin_amdgcn_cvt_scalef32_pk_fp8_bf16(I16x2{}, source, 1.0f, false); + packed = __builtin_bit_cast(uint32_t, result); +#else + // The host audit validates the fixup and boundary algorithm. Actual hardware + packed = uint32_t(from_fp32(bf16_float(lo))) | (uint32_t(from_fp32(bf16_float(hi))) << 8); +#endif + const uint32_t a = (lo & 0x7fffu) > 0x43e8u ? ((lo >> 8) & 0x80u) | 0x7fu : packed & 0xffu; + const uint32_t b = + (hi & 0x7fffu) > 0x43e8u ? ((hi >> 8) & 0x80u) | 0x7fu : (packed >> 8) & 0xffu; + return a | (b << 8); +} +__host__ __device__ __forceinline__ Packed4 pack_loaded16(Packed4 lo, Packed4 hi) +{ + Packed4 result; + result[0] = fast_pair(lo[0]) | (fast_pair(lo[1]) << 16); + result[1] = fast_pair(lo[2]) | (fast_pair(lo[3]) << 16); + result[2] = fast_pair(hi[0]) | (fast_pair(hi[1]) << 16); + result[3] = fast_pair(hi[2]) | (fast_pair(hi[3]) << 16); + return result; +} +template +__device__ __forceinline__ Packed4 pack16(const uint16_t* src, float scale) +{ + Packed4 result; + if constexpr(Fast) + { + const Packed4 lo = *reinterpret_cast(src); + const Packed4 hi = *reinterpret_cast(src + 8); + result = pack_loaded16(lo, hi); + } + else + { +#pragma unroll + for(int j = 0; j < 4; ++j) + result[j] = pack4(src + j * 4, scale); + } + return result; +} +} // namespace q_encode diff --git a/csrc/kernels/mla/fp8/reduce.cuh b/csrc/kernels/mla/fp8/reduce.cuh new file mode 100644 index 00000000000..2946a999234 --- /dev/null +++ b/csrc/kernels/mla/fp8/reduce.cuh @@ -0,0 +1,110 @@ +// SPDX-License-Identifier: MIT +#pragma once +#include "online_common.cuh" +namespace sparse_mla_fp8::online { +template +__global__ __launch_bounds__(64) void reduce_splits(Params p) +{ + constexpr int V = Cols / 16; + int lane = threadIdx.x, head = SingleHead ? int(blockIdx.y) : int(blockIdx.y) * 16 + lane / 4; + int col = + SingleHead ? int(blockIdx.z) * 256 + lane * 4 : int(blockIdx.z) * Cols + (lane % 4) * 4; + int64_t row = int64_t(blockIdx.x) * 64 + head; + float accum[V][4]{}; + float z = 0.f, m = -INFINITY; + if constexpr(Mode == 0) + { +#pragma unroll + for(int s = 0; s < Splits; ++s) + { + float l = p.partial_lse[(int64_t(blockIdx.x) * Splits + s) * 64 + head]; + float d = exp2f(-fabsf((m == -INFINITY && l == -INFINITY) ? 0.f : m - l) * LOG2E); + float alpha = m == -INFINITY ? 0.f : (m >= l ? 1.f : d), + beta = l == -INFINITY ? 0.f : (l >= m ? 1.f : d); + z = z * alpha + beta; + m = fmaxf(m, l); +#pragma unroll + for(int v = 0; v < V; ++v) + { + F4 x = partial_load(p.partial_o + + ((int64_t(blockIdx.x) * Splits + s) * 64 + head) * 512 + col + + v * 16); +#pragma unroll + for(int j = 0; j < 4; ++j) + accum[v][j] = fmaf(x[j], beta, accum[v][j] * alpha); + } + } + } + else if constexpr(Mode == 1) + { + float ls[Splits]; +#pragma unroll + for(int s = 0; s < Splits; ++s) + { + ls[s] = p.partial_lse[(int64_t(blockIdx.x) * Splits + s) * 64 + head]; + m = fmaxf(m, ls[s]); + } +#pragma unroll + for(int s = 0; s < Splits; ++s) + { + float w = m == -INFINITY ? 0.f : exp2f((ls[s] - m) * LOG2E); + z += w; +#pragma unroll + for(int v = 0; v < V; ++v) + { + F4 x = partial_load(p.partial_o + + ((int64_t(blockIdx.x) * Splits + s) * 64 + head) * 512 + col + + v * 16); +#pragma unroll + for(int j = 0; j < 4; ++j) + accum[v][j] = fmaf(w > 0 ? x[j] : 0.f, w, accum[v][j]); + } + } + } + if constexpr(Mode >= 2) + { +#pragma unroll + for(int s = 0; s < Splits; ++s) + m = fmaxf(m, p.partial_lse[(int64_t(blockIdx.x) * Splits + s) * 64 + head]); +#pragma unroll(Mode == 3 ? 4 : 1) + for(int s = 0; s < Splits; ++s) + { + float l = p.partial_lse[(int64_t(blockIdx.x) * Splits + s) * 64 + head]; + float w = m == -INFINITY ? 0.f : exp2f((l - m) * LOG2E); + z += w; +#pragma unroll + for(int v = 0; v < V; ++v) + { + F4 x = partial_load(p.partial_o + + ((int64_t(blockIdx.x) * Splits + s) * 64 + head) * 512 + col + + v * 16); +#pragma unroll + for(int j = 0; j < 4; ++j) + accum[v][j] = fmaf(w > 0 ? x[j] : 0.f, w, accum[v][j]); + } + } + } + float inv = z > 0 ? 1.f / (z * (OutputFP8 ? 1.f : 1.f)) : 0.f; +#pragma unroll + for(int v = 0; v < V; ++v) + { + F4 x; + for(int j = 0; j < 4; ++j) + { + x[j] = accum[v][j] * inv; + } + if constexpr(OutputFP8) + *reinterpret_cast(p.out + row * 512 + col + v * 16) = pack4(x); + else + { + uint32_t lo, hi; + asm("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(lo) : "v"(x[0]), "v"(x[1])); + asm("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(hi) : "v"(x[2]), "v"(x[3])); + *reinterpret_cast(reinterpret_cast(p.out) + row * 512 + col + v * 16) = + I2{int(lo), int(hi)}; + } + } + if(p.lse && blockIdx.z == 0 && (SingleHead ? lane == 0 : lane % 4 == 0)) + p.lse[row] = z > 0 ? m + log2f(z) * LN2 : -INFINITY; +} +} // namespace sparse_mla_fp8::online diff --git a/csrc/kernels/mla/fp8/rotating.cu b/csrc/kernels/mla/fp8/rotating.cu new file mode 100644 index 00000000000..b3adf86e0ee --- /dev/null +++ b/csrc/kernels/mla/fp8/rotating.cu @@ -0,0 +1,365 @@ +// SPDX-License-Identifier: MIT +#include "online_common.cuh" +#include "reduce.cuh" +#define PIPE_MODE 3 +namespace sparse_mla_fp8::online { +__device__ __forceinline__ void rotate_stage(Params p, uint8_t* kv, const int* selected) +{ + const int lane = threadIdx.x & 63, wave = threadIdx.x / 64; + const int id = selected[wave * 16 + lane / 4]; +#pragma unroll + for(int i = 0; i < 8; ++i) + { + const int chunk = wave + i * 8, r = (chunk % 8) * 16 + lane / 4; + const int col = (chunk / 8) * 64 + ((lane % 4) * 16 ^ ((r & 3) * 16)); + if(id >= 0) + direct_load(kv + chunk * KV_CHUNK_STRIDE, p.kv + int64_t(id) * 512 + col); + else + *reinterpret_cast(kv + chunk * KV_CHUNK_STRIDE + lane * 16) = U4{}; + } +} +template +__global__ __launch_bounds__(Waves * 64) void rotating_attention(Params p) +{ + + constexpr bool Pipeline = PIPE_MODE >= 2; + constexpr bool QAlias = PIPE_MODE == 1 || PIPE_MODE >= 3 || HeadRows == 64; + constexpr int Buffers = Pipeline ? 2 : 1; + constexpr int InitialTiles = Pipeline ? (QAlias ? 1 : 2) : 0; + constexpr int StartBuffer = Pipeline && QAlias ? 1 : 0; + constexpr int MaxTiles = (17 + Splits - 1) / Splits; + constexpr int WPH = Waves * 16 / HeadRows, Fragments = BlockK / (WPH * 16), PVTiles = 32 / WPH; + static_assert((HeadRows == 32 || HeadRows == 64) && (WPH == 1 || WPH == 2 || WPH == 4)); + static_assert(BlockK == 128 || (BlockK == 64 && (Waves == 4 || Waves == 8))); + __shared__ __align__(16) uint8_t kv_storage[Buffers][(BlockK / 2) * KV_CHUNK_STRIDE]; + __shared__ __align__(16) uint8_t ps[HeadRows * BlockK]; + __shared__ __align__(16) uint8_t qs[QAlias ? 1 : HeadRows * 512]; + __shared__ float maxs[WPH >= 2 ? Waves * 16 : 1], sums[WPH >= 2 ? Waves * 16 : 1]; + __shared__ int selected_storage[MaxTiles][BlockK]; + const int tid = threadIdx.x, lane = tid & 63, wave = tid / 64, head = lane & 15, + group = lane / 16; + const auto work = order_ids( + blockIdx.x, blockIdx.y, blockIdx.z, p.queries); + const int query = work.query, part = work.part, half = wave % WPH, hgroup = wave / WPH; + const int hbase = work.head * HeadRows + hgroup * 16; + const int start = (query * p.capacity), end = ((query + 1) * p.capacity); + const int total_tiles = (end - start + BlockK - 1) / BlockK, quotient = total_tiles / Splits, + remainder = total_tiles % Splits; + const int tile_begin = part * quotient + (part < remainder ? part : remainder), + tile_count = quotient + (part < remainder ? 1 : 0); + const int64_t qoff = + int64_t(query) * p.q_stride_row + int64_t(hbase + head) * p.q_stride_head + group * 32; + static_assert(Waves == 8 && BlockK == 128 && QMask == 0 && !DoubleBuffer); + static_assert(PIPE_MODE >= 1 && PIPE_MODE <= 4); + // The public serving Call limits each selection row to capacity<=2176 (17 tiles). + // Load the complete split's IDs once; they stay immutable through the loop. + for(int i = tid; i < tile_count * BlockK; i += Waves * 64) + { + const int at = start + tile_begin * BlockK + i; + const int id = at < end ? p.indices[at] : -1; + selected_storage[i / BlockK][i % BlockK] = (id >= 0 && id < p.slots) ? id : -1; + } + __syncthreads(); + if constexpr(Pipeline) + { +#pragma unroll + for(int t = 0; t < InitialTiles; ++t) + if(t < tile_count) + rotate_stage(p, kv_storage[(t + StartBuffer) & 1], selected_storage[t]); + } + // H64 and alias variants use the initially unoccupied KV buffer for Q. + uint8_t* qstage = QAlias ? kv_storage[0] : qs; + constexpr int QChunks = HeadRows / 16 * 8; +#pragma unroll + for(int chunk = wave; chunk < QChunks; chunk += Waves) + { + const int row = (chunk % (HeadRows / 16)) * 16 + lane / 4; + const int col = (chunk / (HeadRows / 16)) * 64 + ((lane % 4) * 16 ^ ((row & 3) * 16)); + const int64_t off = int64_t(query) * p.q_stride_row + + int64_t(work.head * HeadRows + row) * p.q_stride_head + col; + U4 value; + if constexpr(QBF16) + value = q_encode::pack16(reinterpret_cast(p.q) + off, + p.q_scale); + else + value = *reinterpret_cast(p.q + off); + *reinterpret_cast(qstage + chunk * 1024 + lane * 16) = value; + } + asm volatile("s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)" : : : "memory"); + __syncthreads(); + I8 retained[4]; +#pragma unroll + for(int k = 0; k < 4; ++k) + { + const int col = k * 128 + group * 32, row = hgroup * 16 + head; + U4 a, b; + const uint32_t lo = + uint32_t(reinterpret_cast(qstage + kv_address(row, col))); + const uint32_t hi = + uint32_t(reinterpret_cast(qstage + kv_address(row, col + 16))); + asm volatile("ds_read_b128 %0, %2\n\tds_read_b128 %1, %3" + : "=&v"(a), "=&v"(b) + : "v"(lo), "v"(hi) + : "memory"); +#pragma unroll + for(int j = 0; j < 4; ++j) + { + retained[k][j] = a[j]; + retained[k][j + 4] = b[j]; + } + } + asm volatile("s_waitcnt lgkmcnt(0)" + : "+v"(retained[0]), "+v"(retained[1]), "+v"(retained[2]), "+v"(retained[3]) + : + : "memory"); + if constexpr(QAlias) + __syncthreads(); // all retained Q reads precede KV reuse + F4 merged[PVTiles]{}; + float merge_max = -INFINITY, merge_sum = 0.f; +#pragma unroll 1 + for(int iteration = 0; iteration < tile_count; ++iteration) + { + const int split = tile_begin + iteration; + const int current = Pipeline ? ((iteration + StartBuffer) & 1) : 0; + uint8_t* kvs = kv_storage[current]; + const int* selected = selected_storage[iteration]; + if constexpr(!Pipeline) + { + rotate_stage(p, kvs, selected); + asm volatile("s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)" : : : "memory"); + __syncthreads(); + } + if constexpr(Pipeline && PIPE_MODE != 4) + { + if(iteration + 1 < tile_count && (QAlias || iteration > 0)) + rotate_stage(p, kv_storage[current ^ 1], selected_storage[iteration + 1]); + } + F4 score[Fragments]{}; +#define TUNE_QK(K) \ + qk_step( \ + kvs, \ + half * (BlockK / WPH) + head, \ + group * 32, \ + q_fragment( \ + p, retained, qs, qoff, iteration, hgroup * 16 + head, group * 32, half), \ + score) + TUNE_QK(0); + TUNE_QK(1); + TUNE_QK(2); + TUNE_QK(3); +#undef TUNE_QK + // Inline MFMA does not expose its result latency to the scheduler. + // Fence every H32 score operand before non-MFMA consumers; otherwise + // a different compiler schedule may read it before the 12-cycle delay. + if constexpr(Fragments == 2) + asm volatile("s_nop 7\n\ts_nop 3" : "+v"(score[0]), "+v"(score[1]) : : "memory"); + if constexpr(Pipeline && PIPE_MODE == 4) + { + if(iteration + 1 < tile_count) + rotate_stage(p, kv_storage[current ^ 1], selected_storage[iteration + 1]); + } + float m = -INFINITY; + const float scale = (p.q_scale * p.kv_scale * p.attention_scale) * LOG2E; +#pragma unroll + for(int n = 0; n < Fragments; ++n) + { +#pragma unroll + for(int j = 0; j < 4; ++j) + { + const int key = half * (BlockK / WPH) + n * 16 + group * 4 + j; + score[n][j] = selected[key] >= 0 ? score[n][j] * scale : -INFINITY; + m = fmaxf(m, score[n][j]); + const int slot = split * BlockK + key; + } + } + m = softmax_four_groups(m); + if constexpr(WPH >= 2) + { + if(group == 0) + maxs[wave * 16 + head] = m; + __syncthreads(); + if constexpr(WPH == 4) + { + const int first = (wave & ~3) * 16 + head; + m = fmaxf(fmaxf(maxs[first], maxs[first + 16]), + fmaxf(maxs[first + 32], maxs[first + 48])); + } + else + m = fmaxf(m, maxs[(wave ^ 1) * 16 + head]); + } + const float newmax = fmaxf(merge_max, m); + const float alpha = merge_max == -INFINITY ? 0.f : domain_exp2(merge_max - newmax); + merge_max = newmax; + m = newmax; + float zz[Fragments / 2]{}; +#pragma unroll + for(int pair = 0; pair < Fragments / 2; ++pair) + { +#pragma unroll + for(int j = 0; j < 4; ++j) + { + score[2 * pair][j] = + m == -INFINITY ? 0.f : domain_exp2(score[2 * pair][j] - m); + score[2 * pair + 1][j] = + m == -INFINITY ? 0.f : domain_exp2(score[2 * pair + 1][j] - m); + zz[pair] += score[2 * pair][j] + score[2 * pair + 1][j]; + } + zz[pair] = softmax_four_groups(zz[pair]); + } + float local_z = zz[0]; + if constexpr(Fragments >= 4) + local_z += zz[1]; + if constexpr(Fragments == 8) + local_z += zz[2] + zz[3]; + if constexpr(WPH >= 2) + if(group == 0) + sums[wave * 16 + head] = local_z; +// Publish P together with the sum exchange, saving the separate barrier. +#pragma unroll + for(int n = 0; n < Fragments; ++n) + { + const int key = half * (BlockK / WPH) + n * 16 + group * 4; + uint32_t prob = pack4(score[n] * (1.f / p.p_scale)); + *reinterpret_cast(ps + hgroup * 16 * BlockK + p_address(head, key)) = + prob; + const int slot = split * BlockK + key; + } + float z = local_z; + if constexpr(WPH >= 2) + { + __syncthreads(); + if constexpr(WPH == 4) + { + // Same tree as two pairs inside each original WPH2 wave, + // followed by the original cross-wave sum: (z0+z1)+(z2+z3). + const int first = (wave & ~3) * 16 + head; + z = (sums[first] + sums[first + 16]) + (sums[first + 32] + sums[first + 48]); + } + else + { + const float other = sums[(wave ^ 1) * 16 + head]; + z = half == 0 ? local_z + other : other + local_z; + } + } + else + { + // Each wave owns sixteen heads; no other wave reads its P region. + asm volatile("s_waitcnt lgkmcnt(0)" : : : "memory"); + } + merge_sum = merge_sum * alpha + z; + if(group == 0 && half == 0 && (BlockK == 128 || (split % 2) == 0)) + { + const int tile128 = (split * BlockK) / 128; + const int64_t off = (int64_t(query) * 17 + tile128) * 64 + hbase + head; + } + const uint8_t* pvbase = kvs + half * (8 / WPH) * (BlockK / 16) * KV_CHUNK_STRIDE; + if constexpr(PVMode >= 3) + { + static_assert(BlockK == 64 && (PVMode == 3 || PVMode == 4)); + I2 p0 = p32<0>(ps + hgroup * 16 * BlockK, head, group); + I2 p1 = p32<1>(ps + hgroup * 16 * BlockK, head, group); + asm volatile("s_waitcnt lgkmcnt(0)" : "+v"(p0), "+v"(p1) : : "memory"); + pv32<0, PVTiles, PVMode == 4>(pvbase, p0, p1, merged, alpha); + } + else + { + I8 prob = p_fragment(ps + hgroup * 16 * BlockK, head, group * 32); + asm volatile("s_waitcnt lgkmcnt(0)" : "+v"(prob) : : "memory"); + if constexpr(PVMode == 2) + { + static_assert(BlockK == 128); + I8 first = v_read<0, 128>(pvbase); + pv_pipeline<0, PVTiles>(pvbase, prob, merged, alpha, first); + } + else + pv<0, PVTiles, BlockK, PVMode>(pvbase, prob, merged, alpha); + } + if(iteration + 1 < tile_count) + { + // Publish incoming KV, finish all current LDS readers, and protect + // P/max/sum before another wave begins the following iteration. + asm volatile("s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)" : : : "memory"); + __syncthreads(); + } + } + const float inv = merge_sum > 0 ? p.p_scale * p.kv_scale / merge_sum : 0.f; +#pragma unroll + for(int tile = 0; tile < PVTiles; ++tile) + { + F4 x = merged[tile] * inv; + const int col = half * (512 / WPH) + tile * 16 + group * 4; + if constexpr(Splits == 1) + { + const int64_t off = (int64_t(query) * 64 + hbase + head) * 512 + col; + + uint32_t lo, hi; + asm("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(lo) : "v"(x[0]), "v"(x[1])); + asm("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(hi) : "v"(x[2]), "v"(x[3])); + *reinterpret_cast(reinterpret_cast(p.out) + off) = I2{int(lo), int(hi)}; + } + else + { + const int64_t off = ((int64_t(query) * Splits + part) * 64 + hbase + head) * 512 + col; + partial_store(p.partial_o + off, x); + } + } + if(group == 0 && half == 0) + { + const float l = merge_sum > 0 ? (merge_max + log2f(merge_sum)) * LN2 : -INFINITY; + if constexpr(Splits == 1) + { + if(p.lse) + p.lse[int64_t(query) * 64 + hbase + head] = l; + } + else + p.partial_lse[(int64_t(query) * Splits + part) * 64 + hbase + head] = l; + } +} + +template +hipError_t rotating_typed(Params p, hipStream_t stream) +{ + hipLaunchKernelGGL((rotating_attention), + dim3(p.queries, S, 64 / H), + dim3(512), + 0, + stream, + p); + auto status = hipGetLastError(); + if(status != hipSuccess) + return status; + hipLaunchKernelGGL( + (reduce_splits), dim3(p.queries, 64, 2), dim3(64), 0, stream, p); + return hipGetLastError(); +} +template +hipError_t rotating(Params p, hipStream_t stream) +{ + if(p.queries <= 21) + return rotating_typed<6, 32, BF16, Unit>(p, stream); + if(p.queries <= 23) + return rotating_typed<4, 32, BF16, Unit>(p, stream); + if(p.queries <= 42) + return rotating_typed<6, 64, BF16, Unit>(p, stream); + return rotating_typed<4, 64, BF16, Unit>(p, stream); +} +} // namespace sparse_mla_fp8::online +namespace sparse_mla_fp8 { +hipError_t launch_rotating(Params p, hipStream_t stream) +{ + if(!p.bf16_q) + return online::rotating(p, stream); + if(p.q_scale == 1.f) + return online::rotating(p, stream); + return online::rotating(p, stream); +} +} // namespace sparse_mla_fp8 diff --git a/csrc/kernels/mla/fp8/small.cu b/csrc/kernels/mla/fp8/small.cu new file mode 100644 index 00000000000..20855b897df --- /dev/null +++ b/csrc/kernels/mla/fp8/small.cu @@ -0,0 +1,847 @@ +// SPDX-License-Identifier: MIT +#include "params.cuh" +#include "partial_io.cuh" +#include "q_encode.cuh" +#define PV_PIPE_WAIT 1 +namespace sparse_mla_fp8::small { + +#ifndef CACHEX_READ_NT +#define CACHEX_READ_NT 0 +#endif +#ifndef CACHEX_STORE_NT +#define CACHEX_STORE_NT 0 +#endif +#ifndef LOAD_MODE +#define LOAD_MODE 0 +#endif +#ifndef IDS_MODE +#define IDS_MODE 0 +#endif +// SPDX-License-Identifier: MIT +using I8 = int __attribute__((ext_vector_type(8))); +using U4 = uint32_t __attribute__((ext_vector_type(4))); +using F4 = float __attribute__((ext_vector_type(4))); +using Params = sparse_mla_fp8::Params; +constexpr float LOG2E = 1.4426950408889634f, LN2 = .6931471805599453f; +__device__ __forceinline__ F4 mma(I8 a, I8 b, F4 c) +{ + return __builtin_amdgcn_mfma_scale_f32_16x16x128_f8f6f4(a, b, c, 0, 0, 0, 127, 0, 127); +} +// unchanged. Keep producer/consumer hazards inside the same assembly block: +// Table38 non-DL VALU->MFMA 2 waits; XDL result->VALU 12 waits for this form. +template +__device__ __forceinline__ F4 online_pv_mma(I8 a, I8 b, F4 c) +{ + if constexpr(!VGPR) + return mma(a, b, c); + else + { + const int scale = 127; + asm volatile( + "s_nop 1\n\t" + "v_mfma_scale_f32_16x16x128_f8f6f4 %0, %1, %2, %0, %3, %3 op_sel_hi:[0,0,0]\n\t" + "s_nop 7\n\ts_nop 3" + : "+&v"(c) + : "v"(a), "v"(b), "v"(scale) + : "memory"); + return c; + } +} +__device__ __forceinline__ uint32_t pack4(F4 x) +{ +#pragma unroll + for(int j = 0; j < 4; ++j) + x[j] = fminf(448.f, fmaxf(-448.f, x[j])); + int w = __builtin_amdgcn_cvt_pk_fp8_f32(x[0], x[1], 0, 0); + return __builtin_amdgcn_cvt_pk_fp8_f32(x[2], x[3], w, 1); +} +// Each 1 KiB chunk is 16 rows x 64 bytes. XOR preserves 16-byte vectors. +#include "kv_layout.cuh" +template +__device__ __forceinline__ I8 fragment(const uint8_t* base, int row, int col) +{ + I8 x; + const U4 a = *reinterpret_cast(base + address(row, col)); + const U4 b = *reinterpret_cast(base + address(row, col + 16)); +#pragma unroll + for(int j = 0; j < 4; ++j) + { + x[j] = a[j]; + x[j + 4] = b[j]; + } + return x; +} +__device__ __forceinline__ int p_address(int head, int key) +{ + return head * 128 + (key ^ ((head & 7) * 16)); +} +__device__ __forceinline__ I8 p_fragment(const uint8_t* base, int head, int col) +{ + I8 x; + U4 a = *reinterpret_cast(base + p_address(head, col)); + U4 b = *reinterpret_cast(base + p_address(head, col + 16)); +#pragma unroll + for(int j = 0; j < 4; ++j) + { + x[j] = a[j]; + x[j + 4] = b[j]; + } + return x; +} +__device__ __forceinline__ I8 v_fragment(const uint8_t* base, int col, int key) +{ + I8 x{}; +#pragma unroll + for(int j = 0; j < 8; ++j) + { + uint32_t w = 0; +#pragma unroll + for(int b = 0; b < 4; ++b) + w |= uint32_t(base[address<128>(key + j * 4 + b, col)]) << (8 * b); + x[j] = w; + } + return x; +} +using I2 = int __attribute__((ext_vector_type(2))); +__device__ __forceinline__ I8 v_fragment_tr(const uint8_t* base, int col_base) +{ + const int lane = threadIdx.x & 63; + I8 x; +#pragma unroll + for(int i = 0; i < 4; ++i) + { + int row = (lane / 16) * 32 + i * 8 + (lane % 16) / 2; + int col = col_base + (lane % 2) * 8; + uint32_t addr = uint32_t(reinterpret_cast(base + address<128>(row, col))); + I2 r; + asm volatile("ds_read_b64_tr_b8 %0, %1" : "=v"(r) : "v"(addr) : "memory"); + x[i * 2] = r[0]; + x[i * 2 + 1] = r[1]; + } + // Tie the loaded operand to the wait. A memory-only barrier can move + // after the last register-only MFMA when its results stay in registers. + asm volatile("s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)" : "+v"(x) : : "memory"); + return x; +} +__device__ __forceinline__ void direct_load(uint8_t* dest, const uint8_t* src) +{ + uint32_t lds = __builtin_amdgcn_readfirstlane(uint32_t(reinterpret_cast(dest))); + uint64_t addr = reinterpret_cast(src); + asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tglobal_load_lds_dwordx4 %1, off" + : + : "s"(lds), "v"(addr) + : "m0", "memory"); +} +__host__ __device__ constexpr bool narrow_kv(int slots) +{ + return slots >= 0 && uint64_t(slots) * 512 <= 0x7fffffffULL; +} +template +__device__ __forceinline__ void +direct_kv(uint8_t* dest, const uint8_t* base, int id, int col, int slots) +{ + if constexpr(Narrow && LOAD_MODE) + { + uint32_t lds = __builtin_amdgcn_readfirstlane(uint32_t(reinterpret_cast(dest))); + uint64_t addr = reinterpret_cast(base); + uint32_t offset = (uint32_t(id) << 9) + uint32_t(col); +#if LOAD_MODE == 1 + U4 resource = {uint32_t(addr), uint32_t(addr >> 32), uint32_t(slots) * 512, 0x00020000}; + asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tbuffer_load_dwordx4 %1, %2, 0 offen lds" + : + : "s"(lds), "v"(offset), "s"(resource) + : "m0", "memory"); +#elif LOAD_MODE == 2 + asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tglobal_load_lds_dwordx4 %1, %2" + : + : "s"(lds), "v"(offset), "s"(addr) + : "m0", "memory"); +#endif + } + else + direct_load(dest, base + int64_t(id) * 512 + col); +} + +#include "partial_layout.cuh" +#include "pv_pipe.cuh" +#include "softmax_lane.cuh" +// Pad16 addresses keep the byte payload and arithmetic unchanged. Offsets are +// compile-time instruction immediates; no array of32 persistent PV addresses. +#include "pad16_factored.cuh" + +template +__global__ __launch_bounds__(Waves * 64) void attention(Params p) +{ + constexpr int Tiles = (17 + Splits - 1) / Splits; + constexpr int Pitch = PadFactored ? 1040 : KV_STRIDE; + static_assert(!PadFactored || (PVGroup == 1 && IdsMode == 0)); + constexpr bool Transpose = true, Direct = true, LocalReduce = true, LoadSchedule = true; + __shared__ __align__(16) uint8_t kvs[64 * Pitch]; + __shared__ __align__(16) uint8_t qs[8192]; + __shared__ __align__(16) uint8_t ps[2048]; + __shared__ float exchange[64], row_max[16], row_inv[16]; + __shared__ float exchange_sum[LocalReduce ? 64 : 1]; + __shared__ int selected[IdsMode == 0 ? 128 : 1]; + const int tid = threadIdx.x, lane = tid & 63, wave = tid / 64, head = lane & 15, + group = lane / 16; + static_assert(Order == 0 && (PVGroup == 1 || PVGroup == 4)); + CtaTile work; + if constexpr(Remap) + { + const auto t = cta_tile<4, Splits, true>(); + work = {t.q, t.c, t.b}; + } + else + work = {int(blockIdx.x), int(blockIdx.y), int(blockIdx.z)}; + const int query = work.q, part = work.b, hbase = work.c * 16; + const int start = (query * p.capacity), end = ((query + 1) * p.capacity); + // Uniform within the CTA. Distribute the seventeenth tail tile without + // computing ceil(17/S) tiles for every split. No host reads of selection values. + const int total_tiles = (end - start + 127) / 128; + const int quotient = total_tiles / Splits, remainder = total_tiles % Splits; + const int tile_begin = part * quotient + (part < remainder ? part : remainder); + const int tile_count = quotient + (part < remainder ? 1 : 0); + if(wave < 4) + { +#pragma unroll + for(int i = 0; i < 2; ++i) + { + int chunk = wave + i * 4; + int r = lane / 4, c = chunk * 64 + ((lane % 4) * 16 ^ ((r & 3) * 16)); + const int64_t off = + int64_t(query) * p.q_stride_row + int64_t(hbase + r) * p.q_stride_head + c; + if constexpr(QBF16) + { + const auto* src = reinterpret_cast(p.q) + off; + U4 encoded = q_encode::pack16(src, p.q_scale); + *reinterpret_cast(qs + chunk * 1024 + lane * 16) = encoded; + } + else if constexpr(Direct) + direct_load(qs + chunk * 1024, p.q + off); + else + *reinterpret_cast(qs + chunk * 1024 + lane * 16) = + *reinterpret_cast(p.q + off); + } + } + F4 merged[8]{}; + float merge_max = -INFINITY, merge_sum = 0.f; + int next_id = -1, next_mask_id = -1, next_id0 = -1, next_id1 = -1; +// Q is staged once per grouped CTA. KV/P shared storage is reused per tile. +#pragma unroll 1 + for(int iteration = 0; iteration < tile_count; ++iteration) + { + const int split = tile_begin + iteration; + uint32_t valid_bits = 0; + int reg_id0 = -1, reg_id1 = -1; + if constexpr(IdsMode == 0) + { + if(tid < 128) + { + int id; + if constexpr(Prefetch) + { + if(iteration == 0) + { + int at = start + split * 128 + tid; + id = at < end ? p.indices[at] : -1; + } + else + id = next_id; + } + else + { + int at = start + split * 128 + tid; + id = at < end ? p.indices[at] : -1; + } + selected[tid] = (id >= 0 && id < p.slots) ? id : -1; + } + // Also protects the previous tile's KV/P consumers when TailBarrier=false. + __syncthreads(); + if constexpr(Prefetch) + { + if(tid < 128 && iteration + 1 < tile_count) + { + int at = start + (split + 1) * 128 + tid; + next_id = at < end ? p.indices[at] : -1; + } + } + } + else + { + // First tile needs no ID exchange. Later tiles MUST wait for all + // preceding KV/P consumers before overwriting shared storage. + if constexpr(IdsMode == 3) + __syncthreads(); + else if(iteration > 0 && !TailBarrier) + __syncthreads(); + if(wave < 4) + { + int at = start + split * 128 + wave * 32 + (lane & 31); + int id; + if constexpr(IdsMode == 2 && Prefetch) + id = iteration ? next_mask_id : (at < end ? p.indices[at] : -1); + else + id = at < end ? p.indices[at] : -1; + valid_bits = uint32_t(__ballot(id >= 0 && id < p.slots)); + } + if(Waves == 4 || wave >= 4) + { + int loader = Waves == 4 ? wave : wave - 4; + int at0 = start + split * 128 + loader * 16 + lane / 4, at1 = at0 + 64; + if constexpr(IdsMode == 2 && Prefetch) + { + reg_id0 = iteration ? next_id0 : (at0 < end ? p.indices[at0] : -1); + reg_id1 = iteration ? next_id1 : (at1 < end ? p.indices[at1] : -1); + } + else + { + reg_id0 = at0 < end ? p.indices[at0] : -1; + reg_id1 = at1 < end ? p.indices[at1] : -1; + } + reg_id0 = (reg_id0 >= 0 && reg_id0 < p.slots) ? reg_id0 : -1; + reg_id1 = (reg_id1 >= 0 && reg_id1 < p.slots) ? reg_id1 : -1; + } + if constexpr(IdsMode == 2 && Prefetch) + { + if(iteration + 1 < tile_count) + { + if(wave < 4) + { + int at = start + (split + 1) * 128 + wave * 32 + (lane & 31); + next_mask_id = at < end ? p.indices[at] : -1; + } + if(Waves == 4 || wave >= 4) + { + int loader = Waves == 4 ? wave : wave - 4; + int at0 = start + (split + 1) * 128 + loader * 16 + lane / 4, + at1 = at0 + 64; + next_id0 = at0 < end ? p.indices[at0] : -1; + next_id1 = at1 < end ? p.indices[at1] : -1; + } + } + } + } + if(Waves == 4 || wave >= 4) + { + int loader = Waves == 4 ? wave : wave - 4; + if constexpr(LoadSchedule) + { + int id0, id1; + if constexpr(IdsMode == 0) + { + id0 = selected[loader * 16 + lane / 4]; + id1 = selected[(loader + 4) * 16 + lane / 4]; + } + else + { + id0 = reg_id0; + id1 = reg_id1; + } +#pragma unroll + for(int i = 0; i < 16; ++i) + { + int chunk = loader + i * 4, r = (chunk % 8) * 16 + lane / 4; + int c = (chunk / 8) * 64 + ((lane % 4) * 16 ^ kv_xor(r)); + int id = (i & 1) ? id1 : id0; + if(id >= 0) + direct_kv(kvs + chunk * Pitch, p.kv, id, c, p.slots); + else + *reinterpret_cast(kvs + chunk * Pitch + lane * 16) = U4{}; + } + } + else + { +#pragma unroll + for(int chunk = loader; chunk < 64; chunk += 4) + { + int r = (chunk % 8) * 16 + lane / 4, + c = (chunk / 8) * 64 + ((lane % 4) * 16 ^ kv_xor(r)); + int id; + if constexpr(IdsMode == 0) + id = selected[r]; + else + { + int at = start + split * 128 + r; + id = at < end ? p.indices[at] : -1; + id = (id >= 0 && id < p.slots) ? id : -1; + } + if constexpr(Direct) + { + if(id >= 0) + direct_kv(kvs + chunk * Pitch, p.kv, id, c, p.slots); + else + *reinterpret_cast(kvs + chunk * Pitch + lane * 16) = U4{}; + } + else + { + U4 x{}; + if(id >= 0) + x = *reinterpret_cast(p.kv + int64_t(id) * 512 + c); + *reinterpret_cast(kvs + chunk * Pitch + lane * 16) = x; + } + } + } + } + if constexpr(Direct) + asm volatile("s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)" : : : "memory"); + __syncthreads(); + F4 a{}, b{}; + if(wave < 4) + { + if constexpr(PadFactored) + { + padded_qk_step<0>(qs, kvs, head, group, wave, a, b); + } + else + { +#pragma unroll + for(int k = 0; k < 4; ++k) + { + I8 q = fragment<16>(qs, head, k * 128 + group * 32); + a = mma(fragment<128>(kvs, wave * 32 + head, k * 128 + group * 32), q, a); + b = mma(fragment<128>(kvs, wave * 32 + 16 + head, k * 128 + group * 32), q, b); + } + } + float scale = p.q_scale * p.kv_scale * p.attention_scale; +#pragma unroll + for(int j = 0; j < 4; ++j) + { + if constexpr(IdsMode == 0) + { + a[j] = selected[wave * 32 + group * 4 + j] >= 0 ? a[j] * scale : -INFINITY; + b[j] = selected[wave * 32 + 16 + group * 4 + j] >= 0 ? b[j] * scale : -INFINITY; + } + else + { + a[j] = ((valid_bits >> (group * 4 + j)) & 1) ? a[j] * scale : -INFINITY; + b[j] = ((valid_bits >> (16 + group * 4 + j)) & 1) ? b[j] * scale : -INFINITY; + } + } + float m = -INFINITY; +#pragma unroll + for(int j = 0; j < 4; ++j) + m = fmaxf(m, fmaxf(a[j], b[j])); +#if SM_LANE_MODE & 1 + if constexpr(Splits == 8 && Order == 0 && Prefetch && !TailBarrier && Waves == 4) + m = softmax_four_groups(m); + else + { + m = fmaxf(m, __shfl_xor(m, 16)); + m = fmaxf(m, __shfl_xor(m, 32)); + } +#else + m = fmaxf(m, __shfl_xor(m, 16)); + m = fmaxf(m, __shfl_xor(m, 32)); +#endif + if(group == 0) + exchange[wave * 16 + head] = m; + } + __syncthreads(); + float m = 0.f, inverse = 0.f, tile_lse = -INFINITY, online_alpha = 0.f; + if constexpr(LocalReduce) + { + if(wave < 4) + m = fmaxf(fmaxf(exchange[head], exchange[16 + head]), + fmaxf(exchange[32 + head], exchange[48 + head])); + } + else + { + if(wave == (Waves == 8 ? 4 : 0) && lane < 16) + row_max[lane] = fmaxf(fmaxf(exchange[lane], exchange[16 + lane]), + fmaxf(exchange[32 + lane], exchange[48 + lane])); + __syncthreads(); + } + if(wave < 4) + { + if constexpr(!LocalReduce) + m = row_max[head]; + if constexpr(Online) + { + const float newmax = fmaxf(merge_max, m); + online_alpha = merge_max == -INFINITY ? 0.f : exp2f((merge_max - newmax) * LOG2E); + merge_max = newmax; + m = newmax; + } + float z = 0.f; +#pragma unroll + for(int j = 0; j < 4; ++j) + { + a[j] = m == -INFINITY ? 0.f : exp2f((a[j] - m) * LOG2E); + b[j] = m == -INFINITY ? 0.f : exp2f((b[j] - m) * LOG2E); + z += a[j] + b[j]; + } +#if SM_LANE_MODE & 2 + if constexpr(Splits == 8 && Order == 0 && Prefetch && !TailBarrier && Waves == 4) + z = softmax_four_groups(z); + else + { + z += __shfl_xor(z, 16); + z += __shfl_xor(z, 32); + } +#else + z += __shfl_xor(z, 16); + z += __shfl_xor(z, 32); +#endif + if(group == 0) + { + if constexpr(LocalReduce) + exchange_sum[wave * 16 + head] = z; + else + exchange[wave * 16 + head] = z; + } + } + __syncthreads(); + if constexpr(LocalReduce) + { + if(wave < 4) + { + float z = (exchange_sum[head] + exchange_sum[16 + head]) + + (exchange_sum[32 + head] + exchange_sum[48 + head]); + if constexpr(Online) + merge_sum = merge_sum * online_alpha + z; + else + { + inverse = z > 0 ? 1.f / z : 0.f; + tile_lse = z > 0 ? m + log2f(z) * LN2 : -INFINITY; + } + } + } + else + { + if(wave == (Waves == 8 ? 4 : 0) && lane < 16) + { + float z = (exchange[lane] + exchange[16 + lane]) + + (exchange[32 + lane] + exchange[48 + lane]); + row_inv[lane] = z > 0 ? 1.f / z : 0.f; + p.partial_lse[(int64_t(query) * 17 + split) * 64 + hbase + lane] = + z > 0 ? row_max[lane] + log2f(z) * LN2 : -INFINITY; + } + __syncthreads(); + } + if(wave < 4) + { + float inv = (Online ? 1.f : (LocalReduce ? inverse : row_inv[head])) / p.p_scale; + uint32_t pa = pack4(a * inv), pb = pack4(b * inv); + *reinterpret_cast(ps + p_address(head, wave * 32 + group * 4)) = pa; + *reinterpret_cast(ps + p_address(head, wave * 32 + 16 + group * 4)) = pb; + } + __syncthreads(); + if(wave < 4) + { + float alpha = 0.f, beta = 0.f; + if constexpr(!Online) + { + if constexpr(Tiles > 1) + { + if(iteration == 0) + { + merge_max = tile_lse; + merge_sum = tile_lse == -INFINITY ? 0.f : 1.f; + } + else + { + float delta = exp2f(-fabsf((merge_max == -INFINITY && tile_lse == -INFINITY) + ? 0.f + : merge_max - tile_lse) * + LOG2E); + alpha = + merge_max == -INFINITY ? 0.f : (merge_max >= tile_lse ? 1.f : delta); + beta = tile_lse == -INFINITY ? 0.f : (tile_lse >= merge_max ? 1.f : delta); + merge_sum = merge_sum * alpha + beta; + merge_max = fmaxf(merge_max, tile_lse); + } + } + else + { + merge_max = tile_lse; + merge_sum = tile_lse == -INFINITY ? 0.f : 1.f; + } + } + I8 prob = p_fragment(ps, head, group * 32); + if constexpr(Online) + { + static_assert(PVGroup == 1 && LocalReduce); + if constexpr(PadFactored) + padded_online_pv_step<0, PVVGPR>(kvs, prob, wave, online_alpha, merged); + else + { +#pragma unroll + for(int tile = 0; tile < 8; ++tile) + { + I8 v = v_fragment_tr(kvs, wave * 128 + tile * 16); + merged[tile] = online_pv_mma(v, prob, merged[tile] * online_alpha); + } + } + } + else if constexpr(PVGroup == 4) + { +#pragma unroll + for(int base_tile = 0; base_tile < 8; base_tile += 4) + { + I8 values[4]; +#pragma unroll + for(int j = 0; j < 4; ++j) + values[j] = pv_pipe_fragment(kvs, wave * 128 + (base_tile + j) * 16); + // Explicit register inputs/outputs prevent wait moving past MFMA. + pv_pipe_wait(values); + F4 products[4]; +#pragma unroll + for(int j = 0; j < 4; ++j) + products[j] = mma(values[j], prob, F4{}); + asm volatile("" + : "+v"(products[0]), + "+v"(products[1]), + "+v"(products[2]), + "+v"(products[3])); +#pragma unroll + for(int j = 0; j < 4; ++j) + { + const int tile = base_tile + j; + F4 value = products[j] * (p.p_scale * p.kv_scale); + if(iteration == 0) + merged[tile] = value; + else + { +#pragma unroll + for(int k = 0; k < 4; ++k) + merged[tile][k] = fmaf(value[k], beta, merged[tile][k] * alpha); + } + } + } + } + else + { + if constexpr(PadFactored) + { + padded_pv_step<0>( + kvs, prob, wave, iteration, alpha, beta, p.p_scale * p.kv_scale, merged); + } + else + { +#pragma unroll + for(int tile = 0; tile < 8; ++tile) + { + I8 v; + if constexpr(Transpose) + v = v_fragment_tr(kvs, wave * 128 + tile * 16); + else + v = v_fragment(kvs, wave * 128 + tile * 16 + head, group * 32); + F4 out = mma(v, prob, F4{}); + out *= p.p_scale * p.kv_scale; + if(iteration == 0) + merged[tile] = out; + else + { +#pragma unroll + for(int j = 0; j < 4; ++j) + merged[tile][j] = fmaf(out[j], beta, merged[tile][j] * alpha); + } + } + } + } + } + if constexpr(TailBarrier && Tiles > 1) + __syncthreads(); + } + if(wave < 4) + { + float inv = merge_sum > 0 ? (Online ? p.p_scale * p.kv_scale : 1.f) / merge_sum : 0.f; +#pragma unroll + for(int tile = 0; tile < 8; ++tile) + { + int64_t off = + partial_address(((int64_t(query) * Splits + part) * 64 + hbase + head), + wave * 128 + tile * 16 + group * 4); + partial_store(p.partial_o + off, merged[tile] * inv); + } + if(wave == 0 && group == 0) + p.partial_lse[(int64_t(query) * Splits + part) * 64 + hbase + head] = + merge_sum > 0 ? merge_max + log2f(merge_sum) * LN2 : -INFINITY; + } +} + +template +__device__ __forceinline__ float split_reduce(float x) +{ + if constexpr(Permute) + { + // Rotate each 16-lane row by eight == lane xor 8. These builtins + // expose dependencies to LLVM's scheduling/hazard handling. + float a = __builtin_bit_cast( + float, __builtin_amdgcn_mov_dpp(__builtin_bit_cast(int, x), 0x128, 0xf, 0xf, false)); + x = Max ? fmaxf(x, a) : x + a; + // reducer. Pair-valued builtins lost a half under this compiler. + uint32_t seed16, bits16 = __builtin_bit_cast(uint32_t, x); + asm("v_mov_b32_e32 %0, %1" : "=v"(seed16) : "v"(bits16)); + asm("s_nop 2\n\tv_permlane16_swap_b32 %0, %1" : "+v"(seed16), "+v"(bits16)); + float l16 = __builtin_bit_cast(float, seed16), r16 = __builtin_bit_cast(float, bits16); + x = Max ? fmaxf(l16, r16) : l16 + r16; + // Seed both halves: one result is self, the other is xor-32 partner. + uint32_t seed32, bits32 = __builtin_bit_cast(uint32_t, x); + asm("v_mov_b32_e32 %0, %1" : "=v"(seed32) : "v"(bits32)); + asm("s_nop 2\n\tv_permlane32_swap_b32 %0, %1" : "+v"(seed32), "+v"(bits32)); + float l = __builtin_bit_cast(float, seed32), r = __builtin_bit_cast(float, bits32); + return Max ? fmaxf(l, r) : l + r; + } + else + { +#pragma unroll + for(int d = 8; d < 64; d *= 2) + { + float a = __shfl_xor(x, d, 64); + x = Max ? fmaxf(x, a) : x + a; + } + return x; + } +} + +// One common BF16 merge body; Parallel changes the split reduction tree only. +__device__ __forceinline__ void store_bf16(uint8_t* out, int64_t offset, F4 value) +{ + + uint32_t lo, hi; + asm("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(lo) : "v"(value[0]), "v"(value[1])); + asm("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(hi) : "v"(value[2]), "v"(value[3])); + *reinterpret_cast(reinterpret_cast(out) + offset) = I2{int(lo), int(hi)}; +} +template +__device__ __forceinline__ F4 load_partial(const Params& p, int64_t row, int col) +{ + return partial_load(p.partial_o + partial_address(row, col)); +} +template +__global__ __launch_bounds__((Parallel ? 4 : 1) * 64) void merge(Params p) +{ + const int lane = threadIdx.x & 63, wave = threadIdx.x / 64; + if constexpr(Parallel) + { + constexpr int R = (Splits + 7) / 8; + const CtaTile work = cta_tile<64, 4, Remap>(); + const int split = lane / 8, col = work.c * 128 + wave * 32 + (lane % 8) * 4, head = work.b; + const int64_t row = int64_t(work.q) * 64 + head; + float ls[R]; + F4 values[R]; + float m = -INFINITY; +#pragma unroll + for(int r = 0; r < R; ++r) + { + const int s = split + r * 8; + const int64_t stat = (int64_t(work.q) * Splits + s) * 64 + head; + ls[r] = s < Splits ? p.partial_lse[stat] : -INFINITY; + values[r] = s < Splits ? load_partial(p, stat, col) : F4{}; + m = fmaxf(m, ls[r]); + } + m = split_reduce(m); + float weights[R]; + float z = 0.f; +#pragma unroll + for(int r = 0; r < R; ++r) + { + weights[r] = ls[r] == -INFINITY ? 0.f : exp2f((ls[r] - m) * LOG2E); + z += weights[r]; + } + z = split_reduce(z); + F4 out{}; + const float inv = z > 0 ? 1.f / z : 0.f; +#pragma unroll + for(int j = 0; j < 4; ++j) + { + float x = 0.f; +#pragma unroll + for(int r = 0; r < R; ++r) + x += weights[r] > 0 ? values[r][j] * weights[r] : 0.f; + out[j] = split_reduce(x) * inv; + } + if(split == 0) + store_bf16(p.out, row * 512 + col, out); + if(p.lse && work.c == 0 && threadIdx.x == 0) + p.lse[row] = z > 0 ? m + log2f(z) * LN2 : -INFINITY; + } + else + { + const CtaTile work = cta_tile<4, 32, Remap>(); + const int head = work.b * 16 + lane / 4, col = work.c * 16 + (lane % 4) * 4; + const int64_t row = int64_t(work.q) * 64 + head; + float ls[Splits], m = -INFINITY, z = 0.f; + F4 accum{}; +#pragma unroll + for(int s = 0; s < Splits; ++s) + { + ls[s] = p.partial_lse[(int64_t(work.q) * Splits + s) * 64 + head]; + m = fmaxf(m, ls[s]); + } +#pragma unroll + for(int s = 0; s < Splits; ++s) + { + float w = m == -INFINITY ? 0.f : exp2f((ls[s] - m) * LOG2E); + z += w; + F4 value = + load_partial(p, (int64_t(work.q) * Splits + s) * 64 + head, col); +#pragma unroll + for(int j = 0; j < 4; ++j) + accum[j] = fmaf(w > 0 ? value[j] : 0.f, w, accum[j]); + } + float inv = z > 0 ? 1.f / z : 0.f; + store_bf16(p.out, row * 512 + col, accum * inv); + if(p.lse && work.c == 0 && lane % 4 == 0) + p.lse[row] = z > 0 ? m + log2f(z) * LN2 : -INFINITY; + } +} + +template +hipError_t launch(Params p, hipStream_t stream) +{ + hipLaunchKernelGGL((attention<17, + 0, + true, + false, + 4, + false, + BF16, + Unit, + 16, + true, + true, + 1, + 0, + false, + true, + true>), + dim3(p.queries, 17, 4), + dim3(256), + 0, + stream, + p); + auto status = hipGetLastError(); + if(status != hipSuccess) + return status; + hipLaunchKernelGGL( + (merge<17, 16, true, true, false>), dim3(p.queries, 4, 32), dim3(64), 0, stream, p); + return hipGetLastError(); +} +} // namespace sparse_mla_fp8::small +namespace sparse_mla_fp8 { +hipError_t launch_small(Params p, hipStream_t stream) +{ + if(!p.bf16_q) + return small::launch(p, stream); + if(p.q_scale == 1.f) + return small::launch(p, stream); + return small::launch(p, stream); +} +} // namespace sparse_mla_fp8 diff --git a/csrc/kernels/mla/fp8/softmax_lane.cuh b/csrc/kernels/mla/fp8/softmax_lane.cuh new file mode 100644 index 00000000000..4bb33382294 --- /dev/null +++ b/csrc/kernels/mla/fp8/softmax_lane.cuh @@ -0,0 +1,21 @@ +// SPDX-License-Identifier: MIT +#pragma once +#ifndef SM_LANE_MODE +#define SM_LANE_MODE 0 +#endif +// Preserve native xor16 then xor32 arithmetic order. Gluon uses these +// opcodes but its reduction tree/FP8 arithmetic contract is different. +template +__device__ __forceinline__ float softmax_four_groups(float x) +{ + uint32_t seed16, bits16 = __builtin_bit_cast(uint32_t, x); + asm("v_mov_b32_e32 %0, %1" : "=v"(seed16) : "v"(bits16)); + asm("s_nop 2\n\tv_permlane16_swap_b32 %0, %1" : "+v"(seed16), "+v"(bits16)); + float l16 = __builtin_bit_cast(float, seed16), r16 = __builtin_bit_cast(float, bits16); + x = Max ? fmaxf(l16, r16) : l16 + r16; + uint32_t seed32, bits32 = __builtin_bit_cast(uint32_t, x); + asm("v_mov_b32_e32 %0, %1" : "=v"(seed32) : "v"(bits32)); + asm("s_nop 2\n\tv_permlane32_swap_b32 %0, %1" : "+v"(seed32), "+v"(bits32)); + float l32 = __builtin_bit_cast(float, seed32), r32 = __builtin_bit_cast(float, bits32); + return Max ? fmaxf(l32, r32) : l32 + r32; +} diff --git a/csrc/kernels/mla/fp8/streaming.cu b/csrc/kernels/mla/fp8/streaming.cu new file mode 100644 index 00000000000..1408e2f0914 --- /dev/null +++ b/csrc/kernels/mla/fp8/streaming.cu @@ -0,0 +1,303 @@ +// SPDX-License-Identifier: MIT +#include "online_common.cuh" +#include "reduce.cuh" +namespace sparse_mla_fp8::online { +template +__global__ __launch_bounds__(Waves * 64) void streaming_attention(Params p) +{ + + static_assert(DoubleBuffer && !Gradual); + constexpr int WPH = Waves * 16 / HeadRows, Fragments = BlockK / (WPH * 16), PVTiles = 32 / WPH; + static_assert((HeadRows == 32 || HeadRows == 64) && (WPH == 1 || WPH == 2)); + static_assert(BlockK == 128 || (BlockK == 64 && (Waves == 4 || Waves == 8))); + __shared__ __align__(16) + uint8_t kv_storage[DoubleBuffer ? 2 : 1][(BlockK / 2) * KV_CHUNK_STRIDE]; + __shared__ __align__(16) uint8_t ps[HeadRows * BlockK]; + __shared__ __align__(16) uint8_t qs[QMask ? HeadRows * 128 * __builtin_popcount(QMask) : 1]; + __shared__ float maxs[WPH == 2 ? Waves * 16 : 1], sums[WPH == 2 ? Waves * 16 : 1]; + __shared__ int selected_storage[DoubleBuffer ? 2 : 1][BlockK]; + const int tid = threadIdx.x, lane = tid & 63, wave = tid / 64, head = lane & 15, + group = lane / 16; + const auto work = order_ids( + blockIdx.x, blockIdx.y, blockIdx.z, p.queries); + const int query = work.query, part = work.part, half = wave % WPH, hgroup = wave / WPH; + const int hbase = work.head * HeadRows + hgroup * 16; + const int start = (query * p.capacity), end = ((query + 1) * p.capacity); + const int total_tiles = (end - start + BlockK - 1) / BlockK, quotient = total_tiles / Splits, + remainder = total_tiles % Splits; + const int tile_begin = part * quotient + (part < remainder ? part : remainder), + tile_count = quotient + (part < remainder ? 1 : 0); + const int64_t qoff = + int64_t(query) * p.q_stride_row + int64_t(hbase + head) * p.q_stride_head + group * 32; + I8 retained[4]; +#pragma unroll + for(int k = 0; k < 4; ++k) + if(!(QMask & (1 << k))) + retained[k] = q_load(p, qoff + k * 128); + F4 merged[PVTiles]{}; + float merge_max = -INFINITY, merge_sum = 0.f; + int next_id = -1; + if constexpr(DoubleBuffer) + { + static_assert(BlockK == 128 && HeadRows == 64 && Waves == 8 && QMask == 0 && PVMode <= 2); + if(tile_count > 0) + { + stage_iddirect( + p, kv_storage[0], selected_storage[0], start, end, tile_begin); + asm volatile("s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)" : : : "memory"); + __syncthreads(); + } + } +#pragma unroll 1 + for(int iteration = 0; iteration < tile_count; ++iteration) + { + const int split = tile_begin + iteration; + uint8_t* kvs = kv_storage[DoubleBuffer ? (iteration & 1) : 0]; + int* selected = selected_storage[DoubleBuffer ? (iteration & 1) : 0]; + if constexpr(DoubleBuffer) + { + if(iteration + 1 < tile_count) + stage_iddirect(p, + kv_storage[(iteration + 1) & 1], + selected_storage[(iteration + 1) & 1], + start, + end, + split + 1); + } + else + { + if(tid < BlockK) + { + const int at = start + split * BlockK + tid; + const int id = iteration == 0 ? (at < end ? p.indices[at] : -1) : next_id; + selected[tid] = (id >= 0 && id < p.slots) ? id : -1; + } + __syncthreads(); + if(tid < BlockK && iteration + 1 < tile_count) + { + const int at = start + (split + 1) * BlockK + tid; + next_id = at < end ? p.indices[at] : -1; + } + const int id0 = selected[(wave % (BlockK / 16)) * 16 + lane / 4]; + const int id1 = + (Waves == 4 && BlockK == 128) ? selected[(wave + 4) * 16 + lane / 4] : -1; +#pragma unroll + for(int i = 0; i < (BlockK / 2) / Waves; ++i) + { + const int chunk = wave + i * Waves, r = (chunk % (BlockK / 16)) * 16 + lane / 4; + const int c = (chunk / (BlockK / 16)) * 64 + ((lane % 4) * 16 ^ ((r & 3) * 16)); + const int id = (Waves == 4 && BlockK == 128 && (i & 1)) ? id1 : id0; + if(id >= 0) + direct_load(kvs + chunk * KV_CHUNK_STRIDE, p.kv + int64_t(id) * 512 + c); + else + *reinterpret_cast(kvs + chunk * KV_CHUNK_STRIDE + lane * 16) = U4{}; + } + asm volatile("s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)" : : : "memory"); + __syncthreads(); + } + F4 score[Fragments]{}; +#define TUNE_QK(K) \ + qk_step( \ + kvs, \ + half * (BlockK / WPH) + head, \ + group * 32, \ + q_fragment( \ + p, retained, qs, qoff, iteration, hgroup * 16 + head, group * 32, half), \ + score) + TUNE_QK(0); + TUNE_QK(1); + TUNE_QK(2); + TUNE_QK(3); +#undef TUNE_QK + // Inline MFMA does not expose its result latency to the scheduler. + // Fence every H32 score operand before non-MFMA consumers; otherwise + // a different compiler schedule may read it before the 12-cycle delay. + if constexpr(Fragments == 2) + asm volatile("s_nop 7\n\ts_nop 3" : "+v"(score[0]), "+v"(score[1]) : : "memory"); + float m = -INFINITY; + const float scale = (p.q_scale * p.kv_scale * p.attention_scale) * LOG2E; +#pragma unroll + for(int n = 0; n < Fragments; ++n) + { +#pragma unroll + for(int j = 0; j < 4; ++j) + { + const int key = half * (BlockK / WPH) + n * 16 + group * 4 + j; + score[n][j] = selected[key] >= 0 ? score[n][j] * scale : -INFINITY; + m = fmaxf(m, score[n][j]); + const int slot = split * BlockK + key; + } + } + m = fmaxf(m, __shfl_xor(m, 16)); + m = fmaxf(m, __shfl_xor(m, 32)); + if constexpr(WPH == 2) + { + if(group == 0) + maxs[wave * 16 + head] = m; + __syncthreads(); + m = fmaxf(m, maxs[(wave ^ 1) * 16 + head]); + } + const float newmax = fmaxf(merge_max, m); + const float alpha = merge_max == -INFINITY ? 0.f : domain_exp2(merge_max - newmax); + merge_max = newmax; + m = newmax; + float zz[Fragments / 2]{}; +#pragma unroll + for(int pair = 0; pair < Fragments / 2; ++pair) + { +#pragma unroll + for(int j = 0; j < 4; ++j) + { + score[2 * pair][j] = + m == -INFINITY ? 0.f : domain_exp2(score[2 * pair][j] - m); + score[2 * pair + 1][j] = + m == -INFINITY ? 0.f : domain_exp2(score[2 * pair + 1][j] - m); + zz[pair] += score[2 * pair][j] + score[2 * pair + 1][j]; + } + zz[pair] += __shfl_xor(zz[pair], 16); + zz[pair] += __shfl_xor(zz[pair], 32); + } + float local_z = zz[0]; + if constexpr(Fragments >= 4) + local_z += zz[1]; + if constexpr(Fragments == 8) + local_z += zz[2] + zz[3]; + if constexpr(WPH == 2) + if(group == 0) + sums[wave * 16 + head] = local_z; +// Publish P together with the sum exchange, saving the separate barrier. +#pragma unroll + for(int n = 0; n < Fragments; ++n) + { + const int key = half * (BlockK / WPH) + n * 16 + group * 4; + uint32_t prob = pack4(score[n] * (1.f / p.p_scale)); + *reinterpret_cast(ps + hgroup * 16 * BlockK + p_address(head, key)) = + prob; + const int slot = split * BlockK + key; + } + float z = local_z; + if constexpr(WPH == 2) + { + __syncthreads(); + const float other = sums[(wave ^ 1) * 16 + head]; + z = half == 0 ? local_z + other : other + local_z; + } + else + { + // Each wave owns sixteen heads; no other wave reads its P region. + asm volatile("s_waitcnt lgkmcnt(0)" : : : "memory"); + } + merge_sum = merge_sum * alpha + z; + if(group == 0 && half == 0 && (BlockK == 128 || (split % 2) == 0)) + { + const int tile128 = (split * BlockK) / 128; + const int64_t off = (int64_t(query) * 17 + tile128) * 64 + hbase + head; + } + const uint8_t* pvbase = kvs + half * (8 / WPH) * (BlockK / 16) * KV_CHUNK_STRIDE; + if constexpr(PVMode >= 3) + { + static_assert(BlockK == 64 && (PVMode == 3 || PVMode == 4)); + I2 p0 = p32<0>(ps + hgroup * 16 * BlockK, head, group); + I2 p1 = p32<1>(ps + hgroup * 16 * BlockK, head, group); + asm volatile("s_waitcnt lgkmcnt(0)" : "+v"(p0), "+v"(p1) : : "memory"); + pv32<0, PVTiles, PVMode == 4>(pvbase, p0, p1, merged, alpha); + } + else + { + I8 prob = p_fragment(ps + hgroup * 16 * BlockK, head, group * 32); + asm volatile("s_waitcnt lgkmcnt(0)" : "+v"(prob) : : "memory"); + if constexpr(PVMode == 2) + { + static_assert(BlockK == 128); + I8 first = v_read<0, 128>(pvbase); + pv_pipeline<0, PVTiles>(pvbase, prob, merged, alpha, first); + } + else + pv<0, PVTiles, BlockK, PVMode>(pvbase, prob, merged, alpha); + } + if constexpr(DoubleBuffer) + { + // Wait for next KV writes AND this wave's current LDS readers. + // All waves finish before the old current buffer can be reused. + asm volatile("s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)" : : : "memory"); + __syncthreads(); + } + } + const float inv = merge_sum > 0 ? p.p_scale * p.kv_scale / merge_sum : 0.f; +#pragma unroll + for(int tile = 0; tile < PVTiles; ++tile) + { + F4 x = merged[tile] * inv; + const int col = half * (512 / WPH) + tile * 16 + group * 4; + if constexpr(Splits == 1) + { + const int64_t off = (int64_t(query) * 64 + hbase + head) * 512 + col; + + uint32_t lo, hi; + asm("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(lo) : "v"(x[0]), "v"(x[1])); + asm("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(hi) : "v"(x[2]), "v"(x[3])); + *reinterpret_cast(reinterpret_cast(p.out) + off) = I2{int(lo), int(hi)}; + } + else + { + const int64_t off = ((int64_t(query) * Splits + part) * 64 + hbase + head) * 512 + col; + partial_store(p.partial_o + off, x); + } + } + if(group == 0 && half == 0) + { + const float l = merge_sum > 0 ? (merge_max + log2f(merge_sum)) * LN2 : -INFINITY; + if constexpr(Splits == 1) + { + if(p.lse) + p.lse[int64_t(query) * 64 + hbase + head] = l; + } + else + p.partial_lse[(int64_t(query) * Splits + part) * 64 + hbase + head] = l; + } +} + +template +hipError_t streaming_typed(Params p, hipStream_t stream) +{ + hipLaunchKernelGGL((streaming_attention), + dim3(p.queries, S, 1), + dim3(512), + 0, + stream, + p); + auto status = hipGetLastError(); + if(status != hipSuccess) + return status; + if constexpr(S > 1) + hipLaunchKernelGGL( + (reduce_splits), dim3(p.queries, 64, 2), dim3(64), 0, stream, p); + return hipGetLastError(); +} +template +hipError_t streaming(Params p, hipStream_t stream) +{ + if(p.queries <= 128) + return streaming_typed<2, BF16, Unit>(p, stream); + return streaming_typed<1, BF16, Unit>(p, stream); +} +} // namespace sparse_mla_fp8::online +namespace sparse_mla_fp8 { +hipError_t launch_streaming(Params p, hipStream_t stream) +{ + if(!p.bf16_q) + return online::streaming(p, stream); + if(p.q_scale == 1.f) + return online::streaming(p, stream); + return online::streaming(p, stream); +} +} // namespace sparse_mla_fp8 diff --git a/csrc/kernels/mla/hk/sparse_mla_bf16_common.cuh b/csrc/kernels/mla/hk/sparse_mla_bf16_common.cuh deleted file mode 100644 index 07a41673df9..00000000000 --- a/csrc/kernels/mla/hk/sparse_mla_bf16_common.cuh +++ /dev/null @@ -1,878 +0,0 @@ -// SPDX-License-Identifier: MIT -// Derived from ROCm/aiter PR #3459 (e961cb0c), HipKittens d3cd9b31. -// Sparse MLA with a shared BF16 K/V latent: D512, no appended RoPE, H16/H64. -#pragma once - -#include "hk_mla_v40_fwd_decode_gen1_common.cuh" -#include -#include - -namespace sparse_mla_bf16 { - -template -struct Traits -{ - static_assert(Heads == 16 || Heads == 64); - static constexpr uint32_t kBlockM = Heads, kNumWarps = Heads / 16; - static constexpr uint32_t kBlockN = 32, kBlockK = 32, kTileM = 16; - static constexpr uint32_t kQkHeadDim = 512, kVoHeadDim = 512; - static constexpr uint32_t kQkNopeHeadDim = 512, kQkRopeHeadDim = 0; - static constexpr uint32_t kNumThreads = 64 * kNumWarps; - static constexpr uint32_t kRoundMode = 0; // nearest even, including BF16 ties - // H16 keeps one KV tile for four LDS-limited workgroups/CU; H64 uses two - // tiles to overlap next-KV DMA with current MFMA. Q staging overlays this - // allocation only in the prologue; terminal O bounce overlays dead KV0. - static constexpr uint32_t kKvBytes = 32 * 512 * 2; - static constexpr uint32_t kKvBuffers = Heads == 64 ? 2u : 1u; - static constexpr uint32_t kIndexOffset = kKvBuffers * kKvBytes; - static constexpr uint32_t kLdsBytes = kIndexOffset + 256; - static_assert(kLdsBytes <= 80 * 1024); - static_assert(kNumWarps * 4352 <= kKvBytes); -}; - -struct Params -{ - const uint16_t* q; - const uint16_t* kv; - const int32_t* indptr; - const int32_t* indices; - uint16_t* out; - float* lse; - float* partial_o; - float* partial_lse; - int32_t queries; - int32_t splits; - int64_t slots; - int64_t q_stride0; - int64_t q_stride1; - int64_t kv_stride0; - float scale; -}; - -// These BF16 LDS readers are copied from PR #3459 unchanged. They intentionally -// have no FP8/scale/RoPE loaders or traits: the new writer supplies their exact -// sb8-permuted, row-bank-swizzled 32x512 BF16 layout. -template -struct Bf16KvManager -{ - static constexpr uint32_t kLoadBlockN = 32, kSubBlockRows = 16; - static constexpr uint32_t kSubBlockCols = 32, kSubBlockBytes = 1024; - static constexpr uint32_t kNumRowTiles = 2; - __device__ __forceinline__ static constexpr uint32_t sub_block_byte_offset(uint32_t row_tile, - uint32_t col_tile) - { - return (col_tile * kNumRowTiles + row_tile) * kSubBlockBytes; - } - // ---- LDS -> VGPR readout for QK / PV mfmas ----------------------------- - // QK A-tile load: ds_read_b128 of one 16 x 32 bf16 sub-block into 4 vgprs. - // (kRowOffset, kColOffset) selects which (row_tile, col_tile) of the pong; - // the per-lane offset within the sub-block follows the mfma_f32_16x16x32_bf16 - // A-operand layout (lane = (row_in_tile, group_in_row)). - // kRowOffset spans the FULL logical tile (0/16 for kBlockN=32; 0/16/32/48 for - // kBlockN=64). The sub-tile (kRowOffset/32) is folded into the ds_read imm as - // kSubIdx*kSubPongBytes so A and B share the same base ptr -- avoids a per-call - // base+kSubPong address VGPR that CSE holds across QK and overlaps a pinned reg. - template - __device__ __forceinline__ static void load_k_to_gpr(RT& dst, const uintptr_t p_lds_kv) - { - static_assert(kRowOffset % kSubBlockRows == 0, - "load_k_to_gpr: kRowOffset must be a multiple of 16."); - static_assert(kColOffset % kSubBlockCols == 0, - "load_k_to_gpr: kColOffset must be a multiple of 32."); - static_assert(kRowOffset < T::kBlockN, "load_k_to_gpr: kRowOffset out of range."); - static_assert(kColOffset < T::kQkHeadDim, "load_k_to_gpr: kColOffset out of range."); - // Sub-tile B (kRowOffset >= 32) -> +kSubPongBytes in the ds_read imm. - constexpr uint32_t kSubIdx = kRowOffset / kLoadBlockN; - constexpr uint32_t kRowInSub = kRowOffset % kLoadBlockN; - constexpr uint32_t kSubPongBytes = kLoadBlockN * T::kQkHeadDim * sizeof(hk::bf16); - - // mfma_f32_16x16x32_bf16 A-operand layout: lane t holds 8 bf16 from - // row r = lane%16, cols [c, c+8) where c = (lane/16) * 8. - // 8 bf16 = 4 dwords -> ds_read_b128. - constexpr uint32_t kMfmaRows = 16; - constexpr uint32_t kMfmaElemPerThr = 8; - - const uint32_t lane_idx = opus::lane_id(); - const uint32_t row = lane_idx % kMfmaRows; - const uint32_t col = (lane_idx / kMfmaRows) * kMfmaElemPerThr; - - // Un-swizzle: writer XORs intra-sub-block byte position by 32 on - // sub-tile-rows 1 & 3 (rows 4..7 and 12..15) to break the 2-way bank - // conflict; the reader applies the same XOR on the col-byte component. - const uint32_t row_bank_swap = ((row >> 2) & 1u) << 5; - const uint32_t in_sb_byte = - row * (kSubBlockCols * sizeof(hk::bf16)) + ((col * sizeof(hk::bf16)) ^ row_bank_swap); - - // Constexpr sub-block selector (compiles to immediate offset); the sub-tile - // (A/B) contributes kSubIdx*kSubPongBytes so the base ptr stays p_lds_kv. - constexpr uint32_t kFixedOffset = - sub_block_byte_offset(kRowInSub / kSubBlockRows, kColOffset / kSubBlockCols) + - kSubIdx * kSubPongBytes; - - // RT must hold a single 4-vgpr range (16 bf16 mfma A-tile = 4 vgprs/lane). - using range_type = hkdart::get_nth_range_t; - static_assert(range_type::lo + 3 == range_type::hi, - "ds_read_b128 requires 4 consecutive registers"); - - const uintptr_t p_lds_kv_lane = p_lds_kv + in_sb_byte; - hkm::ds_read_b128(static_cast(p_lds_kv_lane), kFixedOffset); - } - - // PV A-tile load: ds_read_b64_tr_b16 (bf16 transpose-read) of one 16-row x 16-col - // bf16 patch from a sub-block, results land in 2 dwords/lane in (GPR, GPR+1). - // - // PV math is V^T @ P^T = O^T computed via mma_ABt(oaccu, kv, p_mfma) (= kv @ p_mfma^T, - // matching the QK convention of K^T @ Q^T = P^T). So `kv` is the A operand of - // v_mfma_f32_16x16x32_bf16, holding V^T values reorganized into the mfma A layout. - // - // Within each 16-lane group, lane t's 4 bf16 (4 bf16 = 1 b64 = 2 dwords/lane) are - // (after HW transpose): - // output_lane[g*16 + l] holds V[g*4+0..g*4+3, kColOffset + l] - // for g = lane_group_idx (0..3), l = lane_in_group (0..15). I.e. each lane gets - // 4 K-rows of one V-col. Caller stitches two row halves (kRowOffset = 0, then 16) - // into a single mfma A operand spanning 8 K-rows (= mfma K = 0..7). - // - // Per-lane source address (within the selected 16x32 sub-block): - // in_sb_byte = (lane >> 2) * (kSubBlockCols * sizeof(bf16)) + (lane & 3) * 8 - // = lane_row * 64 + lane_col_quad * 8 - // (row stride = 32 bf16 cols * 2 B = 64 B; each "col_quad" = 4 bf16 = 8 B.) - // - // Compile-time fixed_offset selects: - // * the sub-block (row_tile = kRowOffset/16, col_tile = kColOffset/32), and - // * the 16-col half within that 32-col sub-block: kColOffset%32 -> +0 or +32 B. - // - // Un-swizzle: writer XORs intra-sub-block byte position by 32 on rows - // whose sub-tile-row index is odd (rows 4..7 and 12..15 within the 16-row - // sub-block). The reader applies the same XOR. With this swizzle both - // cycles of ds_read_b64_tr_b16 (lanes 0..31 covering rows 0..7, lanes - // 32..63 covering rows 8..15) hit 32 distinct conflict slots -- fully - // conflict-free for this per-wave instruction mapping. - // kRowOffset spans the FULL logical tile (0/16 = sub-tile A, 32/48 = sub-tile B). - // The sub-tile (kRowOffset/32) is folded into the ds_read imm as - // kSubIdx*kSubPongBytes so the base ptr stays p_lds_v (no +kSubPong VGPR). - template - __device__ __forceinline__ static void load_transposed_v_to_gpr(const uintptr_t p_lds_v) - { - static_assert((kRowOffset % kSubBlockRows == 0u) && (kRowOffset < T::kBlockN), - "load_transposed_v_to_gpr: kRowOffset must be 0/16 (A) or 32/48 (B)."); - static_assert( - (kColOffset % 16u == 0u) && (kColOffset < T::kVoHeadDim), - "load_transposed_v_to_gpr: kColOffset must be a multiple of 16, < kVoHeadDim."); - - constexpr uint32_t kSubIdx = kRowOffset / kLoadBlockN; // 0 (A) or 1 (B) - constexpr uint32_t kRowInSub = kRowOffset % kLoadBlockN; // 0 or 16 - constexpr uint32_t kSubPongBytes = kLoadBlockN * T::kQkHeadDim * sizeof(hk::bf16); - constexpr uint32_t kRowTile = kRowInSub / kSubBlockRows; // 0 or 1 - constexpr uint32_t kColTile = kColOffset / kSubBlockCols; // 0..15 - constexpr uint32_t kColInSbBytes = (kColOffset % kSubBlockCols) * sizeof(hk::bf16); - // Bank-swizzle re-expressed as a conditional ±32 delta so that - // (kFixedOffset + kColInSbBytes) stays fully constexpr in the - // ds_read_b64_tr_b16 immediate offset: XOR-by-32 against a constexpr - // value flips bit 5, equivalent to "+32 if bit was 0, else -32". - // The sign is compile-time (from kColInSbBytes's bit 5); only the - // boolean `is_swz` (1 bit per lane) is runtime. Avoids materialising - // kColInSbBytes as a runtime VGPR (vs. the plain XOR formulation), - // which freed 2 unpinned VGPRs in the audit. - constexpr int32_t kSwzDelta = (kColInSbBytes & 32u) ? -32 : +32; - constexpr uint32_t kFixedOffset = - sub_block_byte_offset(kRowTile, kColTile) + kColInSbBytes + kSubIdx * kSubPongBytes; - - const uint32_t lane_idx = opus::lane_id(); - const uint32_t row_in_sb = lane_idx >> 2; - const uint32_t is_swz = (row_in_sb >> 2) & 1u; - const uint32_t in_sb = row_in_sb * (kSubBlockCols * sizeof(hk::bf16)) + - (lane_idx & 3u) * 8u + is_swz * static_cast(kSwzDelta); - const uint32_t addr = static_cast(p_lds_v) + in_sb; - - hkm::ds_read_b64_tr_b16(addr, kFixedOffset); - } - - // bf16 ds_read_b64_tr_b16 already lands in the mfma A-operand layout -- no - // intra-lane v_swap_b32 fixup needed (V32's fp8 path interleaved cols c and - // c+16 into the same 2 GPRs and required a swap; the b16 transpose does not). - // Kept as a no-op for caller parity with KvManager8bitsV3. - template - __device__ __forceinline__ static void finalize_load_transposed_v_to_gpr() - { - } -}; - -// PR PV schedule, changed only to accept the BF16-only LDS manager. -template -__device__ __forceinline__ void -pv_gemm(Bf16KvManager& kv_manager, const uintptr_t p_lds_v, const float rescale) -{ - using R = HkMlaV40Regs; - using comp_t = typename R::comp_t; - constexpr uint32_t k_o_begin = R::k_o_begin; - constexpr uint32_t k_v0_begin = R::k_v0_begin; - constexpr uint32_t k_k0_begin = R::k_k0_begin; - constexpr uint32_t k_k1_begin = R::k_k1_begin; - constexpr uint32_t k_k2_begin = R::k_k2_begin; - // kBlockN=64 (m16x8 double-tile) contracts BOTH 32-row sub-tiles in one call: - // sub-tile A (rows 0:32) uses p_mfma_a, B (rows 32:64) uses p_mfma_b. kBlockN=32 - // has only A. Picked per-iter by row_base (see below), no caller arg. - typename R::p_mfma_a_t p_mfma_a; - typename R::p_mfma_b_t p_mfma_b; - typename R::pv_v_0_t pv_v_0; - typename R::pv_v_1_t pv_v_1; - typename R::pv_v_2_t pv_v_2; - typename R::pv_v_3_t pv_v_3; - - // D-tiling: each D-iter emits 2 oaccu 16x16 sub-tiles (32 D-cols); 512/32 = 16 - // D-iters per 32-row KV sub-tile. kBlockN=64 runs BOTH sub-tiles in one call, so - // num_pv_iter = kNumKvSub * 16. iter [0,16) = sub-tile A, [16,32) = sub-tile B. - constexpr uint32_t kDIters = T::kVoHeadDim / (2u * T::kTileM); // 16 - constexpr uint32_t num_pv_iter = R::kNumKvSub * kDIters; // 16 or 32 - // Flat V base tiles across all sub-tiles: 32 per sub-tile (S_0..S_31 each). - constexpr uint32_t kVTilesPerSub = 2u * kDIters; // 32 - constexpr uint32_t kNumVTiles = R::kNumKvSub * kVTilesPerSub; // 32 or 64 - - auto pk_mul_pair = [&](float r, auto base_c) { - constexpr uint32_t base = decltype(base_c)::value; - const float2 r2 = {r, r}; - asm volatile("v_pk_mul_f32 v[%0:%1], %2, v[%0:%1]" : : "n"(base), "n"(base + 1), "v"(r2)); - }; - auto mul_pair = [&](float r, auto base_c) { - constexpr uint32_t base = decltype(base_c)::value; - asm volatile("v_mul_f32_e32 v[%0], %1, v[%0]" : : "n"(base), "v"(r)); - asm volatile("v_mul_f32_e32 v[%0], %1, v[%0]" : : "n"(base + 1), "v"(r)); - }; - - // Issue both ds_read_b64 for base tile S_jj (within its sub-tile) into a - // round-robin slot (4 slots). kRowBase selects the KV sub-tile (0=A, 32=B), - // folded into the V ds_read imm by load_transposed_v_to_gpr. - auto load_S = [&]() { - constexpr uint32_t base = (slot == 0u) ? k_v0_begin - : (slot == 1u) ? k_k0_begin - : (slot == 2u) ? k_k1_begin - : k_k2_begin; - kv_manager.template load_transposed_v_to_gpr(p_lds_v); - kv_manager.template load_transposed_v_to_gpr(p_lds_v); - }; - // mfma: oaccu_dst (+)= pv_v_{slot}^T @ p_mfma[sub-tile]. 3-arg init only on the - // FIRST sub-tile of a first-iter call (kRowBase==0); sub-tile B always accumulates - // onto A's oaccu (same oaccu tile, second N-half of P). - auto do_mma = [&](OA& oaccu_dst) { - auto run = [&](auto& v) { - if constexpr(kIsFirstIter && kRowBase == 0u) - hk::mma_ABt(oaccu_dst, v, p_mfma_a); - else if constexpr(kRowBase == 0u) - hk::mma_ABt(oaccu_dst, v, p_mfma_a, oaccu_dst); - else - hk::mma_ABt(oaccu_dst, v, p_mfma_b, oaccu_dst); - }; - if constexpr(slot == 0u) - run(pv_v_0); - else if constexpr(slot == 1u) - run(pv_v_1); - else if constexpr(slot == 2u) - run(pv_v_2); - else - run(pv_v_3); - }; - - if constexpr(kDoRescale) - { - opus::static_for<2>([&](auto s) { - pk_mul_pair(rescale, opus::number{}); - pk_mul_pair(rescale, opus::number{}); - }); - } - - // Prologue: preload the first 3 flat V tiles (S_0,S_1,S_2 of sub-tile A) into - // slots 0,1,2. ONE prologue for the whole call -- sub-tile B's S_0..S_2 are - // prefetched by sub-tile A's tail iters via the flat prefetch index below. - load_S.template operator()<0u, 0u, 0u>(); - load_S.template operator()<0u, 1u, 1u>(); - load_S.template operator()<0u, 2u, 2u>(); - - opus::static_for([&](auto i) { - constexpr uint32_t iter = i.value; - constexpr uint32_t col_idx = iter % kDIters; // 0..15 D-tile within sub-tile - constexpr uint32_t row_base = (iter / kDIters) * 32u; // 0 (A) or 32 (B) - constexpr bool has_next = (iter + 1u) < num_pv_iter; - // Rescale only folds on the FIRST sub-tile pass (row_base==0) and only when a - // NEXT D-tile in that pass exists (col_idx+1 < kDIters) -- else oaccu overflows. - constexpr bool resc_next = kDoRescale && (row_base == 0u) && (col_idx + 1u < kDIters); - constexpr uint32_t next_oaccu_base = k_o_begin + (col_idx + 1u) * 8u; - - // Flat V-tile index across all sub-tiles: prefetch runs on this so sub-tile A's - // tail prefetches sub-tile B's S_0.. (ONE prologue). Consume uses col_idx. - constexpr uint32_t jf_lo = 2u * iter; // flat mfma idx (a) - constexpr uint32_t jf_hi = 2u * iter + 1u; // flat mfma idx (b) - // 4-slot round-robin (flat tile T -> slot T%4). Reading slot j%4 while refilling - // slot (j+3)%4 keeps the refill ds_read's dst != the mfma's src reg. - constexpr uint32_t slot_lo = jf_lo % 4u; - constexpr uint32_t slot_hi = jf_hi % 4u; - - constexpr uint32_t oaccu_base = k_o_begin + col_idx * 8u; - using oaccu_a_r = - hkdart::split_many_t>, - 4>; - using oaccu_b_r = - hkdart::split_many_t>, - 4>; - hk::art oaccu_a; - hk::art oaccu_b; - - // Prefetch S at flat index jf+3 -> its own sub-tile row_base + local jj%32. - auto prefetch = [&]() { - if constexpr(jf < kNumVTiles) - { - constexpr uint32_t pf_row = (jf / kVTilesPerSub) * 32u; - constexpr uint32_t pf_jj = jf % kVTilesPerSub; - load_S.template operator()(); - } - }; - - // ---- mfma_a (flat jf_lo, reads slot_lo) ---- - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(has_next ? 4 : 2, -1)); - prefetch.template operator()(); - do_mma.template operator()(oaccu_a); - if constexpr(resc_next) - { - mul_pair(rescale, opus::number{}); - mul_pair(rescale, opus::number{}); - } - - // ---- mfma_b (flat jf_hi, reads slot_hi) ---- - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(has_next ? 4 : 0, -1)); - prefetch.template operator()(); - do_mma.template operator()(oaccu_b); - if constexpr(resc_next) - { - mul_pair(rescale, opus::number{}); - mul_pair(rescale, opus::number{}); - } - }); -} - -// Pinned-register access is deliberately opaque, as in the upstream hk::art -// helpers. The entry's compiler scratch cap and clobber map are part of its ABI. -template -__device__ __forceinline__ float read_f32() -{ - float x; - asm volatile("v_mov_b32 %0, v[%1]" : "=v"(x) : "n"(Reg)); - return x; -} -template -__device__ __forceinline__ int32_t read_i32() -{ - int32_t x; - asm volatile("v_mov_b32 %0, v[%1]" : "=v"(x) : "n"(Reg)); - return x; -} -template -__device__ __forceinline__ void write_f32(float x) -{ - asm volatile("v_mov_b32 v[%0], %1" : : "n"(Reg), "v"(x)); -} -template -__device__ __forceinline__ void write_u32(uint32_t x) -{ - asm volatile("v_mov_b32 v[%0], %1" : : "n"(Reg), "v"(x)); -} -__device__ __forceinline__ void drain() -{ - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, 0)); - __builtin_amdgcn_sched_barrier(0); -} - -// The common aligned layout moves an entire eight-BF16 MFMA packet in one -// 128-bit load. The ABI still admits only two-byte alignment: padded/misaligned -// tensors use the exact scalar-pair path below. All values are copied bitwise. -template -__device__ __forceinline__ void load_q(const Params& p) -{ - using R = HkMlaV40Regs; - const uint32_t lane = threadIdx.x & 63u; - const uint32_t head = (threadIdx.x >> 6) * 16u + (lane & 15u); - const uint64_t offset = - uint64_t(blockIdx.x) * uint64_t(p.q_stride0) + uint64_t(head) * uint64_t(p.q_stride1); - const uint16_t* row = p.q + offset; - const bool aligned = ((reinterpret_cast(p.q) | (uint64_t(p.q_stride0) * 2u) | - (uint64_t(p.q_stride1) * 2u)) & - 15u) == 0; - opus::static_for<16>([&](auto c_) { - constexpr uint32_t c = c_.value; - const uint32_t col = sb8_inv_perm_col_elems(c * 32u + (lane >> 4) * 8u); - if(aligned) - { - // sb8 changes the packet order, never the eight values inside it. - const v4ui packet = *reinterpret_cast(row + col); - opus::static_for<4>([&](auto j_) { - constexpr uint32_t j = j_.value; - write_u32(packet[j]); - }); - } - else - { - opus::static_for<4>([&](auto j_) { - constexpr uint32_t j = j_.value; - const uint32_t pair = - uint32_t(row[col + j * 2u]) | (uint32_t(row[col + j * 2u + 1u]) << 16); - write_u32(pair); - }); - } - // Do not hoist an entire 512-column Q row into compiler-managed VGPRs. - __builtin_amdgcn_sched_barrier(0); - }); - drain(); -} - -// Prologue-only Q staging reuses the not-yet-live KV ping/pong allocation. -// H64 occupies exactly64KiB (four private16KiB wave regions); H16 occupies16KiB. -// After every wave reads its pinned Q and crosses the terminal barrier, these -// bytes are dead and KV may overwrite them. No extra steady-state LDS is used. -template -__device__ __forceinline__ void load_q_dma(const Params& p, uint16_t* staging) -{ - using R = HkMlaV40Regs; - static_assert(T::kNumWarps * 16u * 512u * 2u <= T::kIndexOffset); - const bool aligned = ((reinterpret_cast(p.q) | (uint64_t(p.q_stride0) * 2u) | - (uint64_t(p.q_stride1) * 2u)) & - 15u) == 0; - if(!aligned) - { - load_q(p); // preserve the original two-byte-alignment ABI - return; - } - const uint32_t lane = threadIdx.x & 63u; - const uint32_t wave = __builtin_amdgcn_readfirstlane(threadIdx.x >> 6); - const uint32_t r = lane >> 2, quad = lane & 3u; - const uint32_t head = wave * 16u + r; - const uint64_t row_offset = - uint64_t(blockIdx.x) * uint64_t(p.q_stride0) + uint64_t(head) * uint64_t(p.q_stride1); - const uint32_t base = __builtin_amdgcn_readfirstlane( - static_cast(reinterpret_cast(staging))) + - wave * 16384u; - // Queue all16 packets without a VMEM wait between chunks. Every source and - // destination is aligned; Q values remain their supplied BF16 bit patterns. - for(uint32_t c = 0; c < 16u; ++c) - { - const uint32_t L = c * 32u + ((quad * 8u) ^ (((r >> 2) & 1u) * 16u)); - const uint32_t col = sb8_inv_perm_col_elems(L); - const uint64_t address = reinterpret_cast(p.q + row_offset + uint64_t(col)); - const uint32_t destination = base + c * 1024u; - asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tglobal_load_lds_dwordx4 %1, off" - : - : "s"(destination), "v"(address) - : "m0", "memory"); - } - drain(); - __syncthreads(); // all Q packets have reached their wave-private LDS region - const uint32_t row = lane & 15u, col = (lane >> 4) * 8u; - const uint32_t read_address = base + row * 64u + ((col * 2u) ^ (((row >> 2) & 1u) * 32u)); - opus::static_for<16>([&](auto c_) { - constexpr uint32_t c = c_.value; - hkm::ds_read_b128(read_address, c * 1024u); - }); - drain(); - __syncthreads(); // Q LDS lifetime ends before any wave starts writing KV -} - -template -__device__ __forceinline__ void -stage_kv(const Params& p, uint16_t* tile, int32_t* selected, int64_t start, int64_t end) -{ - const uint32_t tid = threadIdx.x; - if(tid < 32) - { - // CSR entry and physical slot are both guarded before their load. - int32_t slot = -1; - if(int64_t(start) + tid < int64_t(end)) - slot = p.indices[int64_t(start) + tid]; - selected[tid] = slot; - } - __syncthreads(); - // Keep each load address 64 bit through multiplication. An invalid slot - // does not produce a pointer, including when the entire cache is empty. - const bool aligned = - ((reinterpret_cast(p.kv) | (uint64_t(p.kv_stride0) * 2u)) & 15u) == 0; - if(aligned) - { - // Eight contiguous BF16 values map to one aligned 16-byte LDS packet: - // sb8 permutation and bank XOR preserve all intra-packet bits. This - // reduces address/index work and global/LDS instructions by eight. - // The PR's K ring and P registers are dead during this stage. Reuse - // v44..v59 as four in-flight vector packets while Q64..127 and O128..255 - // remain resident. Compiler scratch stays strictly below v44. - using R = HkMlaV40Regs; - constexpr uint32_t first_reg = R::k_scratch_budget; - constexpr uint32_t packets = 4; - static_assert(first_reg == 44 && first_reg + packets * 4u <= R::k_q_vgpr_begin); - static_assert((32u * 64u) % (packets * T::kNumThreads) == 0); - for(uint32_t first = tid; first < 32u * 64u; first += packets * T::kNumThreads) - { - // Full EXEC initializes every lane before any predicated global - // load. Invalid lanes thus publish zero without forming a pointer. - opus::static_for([&](auto reg) { - asm volatile("v_mov_b32 v[%0], 0" : : "n"(first_reg + reg.value)); - }); - opus::static_for([&](auto i_) { - constexpr uint32_t i = i_.value, reg = first_reg + i * 4u; - uint32_t packet = first + i * T::kNumThreads; - // Prevent four independently strength-reduced address streams - // from staying live across the pinned prefetch and escaping - // compiler scratch v0..43 (the same break-CSE idiom as PR O). - asm volatile("" : "+v"(packet)); - const uint32_t row = packet / 64u, col = (packet % 64u) * 8u; - const int32_t slot = selected[row]; - if(slot >= 0 && int64_t(slot) < p.slots) - { - const uint64_t offset = uint64_t(slot) * uint64_t(p.kv_stride0) + uint64_t(col); - const uint64_t address = reinterpret_cast(p.kv + offset); - // Flat 64-bit address; never truncate to a cache-wide - // buffer descriptor or a 32-bit byte offset. - asm volatile("global_load_dwordx4 v[%0:%1], %2, off" - : - : "n"(reg), "n"(reg + 3u), "v"(address) - : "memory"); - } - }); - // Some packets can issue no VMEM instruction (all lanes invalid), - // so a fixed partial vmcnt threshold would be incorrect here. - drain(); - opus::static_for([&](auto i_) { - constexpr uint32_t i = i_.value, reg = first_reg + i * 4u; - uint32_t packet = first + i * T::kNumThreads; - // Prevent four independently strength-reduced address streams - // from staying live across the pinned prefetch and escaping - // compiler scratch v0..43 (the same break-CSE idiom as PR O). - asm volatile("" : "+v"(packet)); - const uint32_t row = packet / 64u, col = (packet % 64u) * 8u; - const uint32_t L = sb8_perm_col_elems(col); - const uint32_t r = row & 15u; - const uint32_t byte = (L / 32u * 2u + row / 16u) * 1024u + r * 64u + - ((2u * (L & 31u)) ^ (((r >> 2) & 1u) * 32u)); - hkm::ds_write_b128(reinterpret_cast(tile) + byte, 0); - }); - drain(); // retire every LDS writer before its pinned source is reused - } - } - else - { - for(uint32_t linear = tid; linear < 32u * 512u; linear += T::kNumThreads) - { - const uint32_t row = linear / 512u, col = linear % 512u; - const int32_t slot = selected[row]; - uint16_t value = 0; - if(slot >= 0 && int64_t(slot) < p.slots) - value = p.kv[uint64_t(slot) * uint64_t(p.kv_stride0) + uint64_t(col)]; - const uint32_t L = sb8_perm_col_elems(col); - const uint32_t r = row & 15u; - const uint32_t byte = (L / 32u * 2u + row / 16u) * 1024u + r * 64u + - ((2u * (L & 31u)) ^ (((r >> 2) & 1u) * 32u)); - tile[byte / 2u] = value; - } - } - drain(); - __syncthreads(); -} - -// Gfx950 full-64-bit global-to-LDS DMA. M0 supplies a wave-uniform destination -// base; lane t lands sixteen bytes at M0+t*16. The PR subblock mapping undoes -// bank XOR and sb8 on the global source side, making the destination contiguous. -// Counter semantics: DMA completion is VMEM (vmcnt), not LGKM. This helper only -// retires ordinary LDS operations, deliberately leaving valid DMA in flight. -template -__device__ __forceinline__ void -stage_kv_async(const Params& p, uint16_t* tile, int32_t* selected, int64_t start, int64_t end) -{ - const uint32_t tid = threadIdx.x; - if(tid < 32u) - { - int32_t slot = -1; - if(start + tid < end) - slot = p.indices[start + tid]; - selected[tid] = slot; - } - __syncthreads(); // index publication precedes every DMA; none is issued yet - const uint32_t lane = tid & 63u; - const uint32_t wave = __builtin_amdgcn_readfirstlane(tid >> 6); - const uint32_t base = - __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(tile))); - for(uint32_t block = wave; block < 32u; block += T::kNumWarps) - { - const uint32_t r = lane >> 2, quad = lane & 3u; - const uint32_t row = (block & 1u) * 16u + r; - const uint32_t L = (block >> 1) * 32u + ((quad * 8u) ^ (((r >> 2) & 1u) * 16u)); - const uint32_t col = sb8_inv_perm_col_elems(L); - const int32_t slot = selected[row]; - const uint32_t destination = base + block * 1024u; - if(slot >= 0 && int64_t(slot) < p.slots) - { - const uint64_t offset = uint64_t(slot) * uint64_t(p.kv_stride0) + uint64_t(col); - const uint64_t address = reinterpret_cast(p.kv + offset); - // M0->VMEM hazard gap follows LLVM's gfx950 DMA lowering. Source - // addresses are only formed inside the validity predicate. - asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tglobal_load_lds_dwordx4 %1, off" - : - : "s"(destination), "v"(address) - : "m0", "memory"); - } - else - { - // Inactive DMA lanes do not zero memory. Clear exactly their - // private sixteen-byte landing so stale previous-pong data cannot - // enter QK/PV, including all-invalid first/middle/final tiles. - const v4ui zero = {0, 0, 0, 0}; - *reinterpret_cast(reinterpret_cast(tile) + block * 1024u + - lane * 16u) = zero; - } - } - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, -1)); - __builtin_amdgcn_sched_barrier(0); -} - -// The QK ring and issue/wait/MFMA order are the PR m16x4 schedule. K comes from -// the same sb8 LDS layout; Q uses exactly the corresponding D-axis permutation. -template -__device__ __forceinline__ void qk_gemm(uintptr_t tile) -{ - using R = HkMlaV40Regs; - using mfma_ab_t = hk::bf16; - typename R::p_comp_lo_t lo; - typename R::p_comp_hi_t hi; - auto ring = [](uint32_t s) constexpr { - return s % 3u == 0u ? R::k_k0_begin : s % 3u == 1u ? R::k_k1_begin : R::k_k2_begin; - }; - auto issue = [&]() { - constexpr uint32_t base = ring(S); - using kr = hkdart::split_many_t>, 4>; - hk::art k; - Bf16KvManager::template load_k_to_gpr<(S % 2u) * 16u, (S / 2u) * 32u>(k, tile); - }; - auto mma = [&]() { - constexpr uint32_t base = ring(S), qb = R::k_q_vgpr_begin + (S / 2u) * 4u; - using kr = hkdart::split_many_t>, 4>; - using qr = hkdart::split_many_t>, 4>; - hk::art k; - hk::art q; - if constexpr(S == 0) - hk::mma_ABt(lo, k, q); - else if constexpr(S == 1) - hk::mma_ABt(hi, k, q); - else if constexpr((S % 2u) == 0) - hk::mma_ABt(lo, k, q, lo); - else - hk::mma_ABt(hi, k, q, hi); - }; - opus::static_for<3>([&](auto s) { issue.template operator()(); }); - opus::static_for<32>([&](auto s_) { - constexpr uint32_t s = s_.value, remain = 31u - s; - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(remain < 2u ? remain : 2u, -1)); - mma.template operator()(); - if constexpr(s + 3u < 32u) - issue.template operator()(); - }); - // The PR hides the final MFMA-to-VALU RAW latency behind its next-KV - // conversion/staging. Our synchronous BF16 path has no such intervening - // work. LLVM cannot see the numeric pinned-register dependence, and an - // lgkmcnt/vmcnt drain does not wait for MFMA arithmetic. Reserve a - // conservative 16 issue cycles before softmax first reads p56..p63. - // Without this gap the final D-axis slice is missing in one-hot Q tests. - asm volatile("s_nop 7\n\ts_nop 7"); - __builtin_amdgcn_sched_barrier(0); -} - -template -__device__ __forceinline__ void -softmax(const Params& p, const int32_t* selected, float& maximum, float& denominator) -{ - using R = HkMlaV40Regs; - constexpr float log2e = 1.4426950408889634074f; - const uint32_t group = (threadIdx.x & 63u) >> 4; - // QK's K ring is dead after its explicit MFMA latency gap. Prefetch the - // eight mask indices into v44..51 while keeping live scores in v56..63. - // One LGKM drain replaces eight serialized scalar reads; pending next-KV - // DMA (VMEM counter) remains in flight. PV may reuse the K ring afterward. - constexpr uint32_t mask_reg = R::k_scratch_budget; - static_assert(mask_reg + 8u <= R::k_p_comp_begin); - const uint32_t mask_address = - static_cast(reinterpret_cast(selected)) + group * 4u * sizeof(int32_t); - hkm::ds_read_b128(mask_address, 0); - hkm::ds_read_b128(mask_address, 16u * sizeof(int32_t)); - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, -1)); - __builtin_amdgcn_sched_barrier(0); - float local_max = -INFINITY; - opus::static_for<8>([&](auto i_) { - constexpr uint32_t i = i_.value; - const int32_t slot = read_i32(); - float score = read_f32(); - // Scaling is applied after the supplied BF16 QK FP32 accumulation. - score = (score * p.scale) * log2e; - if(slot < 0 || int64_t(slot) >= p.slots) - score = -INFINITY; - write_f32(score); - local_max = fmaxf(local_max, score); - }); - local_max = fmaxf(local_max, __shfl_xor(local_max, 16, 64)); - local_max = fmaxf(local_max, __shfl_xor(local_max, 32, 64)); - const float new_max = fmaxf(maximum, local_max); - const float alpha = denominator > 0.0f ? __builtin_amdgcn_exp2f(maximum - new_max) : 0.0f; - maximum = new_max; - typename R::oaccu_t o; - hk::mul_vgpr(o, o, alpha); - float sum = 0.0f; - opus::static_for<8>([&](auto i_) { - constexpr uint32_t i = i_.value; - const float score = read_f32(); - // All-empty tiles never evaluate exp2(-inf - -inf). - const float weight = score == -INFINITY ? 0.0f : __builtin_amdgcn_exp2f(score - maximum); - write_f32(weight); - sum += weight; - }); - sum += __shfl_xor(sum, 16, 64); - sum += __shfl_xor(sum, 32, 64); - denominator = denominator * alpha + sum; - // Exactly the upstream low-to-high P packing order and BF16 MFMA operands. - pack_2f32_to_bf16_pair_pinned(); - pack_2f32_to_bf16_pair_pinned(); - pack_2f32_to_bf16_pair_pinned(); - pack_2f32_to_bf16_pair_pinned(); -} - -template -__device__ __forceinline__ void sparse_mla_legacy(Params p) -{ - using T = Traits; - using R = HkMlaV40Regs; - static_assert(R::k_scratch_budget == 44); - // Q/O live throughout; K/P ranges intentionally overlay PV operands. - hkdart::clobber(); - hkdart::clobber(); - hkdart::clobber(); - hkdart::clobber(); - hkdart::clobber(); - hkdart::clobber(); - extern __shared__ __align__(16) unsigned char storage[]; - auto* tile = reinterpret_cast(storage); - auto* next_tile = reinterpret_cast(storage + T::kKvBytes); - auto* selected = reinterpret_cast(storage + T::kIndexOffset); - auto* next_selected = selected + 32; - const uintptr_t output_bounce = reinterpret_cast(storage); - const uint32_t query = blockIdx.x, partition = blockIdx.y; - const int32_t row_start = p.indptr[query], row_end = p.indptr[query + 1u]; - const int32_t length = row_end - row_start; - const int32_t begin = row_start + int64_t(length) * partition / p.splits; - const int32_t end = row_start + int64_t(length) * (partition + 1u) / p.splits; - load_q_dma(p, tile); - typename R::oaccu_t o; - hk::zero(o); - float maximum = -INFINITY, denominator = 0.0f; - Bf16KvManager manager; - const bool async_aligned = - T::kKvBuffers == 2u && - ((reinterpret_cast(p.kv) | (uint64_t(p.kv_stride0) * 2u)) & 15u) == 0; - if(async_aligned && begin < end) - { - stage_kv_async(p, tile, selected, begin, end); - drain(); - __syncthreads(); // tile0 prologue: every DMA/zero has completed - } - for(int64_t pos = begin; pos < end; pos += 32) - { - if(async_aligned) - { - if(pos + 32 < end) - stage_kv_async(p, next_tile, next_selected, pos + 32, end); - // No VMEM wait here: next-pong DMA overlaps current-pong compute. - } - else - { - // Preserve the proven BF16-alignment-only path synchronously. - stage_kv(p, tile, selected, pos, end); - } - const uintptr_t tile_addr = reinterpret_cast(tile); - qk_gemm(tile_addr); - softmax(p, selected, maximum, denominator); - // Reuse PR's exact PV four-slot LDS/MFMA schedule. Rescale occurred once - // above, so this call only accumulates into initialized FP32 O registers. - pv_gemm(manager, tile_addr, 1.0f); - drain(); // next DMA complete AND every current LDS reader retired - __syncthreads(); // publish next tile; old current may now become next - if(async_aligned) - { - auto* old_tile = tile; - tile = next_tile; - next_tile = old_tile; - auto* old_selected = selected; - selected = next_selected; - next_selected = old_selected; - } - } - const float inv_sum = denominator > 0.0f ? 1.0f / denominator : 0.0f; - hk::mul_vgpr(o, o, inv_sum); - constexpr float ln2 = 0.6931471805599453094f; - const float lse = - denominator > 0.0f ? (maximum + __builtin_amdgcn_logf(denominator)) * ln2 : -INFINITY; - const uint32_t lane = threadIdx.x & 63u, wave = threadIdx.x >> 6; - const uint64_t slot = Split ? uint64_t(query) * uint64_t(p.splits) + partition : query; - if(lane < 16u) - { - const uint64_t h = uint64_t(wave) * 16u + lane; - if constexpr(Split) - p.partial_lse[slot * Heads + h] = lse; - else if(p.lse) - p.lse[slot * Heads + h] = lse; - } - // Rebase the descriptor to one query/partition (64-bit) before the upstream - // manager's local 32-bit output offsets. This also works above 4 GiB. - if constexpr(Split) - { - OManager32bitsV4Gen1Swizzle output; - float* dst = p.partial_o + slot * Heads * 512u; - opus::static_for<8>([&](auto i) { - output.template output_to_vram_pair( - dst, wave, 0, 0, output_bounce, Heads); - drain(); - }); - } - else - { - OManager16bitsV4Gen1Swizzle output; - auto* dst = reinterpret_cast(p.out + slot * Heads * 512u); - opus::static_for<8>([&](auto i) { - output.template output_to_vram_pair( - dst, wave, 0, 0, output_bounce, Heads); - drain(); - }); - } -} - -// Guarded normalized-FP32 partial reducer. Each thread owns one D element; -// all-empty partitions contribute exactly zero without exp(-inf - -inf). -__global__ __launch_bounds__(256) void reduce(Params p, int heads) -{ - const uint32_t query = blockIdx.x, head = blockIdx.y; - const uint32_t d0 = threadIdx.x; - const uint64_t first = uint64_t(query) * uint64_t(p.splits) * uint64_t(heads) + head; - float maximum = -INFINITY; - for(int s = 0; s < p.splits; ++s) - maximum = fmaxf(maximum, p.partial_lse[first + uint64_t(s) * heads]); - float sum = 0.0f, a0 = 0.0f, a1 = 0.0f; - for(int s = 0; s < p.splits; ++s) - { - const uint64_t row = first + uint64_t(s) * heads; - const float value = p.partial_lse[row]; - const float weight = value == -INFINITY ? 0.0f : expf(value - maximum); - sum += weight; - a0 += weight * p.partial_o[row * 512u + d0]; - a1 += weight * p.partial_o[row * 512u + d0 + 256u]; - } - const float reciprocal = sum > 0.0f ? 1.0f / sum : 0.0f; - const uint64_t out_row = (uint64_t(query) * uint64_t(heads) + head) * 512u; - // Explicit nearest-even conversion, including subnormal and tie semantics. - uint32_t pair; - const float v0 = a0 * reciprocal, v1 = a1 * reciprocal; - asm volatile("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(pair) : "v"(v0), "v"(v1)); - p.out[out_row + d0] = uint16_t(pair); - p.out[out_row + d0 + 256u] = uint16_t(pair >> 16); - if(threadIdx.x == 0 && p.lse) - p.lse[uint64_t(query) * heads + head] = sum > 0.0f ? maximum + logf(sum) : -INFINITY; -} - -} // namespace sparse_mla_bf16 diff --git a/csrc/kernels/mla/hk/sparse_mla_bf16_full_lds.cuh b/csrc/kernels/mla/hk/sparse_mla_bf16_full_lds.cuh deleted file mode 100644 index e578e2debe0..00000000000 --- a/csrc/kernels/mla/hk/sparse_mla_bf16_full_lds.cuh +++ /dev/null @@ -1,374 +0,0 @@ -// SPDX-License-Identifier: MIT -// Full-LDS H64 implementation: Q, two KV tiles and P stay in 139,520 B LDS. -// Included inside namespace sparse_mla_bf16. -struct FullLdsTraits : Traits<64> -{ - static constexpr uint32_t kNumWarps = 8, kNumThreads = 512; - static constexpr uint32_t kIndexOffset = 139264, kLdsBytes = 139520; -}; - -struct FullLdsRegs -{ - using comp_t = float; - template - using ranges = hkdart::split_many_t>, 4>; - template - using bf16_fragment = hk::art>; - // All compiler-managed instructions must remain below v40. QK operands - // occupy40:55; scores60:63. Softmax reuses40:47 for eight scores and48:55 - // for masks. PV uses40:55 for four V slots and60:63 for packed full-N32 P. - static constexpr uint32_t k_o_begin = 64, kNumKvSub = 1; - static constexpr uint32_t k_v0_begin = 40, k_k0_begin = 44; - static constexpr uint32_t k_k1_begin = 48, k_k2_begin = 52; - using pinned_ranges = ranges<40, 88>; - using oaccu_t = hk::art>; - using score_t = hk::art>; - using p_mfma_a_t = bf16_fragment<60>; - using p_mfma_b_t = bf16_fragment<60>; // unused by the N32 PV schedule - using pv_v_0_t = bf16_fragment<40>; - using pv_v_1_t = bf16_fragment<44>; - using pv_v_2_t = bf16_fragment<48>; - using pv_v_3_t = bf16_fragment<52>; -}; - -#include "sparse_mla_bf16_full_lds_pv.cuh" - -// LDS communications use explicit DS operations and a barrier that does not -// implicitly drain VMEM. Next-KV DMA must be allowed to overlap these phases. -__device__ __forceinline__ void full_lds_barrier() -{ - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, -1)); - __builtin_amdgcn_sched_barrier(0); - asm volatile("s_barrier" ::: "memory"); -} - -// Same packet mapping/validity/zeroing as stage_kv_async; H64 row reuse only. -__device__ __forceinline__ void full_lds_stage_kv_async( - const Params& p, uint16_t* tile, int32_t* selected, int64_t start, int64_t end) -{ - using T = FullLdsTraits; - static_assert(T::kNumWarps == 8u); - const uint32_t tid = threadIdx.x; - if(tid < 32u) - { - int32_t slot = -1; - if(start + tid < end) - slot = p.indices[start + tid]; - selected[tid] = slot; - } - __syncthreads(); // index publication precedes every DMA; none is issued yet - const uint32_t lane = tid & 63u; - const uint32_t wave = __builtin_amdgcn_readfirstlane(tid >> 6); - const uint32_t base = - __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(tile))); - // Four D32 packets share one selected row in this eight-wave layout. - const uint32_t r = lane >> 2, quad = lane & 3u; - const uint32_t row = (wave & 1u) * 16u + r; - const int32_t slot = selected[row]; - const bool valid = slot >= 0 && int64_t(slot) < p.slots; - // Each lane's predicate is invariant across all four packets. Restore - // EXEC for every lane before the common LDS drain below; no CTA barrier - // is entered from either divergent branch. - if(valid) - { - const uint64_t offset = uint64_t(slot) * uint64_t(p.kv_stride0); - const uint64_t row_base = reinterpret_cast(p.kv + offset); - opus::static_for<4>([&](auto packet_) { - const uint32_t block = wave + packet_.value * T::kNumWarps; - const uint32_t L = (block >> 1) * 32u + ((quad * 8u) ^ (((r >> 2) & 1u) * 16u)); - const uint32_t col = sb8_inv_perm_col_elems(L); - const uint32_t destination = base + block * 1024u; - const uint64_t address = row_base + uint64_t(col) * sizeof(uint16_t); - asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tglobal_load_lds_dwordx4 %1, off" - : - : "s"(destination), "v"(address) - : "m0", "memory"); - }); - } - else - { - opus::static_for<4>([&](auto packet_) { - const uint32_t block = wave + packet_.value * T::kNumWarps; - // Invalid lanes own the same private 16 B landing as valid lanes. - // Zero through the dead QK operands to avoid hoisted compiler - // constants crossing the pinned register boundary. - asm volatile("v_mov_b32 v40, 0\n\tv_mov_b32 v41, 0\n\t" - "v_mov_b32 v42, 0\n\tv_mov_b32 v43, 0" :: - : "memory"); - hkm::ds_write_b128<40>(static_cast(reinterpret_cast(tile)) + - block * 1024u + lane * 16u, - 0); - }); - } - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, -1)); - __builtin_amdgcn_sched_barrier(0); -} - -__device__ __forceinline__ void full_lds_load_q(const Params& p, uint16_t* q_lds) -{ - const uint32_t tid = threadIdx.x, lane = tid & 63u; - const uint32_t wave = __builtin_amdgcn_readfirstlane(tid >> 6); - const uint32_t base = - __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(q_lds))); - const bool aligned = ((reinterpret_cast(p.q) | (uint64_t(p.q_stride0) * 2u) | - (uint64_t(p.q_stride1) * 2u)) & - 15u) == 0; - if(aligned) - { - // 4 head groups * 16 D32 blocks. Eight waves jointly load Q exactly - // once; paired compute waves later read the same16KiB head-group area. - for(uint32_t block = wave; block < 64u; block += 8u) - { - const uint32_t group = block / 16u, c = block % 16u; - const uint32_t row = lane >> 2, quad = lane & 3u; - const uint32_t col = - sb8_inv_perm_col_elems(c * 32u + ((quad * 8u) ^ (((row >> 2) & 1u) * 16u))); - const uint64_t offset = uint64_t(blockIdx.x) * uint64_t(p.q_stride0) + - uint64_t(group * 16u + row) * uint64_t(p.q_stride1) + col; - const uint64_t address = reinterpret_cast(p.q + offset); - const uint32_t destination = base + block * 1024u; - asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tglobal_load_lds_dwordx4 %1, off" - : - : "s"(destination), "v"(address) - : "m0", "memory"); - } - } - else - { - // Minimum contract alignment is2B, including odd BF16 row/head strides. - for(uint32_t linear = tid; linear < 64u * 512u; linear += 512u) - { - uint32_t item = linear; - asm volatile("" : "+v"(item)); - const uint32_t head = item / 512u, col = item % 512u; - const uint64_t offset = uint64_t(blockIdx.x) * uint64_t(p.q_stride0) + - uint64_t(head) * uint64_t(p.q_stride1) + col; - const uint16_t value = p.q[offset]; - const uint32_t L = sb8_perm_col_elems(col), row = head & 15u; - const uint32_t byte = (head / 16u) * 16384u + (L / 32u) * 1024u + row * 64u + - ((2u * (L & 31u)) ^ (((row >> 2) & 1u) * 32u)); - q_lds[byte / 2u] = value; - } - } - drain(); - __syncthreads(); -} - -template -__device__ __forceinline__ void full_lds_qk(uintptr_t q_base, uintptr_t kv_base) -{ - using T = FullLdsTraits; - const uint32_t lane = threadIdx.x & 63u, row = lane & 15u; - const uint32_t col = (lane >> 4) * 8u; - const uint32_t q_address = static_cast(q_base) + (threadIdx.x >> 7) * 16384u + - row * 64u + ((col * 2u) ^ (((row >> 2) & 1u) * 32u)); - FullLdsRegs::score_t score; - auto issue = [&]() { - constexpr uint32_t qr = 40u + (C % 2u) * 4u, kr = 48u + (C % 2u) * 4u; - typename FullLdsRegs::template bf16_fragment k; - hkm::ds_read_b128(q_address, C * 1024u); - Bf16KvManager::template load_k_to_gpr(k, kv_base); - }; - auto mma = [&]() { - typename FullLdsRegs::template bf16_fragment<40u + (C % 2u) * 4u> q; - typename FullLdsRegs::template bf16_fragment<48u + (C % 2u) * 4u> k; - if constexpr(C == 0u) - hk::mma_ABt(score, k, q); - else - hk::mma_ABt(score, k, q, score); - }; - issue.template operator()<0u>(); - issue.template operator()<1u>(); - opus::static_for<16>([&](auto c_) { - constexpr uint32_t c = c_.value; - // Each future fragment consists of exactly Q+K LDS reads. Wait until - // the oldest pair is ready, leaving the next pair in flight. - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(c < 15u ? 2 : 0, -1)); - mma.template operator()(); - if constexpr(c + 2u < 16u) - issue.template operator()(); - }); - // Memory waits do not scoreboard the opaque numeric MFMA accumulator. - asm volatile("s_nop 7\n\ts_nop 7"); - __builtin_amdgcn_sched_barrier(0); -} - -__device__ __forceinline__ void full_lds_softmax( - const Params& p, const int32_t* selected, uintptr_t scores, float& maximum, float& denominator) -{ - const uint32_t lane = threadIdx.x & 63u, group = lane >> 4; - const uint32_t head = (threadIdx.x >> 7) * 16u + (lane & 15u); - const uint32_t score_address = - static_cast(scores) + (head / 16u) * 2048u + (group * 16u + (head & 15u)) * 16u; - const uint32_t mask_address = - static_cast(reinterpret_cast(selected)) + group * 16u; - // Both paired waves read the same N32 scores and perform the same softmax. - // This deliberately duplicates its arithmetic in the first candidate but - // needs no extra state slab, no BF16 score approximation and no broadcast. - hkm::ds_read_b128<40>(score_address, 0); - hkm::ds_read_b128<44>(score_address, 1024); - hkm::ds_read_b128<48>(mask_address, 0); - hkm::ds_read_b128<52>(mask_address, 64); - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, -1)); - __builtin_amdgcn_sched_barrier(0); - constexpr float log2e = 1.4426950408889634074f; - float local_max = -INFINITY; - opus::static_for<8>([&](auto i_) { - constexpr uint32_t i = i_.value; - const int32_t slot = read_i32<48u + i>(); - float score = (read_f32<40u + i>() * p.scale) * log2e; - if(slot < 0 || int64_t(slot) >= p.slots) - score = -INFINITY; - write_f32<40u + i>(score); - local_max = fmaxf(local_max, score); - }); - local_max = fmaxf(local_max, __shfl_xor(local_max, 16, 64)); - local_max = fmaxf(local_max, __shfl_xor(local_max, 32, 64)); - const float next_max = fmaxf(maximum, local_max); - const float alpha = denominator > 0.0f ? __builtin_amdgcn_exp2f(maximum - next_max) : 0.0f; - maximum = next_max; - FullLdsRegs::oaccu_t o; - hk::mul_vgpr(o, o, alpha); - float sum = 0.0f; - opus::static_for<8>([&](auto i_) { - constexpr uint32_t i = i_.value; - const float score = read_f32<40u + i>(); - const float weight = score == -INFINITY ? 0.0f : __builtin_amdgcn_exp2f(score - maximum); - write_f32<40u + i>(weight); - sum += weight; - }); - sum += __shfl_xor(sum, 16, 64); - sum += __shfl_xor(sum, 32, 64); - denominator = denominator * alpha + sum; - pack_2f32_to_bf16_pair_pinned<60, 40>(); - pack_2f32_to_bf16_pair_pinned<61, 42>(); - pack_2f32_to_bf16_pair_pinned<62, 44>(); - pack_2f32_to_bf16_pair_pinned<63, 46>(); -} - -template -__device__ __forceinline__ void full_lds_output(const Params& p, uint64_t slot, uintptr_t bounce) -{ - using T = FullLdsTraits; // physical output pitch remains D512 - const uint32_t head_group = threadIdx.x >> 7; - // Preserve independent LDS bounce regions for both waves of each pair. - // Float output needs8*4352=34816B, which fits dead Q64KiB but not one KV32KiB. - const uintptr_t half_bounce = bounce + Half * 4u * 4352u; - if constexpr(Split) - { - OManager32bitsV4Gen1Swizzle output; - float* dst = p.partial_o + slot * 64u * 512u; - opus::static_for<4>([&](auto i) { - output.template output_to_vram_pair<64u + i.value * 16u, - Half * 256u + i.value * 64u, - false>(dst, head_group, 0, 0, half_bounce, 64); - drain(); - }); - } - else - { - OManager16bitsV4Gen1Swizzle output; - auto* dst = reinterpret_cast(p.out + slot * 64u * 512u); - opus::static_for<4>([&](auto i) { - output.template output_to_vram_pair<64u + i.value * 16u, - Half * 256u + i.value * 64u, - false>(dst, head_group, 0, 0, half_bounce, 64); - drain(); - }); - } -} - -template -__device__ __forceinline__ void full_lds_h64(Params p) -{ - using T = FullLdsTraits; - static_assert(T::kNumWarps == 8u && T::kLdsBytes == 139520u); - hkdart::clobber(); - extern __shared__ __align__(16) unsigned char storage[]; - auto* q_lds = reinterpret_cast(storage); - auto* tile = reinterpret_cast(storage + 65536u); - auto* next_tile = reinterpret_cast(storage + 98304u); - const uintptr_t scores = reinterpret_cast(storage + 131072u); - auto* selected = reinterpret_cast(storage + 139264u); - auto* next_selected = selected + 32u; - const uint32_t query = blockIdx.x, partition = blockIdx.y; - const int32_t row_start = p.indptr[query], row_end = p.indptr[query + 1u]; - const int32_t length = row_end - row_start; - const int64_t begin = int64_t(row_start) + int64_t(length) * partition / p.splits; - const int64_t end = int64_t(row_start) + int64_t(length) * (partition + 1u) / p.splits; - full_lds_load_q(p, q_lds); - FullLdsRegs::oaccu_t o; - hk::zero(o); - float maximum = -INFINITY, denominator = 0.0f; - const uint32_t lane = threadIdx.x & 63u, wave = threadIdx.x >> 6; - const bool aligned = - ((reinterpret_cast(p.kv) | (uint64_t(p.kv_stride0) * 2u)) & 15u) == 0; - Bf16KvManager manager; - if(aligned && begin < end) - { - full_lds_stage_kv_async(p, tile, selected, begin, end); - drain(); - __syncthreads(); - } - for(int64_t pos = begin; pos < end; pos += 32) - { - if(aligned) - { - if(pos + 32 < end) - full_lds_stage_kv_async(p, next_tile, next_selected, pos + 32, end); - } - else - { - // v44:59 staging payloads are dead Q/K slots in this map; live - // O64:127 remains untouched. Retains minimum2B input alignment. - stage_kv(p, tile, selected, pos, end); - } - const uintptr_t kv_address = reinterpret_cast(tile); - if(wave & 1u) - full_lds_qk<1u>(reinterpret_cast(q_lds), kv_address); - else - full_lds_qk<0u>(reinterpret_cast(q_lds), kv_address); - const uint32_t head = (wave / 2u) * 16u + (lane & 15u); - const uint32_t key = (wave & 1u) * 16u + (lane >> 4) * 4u; - hkm::ds_write_b128<60>(static_cast(scores) + (head / 16u) * 2048u + - ((key / 4u) * 16u + (head & 15u)) * 16u, - 0); - full_lds_barrier(); // both N16 halves published; next KV stays pending - full_lds_softmax(p, selected, scores, maximum, denominator); - if(wave & 1u) - full_lds_pv_gemm<256u, false, false, T>(manager, kv_address, 1.0f); - else - full_lds_pv_gemm<0u, false, false, T>(manager, kv_address, 1.0f); - drain(); - __syncthreads(); // all score/current reads retired and next DMA complete - if(aligned) - { - auto* old_tile = tile; - tile = next_tile; - next_tile = old_tile; - auto* old_selected = selected; - selected = next_selected; - next_selected = old_selected; - } - } - const float inv_sum = denominator > 0.0f ? 1.0f / denominator : 0.0f; - hk::mul_vgpr(o, o, inv_sum); - constexpr float ln2 = 0.6931471805599453094f; - const float lse = - denominator > 0.0f ? (maximum + __builtin_amdgcn_logf(denominator)) * ln2 : -INFINITY; - const uint64_t slot = Split ? uint64_t(query) * uint64_t(p.splits) + partition : query; - if((wave & 1u) == 0 && lane < 16u) - { - const uint64_t head = uint64_t(wave / 2u) * 16u + lane; - if constexpr(Split) - p.partial_lse[slot * 64u + head] = lse; - else if(p.lse) - p.lse[slot * 64u + head] = lse; - } - // No Q read survives the terminal loop barrier (or prologue barrier when - // empty), so its64KiB region is now safe for the independent O bounce slots. - if(wave & 1u) - full_lds_output(p, slot, reinterpret_cast(storage)); - else - full_lds_output(p, slot, reinterpret_cast(storage)); -} diff --git a/csrc/kernels/mla/hk/sparse_mla_bf16_full_lds_pv.cuh b/csrc/kernels/mla/hk/sparse_mla_bf16_full_lds_pv.cuh deleted file mode 100644 index e64bf00ade5..00000000000 --- a/csrc/kernels/mla/hk/sparse_mla_bf16_full_lds_pv.cuh +++ /dev/null @@ -1,155 +0,0 @@ -// SPDX-License-Identifier: MIT -// Included inside namespace sparse_mla_bf16, after FullLdsRegs. -// PR PV four-slot schedule, compact register map and D256 per wave. -// ColBase chooses the sb8-preserving physical output half. -template -__device__ __forceinline__ void -full_lds_pv_gemm(Bf16KvManager& kv_manager, const uintptr_t p_lds_v, const float rescale) -{ - using R = FullLdsRegs; - using comp_t = typename R::comp_t; - constexpr uint32_t k_o_begin = R::k_o_begin; - constexpr uint32_t k_v0_begin = R::k_v0_begin; - constexpr uint32_t k_k0_begin = R::k_k0_begin; - constexpr uint32_t k_k1_begin = R::k_k1_begin; - constexpr uint32_t k_k2_begin = R::k_k2_begin; - // kBlockN=64 (m16x8 double-tile) contracts BOTH 32-row sub-tiles in one call: - // sub-tile A (rows 0:32) uses p_mfma_a, B (rows 32:64) uses p_mfma_b. kBlockN=32 - // has only A. Picked per-iter by row_base (see below), no caller arg. - typename R::p_mfma_a_t p_mfma_a; - typename R::p_mfma_b_t p_mfma_b; - typename R::pv_v_0_t pv_v_0; - typename R::pv_v_1_t pv_v_1; - typename R::pv_v_2_t pv_v_2; - typename R::pv_v_3_t pv_v_3; - - // D-tiling: each D-iter emits 2 oaccu 16x16 sub-tiles (32 D-cols); 512/32 = 16 - // D-iters per 32-row KV sub-tile. kBlockN=64 runs BOTH sub-tiles in one call, so - // num_pv_iter = kNumKvSub * 16. iter [0,16) = sub-tile A, [16,32) = sub-tile B. - constexpr uint32_t kDIters = 256u / (2u * T::kTileM); // 8 per D256 wave - constexpr uint32_t num_pv_iter = R::kNumKvSub * kDIters; // 16 or 32 - // Flat V base tiles across all sub-tiles: 32 per sub-tile (S_0..S_31 each). - constexpr uint32_t kVTilesPerSub = 2u * kDIters; // 32 - constexpr uint32_t kNumVTiles = R::kNumKvSub * kVTilesPerSub; // 32 or 64 - - auto pk_mul_pair = [&](float r, auto base_c) { - constexpr uint32_t base = decltype(base_c)::value; - const float2 r2 = {r, r}; - asm volatile("v_pk_mul_f32 v[%0:%1], %2, v[%0:%1]" : : "n"(base), "n"(base + 1), "v"(r2)); - }; - auto mul_pair = [&](float r, auto base_c) { - constexpr uint32_t base = decltype(base_c)::value; - asm volatile("v_mul_f32_e32 v[%0], %1, v[%0]" : : "n"(base), "v"(r)); - asm volatile("v_mul_f32_e32 v[%0], %1, v[%0]" : : "n"(base + 1), "v"(r)); - }; - - // Issue both ds_read_b64 for base tile S_jj (within its sub-tile) into a - // round-robin slot (4 slots). kRowBase selects the KV sub-tile (0=A, 32=B), - // folded into the V ds_read imm by load_transposed_v_to_gpr. - auto load_S = [&]() { - constexpr uint32_t base = (slot == 0u) ? k_v0_begin - : (slot == 1u) ? k_k0_begin - : (slot == 2u) ? k_k1_begin - : k_k2_begin; - kv_manager.template load_transposed_v_to_gpr( - p_lds_v); - kv_manager.template load_transposed_v_to_gpr( - p_lds_v); - }; - // mfma: oaccu_dst (+)= pv_v_{slot}^T @ p_mfma[sub-tile]. 3-arg init only on the - // FIRST sub-tile of a first-iter call (kRowBase==0); sub-tile B always accumulates - // onto A's oaccu (same oaccu tile, second N-half of P). - auto do_mma = [&](OA& oaccu_dst) { - auto run = [&](auto& v) { - if constexpr(kIsFirstIter && kRowBase == 0u) - hk::mma_ABt(oaccu_dst, v, p_mfma_a); - else if constexpr(kRowBase == 0u) - hk::mma_ABt(oaccu_dst, v, p_mfma_a, oaccu_dst); - else - hk::mma_ABt(oaccu_dst, v, p_mfma_b, oaccu_dst); - }; - if constexpr(slot == 0u) - run(pv_v_0); - else if constexpr(slot == 1u) - run(pv_v_1); - else if constexpr(slot == 2u) - run(pv_v_2); - else - run(pv_v_3); - }; - - if constexpr(kDoRescale) - { - opus::static_for<2>([&](auto s) { - pk_mul_pair(rescale, opus::number{}); - pk_mul_pair(rescale, opus::number{}); - }); - } - - // Prologue: preload the first 3 flat V tiles (S_0,S_1,S_2 of sub-tile A) into - // slots 0,1,2. ONE prologue for the whole call -- sub-tile B's S_0..S_2 are - // prefetched by sub-tile A's tail iters via the flat prefetch index below. - load_S.template operator()<0u, 0u, 0u>(); - load_S.template operator()<0u, 1u, 1u>(); - load_S.template operator()<0u, 2u, 2u>(); - - opus::static_for([&](auto i) { - constexpr uint32_t iter = i.value; - constexpr uint32_t col_idx = iter % kDIters; // 0..15 D-tile within sub-tile - constexpr uint32_t row_base = (iter / kDIters) * 32u; // 0 (A) or 32 (B) - constexpr bool has_next = (iter + 1u) < num_pv_iter; - // Rescale only folds on the FIRST sub-tile pass (row_base==0) and only when a - // NEXT D-tile in that pass exists (col_idx+1 < kDIters) -- else oaccu overflows. - constexpr bool resc_next = kDoRescale && (row_base == 0u) && (col_idx + 1u < kDIters); - constexpr uint32_t next_oaccu_base = k_o_begin + (col_idx + 1u) * 8u; - - // Flat V-tile index across all sub-tiles: prefetch runs on this so sub-tile A's - // tail prefetches sub-tile B's S_0.. (ONE prologue). Consume uses col_idx. - constexpr uint32_t jf_lo = 2u * iter; // flat mfma idx (a) - constexpr uint32_t jf_hi = 2u * iter + 1u; // flat mfma idx (b) - // 4-slot round-robin (flat tile T -> slot T%4). Reading slot j%4 while refilling - // slot (j+3)%4 keeps the refill ds_read's dst != the mfma's src reg. - constexpr uint32_t slot_lo = jf_lo % 4u; - constexpr uint32_t slot_hi = jf_hi % 4u; - - constexpr uint32_t oaccu_base = k_o_begin + col_idx * 8u; - using oaccu_a_r = - hkdart::split_many_t>, - 4>; - using oaccu_b_r = - hkdart::split_many_t>, - 4>; - hk::art oaccu_a; - hk::art oaccu_b; - - // Prefetch S at flat index jf+3 -> its own sub-tile row_base + local jj%32. - auto prefetch = [&]() { - if constexpr(jf < kNumVTiles) - { - constexpr uint32_t pf_row = (jf / kVTilesPerSub) * 32u; - constexpr uint32_t pf_jj = jf % kVTilesPerSub; - load_S.template operator()(); - } - }; - - // ---- mfma_a (flat jf_lo, reads slot_lo) ---- - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(has_next ? 4 : 2, -1)); - prefetch.template operator()(); - do_mma.template operator()(oaccu_a); - if constexpr(resc_next) - { - mul_pair(rescale, opus::number{}); - mul_pair(rescale, opus::number{}); - } - - // ---- mfma_b (flat jf_hi, reads slot_hi) ---- - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(has_next ? 4 : 0, -1)); - prefetch.template operator()(); - do_mma.template operator()(oaccu_b); - if constexpr(resc_next) - { - mul_pair(rescale, opus::number{}); - mul_pair(rescale, opus::number{}); - } - }); -} diff --git a/csrc/kernels/mla/hk/sparse_mla_bf16_headtiles.cuh b/csrc/kernels/mla/hk/sparse_mla_bf16_headtiles.cuh deleted file mode 100644 index f2b0da8ce3d..00000000000 --- a/csrc/kernels/mla/hk/sparse_mla_bf16_headtiles.cuh +++ /dev/null @@ -1,502 +0,0 @@ -// SPDX-License-Identifier: MIT -// Experimental H16/H64 head tiling: each head16 CTA uses QK N64 / PV D128 per wave. -// Reuses PR #3459 HipKittens BF16 LDS readers, four-slot PV schedule, O managers. -#pragma once -#include "sparse_mla_bf16_common.cuh" - -namespace sparse_mla_bf16::headtiles { -struct Traits -{ - static constexpr uint32_t kBlockN = 64, kBlockK = 32, kTileM = 16; - static constexpr uint32_t kQkHeadDim = 512, kVoHeadDim = 512; - static constexpr uint32_t kNumWarps = 4, kNumThreads = 256, kRoundMode = 0; - static constexpr uint32_t kKvBytes = 65536, kPBytes = 2048; - static constexpr uint32_t kPOffset = kKvBytes; - static constexpr uint32_t kReduceOffset = kPOffset + kPBytes; - static constexpr uint32_t kIndexOffset = kReduceOffset + 512; - static constexpr uint32_t kLdsBytes = kIndexOffset + 256; - static_assert(kLdsBytes == 68352); -}; -template -using ranges = hkdart::split_many_t>, 4>; -struct Regs -{ - using comp_t = float; - static constexpr uint32_t kNumKvSub = 2; - static constexpr uint32_t k_o_begin = 128, k_v0_begin = 52; - static constexpr uint32_t k_k0_begin = 40, k_k1_begin = 44, k_k2_begin = 48; - using all_ranges = ranges<40, 159>; - using o_ranges = ranges<128, 159>; - using oaccu_t = hk::art; - using score_t = hk::art>; - using p_mfma_a_t = hk::art>; - using p_mfma_b_t = hk::art>; - using pv_v_0_t = hk::art>; - using pv_v_1_t = hk::art>; - using pv_v_2_t = hk::art>; - using pv_v_3_t = hk::art>; -}; - -template -__device__ __forceinline__ void -pv_gemm(Bf16KvManager& kv_manager, const uintptr_t p_lds_v, const float rescale) -{ - using R = Regs; - using comp_t = typename R::comp_t; - constexpr uint32_t k_o_begin = R::k_o_begin; - constexpr uint32_t k_v0_begin = R::k_v0_begin; - constexpr uint32_t k_k0_begin = R::k_k0_begin; - constexpr uint32_t k_k1_begin = R::k_k1_begin; - constexpr uint32_t k_k2_begin = R::k_k2_begin; - // kBlockN=64 (m16x8 double-tile) contracts BOTH 32-row sub-tiles in one call: - // sub-tile A (rows 0:32) uses p_mfma_a, B (rows 32:64) uses p_mfma_b. kBlockN=32 - // has only A. Picked per-iter by row_base (see below), no caller arg. - typename R::p_mfma_a_t p_mfma_a; - typename R::p_mfma_b_t p_mfma_b; - typename R::pv_v_0_t pv_v_0; - typename R::pv_v_1_t pv_v_1; - typename R::pv_v_2_t pv_v_2; - typename R::pv_v_3_t pv_v_3; - - // Each wave owns D128: four D32 iterations per N32 half, eight total. - // Iterations0:4 use keys0:32; iterations4:8 use keys32:64. - constexpr uint32_t kDIters = 128u / (2u * T::kTileM); // 4 - constexpr uint32_t num_pv_iter = R::kNumKvSub * kDIters; // 8 - // Eight V base tiles per N32 half, sixteen total. - constexpr uint32_t kVTilesPerSub = 2u * kDIters; // 8 - constexpr uint32_t kNumVTiles = R::kNumKvSub * kVTilesPerSub; // 16 - - auto pk_mul_pair = [&](float r, auto base_c) { - constexpr uint32_t base = decltype(base_c)::value; - const float2 r2 = {r, r}; - asm volatile("v_pk_mul_f32 v[%0:%1], %2, v[%0:%1]" : : "n"(base), "n"(base + 1), "v"(r2)); - }; - auto mul_pair = [&](float r, auto base_c) { - constexpr uint32_t base = decltype(base_c)::value; - asm volatile("v_mul_f32_e32 v[%0], %1, v[%0]" : : "n"(base), "v"(r)); - asm volatile("v_mul_f32_e32 v[%0], %1, v[%0]" : : "n"(base + 1), "v"(r)); - }; - - // Issue both ds_read_b64 for base tile S_jj (within its sub-tile) into a - // round-robin slot (4 slots). kRowBase selects the KV sub-tile (0=A, 32=B), - // folded into the V ds_read imm by load_transposed_v_to_gpr. - auto load_S = [&]() { - constexpr uint32_t base = (slot == 0u) ? k_v0_begin - : (slot == 1u) ? k_k0_begin - : (slot == 2u) ? k_k1_begin - : k_k2_begin; - kv_manager.template load_transposed_v_to_gpr(p_lds_v); - kv_manager.template load_transposed_v_to_gpr(p_lds_v); - }; - // mfma: oaccu_dst (+)= pv_v_{slot}^T @ p_mfma[sub-tile]. 3-arg init only on the - // FIRST sub-tile of a first-iter call (kRowBase==0); sub-tile B always accumulates - // onto A's oaccu (same oaccu tile, second N-half of P). - auto do_mma = [&](OA& oaccu_dst) { - auto run = [&](auto& v) { - if constexpr(kIsFirstIter && kRowBase == 0u) - hk::mma_ABt(oaccu_dst, v, p_mfma_a); - else if constexpr(kRowBase == 0u) - hk::mma_ABt(oaccu_dst, v, p_mfma_a, oaccu_dst); - else - hk::mma_ABt(oaccu_dst, v, p_mfma_b, oaccu_dst); - }; - if constexpr(slot == 0u) - run(pv_v_0); - else if constexpr(slot == 1u) - run(pv_v_1); - else if constexpr(slot == 2u) - run(pv_v_2); - else - run(pv_v_3); - }; - - if constexpr(kDoRescale) - { - opus::static_for<2>([&](auto s) { - pk_mul_pair(rescale, opus::number{}); - pk_mul_pair(rescale, opus::number{}); - }); - } - - // Prologue: preload the first 3 flat V tiles (S_0,S_1,S_2 of sub-tile A) into - // slots 0,1,2. ONE prologue for the whole call -- sub-tile B's S_0..S_2 are - // prefetched by sub-tile A's tail iters via the flat prefetch index below. - load_S.template operator()<0u, 0u, 0u>(); - load_S.template operator()<0u, 1u, 1u>(); - load_S.template operator()<0u, 2u, 2u>(); - - opus::static_for([&](auto i) { - constexpr uint32_t iter = i.value; - constexpr uint32_t col_idx = iter % kDIters; // 0..3 D-tile within sub-tile - constexpr uint32_t row_base = (iter / kDIters) * 32u; // 0 (A) or 32 (B) - constexpr bool has_next = (iter + 1u) < num_pv_iter; - // Rescale only folds on the FIRST sub-tile pass (row_base==0) and only when a - // NEXT D-tile in that pass exists (col_idx+1 < kDIters) -- else oaccu overflows. - constexpr bool resc_next = kDoRescale && (row_base == 0u) && (col_idx + 1u < kDIters); - constexpr uint32_t next_oaccu_base = k_o_begin + (col_idx + 1u) * 8u; - - // Flat V-tile index across all sub-tiles: prefetch runs on this so sub-tile A's - // tail prefetches sub-tile B's S_0.. (ONE prologue). Consume uses col_idx. - constexpr uint32_t jf_lo = 2u * iter; // flat mfma idx (a) - constexpr uint32_t jf_hi = 2u * iter + 1u; // flat mfma idx (b) - // 4-slot round-robin (flat tile T -> slot T%4). Reading slot j%4 while refilling - // slot (j+3)%4 keeps the refill ds_read's dst != the mfma's src reg. - constexpr uint32_t slot_lo = jf_lo % 4u; - constexpr uint32_t slot_hi = jf_hi % 4u; - - constexpr uint32_t oaccu_base = k_o_begin + col_idx * 8u; - using oaccu_a_r = - hkdart::split_many_t>, - 4>; - using oaccu_b_r = - hkdart::split_many_t>, - 4>; - hk::art oaccu_a; - hk::art oaccu_b; - - // Prefetch S at flat index jf+3 -> its own sub-tile row_base + local jj%8. - auto prefetch = [&]() { - if constexpr(jf < kNumVTiles) - { - constexpr uint32_t pf_row = (jf / kVTilesPerSub) * 32u; - constexpr uint32_t pf_jj = jf % kVTilesPerSub; - load_S.template operator()(); - } - }; - - // ---- mfma_a (flat jf_lo, reads slot_lo) ---- - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(has_next ? 4 : 2, -1)); - prefetch.template operator()(); - do_mma.template operator()(oaccu_a); - if constexpr(resc_next) - { - mul_pair(rescale, opus::number{}); - mul_pair(rescale, opus::number{}); - } - - // ---- mfma_b (flat jf_hi, reads slot_hi) ---- - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(has_next ? 4 : 0, -1)); - prefetch.template operator()(); - do_mma.template operator()(oaccu_b); - if constexpr(resc_next) - { - mul_pair(rescale, opus::number{}); - mul_pair(rescale, opus::number{}); - } - }); -} - -// All waves need the same head16 Q block selected by grid.z. Only one cooperative 16KiB global/LDS -// copy occurs; the four wave-private register copies are populated after the barrier. -__device__ __forceinline__ void load_q(const Params& p, uint16_t* tile) -{ - const uint32_t lane = threadIdx.x & 63u; - const uint32_t wave = __builtin_amdgcn_readfirstlane(threadIdx.x >> 6); - const bool aligned = ((reinterpret_cast(p.q) | uint64_t(p.q_stride0) * 2u | - uint64_t(p.q_stride1) * 2u) & - 15u) == 0; - if(aligned) - { - const uint32_t row = lane >> 2, quad = lane & 3u; - const uint64_t offset = uint64_t(blockIdx.x) * uint64_t(p.q_stride0) + - uint64_t(blockIdx.z * 16u + row) * uint64_t(p.q_stride1); - const uint32_t base = __builtin_amdgcn_readfirstlane( - static_cast(reinterpret_cast(tile))); - for(uint32_t c = wave; c < 16u; c += 4u) - { - const uint32_t L = c * 32u + ((quad * 8u) ^ (((row >> 2) & 1u) * 16u)); - const uint32_t col = sb8_inv_perm_col_elems(L); - const uint64_t address = reinterpret_cast(p.q + offset + col); - const uint32_t destination = base + c * 1024u; - asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tglobal_load_lds_dwordx4 %1, off" - : - : "s"(destination), "v"(address) - : "m0", "memory"); - } - drain(); - __syncthreads(); - const uint32_t r = lane & 15u, col = (lane >> 4) * 8u; - const uint32_t address = base + r * 64u + ((col * 2u) ^ (((r >> 2) & 1u) * 32u)); - opus::static_for<16>( - [&](auto c) { hkm::ds_read_b128<64u + c.value * 4u>(address, c.value * 1024u); }); - } - else - { - const uint32_t head = lane & 15u; - const uint64_t offset = uint64_t(blockIdx.x) * uint64_t(p.q_stride0) + - uint64_t(blockIdx.z * 16u + head) * uint64_t(p.q_stride1); - opus::static_for<16>([&](auto c) { - const uint32_t col = sb8_inv_perm_col_elems(c.value * 32u + (lane >> 4) * 8u); - opus::static_for<4>([&](auto j) { - const uint32_t pair = uint32_t(p.q[offset + col + j.value * 2u]) | - (uint32_t(p.q[offset + col + j.value * 2u + 1u]) << 16); - write_u32<64u + c.value * 4u + j.value>(pair); - }); - __builtin_amdgcn_sched_barrier(0); - }); - } - drain(); - __syncthreads(); // every replicated Q register read retires before KV overwrite -} - -// Two PR 32-row subtiles, each preserving the exact sb8 + row-bank swizzle. -// This initial N64 geometry intentionally uses a single synchronous KV buffer. -__device__ __forceinline__ void -stage_kv(const Params& p, uint16_t* tile, int32_t* selected, int64_t start, int64_t end) -{ - const uint32_t tid = threadIdx.x, lane = tid & 63u; - if(tid < 64u) - { - int32_t slot = -1; - if(start + tid < end) - slot = p.indices[start + tid]; - selected[tid] = slot; - } - __syncthreads(); - const bool aligned = - ((reinterpret_cast(p.kv) | uint64_t(p.kv_stride0) * 2u) & 15u) == 0; - if(aligned) - { - const uint32_t wave = __builtin_amdgcn_readfirstlane(tid >> 6); - const uint32_t base = __builtin_amdgcn_readfirstlane( - static_cast(reinterpret_cast(tile))); - for(uint32_t block = wave; block < 64u; block += 4u) - { - const uint32_t local = block & 31u, r = lane >> 2, quad = lane & 3u; - const uint32_t row = (block >> 5) * 32u + (local & 1u) * 16u + r; - const uint32_t L = (local >> 1) * 32u + ((quad * 8u) ^ (((r >> 2) & 1u) * 16u)); - const uint32_t col = sb8_inv_perm_col_elems(L); - const int32_t slot = selected[row]; - const uint32_t destination = base + block * 1024u; - if(slot >= 0 && int64_t(slot) < p.slots) - { - const uint64_t offset = uint64_t(slot) * uint64_t(p.kv_stride0) + uint64_t(col); - const uint64_t address = reinterpret_cast(p.kv + offset); - asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tglobal_load_lds_dwordx4 %1, off" - : - : "s"(destination), "v"(address) - : "m0", "memory"); - } - else - { - const v4ui zero = {0, 0, 0, 0}; - *reinterpret_cast(reinterpret_cast(tile) + block * 1024u + - lane * 16u) = zero; - } - } - } - else - { - for(uint32_t linear = tid; linear < 64u * 512u; linear += 256u) - { - const uint32_t row = linear / 512u, col = linear % 512u; - const int32_t slot = selected[row]; - uint16_t value = 0; - if(slot >= 0 && int64_t(slot) < p.slots) - value = p.kv[uint64_t(slot) * uint64_t(p.kv_stride0) + uint64_t(col)]; - const uint32_t L = sb8_perm_col_elems(col), r = row & 15u; - const uint32_t byte = (row / 32u) * 32768u + - (L / 32u * 2u + (row % 32u) / 16u) * 1024u + r * 64u + - ((2u * (L & 31u)) ^ (((r >> 2) & 1u) * 32u)); - tile[byte / 2u] = value; - } - } - drain(); - __syncthreads(); -} - -__device__ __forceinline__ void qk(uintptr_t tile) -{ - // Wave w computes its 16 selected tokens against all 16 query heads. - const uint32_t wave = __builtin_amdgcn_readfirstlane(threadIdx.x >> 6); - const uintptr_t wave_tile = tile + (wave / 2u) * 32768u + (wave % 2u) * 1024u; - Regs::score_t score; - auto issue = [&]() { - constexpr uint32_t base = 40u + (S % 3u) * 4u; - hk::art> k; - Bf16KvManager::template load_k_to_gpr<0u, S * 32u>(k, wave_tile); - }; - auto mma = [&]() { - constexpr uint32_t base = 40u + (S % 3u) * 4u, qb = 64u + S * 4u; - hk::art> k; - hk::art> q; - if constexpr(S == 0) - hk::mma_ABt(score, k, q); - else - hk::mma_ABt(score, k, q, score); - }; - opus::static_for<3>([&](auto s) { issue.template operator()(); }); - opus::static_for<16>([&](auto s) { - constexpr uint32_t remain = 15u - s.value; - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(remain < 2u ? remain : 2u, -1)); - mma.template operator()(); - if constexpr(s.value + 3u < 16u) - issue.template operator()(); - }); - asm volatile("s_nop 7\n\ts_nop 7"); // opaque MFMA score dependency is not compiler-visible - __builtin_amdgcn_sched_barrier(0); -} - -__device__ __forceinline__ void softmax(const Params& p, - const int32_t* selected, - float* exchange, - uint16_t* p_lds, - float& maximum, - float& denominator) -{ - const uint32_t lane = threadIdx.x & 63u, wave = threadIdx.x >> 6; - const uint32_t head = lane & 15u, group = lane >> 4; - const uint32_t first_key = wave * 16u + group * 4u; - const uint32_t index_address = - static_cast(reinterpret_cast(selected)) + first_key * sizeof(int32_t); - hkm::ds_read_b128<40>(index_address, 0); - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, -1)); - float local_max = -INFINITY; - opus::static_for<4>([&](auto i) { - const int32_t slot = read_i32<40u + i.value>(); - float x = (read_f32<56u + i.value>() * p.scale) * 1.4426950408889634074f; - if(slot < 0 || int64_t(slot) >= p.slots) - x = -INFINITY; - write_f32<56u + i.value>(x); - local_max = fmaxf(local_max, x); - }); - local_max = fmaxf(local_max, __shfl_xor(local_max, 16, 64)); - local_max = fmaxf(local_max, __shfl_xor(local_max, 32, 64)); - if(lane < 16u) - exchange[wave * 16u + head] = local_max; - __syncthreads(); - float all_max = -INFINITY; - opus::static_for<4>([&](auto w) { all_max = fmaxf(all_max, exchange[w.value * 16u + head]); }); - const float new_max = fmaxf(maximum, all_max); - const float alpha = denominator > 0.0f ? __builtin_amdgcn_exp2f(maximum - new_max) : 0.0f; - maximum = new_max; - Regs::oaccu_t o; - hk::mul_vgpr(o, o, alpha); - float local_sum = 0.0f; - opus::static_for<4>([&](auto i) { - const float x = read_f32<56u + i.value>(); - const float weight = x == -INFINITY ? 0.0f : __builtin_amdgcn_exp2f(x - maximum); - write_f32<56u + i.value>(weight); - local_sum += weight; - }); - local_sum += __shfl_xor(local_sum, 16, 64); - local_sum += __shfl_xor(local_sum, 32, 64); - if(lane < 16u) - exchange[64u + wave * 16u + head] = local_sum; - // Each wave owns a disjoint N16 slab of BF16 P. Normal [head,key] LDS layout - // makes the subsequent full-N64 operand copies explicit and bit-preserving. - v2ui packed; - const float p0 = read_f32<56>(), p1 = read_f32<57>(); - const float p2 = read_f32<58>(), p3 = read_f32<59>(); - asm volatile("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(packed[0]) : "v"(p0), "v"(p1)); - asm volatile("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(packed[1]) : "v"(p2), "v"(p3)); - const uintptr_t address = reinterpret_cast(p_lds) + (head * 64u + first_key) * 2u; - hkm::ds_write_b64(packed, static_cast(address), 0); - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, -1)); - __syncthreads(); // all P and four partial denominators are visible - float all_sum = 0.0f; - opus::static_for<4>([&](auto w) { all_sum += exchange[64u + w.value * 16u + head]; }); - denominator = denominator * alpha + all_sum; - const uint32_t p_address = - static_cast(reinterpret_cast(p_lds)) + (head * 64u + group * 4u) * 2u; - // PR's transposed-V operand orders the two N16 halves as group*4 and +16. - hkm::ds_read_b64<56>(p_address, 0); - hkm::ds_read_b64<58>(p_address, 32); - hkm::ds_read_b64<60>(p_address, 64); - hkm::ds_read_b64<62>(p_address, 96); - drain(); - __syncthreads(); // denominator/max consumers retire before next iteration reuses exchange -} - -template -__device__ __forceinline__ void run(Params p) -{ - hkdart::clobber(); - extern __shared__ __align__(16) unsigned char storage[]; - auto* tile = reinterpret_cast(storage); - auto* p_lds = reinterpret_cast(storage + Traits::kPOffset); - auto* exchange = reinterpret_cast(storage + Traits::kReduceOffset); - auto* selected = reinterpret_cast(storage + Traits::kIndexOffset); - const uint32_t query = blockIdx.x, partition = blockIdx.y; - const uint32_t head_base = blockIdx.z * 16u; - static_assert(TotalHeads == 16 || TotalHeads == 64); - const int32_t first = p.indptr[query], last = p.indptr[query + 1u], length = last - first; - const int32_t begin = first + int64_t(length) * partition / p.splits; - const int32_t end = first + int64_t(length) * (partition + 1u) / p.splits; - load_q(p, tile); - Regs::oaccu_t o; - hk::zero(o); - float maximum = -INFINITY, denominator = 0.0f; - Bf16KvManager manager; - const uint32_t wave = __builtin_amdgcn_readfirstlane(threadIdx.x >> 6); - for(int64_t pos = begin; pos < end; pos += 64) - { - stage_kv(p, tile, selected, pos, end); - qk(reinterpret_cast(tile)); - softmax(p, selected, exchange, p_lds, maximum, denominator); - // D128/wave is a multiple of the sb8 permutation's 64-element period. - const uintptr_t v_tile = reinterpret_cast(tile) + wave * 8192u; - sparse_mla_bf16::headtiles::pv_gemm(manager, v_tile, 1.0f); - asm volatile("s_nop 7\n\ts_nop 7"); - drain(); - __syncthreads(); - } - const float inverse = denominator > 0.0f ? 1.0f / denominator : 0.0f; - hk::mul_vgpr(o, o, inverse); - const float lse = denominator > 0.0f - ? (maximum + __builtin_amdgcn_logf(denominator)) * 0.6931471805599453094f - : -INFINITY; - const uint32_t lane = threadIdx.x & 63u; - const uint64_t output_slot = Split ? uint64_t(query) * uint64_t(p.splits) + partition : query; - if(wave == 0u && lane < 16u) - { - if constexpr(Split) - p.partial_lse[output_slot * TotalHeads + head_base + lane] = lse; - else if(p.lse) - p.lse[output_slot * TotalHeads + head_base + lane] = lse; - } - // PR O manager preserves global head stride512. Offset its descriptor to - // this wave's D128, set head-wave id to zero, and give each wave a private - // dead-KV LDS bounce so its data cannot race another output wave. - const uintptr_t bounce = reinterpret_cast(tile) + wave * 4352u; - // grid.z introduces another head-block base. The PR buffer descriptor must - // be an SGPR tuple, so state its per-wave uniformity explicitly in64bits. - auto uniform_pointer = [](uint64_t address) { - const uint32_t lo = __builtin_amdgcn_readfirstlane(uint32_t(address)); - const uint32_t hi = __builtin_amdgcn_readfirstlane(uint32_t(address >> 32)); - return uint64_t(lo) | (uint64_t(hi) << 32); - }; - if constexpr(Split) - { - OManager32bitsV4Gen1Swizzle output; - float* dst = reinterpret_cast(uniform_pointer(reinterpret_cast( - p.partial_o + (output_slot * TotalHeads + head_base) * 512u + wave * 128u))); - opus::static_for<2>([&](auto i) { - output.template output_to_vram_pair<128u + i.value * 16u, i.value * 64u, false>( - dst, 0, 0, 0, bounce, TotalHeads); - drain(); - }); - } - else - { - OManager16bitsV4Gen1Swizzle output; - auto* dst = reinterpret_cast(uniform_pointer(reinterpret_cast( - p.out + (output_slot * TotalHeads + head_base) * 512u + wave * 128u))); - opus::static_for<2>([&](auto i) { - output.template output_to_vram_pair<128u + i.value * 16u, i.value * 64u, false>( - dst, 0, 0, 0, bounce, TotalHeads); - drain(); - }); - } -} -} // namespace sparse_mla_bf16::headtiles - -namespace sparse_mla_bf16 { -template -__global__ __launch_bounds__(256) - __attribute__((amdgpu_num_vgpr(40))) void sparse_mla_headtiles(Params p) -{ - static_assert(Heads == 64); - headtiles::run(p); -} -} // namespace sparse_mla_bf16 diff --git a/csrc/kernels/mla/hk/sparse_mla_bf16_reduce.cuh b/csrc/kernels/mla/hk/sparse_mla_bf16_reduce.cuh deleted file mode 100644 index ad1b3395c80..00000000000 --- a/csrc/kernels/mla/hk/sparse_mla_bf16_reduce.cuh +++ /dev/null @@ -1,107 +0,0 @@ -// SPDX-License-Identifier: MIT -// Fully unrolled, one-wave split-K reducer shared by the v1/v2 implementations. -#pragma once -#include - -namespace sparse_mla_bf16 { - -template -__global__ __launch_bounds__(64) void reduce(Params p) -{ - static_assert(Heads == 16 || Heads == 64); - static_assert(Splits == 2 || Splits == 4 || Splits == 8 || Splits == 16 || Splits == 32); - const uint32_t query = blockIdx.x, head = blockIdx.y; - const uint32_t d = threadIdx.x * 8u; - const uint64_t first = uint64_t(query) * Splits * Heads + head; - float maximum = -INFINITY; -#pragma unroll - for(int s = 0; s < Splits; ++s) - maximum = fmaxf(maximum, p.partial_lse[first + uint64_t(s) * Heads]); - float sum = 0.0f; - float accumulator[8] = {}; - // Every partial row is 2048 B; caller-provided workspace may be unaligned. - // Retain a scalar-load fallback for callers of the raw C ABI. - const bool aligned_partial = (reinterpret_cast(p.partial_o) & 15u) == 0; -#pragma unroll - for(int s = 0; s < Splits; ++s) - { - const uint64_t row = first + uint64_t(s) * Heads; - const float value = p.partial_lse[row]; - // Preserve the baseline summation order and empty-partition semantics. - const float weight = value == -INFINITY ? 0.0f : expf(value - maximum); - sum += weight; - const float* source = p.partial_o + row * 512u + d; - float values[8]; - if(aligned_partial) - { - const float4 lo = *reinterpret_cast(source); - const float4 hi = *reinterpret_cast(source + 4); - values[0] = lo.x; - values[1] = lo.y; - values[2] = lo.z; - values[3] = lo.w; - values[4] = hi.x; - values[5] = hi.y; - values[6] = hi.z; - values[7] = hi.w; - } - else - { -#pragma unroll - for(int i = 0; i < 8; ++i) - values[i] = source[i]; - } -#pragma unroll - for(int i = 0; i < 8; ++i) - accumulator[i] += weight * values[i]; - } - const float reciprocal = sum > 0.0f ? 1.0f / sum : 0.0f; - uint32_t pairs[4]; -#pragma unroll - for(int i = 0; i < 4; ++i) - { - const float lo = accumulator[2 * i] * reciprocal; - const float hi = accumulator[2 * i + 1] * reciprocal; - asm volatile("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(pairs[i]) : "v"(lo), "v"(hi)); - } - uint16_t* destination = p.out + (uint64_t(query) * Heads + head) * 512u + d; - if((reinterpret_cast(destination) & 15u) == 0) - { - const uint4 packed{pairs[0], pairs[1], pairs[2], pairs[3]}; - *reinterpret_cast(destination) = packed; - } - else - { - // The API permits an output storage offset of one BF16 (2 B alignment). -#pragma unroll - for(int i = 0; i < 4; ++i) - { - destination[2 * i] = uint16_t(pairs[i]); - destination[2 * i + 1] = uint16_t(pairs[i] >> 16); - } - } - if(threadIdx.x == 0 && p.lse) - p.lse[uint64_t(query) * Heads + head] = sum > 0.0f ? maximum + logf(sum) : -INFINITY; -} - -template -inline hipError_t launch_reduce64(Params p, hipStream_t stream) -{ -#define SPARSE_MLA_REDUCE_CASE(S) \ - case S: \ - hipLaunchKernelGGL((reduce), dim3(p.queries, Heads), dim3(64), 0, stream, p); \ - break - switch(p.splits) - { - SPARSE_MLA_REDUCE_CASE(2); - SPARSE_MLA_REDUCE_CASE(4); - SPARSE_MLA_REDUCE_CASE(8); - SPARSE_MLA_REDUCE_CASE(16); - SPARSE_MLA_REDUCE_CASE(32); - default: return hipErrorInvalidValue; - } -#undef SPARSE_MLA_REDUCE_CASE - return hipGetLastError(); -} - -} // namespace sparse_mla_bf16 diff --git a/csrc/kernels/mla/hk/sparse_mla_bf16_wave4.cuh b/csrc/kernels/mla/hk/sparse_mla_bf16_wave4.cuh deleted file mode 100644 index fae437157e7..00000000000 --- a/csrc/kernels/mla/hk/sparse_mla_bf16_wave4.cuh +++ /dev/null @@ -1,493 +0,0 @@ -// SPDX-License-Identifier: MIT -// Experimental H16 four-wave geometry: QK partitions N64, PV partitions D512. -// Reuses PR #3459 HipKittens BF16 LDS readers, four-slot PV schedule, O managers. -#pragma once -#include "sparse_mla_bf16_common.cuh" - -namespace sparse_mla_bf16::wave4 { -struct Traits -{ - static constexpr uint32_t kBlockN = 64, kBlockK = 32, kTileM = 16; - static constexpr uint32_t kQkHeadDim = 512, kVoHeadDim = 512; - static constexpr uint32_t kNumWarps = 4, kNumThreads = 256, kRoundMode = 0; - static constexpr uint32_t kKvBytes = 65536, kPBytes = 2048; - static constexpr uint32_t kPOffset = kKvBytes; - static constexpr uint32_t kReduceOffset = kPOffset + kPBytes; - static constexpr uint32_t kIndexOffset = kReduceOffset + 512; - static constexpr uint32_t kLdsBytes = kIndexOffset + 256; - static_assert(kLdsBytes == 68352); -}; -template -using ranges = hkdart::split_many_t>, 4>; -struct Regs -{ - using comp_t = float; - static constexpr uint32_t kNumKvSub = 2; - static constexpr uint32_t k_o_begin = 128, k_v0_begin = 52; - static constexpr uint32_t k_k0_begin = 40, k_k1_begin = 44, k_k2_begin = 48; - using all_ranges = ranges<40, 159>; - using o_ranges = ranges<128, 159>; - using oaccu_t = hk::art; - using score_t = hk::art>; - using p_mfma_a_t = hk::art>; - using p_mfma_b_t = hk::art>; - using pv_v_0_t = hk::art>; - using pv_v_1_t = hk::art>; - using pv_v_2_t = hk::art>; - using pv_v_3_t = hk::art>; -}; - -template -__device__ __forceinline__ void -pv_gemm(Bf16KvManager& kv_manager, const uintptr_t p_lds_v, const float rescale) -{ - using R = Regs; - using comp_t = typename R::comp_t; - constexpr uint32_t k_o_begin = R::k_o_begin; - constexpr uint32_t k_v0_begin = R::k_v0_begin; - constexpr uint32_t k_k0_begin = R::k_k0_begin; - constexpr uint32_t k_k1_begin = R::k_k1_begin; - constexpr uint32_t k_k2_begin = R::k_k2_begin; - // kBlockN=64 (m16x8 double-tile) contracts BOTH 32-row sub-tiles in one call: - // sub-tile A (rows 0:32) uses p_mfma_a, B (rows 32:64) uses p_mfma_b. kBlockN=32 - // has only A. Picked per-iter by row_base (see below), no caller arg. - typename R::p_mfma_a_t p_mfma_a; - typename R::p_mfma_b_t p_mfma_b; - typename R::pv_v_0_t pv_v_0; - typename R::pv_v_1_t pv_v_1; - typename R::pv_v_2_t pv_v_2; - typename R::pv_v_3_t pv_v_3; - - // D-tiling: each D-iter emits 2 oaccu 16x16 sub-tiles (32 D-cols); 512/32 = 16 - // D-iters per 32-row KV sub-tile. kBlockN=64 runs BOTH sub-tiles in one call, so - // num_pv_iter = kNumKvSub * 16. iter [0,16) = sub-tile A, [16,32) = sub-tile B. - constexpr uint32_t kDIters = 128u / (2u * T::kTileM); // 16 - constexpr uint32_t num_pv_iter = R::kNumKvSub * kDIters; // 16 or 32 - // Flat V base tiles across all sub-tiles: 32 per sub-tile (S_0..S_31 each). - constexpr uint32_t kVTilesPerSub = 2u * kDIters; // 32 - constexpr uint32_t kNumVTiles = R::kNumKvSub * kVTilesPerSub; // 32 or 64 - - auto pk_mul_pair = [&](float r, auto base_c) { - constexpr uint32_t base = decltype(base_c)::value; - const float2 r2 = {r, r}; - asm volatile("v_pk_mul_f32 v[%0:%1], %2, v[%0:%1]" : : "n"(base), "n"(base + 1), "v"(r2)); - }; - auto mul_pair = [&](float r, auto base_c) { - constexpr uint32_t base = decltype(base_c)::value; - asm volatile("v_mul_f32_e32 v[%0], %1, v[%0]" : : "n"(base), "v"(r)); - asm volatile("v_mul_f32_e32 v[%0], %1, v[%0]" : : "n"(base + 1), "v"(r)); - }; - - // Issue both ds_read_b64 for base tile S_jj (within its sub-tile) into a - // round-robin slot (4 slots). kRowBase selects the KV sub-tile (0=A, 32=B), - // folded into the V ds_read imm by load_transposed_v_to_gpr. - auto load_S = [&]() { - constexpr uint32_t base = (slot == 0u) ? k_v0_begin - : (slot == 1u) ? k_k0_begin - : (slot == 2u) ? k_k1_begin - : k_k2_begin; - kv_manager.template load_transposed_v_to_gpr(p_lds_v); - kv_manager.template load_transposed_v_to_gpr(p_lds_v); - }; - // mfma: oaccu_dst (+)= pv_v_{slot}^T @ p_mfma[sub-tile]. 3-arg init only on the - // FIRST sub-tile of a first-iter call (kRowBase==0); sub-tile B always accumulates - // onto A's oaccu (same oaccu tile, second N-half of P). - auto do_mma = [&](OA& oaccu_dst) { - auto run = [&](auto& v) { - if constexpr(kIsFirstIter && kRowBase == 0u) - hk::mma_ABt(oaccu_dst, v, p_mfma_a); - else if constexpr(kRowBase == 0u) - hk::mma_ABt(oaccu_dst, v, p_mfma_a, oaccu_dst); - else - hk::mma_ABt(oaccu_dst, v, p_mfma_b, oaccu_dst); - }; - if constexpr(slot == 0u) - run(pv_v_0); - else if constexpr(slot == 1u) - run(pv_v_1); - else if constexpr(slot == 2u) - run(pv_v_2); - else - run(pv_v_3); - }; - - if constexpr(kDoRescale) - { - opus::static_for<2>([&](auto s) { - pk_mul_pair(rescale, opus::number{}); - pk_mul_pair(rescale, opus::number{}); - }); - } - - // Prologue: preload the first 3 flat V tiles (S_0,S_1,S_2 of sub-tile A) into - // slots 0,1,2. ONE prologue for the whole call -- sub-tile B's S_0..S_2 are - // prefetched by sub-tile A's tail iters via the flat prefetch index below. - load_S.template operator()<0u, 0u, 0u>(); - load_S.template operator()<0u, 1u, 1u>(); - load_S.template operator()<0u, 2u, 2u>(); - - opus::static_for([&](auto i) { - constexpr uint32_t iter = i.value; - constexpr uint32_t col_idx = iter % kDIters; // 0..15 D-tile within sub-tile - constexpr uint32_t row_base = (iter / kDIters) * 32u; // 0 (A) or 32 (B) - constexpr bool has_next = (iter + 1u) < num_pv_iter; - // Rescale only folds on the FIRST sub-tile pass (row_base==0) and only when a - // NEXT D-tile in that pass exists (col_idx+1 < kDIters) -- else oaccu overflows. - constexpr bool resc_next = kDoRescale && (row_base == 0u) && (col_idx + 1u < kDIters); - constexpr uint32_t next_oaccu_base = k_o_begin + (col_idx + 1u) * 8u; - - // Flat V-tile index across all sub-tiles: prefetch runs on this so sub-tile A's - // tail prefetches sub-tile B's S_0.. (ONE prologue). Consume uses col_idx. - constexpr uint32_t jf_lo = 2u * iter; // flat mfma idx (a) - constexpr uint32_t jf_hi = 2u * iter + 1u; // flat mfma idx (b) - // 4-slot round-robin (flat tile T -> slot T%4). Reading slot j%4 while refilling - // slot (j+3)%4 keeps the refill ds_read's dst != the mfma's src reg. - constexpr uint32_t slot_lo = jf_lo % 4u; - constexpr uint32_t slot_hi = jf_hi % 4u; - - constexpr uint32_t oaccu_base = k_o_begin + col_idx * 8u; - using oaccu_a_r = - hkdart::split_many_t>, - 4>; - using oaccu_b_r = - hkdart::split_many_t>, - 4>; - hk::art oaccu_a; - hk::art oaccu_b; - - // Prefetch S at flat index jf+3 -> its own sub-tile row_base + local jj%32. - auto prefetch = [&]() { - if constexpr(jf < kNumVTiles) - { - constexpr uint32_t pf_row = (jf / kVTilesPerSub) * 32u; - constexpr uint32_t pf_jj = jf % kVTilesPerSub; - load_S.template operator()(); - } - }; - - // ---- mfma_a (flat jf_lo, reads slot_lo) ---- - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(has_next ? 4 : 2, -1)); - prefetch.template operator()(); - do_mma.template operator()(oaccu_a); - if constexpr(resc_next) - { - mul_pair(rescale, opus::number{}); - mul_pair(rescale, opus::number{}); - } - - // ---- mfma_b (flat jf_hi, reads slot_hi) ---- - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(has_next ? 4 : 0, -1)); - prefetch.template operator()(); - do_mma.template operator()(oaccu_b); - if constexpr(resc_next) - { - mul_pair(rescale, opus::number{}); - mul_pair(rescale, opus::number{}); - } - }); -} - -// All waves need the same H16 Q. Only one cooperative 16KiB global/LDS copy -// occurs; the four wave-private register copies are populated after the barrier. -__device__ __forceinline__ void load_q(const Params& p, uint16_t* tile) -{ - const uint32_t lane = threadIdx.x & 63u; - const uint32_t wave = __builtin_amdgcn_readfirstlane(threadIdx.x >> 6); - const bool aligned = ((reinterpret_cast(p.q) | uint64_t(p.q_stride0) * 2u | - uint64_t(p.q_stride1) * 2u) & - 15u) == 0; - if(aligned) - { - const uint32_t row = lane >> 2, quad = lane & 3u; - const uint64_t offset = - uint64_t(blockIdx.x) * uint64_t(p.q_stride0) + uint64_t(row) * uint64_t(p.q_stride1); - const uint32_t base = __builtin_amdgcn_readfirstlane( - static_cast(reinterpret_cast(tile))); - for(uint32_t c = wave; c < 16u; c += 4u) - { - const uint32_t L = c * 32u + ((quad * 8u) ^ (((row >> 2) & 1u) * 16u)); - const uint32_t col = sb8_inv_perm_col_elems(L); - const uint64_t address = reinterpret_cast(p.q + offset + col); - const uint32_t destination = base + c * 1024u; - asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tglobal_load_lds_dwordx4 %1, off" - : - : "s"(destination), "v"(address) - : "m0", "memory"); - } - drain(); - __syncthreads(); - const uint32_t r = lane & 15u, col = (lane >> 4) * 8u; - const uint32_t address = base + r * 64u + ((col * 2u) ^ (((r >> 2) & 1u) * 32u)); - opus::static_for<16>( - [&](auto c) { hkm::ds_read_b128<64u + c.value * 4u>(address, c.value * 1024u); }); - } - else - { - const uint32_t head = lane & 15u; - const uint64_t offset = - uint64_t(blockIdx.x) * uint64_t(p.q_stride0) + uint64_t(head) * uint64_t(p.q_stride1); - opus::static_for<16>([&](auto c) { - const uint32_t col = sb8_inv_perm_col_elems(c.value * 32u + (lane >> 4) * 8u); - opus::static_for<4>([&](auto j) { - const uint32_t pair = uint32_t(p.q[offset + col + j.value * 2u]) | - (uint32_t(p.q[offset + col + j.value * 2u + 1u]) << 16); - write_u32<64u + c.value * 4u + j.value>(pair); - }); - __builtin_amdgcn_sched_barrier(0); - }); - } - drain(); - __syncthreads(); // every replicated Q register read retires before KV overwrite -} - -// Two PR 32-row subtiles, each preserving the exact sb8 + row-bank swizzle. -// This initial N64 geometry intentionally uses a single synchronous KV buffer. -__device__ __forceinline__ void -stage_kv(const Params& p, uint16_t* tile, int32_t* selected, int64_t start, int64_t end) -{ - const uint32_t tid = threadIdx.x, lane = tid & 63u; - if(tid < 64u) - { - int32_t slot = -1; - if(start + tid < end) - slot = p.indices[start + tid]; - selected[tid] = slot; - } - __syncthreads(); - const bool aligned = - ((reinterpret_cast(p.kv) | uint64_t(p.kv_stride0) * 2u) & 15u) == 0; - if(aligned) - { - const uint32_t wave = __builtin_amdgcn_readfirstlane(tid >> 6); - const uint32_t base = __builtin_amdgcn_readfirstlane( - static_cast(reinterpret_cast(tile))); - for(uint32_t block = wave; block < 64u; block += 4u) - { - const uint32_t local = block & 31u, r = lane >> 2, quad = lane & 3u; - const uint32_t row = (block >> 5) * 32u + (local & 1u) * 16u + r; - const uint32_t L = (local >> 1) * 32u + ((quad * 8u) ^ (((r >> 2) & 1u) * 16u)); - const uint32_t col = sb8_inv_perm_col_elems(L); - const int32_t slot = selected[row]; - const uint32_t destination = base + block * 1024u; - if(slot >= 0 && int64_t(slot) < p.slots) - { - const uint64_t offset = uint64_t(slot) * uint64_t(p.kv_stride0) + uint64_t(col); - const uint64_t address = reinterpret_cast(p.kv + offset); - asm volatile("s_mov_b32 m0, %0\n\ts_nop 0\n\tglobal_load_lds_dwordx4 %1, off" - : - : "s"(destination), "v"(address) - : "m0", "memory"); - } - else - { - const v4ui zero = {0, 0, 0, 0}; - *reinterpret_cast(reinterpret_cast(tile) + block * 1024u + - lane * 16u) = zero; - } - } - } - else - { - for(uint32_t linear = tid; linear < 64u * 512u; linear += 256u) - { - const uint32_t row = linear / 512u, col = linear % 512u; - const int32_t slot = selected[row]; - uint16_t value = 0; - if(slot >= 0 && int64_t(slot) < p.slots) - value = p.kv[uint64_t(slot) * uint64_t(p.kv_stride0) + uint64_t(col)]; - const uint32_t L = sb8_perm_col_elems(col), r = row & 15u; - const uint32_t byte = (row / 32u) * 32768u + - (L / 32u * 2u + (row % 32u) / 16u) * 1024u + r * 64u + - ((2u * (L & 31u)) ^ (((r >> 2) & 1u) * 32u)); - tile[byte / 2u] = value; - } - } - drain(); - __syncthreads(); -} - -__device__ __forceinline__ void qk(uintptr_t tile) -{ - // Wave w computes its 16 selected tokens against all 16 query heads. - const uint32_t wave = __builtin_amdgcn_readfirstlane(threadIdx.x >> 6); - const uintptr_t wave_tile = tile + (wave / 2u) * 32768u + (wave % 2u) * 1024u; - Regs::score_t score; - auto issue = [&]() { - constexpr uint32_t base = 40u + (S % 3u) * 4u; - hk::art> k; - Bf16KvManager::template load_k_to_gpr<0u, S * 32u>(k, wave_tile); - }; - auto mma = [&]() { - constexpr uint32_t base = 40u + (S % 3u) * 4u, qb = 64u + S * 4u; - hk::art> k; - hk::art> q; - if constexpr(S == 0) - hk::mma_ABt(score, k, q); - else - hk::mma_ABt(score, k, q, score); - }; - opus::static_for<3>([&](auto s) { issue.template operator()(); }); - opus::static_for<16>([&](auto s) { - constexpr uint32_t remain = 15u - s.value; - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(remain < 2u ? remain : 2u, -1)); - mma.template operator()(); - if constexpr(s.value + 3u < 16u) - issue.template operator()(); - }); - asm volatile("s_nop 7\n\ts_nop 7"); // opaque MFMA score dependency is not compiler-visible - __builtin_amdgcn_sched_barrier(0); -} - -__device__ __forceinline__ void softmax(const Params& p, - const int32_t* selected, - float* exchange, - uint16_t* p_lds, - float& maximum, - float& denominator) -{ - const uint32_t lane = threadIdx.x & 63u, wave = threadIdx.x >> 6; - const uint32_t head = lane & 15u, group = lane >> 4; - const uint32_t first_key = wave * 16u + group * 4u; - const uint32_t index_address = - static_cast(reinterpret_cast(selected)) + first_key * sizeof(int32_t); - hkm::ds_read_b128<40>(index_address, 0); - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, -1)); - float local_max = -INFINITY; - opus::static_for<4>([&](auto i) { - const int32_t slot = read_i32<40u + i.value>(); - float x = (read_f32<56u + i.value>() * p.scale) * 1.4426950408889634074f; - if(slot < 0 || int64_t(slot) >= p.slots) - x = -INFINITY; - write_f32<56u + i.value>(x); - local_max = fmaxf(local_max, x); - }); - local_max = fmaxf(local_max, __shfl_xor(local_max, 16, 64)); - local_max = fmaxf(local_max, __shfl_xor(local_max, 32, 64)); - if(lane < 16u) - exchange[wave * 16u + head] = local_max; - __syncthreads(); - float all_max = -INFINITY; - opus::static_for<4>([&](auto w) { all_max = fmaxf(all_max, exchange[w.value * 16u + head]); }); - const float new_max = fmaxf(maximum, all_max); - const float alpha = denominator > 0.0f ? __builtin_amdgcn_exp2f(maximum - new_max) : 0.0f; - maximum = new_max; - Regs::oaccu_t o; - hk::mul_vgpr(o, o, alpha); - float local_sum = 0.0f; - opus::static_for<4>([&](auto i) { - const float x = read_f32<56u + i.value>(); - const float weight = x == -INFINITY ? 0.0f : __builtin_amdgcn_exp2f(x - maximum); - write_f32<56u + i.value>(weight); - local_sum += weight; - }); - local_sum += __shfl_xor(local_sum, 16, 64); - local_sum += __shfl_xor(local_sum, 32, 64); - if(lane < 16u) - exchange[64u + wave * 16u + head] = local_sum; - // Each wave owns a disjoint N16 slab of BF16 P. Normal [head,key] LDS layout - // makes the subsequent full-N64 operand copies explicit and bit-preserving. - v2ui packed; - const float p0 = read_f32<56>(), p1 = read_f32<57>(); - const float p2 = read_f32<58>(), p3 = read_f32<59>(); - asm volatile("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(packed[0]) : "v"(p0), "v"(p1)); - asm volatile("v_cvt_pk_bf16_f32 %0, %1, %2" : "=v"(packed[1]) : "v"(p2), "v"(p3)); - const uintptr_t address = reinterpret_cast(p_lds) + (head * 64u + first_key) * 2u; - hkm::ds_write_b64(packed, static_cast(address), 0); - __builtin_amdgcn_s_waitcnt(hk_mla::encode_s_waitcnt(0, -1)); - __syncthreads(); // all P and four partial denominators are visible - float all_sum = 0.0f; - opus::static_for<4>([&](auto w) { all_sum += exchange[64u + w.value * 16u + head]; }); - denominator = denominator * alpha + all_sum; - const uint32_t p_address = - static_cast(reinterpret_cast(p_lds)) + (head * 64u + group * 4u) * 2u; - // PR's transposed-V operand orders the two N16 halves as group*4 and +16. - hkm::ds_read_b64<56>(p_address, 0); - hkm::ds_read_b64<58>(p_address, 32); - hkm::ds_read_b64<60>(p_address, 64); - hkm::ds_read_b64<62>(p_address, 96); - drain(); - __syncthreads(); // denominator/max consumers retire before next iteration reuses exchange -} - -template -__device__ __forceinline__ void run(Params p) -{ - hkdart::clobber(); - extern __shared__ __align__(16) unsigned char storage[]; - auto* tile = reinterpret_cast(storage); - auto* p_lds = reinterpret_cast(storage + Traits::kPOffset); - auto* exchange = reinterpret_cast(storage + Traits::kReduceOffset); - auto* selected = reinterpret_cast(storage + Traits::kIndexOffset); - const uint32_t query = blockIdx.x, partition = blockIdx.y; - const int32_t first = p.indptr[query], last = p.indptr[query + 1u], length = last - first; - const int32_t begin = first + int64_t(length) * partition / p.splits; - const int32_t end = first + int64_t(length) * (partition + 1u) / p.splits; - load_q(p, tile); - Regs::oaccu_t o; - hk::zero(o); - float maximum = -INFINITY, denominator = 0.0f; - Bf16KvManager manager; - const uint32_t wave = __builtin_amdgcn_readfirstlane(threadIdx.x >> 6); - for(int64_t pos = begin; pos < end; pos += 64) - { - stage_kv(p, tile, selected, pos, end); - qk(reinterpret_cast(tile)); - softmax(p, selected, exchange, p_lds, maximum, denominator); - // D128/wave is a multiple of the sb8 permutation's 64-element period. - const uintptr_t v_tile = reinterpret_cast(tile) + wave * 8192u; - sparse_mla_bf16::wave4::pv_gemm(manager, v_tile, 1.0f); - asm volatile("s_nop 7\n\ts_nop 7"); - drain(); - __syncthreads(); - } - const float inverse = denominator > 0.0f ? 1.0f / denominator : 0.0f; - hk::mul_vgpr(o, o, inverse); - const float lse = denominator > 0.0f - ? (maximum + __builtin_amdgcn_logf(denominator)) * 0.6931471805599453094f - : -INFINITY; - const uint32_t lane = threadIdx.x & 63u; - const uint64_t output_slot = Split ? uint64_t(query) * uint64_t(p.splits) + partition : query; - if(wave == 0u && lane < 16u) - { - if constexpr(Split) - p.partial_lse[output_slot * 16u + lane] = lse; - else if(p.lse) - p.lse[output_slot * 16u + lane] = lse; - } - // PR O manager preserves global head stride512. Offset its descriptor to - // this wave's D128, set head-wave id to zero, and give each wave a private - // dead-KV LDS bounce so its data cannot race another output wave. - const uintptr_t bounce = reinterpret_cast(tile) + wave * 4352u; - if constexpr(Split) - { - OManager32bitsV4Gen1Swizzle output; - float* dst = p.partial_o + output_slot * 16u * 512u + wave * 128u; - opus::static_for<2>([&](auto i) { - output.template output_to_vram_pair<128u + i.value * 16u, i.value * 64u, false>( - dst, 0, 0, 0, bounce, 16); - drain(); - }); - } - else - { - OManager16bitsV4Gen1Swizzle output; - auto* dst = reinterpret_cast(p.out + output_slot * 16u * 512u + wave * 128u); - opus::static_for<2>([&](auto i) { - output.template output_to_vram_pair<128u + i.value * 16u, i.value * 64u, false>( - dst, 0, 0, 0, bounce, 16); - drain(); - }); - } -} -} // namespace sparse_mla_bf16::wave4 - -namespace sparse_mla_bf16 { -template -__global__ __launch_bounds__(256) __attribute__((amdgpu_num_vgpr(40))) void sparse_mla_v2(Params p) -{ - if constexpr(Heads == 16) - wave4::run(p); - else - sparse_mla_legacy(p); -} -} // namespace sparse_mla_bf16 diff --git a/csrc/kernels/mla/sparse_mla_bf16.cu b/csrc/kernels/mla/sparse_mla_bf16.cu deleted file mode 100644 index ac2050105ea..00000000000 --- a/csrc/kernels/mla/sparse_mla_bf16.cu +++ /dev/null @@ -1,218 +0,0 @@ -// SPDX-License-Identifier: MIT -// BF16, D512 sparse MLA for gfx950. v1: full LDS; v2: hybrid + unrolled reducer. -// aiter_ctypes_error.h requires the common declarations supplied by aiter_tensor.h. -// clang-format off -#include "aiter_tensor.h" -#include "aiter_ctypes_error.h" -#include "hk/sparse_mla_bf16_wave4.cuh" -#include "hk/sparse_mla_bf16_headtiles.cuh" -#include "hk/sparse_mla_bf16_reduce.cuh" -// clang-format on -#include -#include - -AITER_CTYPES_ERROR_DEF - -namespace sparse_mla_bf16 { -#include "hk/sparse_mla_bf16_full_lds.cuh" - -template -__global__ __launch_bounds__(Heads == 64 ? 512 : 64) - __attribute__((amdgpu_num_vgpr(40))) void sparse_mla_v1(Params p) -{ - if constexpr(Heads == 64) - full_lds_h64(p); - else - sparse_mla_legacy(p); -} - -template -hipError_t launch_main(Params p, int version, hipStream_t stream) -{ - const dim3 grid(p.queries, p.splits); - if(version == 1) - { - constexpr int threads = Heads == 64 ? 512 : 64; - constexpr int lds = Heads == 64 ? FullLdsTraits::kLdsBytes : Traits<16>::kLdsBytes; - hipLaunchKernelGGL((sparse_mla_v1), grid, dim3(threads), lds, stream, p); - } - else if(Heads == 64 && uint64_t(p.queries) * p.splits <= 128u) - { - hipLaunchKernelGGL((sparse_mla_headtiles<64, Split>), - dim3(p.queries, p.splits, 4), - dim3(256), - headtiles::Traits::kLdsBytes, - stream, - p); - } - else - { - constexpr int lds = Heads == 16 ? wave4::Traits::kLdsBytes : Traits<64>::kLdsBytes; - hipLaunchKernelGGL((sparse_mla_v2), grid, dim3(256), lds, stream, p); - } - return hipGetLastError(); -} - -template -hipError_t launch(Params p, int version, hipStream_t stream) -{ - if(p.splits == 1) - return launch_main(p, version, stream); - hipError_t status = launch_main(p, version, stream); - if(status != hipSuccess) - return status; - if constexpr(Heads == 16) - { - if(version == 1) - { - hipLaunchKernelGGL(reduce, dim3(p.queries, Heads), dim3(256), 0, stream, p, Heads); - return hipGetLastError(); - } - } - return launch_reduce64(p, stream); -} -} // namespace sparse_mla_bf16 - -namespace { -void check_tensor(const aiter_tensor_t* t, const char* name, AiterDtype dtype, int device) -{ - AITER_CHECK(t && t->is_gpu() && t->device_id == device, name, ": expected the query device"); - AITER_CHECK(t->dtype() == dtype, name, ": unexpected dtype"); - AITER_CHECK(t->numel() == 0 || t->data_ptr(), name, ": null storage"); -} - -void check_shape(const aiter_tensor_t* t, const char* name, std::initializer_list shape) -{ - AITER_CHECK(t->dim() == int(shape.size()), name, ": unexpected rank"); - int d = 0; - for(int64_t extent : shape) - AITER_CHECK(t->size(d++) == extent, name, ": unexpected shape"); -} - -// Conservatively reject overlapping storage spans. Inputs can alias each other; -// writable tensors must not alias inputs or another writable tensor. -uintptr_t end_address(const aiter_tensor_t* t) -{ - if(!t || !t->numel()) - return 0; - size_t span = 1; - for(int d = 0; d < t->dim(); ++d) - span += (t->size(d) - 1) * t->stride(d); - return reinterpret_cast(t->data_ptr()) + span * t->element_size(); -} -void check_no_overlap(const aiter_tensor_t* a, const aiter_tensor_t* b) -{ - if(!a || !b || !a->numel() || !b->numel()) - return; - AITER_CHECK(reinterpret_cast(a->data_ptr()) >= end_address(b) || - reinterpret_cast(b->data_ptr()) >= end_address(a), - "sparse_mla_bf16: output/workspace storage must not overlap"); -} -} // namespace - -AITER_CTYPES_DEFINE_ENTRYPOINT_VOID(sparse_mla_bf16_fwd_out, - (aiter_tensor_t * q, - aiter_tensor_t* kv_buffer, - aiter_tensor_t* kv_indptr, - aiter_tensor_t* kv_indices, - aiter_tensor_t* out, - aiter_tensor_t* lse, - aiter_tensor_t* partial_o, - aiter_tensor_t* partial_lse, - float softmax_scale, - int64_t kv_splits, - int64_t version, - hipStream_t stream), - (q, - kv_buffer, - kv_indptr, - kv_indices, - out, - lse, - partial_o, - partial_lse, - softmax_scale, - kv_splits, - version, - stream)) -{ - AITER_CHECK(q && q->is_gpu(), "q must be on a HIP device"); - const int device = q->device_id; - HipDeviceGuard guard(device); - AITER_CHECK(get_gpu_arch() == "gfx950", "sparse_mla_bf16 requires gfx950"); - AITER_CHECK(version == 1 || version == 2, "version must be 1 or 2"); - AITER_CHECK(kv_splits >= 1 && kv_splits <= 32 && (kv_splits & (kv_splits - 1)) == 0, - "kv_splits must be one of 1, 2, 4, 8, 16, 32"); - AITER_CHECK(std::isfinite(softmax_scale) && softmax_scale > 0, "invalid softmax_scale"); - check_tensor(q, "q", AITER_DTYPE_bf16, device); - AITER_CHECK(q->dim() == 3 && q->size(2) == 512 && (q->size(1) == 16 || q->size(1) == 64), - "expected q=[Q,16|64,512]"); - const int64_t queries = q->size(0), heads = q->size(1); - AITER_CHECK(queries <= INT_MAX, "too many query rows"); - AITER_CHECK(q->stride(2) == 1 && q->stride(1) >= 512 && - q->stride(0) >= (heads - 1) * q->stride(1) + 512, - "q must have non-overlapping, unit-inner-stride rows"); - check_tensor(kv_buffer, "kv_buffer", AITER_DTYPE_bf16, device); - AITER_CHECK(kv_buffer->dim() == 2 && kv_buffer->size(1) == 512 && kv_buffer->stride(1) == 1 && - kv_buffer->stride(0) >= 512, - "expected kv_buffer=[slots,512] with unit inner stride"); - check_tensor(kv_indptr, "kv_indptr", AITER_DTYPE_i32, device); - check_tensor(kv_indices, "kv_indices", AITER_DTYPE_i32, device); - check_shape(kv_indptr, "kv_indptr", {queries + 1}); - AITER_CHECK(kv_indices->dim() == 1 && kv_indptr->is_contiguous() && kv_indices->is_contiguous(), - "CSR tensors must be contiguous vectors"); - AITER_CHECK(kv_indices->numel() <= INT_MAX, "CSR offsets must fit int32"); - check_tensor(out, "out", AITER_DTYPE_bf16, device); - check_shape(out, "out", {queries, heads, 512}); - AITER_CHECK(out->is_contiguous(), "out must be contiguous"); - if(lse) - { - check_tensor(lse, "lse", AITER_DTYPE_fp32, device); - check_shape(lse, "lse", {queries, heads}); - AITER_CHECK(lse->is_contiguous(), "lse must be contiguous"); - } - if(kv_splits > 1) - { - check_tensor(partial_o, "partial_o", AITER_DTYPE_fp32, device); - check_tensor(partial_lse, "partial_lse", AITER_DTYPE_fp32, device); - check_shape(partial_o, "partial_o", {queries, kv_splits, heads, 512}); - check_shape(partial_lse, "partial_lse", {queries, kv_splits, heads}); - AITER_CHECK(partial_o->is_contiguous() && partial_lse->is_contiguous(), - "workspace must be contiguous"); - } - else - { - partial_o = partial_lse = nullptr; - } - const aiter_tensor_t* writes[] = {out, lse, partial_o, partial_lse}; - for(int i = 0; i < 4; ++i) - { - for(auto input : {q, kv_buffer, kv_indptr, kv_indices}) - check_no_overlap(writes[i], input); - for(int j = 0; j < i; ++j) - check_no_overlap(writes[i], writes[j]); - } - // CSR offsets are caller-owned device data: monotone, starting at zero and - // ending at indices.numel(). No D2H reads or synchronization during capture. - // Out-of-range slot ids (including -1) are masked in the kernel. - if(!queries) - return; - sparse_mla_bf16::Params p{static_cast(q->data_ptr()), - static_cast(kv_buffer->data_ptr()), - static_cast(kv_indptr->data_ptr()), - static_cast(kv_indices->data_ptr()), - static_cast(out->data_ptr()), - lse ? static_cast(lse->data_ptr()) : nullptr, - partial_o ? static_cast(partial_o->data_ptr()) : nullptr, - partial_lse ? static_cast(partial_lse->data_ptr()) : nullptr, - int32_t(queries), - int32_t(kv_splits), - kv_buffer->size(0), - q->stride(0), - q->stride(1), - kv_buffer->stride(0), - softmax_scale}; - const auto status = heads == 16 ? sparse_mla_bf16::launch<16>(p, version, stream) - : sparse_mla_bf16::launch<64>(p, version, stream); - HIP_CALL_LAUNCH(status); -} diff --git a/csrc/kernels/mla/sparse_mla_fp8.cu b/csrc/kernels/mla/sparse_mla_fp8.cu new file mode 100644 index 00000000000..3d6d368f78e --- /dev/null +++ b/csrc/kernels/mla/sparse_mla_fp8.cu @@ -0,0 +1,162 @@ +// SPDX-License-Identifier: MIT +// FP8 sparse MLA with fused BF16 query conversion and BF16 output. +// clang-format off +#include "aiter_tensor.h" +#include "aiter_ctypes_error.h" +#include "fp8/params.cuh" +// clang-format on +#include +#include + +AITER_CTYPES_ERROR_DEF + +namespace { +void check_tensor(const aiter_tensor_t* t, const char* name, AiterDtype dtype, int device) +{ + AITER_CHECK(t && t->is_gpu() && t->device_id == device, name, ": expected the query device"); + AITER_CHECK(t->dtype() == dtype, name, ": unexpected dtype"); + AITER_CHECK(t->numel() == 0 || t->data_ptr(), name, ": null storage"); +} +void check_shape(const aiter_tensor_t* t, const char* name, std::initializer_list shape) +{ + AITER_CHECK(t->dim() == int(shape.size()), name, ": unexpected rank"); + int d = 0; + for(int64_t extent : shape) + AITER_CHECK(t->size(d++) == extent, name, ": unexpected shape"); +} +void check_aligned(const aiter_tensor_t* t, const char* name, size_t alignment) +{ + AITER_CHECK(!t->numel() || reinterpret_cast(t->data_ptr()) % alignment == 0, + name, + ": unaligned storage"); +} +uintptr_t end_address(const aiter_tensor_t* t) +{ + if(!t || !t->numel()) + return 0; + size_t span = 1; + for(int d = 0; d < t->dim(); ++d) + span += (t->size(d) - 1) * t->stride(d); + return reinterpret_cast(t->data_ptr()) + span * t->element_size(); +} +void check_no_overlap(const aiter_tensor_t* a, const aiter_tensor_t* b) +{ + if(!a || !b || !a->numel() || !b->numel()) + return; + AITER_CHECK(reinterpret_cast(a->data_ptr()) >= end_address(b) || + reinterpret_cast(b->data_ptr()) >= end_address(a), + "sparse_mla_fp8: output/workspace storage must not overlap"); +} +} // namespace + +AITER_CTYPES_DEFINE_ENTRYPOINT_VOID(sparse_mla_fp8_fwd_out, + (aiter_tensor_t * q, + aiter_tensor_t* kv_buffer, + aiter_tensor_t* kv_indices, + aiter_tensor_t* out, + aiter_tensor_t* lse, + aiter_tensor_t* partial_o, + aiter_tensor_t* partial_lse, + float softmax_scale, + float q_scale, + float kv_scale, + hipStream_t stream), + (q, + kv_buffer, + kv_indices, + out, + lse, + partial_o, + partial_lse, + softmax_scale, + q_scale, + kv_scale, + stream)) +{ + AITER_CHECK(q && q->is_gpu(), "q must be on a HIP device"); + const int device = q->device_id; + HipDeviceGuard guard(device); + AITER_CHECK(get_gpu_arch() == "gfx950", "sparse_mla_fp8 requires gfx950"); + for(float scale : {softmax_scale, q_scale, kv_scale}) + AITER_CHECK(std::isfinite(scale) && scale > 0, "scales must be finite and positive"); + AITER_CHECK(q->dtype() == AITER_DTYPE_bf16 || q->dtype() == AITER_DTYPE_fp8, + "q must be BF16 or FP8"); + check_tensor(q, "q", q->dtype(), device); + AITER_CHECK(q->dim() == 3 && q->size(1) == 64 && q->size(2) == 512, "expected q=[Q,64,512]"); + const int64_t queries = q->size(0); + AITER_CHECK(queries <= 2048, "Q must be <= 2048"); + AITER_CHECK(q->stride(2) == 1 && q->stride(1) >= 512 && q->stride(0) >= 63 * q->stride(1) + 512, + "q must have non-overlapping, unit-inner-stride rows"); + const int bytes = q->element_size(); + AITER_CHECK((q->stride(0) * bytes) % 16 == 0 && (q->stride(1) * bytes) % 16 == 0, + "q rows must be aligned to 16 bytes"); + check_aligned(q, "q", 16); + check_tensor(kv_buffer, "kv_buffer", AITER_DTYPE_fp8, device); + AITER_CHECK(kv_buffer->dim() == 2 && kv_buffer->size(1) == 512 && kv_buffer->size(0) > 0 && + kv_buffer->size(0) <= INT_MAX && kv_buffer->is_contiguous(), + "expected contiguous kv_buffer=[slots,512]"); + check_aligned(kv_buffer, "kv_buffer", 16); + check_tensor(kv_indices, "kv_indices", AITER_DTYPE_i32, device); + AITER_CHECK(kv_indices->dim() == 2 && kv_indices->size(0) == queries && + kv_indices->size(1) <= 2176 && kv_indices->is_contiguous(), + "expected contiguous kv_indices=[Q,K], K <= 2176"); + check_aligned(kv_indices, "kv_indices", 4); + check_tensor(out, "out", AITER_DTYPE_bf16, device); + check_shape(out, "out", {queries, 64, 512}); + AITER_CHECK(out->is_contiguous(), "out must be contiguous"); + check_aligned(out, "out", 16); + if(lse) + { + check_tensor(lse, "lse", AITER_DTYPE_fp32, device); + check_shape(lse, "lse", {queries, 64}); + AITER_CHECK(lse->is_contiguous(), "lse must be contiguous"); + check_aligned(lse, "lse", 4); + } + const int splits = sparse_mla_fp8::split_count(queries); + const int64_t rows = splits > 1 ? queries * splits * 64 : 0; + check_tensor(partial_o, "partial_o", AITER_DTYPE_bf16, device); + check_tensor(partial_lse, "partial_lse", AITER_DTYPE_fp32, device); + check_shape(partial_o, "partial_o", {rows * 512}); + check_shape(partial_lse, "partial_lse", {rows}); + AITER_CHECK(partial_o->is_contiguous() && partial_lse->is_contiguous(), + "workspace must be contiguous"); + check_aligned(partial_o, "partial_o", 16); + check_aligned(partial_lse, "partial_lse", 4); + const aiter_tensor_t* writes[] = {out, lse, partial_o, partial_lse}; + for(int i = 0; i < 4; ++i) + { + for(auto input : {q, kv_buffer, kv_indices}) + check_no_overlap(writes[i], input); + for(int j = 0; j < i; ++j) + check_no_overlap(writes[i], writes[j]); + } + if(!queries) + return; + sparse_mla_fp8::Params p{}; + p.q = static_cast(q->data_ptr()); + p.kv = static_cast(kv_buffer->data_ptr()); + p.indices = static_cast(kv_indices->data_ptr()); + p.partial_o = static_cast(partial_o->data_ptr()); + p.partial_lse = static_cast(partial_lse->data_ptr()); + p.out = static_cast(out->data_ptr()); + p.lse = lse ? static_cast(lse->data_ptr()) : nullptr; + p.queries = int(queries); + p.slots = int(kv_buffer->size(0)); + p.capacity = int(kv_indices->size(1)); + p.q_scale = q_scale; + p.kv_scale = kv_scale; + p.attention_scale = softmax_scale; + p.q_stride_row = q->stride(0); + p.q_stride_head = q->stride(1); + p.bf16_q = q->dtype() == AITER_DTYPE_bf16; + hipError_t status; + if(queries <= 6) + status = sparse_mla_fp8::launch_small(p, stream); + else if(queries <= 12) + status = sparse_mla_fp8::launch_prefetch(p, stream); + else if(queries <= 64) + status = sparse_mla_fp8::launch_rotating(p, stream); + else + status = sparse_mla_fp8::launch_streaming(p, stream); + HIP_CALL_LAUNCH(status); +} diff --git a/docs/sparse_mla_bf16.md b/docs/sparse_mla_bf16.md deleted file mode 100644 index 4492eccb084..00000000000 --- a/docs/sparse_mla_bf16.md +++ /dev/null @@ -1,125 +0,0 @@ -# BF16 sparse MLA on gfx950 - -`aiter.sparse_mla_bf16_fwd` computes attention over a per-query CSR list of -global cache slots. K and V are the same BF16 latent, D=512, with no appended -RoPE dimensions. Both prefill and decode use this contract. H16 and H64 are -supported; GLM-5.3-Flash is one caller of this geometry. - -This is an explicit HIP entry point. The existing Gluon `sparse_mla_fwd` -continues to handle its broader set of formats and geometries. An engine can -select this entry point after checking its device, dtype, head count and -geometry, and retain its existing implementation for other configurations. - -## API and storage - -```python -import aiter -import torch - -# q: [Q,H,512] BF16; kv: [slots,512] BF16 -# ptr: [Q+1] int32; idx: [nnz] int32 GLOBAL physical slot ids -splits = 8 # choose outside capture; supported: 1,2,4,8,16,32 -out = torch.empty_like(q, memory_format=torch.contiguous_format) -lse = torch.empty(q.shape[:2], device=q.device, dtype=torch.float32) -workspace = ( - torch.empty((q.shape[0], splits, q.shape[1], 512), device=q.device), - torch.empty((q.shape[0], splits, q.shape[1]), device=q.device), -) -result, result_lse = aiter.sparse_mla_bf16_fwd( - q, kv, ptr, idx, softmax_scale, - version=2, kv_splits=splits, - out=out, return_lse=True, lse=lse, workspace=workspace, -) -``` - -O is contiguous BF16 `[Q,H,512]`. LSE is optional, contiguous FP32 `[Q,H]` -and uses the natural logarithm. Split-K workspace contains normalized FP32 -partial O and natural-log partial LSE, contiguous `[Q,S,H,512]` and `[Q,S,H]`. -It is unused for S=1. The convenience wrapper allocates omitted buffers; -supplying all buffers makes the call allocation-free after JIT warmup. -Warm the extension before capturing a graph. Calls use the query device's -current stream and do not synchronize to inspect CSR values. - -Q/KV allow padded outer strides and BF16 storage offsets, with stride one -in the last dimension. KV also accepts flat-view-compatible -`[pages,page_size,512]` and `[slots,1,1,512]` layouts. Paged inputs are viewed, -not copied. Out/workspace storage must not overlap inputs or one another. -An output with a larger physical head stride requires the engine to use a -contiguous latent output and then its existing projection/copy path. - -The caller supplies monotone CSR offsets starting at zero and ending at -`idx.numel()`. Slot ids outside `[0,slots)` are masked; duplicates retain -their multiplicity. Empty or fully masked rows yield O=0, LSE=-inf. -Causal/window selection is expressed through CSR. There is no implicit -top-k truncation, attention sink, cache quantization, or separate RoPE input. -In particular, a 2051-entry row preserves all 2051 entries. - -## Implementations - -| Version | H16 main | H64 main | Split-K reduction | -|---|---|---|---| -| 1 | One wave, Q in VGPR, 33,024 B LDS | Full140: eight waves, Q in LDS, 139,520 B LDS | Fully unrolled one-wave H64; original H16 reducer | -| 2 | Four waves, 68,352 B LDS | Four head tiles when Q×S≤128 (68,352 B/CTA), otherwise four waves (65,792 B/CTA) | Fully unrolled one-wave H16/H64 | - -The version selector keeps both implementations reproducible under one API. -The default version is 2; the default split count is the explicit S=1, -not an autotuned choice. Use the supplied sweep to choose a split count -for the deployment workload. No automatic routing changes to the existing -Gluon API are part of this operator. - -The extension reuses AITER's existing HipKittens dependency and pinned register -helpers from [PR #3459](https://github.com/ROCm/aiter/pull/3459). No new external -dependency is introduced. This module uses AITER's CK-free headers. Like the -other HK modules, it is excluded from the -ordinary all-op prebuild unless `AITER_ENABLE_EXPERIMENTAL=1`; explicit JIT -calls can build it directly. The module preserves FP32 denormals, overriding -AITER's usual flush-to-zero flag. Compiler-managed VGPRs must stay below v40. - -## Reproduce correctness and timing - -```bash -python -m pytest -q op_tests/test_sparse_mla_bf16.py -python -m op_tests.op_benchmarks.bench_sparse_mla_bf16 \ - --verify --output sparse_mla_o_only.jsonl -python -m op_tests.op_benchmarks.bench_sparse_mla_bf16 \ - --verify --return-lse --output sparse_mla_o_lse.jsonl -# Independent repeat with reversed case/variant order and fresh allocations: -python -m op_tests.op_benchmarks.bench_sparse_mla_bf16 \ - --verify --reverse --output sparse_mla_o_only_repeat.jsonl -``` - -The benchmark defaults to H16/H64, Q=1/32/128/256/512/1024/2048, K=2048/2051, -S=1/2/4/8/16/32. Each query owns a disjoint 32K-slot pool. Q=2048 uses a -64 GiB KV allocation, so run one benchmark on an otherwise idle MI355X. -Native v1/v2 share input, final output and workspace addresses. The Gluon -comparison comes from the same checkout and includes explicit splits and -its shipped split policy. Internal Gluon workspaces differ from the native -workspace; their traffic is part of the timed operation. - -Timing uses alternating variant order and 60 graph/event samples per point. -Warm and cold results are separate. Cold flushes 512 MiB before the graph's -start event. JIT and Python allocation/dispatch time are excluded. Every -case checks sampled query rows against independent FP64 before and after -timing, and hashes all inputs before/after. Correctness failures are reported -and excluded from ratios; a native failure aborts the run. - -Report the complete per-shape table, matched splits, best tested splits and -Gluon's shipped policy separately. JSONL includes raw samples, median/p90, -source/binary hashes, logical GB/s and QK+PV TFLOP/s. Logical traffic counts -Q, one logical gather of KV, indices, O and optional LSE; it excludes repeated -reads and partial buffers and is not measured HBM traffic. A roofline derived -from it is a lower-bound model, not proof of a memory bottleneck. - -For execution-path validation, run a small sweep under `rocprofv3 ---kernel-trace --hip-trace` with `--purpose trace --samples 1 --warmup 0`. -Trace durations are diagnostic and must not be used as performance results. -Build `sparse_mla_bf16.cu` with `--save-temps` and check the retained assembly: - -```bash -python op_tests/check_sparse_mla_bf16_isa.py path/to/sparse_mla_bf16-hip-amdgcn-amd-amdhsa-gfx950.s -``` - -Operator correctness and latency are distinct from model quality and serving -performance. An engine integration should validate its logical-to-physical -index mapping, selected-set policy (including any tail fitting), graph buffer -reuse and fallback, then measure the model with a stable baseline. diff --git a/docs/sparse_mla_fp8.md b/docs/sparse_mla_fp8.md new file mode 100644 index 00000000000..c2c85645a29 --- /dev/null +++ b/docs/sparse_mla_fp8.md @@ -0,0 +1,118 @@ +# FP8 sparse MLA on gfx950 + +`sparse_mla_fp8_fwd` implements sparse attention over a shared FP8 latent K/V +pool, with BF16 query conversion inside attention and BF16 output. It targets +H64/D512 on gfx950 (MI355X), without appended RoPE. The API is opt-in; an engine +must explicitly route eligible calls and provide its own fallback. + +```python +import torch +from aiter import sparse_mla_fp8_fwd, allocate_sparse_mla_fp8_workspace + +q = torch.randn((8, 64, 512), device="cuda", dtype=torch.bfloat16) +kv = torch.randn((8192, 512), device="cuda", dtype=torch.bfloat16).to(torch.float8_e4m3fn) +indices = torch.randint(8192, (8, 2112), device="cuda", dtype=torch.int32) +indices[:, 2051:] = -1 +out = torch.empty_like(q) +workspace = allocate_sparse_mla_fp8_workspace(q) + +# Warm the JIT before capture. Supplying storage avoids repeated allocation. +out, _ = sparse_mla_fp8_fwd(q, kv, indices, 1 / 16, out=out, workspace=workspace) +graph = torch.cuda.CUDAGraph() +with torch.cuda.graph(graph): + out, _ = sparse_mla_fp8_fwd(q, kv, indices, 1 / 16, out=out, workspace=workspace) +# Update q, kv and indices in place, then replay for the same shapes. +graph.replay() +``` + +## Input and numerical contract + +| Argument | Contract | +|---|---| +| Q | BF16 or `float8_e4m3fn`, `[Q,64,512]`, Q from 0 through 2048 | +| Q layout | Unit inner stride, non-overlapping rows, pointer and row strides aligned to 16 bytes | +| KV | Contiguous `float8_e4m3fn`, `[slots,512]`, `[pages,page_size,512]`, or `[slots,1,1,512]` | +| Slot IDs | Contiguous int32 `[Q,K]` or `[Q,1,K]`, capacity K from 0 through 2176 | +| Scales | Finite positive Python scalars: `softmax_scale`, `q_scale`, `kv_scale` | +| O | Contiguous BF16 `[Q,64,512]` | +| Optional LSE | Contiguous FP32 `[Q,64]`, natural logarithm | +| Workspace | Opaque flat BF16 partial output and FP32 partial LSE from the allocator | + +Slot IDs address physical rows in the flattened pool. Negative and out-of-range +IDs are masked; duplicate IDs retain their multiplicity. K is the array capacity, +including padding. Empty selections produce O=0 and LSE=-infinity. KV addresses +use 64-bit arithmetic, including for pools larger than 4 GiB. + +BF16 Q is encoded as `FP8(Q.float() / q_scale)`. FP8 Q is already encoded and is +not quantized again. Decoding scales apply to Q, K and V; `kv_scale` is used in +both the score and value paths. The E4M3FN conversion uses round-to-nearest-even, +signed zero and gradual underflow; overflow produces NaN, matching the PyTorch +cast. Choose a scale that keeps model activations representable. P is encoded +as FP8 with decoding scale 1/256. There is one FP8 Q component and one QK +computation. Accumulators, softmax and LSE are FP32; split partial O and final O +are BF16. + +Numerical tests compare against independent FP64 attention on the quantized Q, +with per-head relative RMS <=0.07 and elementwise atol/rtol=0.07. This measures +implementation error for FP8 attention; it does not imply equivalence to BF16 +attention or establish model accuracy. + +The operator uses the current HIP stream. Validation reads tensor metadata, +not GPU values. Graph capture fixes the shape-dependent dispatch and replay +reads updated tensor contents. Outputs and workspace must not overlap any +input or each other. Unsupported devices and metadata raise errors. + +## Dispatch and implementation + +| Q | Heads per CTA | K tile | Splits | Main loading path | +|---|---:|---:|---:|---| +| 1–6 | 16 | 128 | 17 | Small-query head tiles | +| 7–12 | 32 | 128 | 9 | Two KV tiles prefetched around Q conversion | +| 13–21 | 32 | 128 | 6 | Rotating KV buffers | +| 22–23 | 32 | 128 | 4 | Rotating KV buffers | +| 24–42 | 64 | 128 | 6 | Rotating KV buffers | +| 43–64 | 64 | 128 | 4 | Rotating KV buffers | +| 65–128 | 64 | 128 | 2 | Streaming double buffer | +| 129–2048 | 64 | 128 | 1 | Streaming double buffer, direct output | + +Q8 uses head-major partial LSE and a packed FP32 reducer. Other split paths use +their corresponding split-major layout. Workspace is opaque so callers do not +need to reproduce these layouts. No workspace initialization kernel is required. +The 2112-entry selection capacity has a specialization in the prefetch path; +other supported capacities use the general path. + +## Reproduction + +The HIP implementation uses AITER's ctypes JIT and HIP intrinsics/assembly; it +does not require HipKittens. Build with `GPU_ARCHS=gfx950`. Warm the operator +before Graph capture. Run correctness and replay checks with: + +```bash +GPU_ARCHS=gfx950 python -m pytest -q op_tests/test_sparse_mla_fp8.py +``` + +The benchmark accepts BF16 Q for both HIP and Gluon, includes conversion inside +attention, and ends with BF16 O ready. It uses a random 5,608,046,592-byte FP8 KV +pool, 2051 valid slots per query, capacity 2112, and `-1` padding. The default +queries are every power of two from 1 through 2048; `--all-small` additionally +covers Q1–64 individually and Q65. + +```bash +GPU_ARCHS=gfx950 python op_tests/op_benchmarks/bench_sparse_mla_fp8.py \ + --all-small --output /tmp/sparse-mla-fp8-results +``` + +Each of two fresh Q seeds is measured in forward and reverse order. Each phase +uses 30 timing samples, with 20 calls captured per HIP Graph. The reported time +is the median of the four phase medians. Compilation, host allocation and +Python execution are excluded; captured GPU work is included. Raw samples, +launch tile size, configuration and per-phase numerical checks are saved. + +HIP uses fixed Q scale one. Gluon uses dynamic tile quantization, checked against +its source-order quantized-Q reference. That reference does not assert bitwise +identity between Torch division and the hardware reciprocal instruction. + +`speedup = Gluon_us / HIP_us - 1`. Logical KV bandwidth is +`Q * valid_K * 512 / HIP_us / 1e6` TB/s: one selected FP8 KV payload per query. +It excludes repeated head reads, Q/O and intermediate/index traffic, so it is +not measured HBM traffic. Operator timings do not establish serving throughput. diff --git a/op_tests/check_sparse_mla_bf16_isa.py b/op_tests/check_sparse_mla_bf16_isa.py deleted file mode 100644 index 2844873f20a..00000000000 --- a/op_tests/check_sparse_mla_bf16_isa.py +++ /dev/null @@ -1,112 +0,0 @@ -# SPDX-License-Identifier: MIT -"""Gate retained gfx950 compiler assembly before running pinned-VGPR kernels. - -Compile sparse_mla_bf16.cu with --save-temps, then pass its device .s here. -Checks compiler-managed VGPR exclusion, FP32 denormals, spills and private memory. -This is a conservative build check, not a proof of synchronization or liveness. -""" - -import argparse -import hashlib -import json -import re -from pathlib import Path - - -def audit(path): - text = path.read_text() - metadata = {} - for block in re.split(r"(?m)^ - (?=\.)", text.split("\t.amdgpu_metadata", 1)[-1])[ - 1: - ]: - fields = dict(re.findall(r"(?m)^\s*\.(\w+):\s*([^\n]+)", block)) - if "name" in fields: - metadata[fields["name"].strip()] = fields - errors, kernels = [], {} - current, inline = None, False - for number, line in enumerate(text.splitlines(), 1): - function = re.match(r"\s*\.type\s+(\S+),@function", line) - if function: - if inline: - errors.append(f"unclosed inline asm before {number}") - current, inline = function[1], False - if "sparse_mla_bf16" in current and current in metadata: - kernels[current] = {"compiler_vgpr_max": -1, "denorm32": None} - if re.search(r"#(?:ASMSTART|APP)\b", line): - inline = True - continue - if re.search(r"#(?:ASMEND|NO_APP)\b", line): - inline = False - continue - if current not in kernels: - continue - info = kernels[current] - mode = re.search(r"\.amdhsa_float_denorm_mode_32\s+(\d+)", line) - if mode: - info["denorm32"] = int(mode[1]) - instruction = re.match(r"^\s+([a-z][a-z0-9_]*)\s*(.*)$", line) - if not instruction: - continue - if instruction[1].startswith("scratch_"): - errors.append(f"{current}:{number}: scratch instruction") - if not inline: - registers = [ - int(b or a or c, 0) - for a, b, c in re.findall( - r"\bv(?:\[(0x[\da-fA-F]+|\d+)(?::(0x[\da-fA-F]+|\d+))?\]|(\d+))", - line, - ) - ] - info["compiler_vgpr_max"] = max(info["compiler_vgpr_max"], *registers, -1) - main = [name for name in kernels if "sparse_mla_v" in name] - tiles = [name for name in kernels if "sparse_mla_headtiles" in name] - if len(main) != 8 or len(tiles) != 2: - errors.append( - f"expected 8 v1/v2 and 2 head-tiled mains, got {len(main)} and {len(tiles)}" - ) - if len(kernels) != 21: - errors.append(f"expected 21 kernels, got {len(kernels)}") - for name, info in kernels.items(): - if name in main + tiles and info["compiler_vgpr_max"] >= 40: - errors.append(f"{name}: compiler crosses pinned v40 boundary") - if info["denorm32"] != 3: - errors.append(f"{name}: FP32 denormals are not preserved") - for field in ( - "vgpr_count", - "agpr_count", - "sgpr_count", - "vgpr_spill_count", - "sgpr_spill_count", - "private_segment_fixed_size", - "group_segment_fixed_size", - ): - value = metadata[name].get(field, "") - info[field] = int(value) if value.isdecimal() else None - if info[field] is None: - errors.append(f"{name}: missing {field}") - for field in ( - "vgpr_spill_count", - "sgpr_spill_count", - "private_segment_fixed_size", - ): - if info[field] != 0: - errors.append(f"{name}: {field} must be zero") - return { - "passed": not errors, - "errors": errors, - "kernels": kernels, - "assembly_sha256": hashlib.sha256(path.read_bytes()).hexdigest(), - } - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("assembly", type=Path) - parser.add_argument("--output", type=Path) - args = parser.parse_args() - result = audit(args.assembly) - report = json.dumps(result, indent=2) + "\n" - if args.output: - args.output.write_text(report) - print(report) - raise SystemExit(not result["passed"]) diff --git a/op_tests/check_sparse_mla_fp8_isa.py b/op_tests/check_sparse_mla_fp8_isa.py new file mode 100644 index 00000000000..2ebba889e97 --- /dev/null +++ b/op_tests/check_sparse_mla_fp8_isa.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: MIT +"""Audit retained gfx950 device assembly for the FP8 sparse MLA operator. + +Pass the device .s files from an AITER build with --save-temps. This checks +resource limits and arithmetic instructions; runtime correctness/replay tests +are still required. It does not prove all assembly dependency hazards. +""" + +import argparse +import hashlib +import json +from pathlib import Path +import re + + +def audit(path): + text = path.read_text() + metadata = {} + for block in re.split(r"(?m)^ - (?=\.)", text.split("\t.amdgpu_metadata", 1)[-1])[ + 1: + ]: + fields = dict(re.findall(r"(?m)^\s*\.(\w+):\s*([^\n]+)", block)) + if "name" in fields: + metadata[fields["name"].strip()] = fields + errors, kernels = [], {} + starts = list(re.finditer(r"(?m)^\s*\.type\s+(\S+),@function", text)) + for i, match in enumerate(starts): + name = match[1] + if name not in metadata or "sparse_mla_fp8" not in name: + continue + body = text[ + match.start() : starts[i + 1].start() if i + 1 < len(starts) else len(text) + ] + fields = metadata[name] + info = {} + for field in [ + "vgpr_count", + "agpr_count", + "sgpr_count", + "vgpr_spill_count", + "sgpr_spill_count", + "private_segment_fixed_size", + "group_segment_fixed_size", + ]: + value = fields.get(field, "") + info[field] = int(value) if value.isdecimal() else None + if info[field] is None: + errors.append(f"{name}: missing {field}") + for field in [ + "vgpr_spill_count", + "sgpr_spill_count", + "private_segment_fixed_size", + ]: + if info[field] != 0: + errors.append(f"{name}: {field} must be zero") + if ( + info["group_segment_fixed_size"] is not None + and info["group_segment_fixed_size"] > 160 * 1024 + ): + errors.append(f"{name}: exceeds gfx950 LDS capacity") + modes = re.findall(r"\.amdhsa_float_denorm_mode_32\s+(\d+)", body) + # Function descriptors can be emitted separately from the function body; + # every descriptor in the input must preserve FP32 denormals. + info["mfma"] = sorted(set(re.findall(r"\bv_mfma\w+", body))) + if any(op != "v_mfma_scale_f32_16x16x128_f8f6f4" for op in info["mfma"]): + errors.append(f"{name}: unexpected matrix instruction") + if "attention" in name and not info["mfma"]: + errors.append(f"{name}: missing FP8 matrix instructions") + if re.search(r"(?m)^\s+scratch_", body): + errors.append(f"{name}: scratch instructions") + kernels[name] = info + if not kernels: + errors.append("no sparse MLA kernels found") + modes = re.findall(r"\.amdhsa_float_denorm_mode_32\s+(\d+)", text) + if not modes or any(mode != "3" for mode in modes): + errors.append("FP32 denormals must be preserved") + return dict( + passed=not errors, + errors=errors, + kernels=kernels, + assembly_sha256=hashlib.sha256(path.read_bytes()).hexdigest(), + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("assembly", type=Path, nargs="+") + parser.add_argument("--output", type=Path) + args = parser.parse_args() + result = {str(p): audit(p) for p in args.assembly} + report = json.dumps(result, indent=2) + "\n" + if args.output: + args.output.write_text(report) + print(report) + raise SystemExit(not all(v["passed"] for v in result.values())) diff --git a/op_tests/op_benchmarks/bench_sparse_mla_bf16.py b/op_tests/op_benchmarks/bench_sparse_mla_bf16.py deleted file mode 100644 index 8e395ab5771..00000000000 --- a/op_tests/op_benchmarks/bench_sparse_mla_bf16.py +++ /dev/null @@ -1,406 +0,0 @@ -# SPDX-License-Identifier: MIT -"""Compare HIP v1/v2 and same-checkout Gluon on disjoint per-query KV pools. - -python op_tests/op_benchmarks/bench_sparse_mla_bf16.py --verify --output result.jsonl -Events inside captured graphs time main + reducer, excluding Python/JIT overhead. -The cold run flushes 512 MiB before the start event. Logical GB/s is not measured -HBM traffic. Sampled FP64 accuracy is checked before and after every timed case. -""" - -import argparse -import gc -import hashlib -import json -import math -from pathlib import Path -import statistics -import sys - -import torch -import triton - -import aiter -from aiter.jit.core import get_user_jit_dir -from aiter.ops.triton.attention import sparse_mla as gluon -from op_tests.test_sparse_mla_bf16 import check, reference - - -def ints(value): - return [int(x) for x in value.split(",")] - - -def tensor_hash(tensor): - data = tensor.detach().contiguous().view(torch.uint8).flatten() - digest = hashlib.sha256() - for start in range(0, data.numel(), 8 << 20): - digest.update(data[start : start + (8 << 20)].cpu().numpy().tobytes()) - return digest.hexdigest() - - -def file_hash(path): - return hashlib.sha256(path.read_bytes()).hexdigest() - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--heads", type=ints, default=[16, 64]) - parser.add_argument( - "--queries", type=ints, default=[1, 32, 128, 256, 512, 1024, 2048] - ) - parser.add_argument("--topk", type=ints, default=[2048, 2051]) - parser.add_argument("--splits", type=ints, default=[1, 2, 4, 8, 16, 32]) - parser.add_argument("--context", type=int, default=32768) - parser.add_argument("--samples", type=int, default=60) - parser.add_argument("--warmup", type=int, default=10) - parser.add_argument("--check-rows", type=int, default=4) - parser.add_argument("--cache-modes", default="warm,cold") - parser.add_argument("--return-lse", action="store_true") - parser.add_argument("--purpose", choices=["timing", "trace"], default="timing") - parser.add_argument("--reverse", action="store_true") - parser.add_argument("--verify", action="store_true", required=True) - parser.add_argument("--output", type=Path, required=True) - args = parser.parse_args() - if args.output.exists(): - raise ValueError("output already exists; use a fresh run") - assert args.samples >= 1 and args.check_rows >= 1 - assert args.purpose == "trace" or (args.samples >= 30 and args.warmup >= 10) - assert all(q > 0 for q in args.queries) - assert set(args.heads) <= {16, 64} and set(args.splits) <= {1, 2, 4, 8, 16, 32} - assert set(args.cache_modes.split(",")) <= {"warm", "cold"} - assert all(0 < k <= args.context for k in args.topk) - assert torch.version.hip and torch.cuda.device_count() == 1 - torch.cuda.set_device(0) - torch.set_num_threads(4) - props = torch.cuda.get_device_properties(0) - assert props.gcnArchName.split(":")[0] == "gfx950" - - def emit(row): - with args.output.open("a") as output: - output.write(json.dumps(row, allow_nan=False) + "\n") - keys = ( - "type", - "heads", - "queries", - "topk", - "variant", - "splits", - "cache", - "median_us", - "status", - ) - print(json.dumps({k: row[k] for k in keys if k in row}), flush=True) - - source = Path(__file__).resolve() - jit = Path(get_user_jit_dir()) / "module_sparse_mla_bf16.so" - emit( - { - "type": "config", - "command": sys.argv, - "device": str(props), - "torch": str(torch.__version__), - "triton": triton.__version__, - "hip": torch.version.hip, - "benchmark_sha256": file_hash(source), - "aiter": aiter.__file__, - "gluon_source_sha256": file_hash(Path(gluon.__file__)), - "native_library_sha256": file_hash(jit) if jit.exists() else None, - "scope": "captured attention main+reduce; not model/engine latency", - "pool_layout": "disjoint", - "output_contract": "O+LSE" if args.return_lse else "O-only", - "args": vars(args) | {"output": str(args.output)}, - } - ) - stream = torch.cuda.Stream() - flush = torch.empty(512 << 20, device="cuda", dtype=torch.uint8) - cases = [(h, q, k) for h in args.heads for q in args.queries for k in args.topk] - if args.reverse: - cases.reverse() - total = failed = 0 - for heads, queries, topk in cases: - with torch.cuda.stream(stream): - torch.manual_seed(1701 + heads + queries + topk) - q = ( - torch.randn((queries, heads, 512), device="cuda", dtype=torch.bfloat16) - * 0.5 - ) - kv = torch.empty( - (queries * args.context, 512), device="cuda", dtype=torch.bfloat16 - ) - kv.normal_().mul_(0.5) - idx = torch.cat( - [ - torch.randperm(args.context, device="cuda")[:topk].to(torch.int32) - + row * args.context - for row in range(queries) - ] - ) - ptr = torch.arange(queries + 1, device="cuda", dtype=torch.int32) * topk - out = torch.empty_like(q) - lse = torch.empty(q.shape[:2], device="cuda") if args.return_lse else None - max_s = max(args.splits) - po = torch.empty((queries * max_s * heads * 512,), device="cuda") - pl = torch.empty((queries * max_s * heads,), device="cuda") - rows = sorted( - { - round(i * (queries - 1) / max(1, min(queries, args.check_rows) - 1)) - for i in range(min(queries, args.check_rows)) - } - ) - selected = torch.tensor(rows, device="cuda", dtype=torch.long) - oracle_ptr = ( - torch.arange(len(rows) + 1, device="cuda", dtype=torch.int32) * topk - ) - expected = reference( - q[selected], - kv, - oracle_ptr, - idx.view(queries, topk)[selected].flatten(), - 1 / 16, - ) - tensors = {"q": q, "kv": kv, "indptr": ptr, "indices": idx} - hashes = {name: tensor_hash(value) for name, value in tensors.items()} - case = { - "heads": heads, - "queries": queries, - "topk": topk, - "pool_bytes": kv.numel() * 2, - "input_sha256": hashes, - "input_addresses": { - name: value.data_ptr() for name, value in tensors.items() - }, - "out_address": out.data_ptr(), - "checked_rows": rows, - "output_contract": "O+LSE" if args.return_lse else "O-only", - } - variants = [(f"v{v}", s) for s in args.splits for v in (1, 2)] - variants += [("gluon", s) for s in args.splits] + [("gluon_default", None)] - if args.reverse: - variants.reverse() - runners = [] - for variant, splits in variants: - state = [None] - if variant.startswith("v"): - workspace = ( - ( - po[: queries * splits * heads * 512].view( - queries, splits, heads, 512 - ), - pl[: queries * splits * heads].view(queries, splits, heads), - ) - if splits > 1 - else None - ) - - def invoke( - v=int(variant[1]), - s=splits, - workspace=workspace, - state=state, - inputs=(q, kv, ptr, idx, out, lse), - ): - q, kv, ptr, idx, out, lse = inputs - state[0] = aiter.sparse_mla_bf16_fwd( - q, - kv, - ptr, - idx, - 1 / 16, - version=v, - kv_splits=s, - out=out, - return_lse=args.return_lse, - lse=lse, - workspace=workspace, - ) - - actual_splits = splits - else: - - def invoke(s=splits, state=state, inputs=(q, kv, ptr, idx, out)): - q, kv, ptr, idx, out = inputs - state[0] = gluon.sparse_mla_fwd( - q, - kv, - ptr, - idx, - 1 / 16, - kv_lora_rank=512, - qk_rope_head_dim=0, - kv_splits=s, - has_invalid=False, - out=out, - return_lse=args.return_lse, - ) - - actual_splits = splits or gluon._mla_num_splits( - queries, heads // 16, topk - ) - key = {"variant": variant, "splits": actual_splits} - invoke() - torch.cuda.synchronize() - try: - before = check( - ( - state[0][0][selected], - state[0][1][selected] if state[0][1] is not None else None, - ), - expected, - ) - except AssertionError as error: - failed += 1 - emit( - { - "type": "accuracy_failure", - **case, - **key, - "status": "excluded", - "error": str(error), - } - ) - if variant.startswith("v"): - raise - continue - start = torch.cuda.Event(enable_timing=True, external=True) - end = torch.cuda.Event(enable_timing=True, external=True) - start.record() - end.record() - stream.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - start.record() - invoke() - end.record() - graph.replay() - stream.synchronize() - check( - ( - state[0][0][selected], - state[0][1][selected] if state[0][1] is not None else None, - ), - expected, - ) - runners.append( - { - **key, - "graph": graph, - "start": start, - "end": end, - "state": state, - "before_rms": before, - } - ) - for mode in args.cache_modes.split(","): - for i in range(args.warmup): - for r in (runners if i % 2 == 0 else runners[::-1]): - if mode == "cold": - flush.fill_(1) - r["graph"].replay() - stream.synchronize() - measured = [[] for _ in runners] - for i in range(args.samples): - order = ( - range(len(runners)) - if i % 2 == 0 - else reversed(range(len(runners))) - ) - for j in order: - r = runners[j] - if mode == "cold": - flush.fill_(1) - r["graph"].replay() - r["end"].synchronize() - value = r["start"].elapsed_time(r["end"]) * 1000 - assert math.isfinite(value) and value > 0 - measured[j].append(value) - for r, samples in zip(runners, measured): - r["graph"].replay() - stream.synchronize() - output, output_lse = r["state"][0] - after = check( - ( - output[selected], - output_lse[selected] if output_lse is not None else None, - ), - expected, - ) - median = statistics.median(samples) - logical_bytes = ( - queries * (heads * 512 * 4 + topk * (512 * 2 + 4)) - + (queries + 1) * 4 - ) - if args.return_lse: - logical_bytes += queries * heads * 4 - emit( - { - "type": ( - "timing" - if args.purpose == "timing" - else "diagnostic_timing" - ), - **case, - "variant": r["variant"], - "splits": r["splits"], - "cache": mode, - "samples_us": samples, - "median_us": median, - "p90_us": sorted(samples)[ - math.ceil(len(samples) * 0.9) - 1 - ], - "logical_bytes": logical_bytes, - "logical_gbps": logical_bytes / median / 1000, - "qk_pv_tflops": 4 - * queries - * heads - * topk - * 512 - / median - / 1e6, - "before_rms": r["before_rms"], - "after_rms": after, - "status": "passed", - } - ) - total += 1 - assert hashes == { - name: tensor_hash(value) for name, value in tensors.items() - } - emit( - { - "type": "case_complete", - **case, - "status": "passed", - "inputs_unchanged": True, - } - ) - runners.clear() - del ( - q, - kv, - ptr, - idx, - out, - lse, - po, - pl, - tensors, - invoke, - state, - graph, - r, - output, - output_lse, - ) - gc.collect() - torch.cuda.empty_cache() - emit( - { - "type": "complete", - "status": "passed", - "timing_records": total, - "accuracy_exclusions": failed, - } - ) - - -if __name__ == "__main__": - main() diff --git a/op_tests/op_benchmarks/bench_sparse_mla_fp8.py b/op_tests/op_benchmarks/bench_sparse_mla_fp8.py new file mode 100644 index 00000000000..3a9a160918c --- /dev/null +++ b/op_tests/op_benchmarks/bench_sparse_mla_fp8.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: MIT +"""BF16 Q -> BF16 O, FP8 KV: HIP and Gluon on identical random slot IDs.""" + +import argparse +import json +import math +from pathlib import Path +import statistics +import sys + +import torch + +from aiter import sparse_mla_fp8_fwd +from aiter.ops.triton.attention.sparse_mla import sparse_mla_fwd as gluon + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from test_sparse_mla_fp8 import assert_reference + + +class RecordLaunch: + def __init__(self, kernel): + self.kernel = kernel + self.block_m = None + + def __getitem__(self, grid): + def launch(*args, **kwargs): + self.block_m = kwargs["BLOCK_M"] + return self.kernel[grid](*args, **kwargs) + + return launch + + +def graph_for(fn, repetitions): + fn() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(repetitions): + out = fn() + for _ in range(4): + graph.replay() + torch.cuda.synchronize() + return graph, out + + +def measure(calls, repetitions, count): + graphs = {name: graph_for(fn, repetitions) for name, fn in calls.items()} + samples = {name: [] for name in calls} + for repeat in range(count): + for name in list(graphs)[:: 1 if repeat % 2 == 0 else -1]: + begin = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + begin.record() + graphs[name][0].replay() + end.record() + end.synchronize() + samples[name].append(begin.elapsed_time(end) * 1000 / repetitions) + return samples + + +def check_gluon(out, q, kv, ids, block_m): + # Source-order model of Gluon's dynamic quantization. This does not assert + # that Torch division and GPU reciprocal instructions are bitwise equal. + blocks = q.float().reshape(q.shape[0], 64 // block_m, block_m, 512) + amax = blocks.abs().amax(dim=(-2, -1), keepdim=True).clamp_min(1e-30) + encoded = (blocks * (448.0 / amax)).to(torch.float8_e4m3fn).reshape(q.shape) + scales = ( + (amax / 448.0) + .expand(q.shape[0], 64 // block_m, block_m, 1) + .reshape(q.shape[0], 64, 1) + ) + return assert_reference(out, None, encoded, kv, ids, q_scale=scales) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--queries", type=int, nargs="+", default=[2**i for i in range(12)] + ) + parser.add_argument( + "--all-small", action="store_true", help="Q1-64 plus 65/128/.../2048" + ) + parser.add_argument("--pool-bytes", type=int, default=5_608_046_592) + parser.add_argument("--valid-k", type=int, default=2051) + parser.add_argument("--capacity", type=int, default=2112) + parser.add_argument("--samples", type=int, default=30) + parser.add_argument("--graph-repetitions", type=int, default=20) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.all_small: + args.queries = list(range(1, 65)) + [65, 128, 256, 512, 1024, 2048] + assert 0 < args.valid_k <= args.capacity <= 2176 + assert args.pool_bytes > 0 and args.pool_bytes % 512 == 0 + assert args.samples > 0 and args.graph_repetitions > 0 + assert all(1 <= n <= 2048 for n in args.queries) + args.output.mkdir(parents=True, exist_ok=False) + (args.output / "config.json").write_text( + json.dumps(vars(args), default=str, indent=2) + ) + slots = args.pool_bytes // 512 + kv = torch.empty((slots, 512), device="cuda", dtype=torch.float8_e4m3fn) + # Bound temporary storage while initializing the complete random pool. + torch.manual_seed(2041103) + for start in range(0, slots, 65536): + kv[start : start + 65536].copy_( + torch.randn( + kv[start : start + 65536].shape, device="cuda", dtype=torch.bfloat16 + ).to(kv.dtype) + ) + one = torch.ones((1,), device="cuda", dtype=torch.float32) + medians = {n: {"HIP": [], "Gluon": []} for n in args.queries} + with (args.output / "phases.jsonl").open("w") as log: + for phase in range(4): + order = args.queries[:: 1 if phase % 2 == 0 else -1] + for n in order: + torch.manual_seed(2121037 + (phase // 2) * 10007 + n) + q = torch.randn((n, 64, 512), device="cuda", dtype=torch.bfloat16) + torch.manual_seed(2041129 + n) + ids = torch.randint( + slots, (n, args.capacity), device="cuda", dtype=torch.int32 + ) + ids[:, args.valid_k :] = -1 + indptr = ( + torch.arange(n + 1, device="cuda", dtype=torch.int32) + * args.capacity + ) + + def run_gluon(): + return gluon( + q, + kv, + indptr, + ids.reshape(-1), + 1 / 16, + kv_scale=one, + kv_lora_rank=512, + qk_rope_head_dim=0, + has_invalid=True, + dot_precision="fp8", + )[0] + + calls = { + "HIP": lambda: sparse_mla_fp8_fwd(q, kv, ids, 1 / 16)[0], + "Gluon": run_gluon, + } + hip_out = calls["HIP"]() + namespace = gluon.__globals__ + original = namespace["_sparse_mla_gfx950"] + tap = RecordLaunch(original) + namespace["_sparse_mla_gfx950"] = tap + try: + gluon_out = run_gluon() + finally: + namespace["_sparse_mla_gfx950"] = original + errors = { + "HIP": assert_reference(hip_out, None, q, kv, ids), + "Gluon": check_gluon(gluon_out, q, kv, ids, tap.block_m), + } + if phase % 2: + calls = dict(reversed(list(calls.items()))) + samples = measure(calls, args.graph_repetitions, args.samples) + row = dict( + phase=phase, + q=n, + rrms=errors, + gluon_block_m=tap.block_m, + samples_us=samples, + medians_us={k: statistics.median(v) for k, v in samples.items()}, + ) + log.write(json.dumps(row) + "\n") + log.flush() + for k, v in row["medians_us"].items(): + medians[n][k].append(v) + print( + json.dumps({k: v for k, v in row.items() if k != "samples_us"}), + flush=True, + ) + rows = [] + for n in args.queries: + hip, g = (statistics.median(medians[n][k]) for k in ["HIP", "Gluon"]) + rows.append( + dict( + q=n, + hip_us=hip, + gluon_us=g, + speedup_pct=100 * (g / hip - 1), + logical_kv_tbps=n * args.valid_k * 512 / hip / 1e6, + ) + ) + result = dict( + rows=rows, + minimum_speedup_pct=min(r["speedup_pct"] for r in rows), + geometric_mean_speedup_pct=100 + * ( + math.exp( + statistics.mean(math.log(r["gluon_us"] / r["hip_us"]) for r in rows) + ) + - 1 + ), + numerical_checks_passed=True, + endpoint="BF16 Q to BF16 O", + kv_dtype="E4M3FN", + logical_bandwidth_is_hbm_traffic=False, + ) + (args.output / "result.json").write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result), flush=True) + + +if __name__ == "__main__": + main() diff --git a/op_tests/test_sparse_mla_bf16.py b/op_tests/test_sparse_mla_bf16.py deleted file mode 100644 index 405515d4067..00000000000 --- a/op_tests/test_sparse_mla_bf16.py +++ /dev/null @@ -1,260 +0,0 @@ -# SPDX-License-Identifier: MIT -"""BF16 D512 sparse MLA: independent FP64, ragged CSR, strides and graph replay. - -Run: python -m pytest -q op_tests/test_sparse_mla_bf16.py -""" - -import pytest -import torch - -from aiter import sparse_mla_bf16_fwd -from aiter.jit.utils.chip_info import get_gfx_runtime - -pytestmark = pytest.mark.skipif(get_gfx_runtime() != "gfx950", reason="requires gfx950") -SPLITS = [1, 2, 4, 8, 16, 32] - - -def reference(q, kv, indptr, indices, scale): - """Reference the mathematical operation, independently of tile/split layout.""" - output = torch.zeros_like(q, dtype=torch.float64) - lse = torch.full(q.shape[:2], -torch.inf, device=q.device, dtype=torch.float64) - offsets = indptr.cpu().tolist() - for row, (start, end) in enumerate(zip(offsets, offsets[1:])): - selected = indices[start:end].long() - selected = selected[(selected >= 0) & (selected < kv.shape[0])] - if not selected.numel(): - continue - values = kv[selected].double() - scores = q[row].double() @ values.T * scale - output[row] = scores.softmax(-1) @ values - lse[row] = scores.logsumexp(-1) - return output, lse - - -def check(actual, expected): - out, lse = actual - ref_o, ref_lse = expected - torch.testing.assert_close(out.double(), ref_o, atol=1e-2, rtol=1e-2) - rms = (out.double() - ref_o).square().mean(-1).sqrt() - relative = rms / ref_o.square().mean(-1).sqrt().clamp_min(1e-3) - assert (relative <= 1e-2).all(), relative.max().item() - if lse is not None: - torch.testing.assert_close(lse.double(), ref_lse, atol=1e-4, rtol=1e-5) - return relative.max().item() if relative.numel() else 0.0 - - -def fixture(heads, topk, padded=False): - torch.manual_seed(321 + heads + topk) - queries, slots = 4, max(65, topk + 7) - if padded: - # Offset by one BF16, with padded query/head/cache rows. - q_store = torch.randn(queries, heads, 520, device="cuda", dtype=torch.bfloat16) - kv_store = torch.randn(slots, 520, device="cuda", dtype=torch.bfloat16) - q, kv = q_store[..., 1:513], kv_store[:, 1:513] - else: - q = torch.randn(queries, heads, 512, device="cuda", dtype=torch.bfloat16) - kv = torch.randn(slots, 512, device="cuda", dtype=torch.bfloat16) - lengths = [topk, max(0, topk - 3), 0, int(topk > 0)] - ptr = torch.tensor( - [0, *torch.tensor(lengths).cumsum(0).tolist()], device="cuda", dtype=torch.int32 - ) - idx = torch.randint(slots, (sum(lengths),), device="cuda", dtype=torch.int32) - if topk > 1: - idx[:topk:11] = -1 - idx[1:topk:17] = slots + 1 - # Duplicate rows intentionally retain multiplicity. - if topk > 4: - idx[2:4] = 3 - return q, kv, ptr, idx - - -@pytest.mark.parametrize("version", [1, 2]) -@pytest.mark.parametrize("heads", [16, 64]) -@pytest.mark.parametrize("splits", SPLITS) -@pytest.mark.parametrize("topk", [0, 1, 31, 32, 33, 63, 65, 2048, 2051]) -@pytest.mark.parametrize("padded", [False, True]) -def test_fp64(version, heads, splits, topk, padded): - q, kv, ptr, idx = fixture(heads, topk, padded) - q_count = q.shape[0] - out_storage = torch.full((q.numel() + 2,), 123, device=q.device, dtype=q.dtype) - out = out_storage[1:-1].view(q.shape) - lse = torch.empty(q.shape[:2], device=q.device, dtype=torch.float32) - workspace = None - if splits > 1: - count = q_count * splits * heads - # Four-byte-aligned partial O exercises the scalar fallback. - partial = torch.empty(count * 512 + 1, device=q.device, dtype=torch.float32)[1:] - workspace = ( - partial.view(q_count, splits, heads, 512), - torch.empty((q_count, splits, heads), device=q.device), - ) - result = sparse_mla_bf16_fwd( - q, - kv, - ptr, - idx, - 1 / 16, - version=version, - kv_splits=splits, - out=out, - lse=lse, - return_lse=True, - workspace=workspace, - ) - assert result[0].data_ptr() == out.data_ptr() - assert result[1].data_ptr() == lse.data_ptr() - check(result, reference(q, kv, ptr, idx, 1 / 16)) - assert out_storage[0] == 123 and out_storage[-1] == 123 - - -@pytest.mark.parametrize("version", [1, 2]) -@pytest.mark.parametrize("heads", [16, 64]) -@pytest.mark.parametrize("splits", SPLITS) -@pytest.mark.parametrize("return_lse", [False, True]) -def test_graph_current_stream(version, heads, splits, return_lse): - q, kv, ptr, idx = fixture(heads, 65) - out = torch.empty_like(q) - lse = torch.empty(q.shape[:2], device=q.device) if return_lse else None - workspace = ( - ( - torch.empty((4, splits, heads, 512), device=q.device), - torch.empty((4, splits, heads), device=q.device), - ) - if splits > 1 - else None - ) - - def run(): - return sparse_mla_bf16_fwd( - q, - kv, - ptr, - idx, - 1 / 16, - version=version, - kv_splits=splits, - out=out, - lse=lse, - return_lse=return_lse, - workspace=workspace, - ) - - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - for _ in range(3): - run() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - run() - # Replay must consume new device values, not captured query/index data. - q.mul_(0.5) - idx.copy_(idx.roll(1)) - for _ in range(30): - graph.replay() - expected = reference(q, kv, ptr, idx, 1 / 16) - torch.cuda.current_stream().wait_stream(stream) - check((out, lse), expected) - - -@pytest.mark.parametrize("version", [1, 2]) -@pytest.mark.parametrize("heads", [16, 64]) -@pytest.mark.parametrize("splits", [1, 32]) -def test_subnormal_single_key(version, heads, splits): - q = torch.zeros((1, heads, 512), device="cuda", dtype=torch.bfloat16) - bits = torch.tensor( - [1, 2, 127, -32767, -32766, -32641, 0, 0], device="cuda", dtype=torch.int16 - ) - kv = bits.repeat(64).view(torch.bfloat16).view(1, 512) - ptr = torch.tensor([0, 1], device="cuda", dtype=torch.int32) - idx = torch.zeros(1, device="cuda", dtype=torch.int32) - out, lse = sparse_mla_bf16_fwd( - q, kv, ptr, idx, 1 / 16, version=version, kv_splits=splits, return_lse=True - ) - assert torch.equal( - out.view(torch.int16), kv.expand(heads, 512)[None].view(torch.int16) - ) - assert (lse == 0).all() - - -@pytest.mark.parametrize("splits", SPLITS) -@pytest.mark.parametrize("above_boundary", [False, True]) -def test_h64_hybrid_dispatch_boundary(splits, above_boundary): - # Q*S=128 uses head tiles; the next query switches to the four-wave main. - # The small ragged fixture alone stays on the head-tiled branch for all S. - queries = 128 // splits + int(above_boundary) - torch.manual_seed(812) - q = torch.randn((queries, 64, 512), dtype=torch.bfloat16, device="cuda") - kv = torch.randn((97, 512), dtype=torch.bfloat16, device="cuda") - ptr = torch.arange(queries + 1, dtype=torch.int32, device="cuda") * 65 - idx = torch.randint(97, (queries, 65), dtype=torch.int32, device="cuda") - idx[::2, 16:32] = -1 - idx[-1] = -1 - result = sparse_mla_bf16_fwd( - q, - kv, - ptr, - idx.flatten(), - 1 / 16, - version=2, - kv_splits=splits, - return_lse=True, - ) - check(result, reference(q, kv, ptr, idx.flatten(), 1 / 16)) - - -@pytest.mark.parametrize("version", [1, 2]) -@pytest.mark.parametrize("pool_shape", ["flat", "paged", "mla"]) -def test_pool_views_empty_q_and_compile(version, pool_shape): - q, kv, ptr, idx = fixture(16, 33) - # Slice to a multiple of the page size; invalid ids are masked by contract. - kv = kv[:64] - pool = {"flat": kv, "paged": kv.view(8, 8, 512), "mla": kv[:, None, None]}[ - pool_shape - ] - run = torch.compile(sparse_mla_bf16_fwd, backend="eager", fullgraph=True) - actual = run( - q, pool, ptr, idx, 1 / 16, version=version, kv_splits=2, return_lse=True - ) - check(actual, reference(q, kv, ptr, idx, 1 / 16)) - empty = sparse_mla_bf16_fwd( - q[:0], - pool, - ptr[:1], - idx[:0], - 1 / 16, - version=version, - kv_splits=2, - return_lse=True, - ) - assert empty[0].shape == (0, 16, 512) and empty[1].shape == (0, 16) - - -@pytest.mark.parametrize( - "change", - ["dtype", "width", "heads", "splits", "version", "scale", "alias", "workspace"], -) -def test_reject_invalid_metadata(change): - q, kv, ptr, idx = fixture(16, 33) - args = dict(kv_splits=2) - if change == "dtype": - q = q.float() - elif change == "width": - q = q[..., :256] - elif change == "heads": - q = q[:, :8] - elif change == "splits": - args["kv_splits"] = 3 - elif change == "version": - args["version"] = 3 - elif change == "alias": - args["out"] = q - elif change == "workspace": - args["workspace"] = ( - torch.empty(1, device=q.device), - torch.empty(1, device=q.device), - ) - with pytest.raises((RuntimeError, ValueError)): - sparse_mla_bf16_fwd( - q, kv, ptr, idx, float("nan") if change == "scale" else 1 / 16, **args - ) diff --git a/op_tests/test_sparse_mla_fp8.py b/op_tests/test_sparse_mla_fp8.py new file mode 100644 index 00000000000..9c07f470eb5 --- /dev/null +++ b/op_tests/test_sparse_mla_fp8.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: MIT +"""Quantized-Q FP64 checks, metadata validation and changed-input graph replay.""" + +import math + +import pytest +import torch + +from aiter import sparse_mla_fp8_fwd, allocate_sparse_mla_fp8_workspace + +QUERIES = list(range(1, 65)) + [65, 128, 256, 512, 1024, 2048] +CAPACITIES = [0, 1, 63, 64, 127, 128, 129, 2048, 2051, 2111, 2112, 2113, 2176] +TOLERANCE = 0.07 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() + or "gfx950" not in torch.cuda.get_device_properties(0).gcnArchName, + reason="gfx950 required", +) + + +def inputs(n, capacity=2112, *, strided=False, seed=2031071): + torch.manual_seed(seed + n + capacity) + shape = (n, 65, 520) if strided else (n, 64, 512) + q = torch.randn(shape, device="cuda", dtype=torch.bfloat16)[:, :64, :512] + kv = torch.randn((8192, 512), device="cuda", dtype=torch.bfloat16).to( + torch.float8_e4m3fn + ) + ids = torch.randint(len(kv), (n, capacity), device="cuda", dtype=torch.int32) + if capacity: + ids[0].fill_(-1) + ids[:, 2051:] = -1 + if n > 1: + ids[1, : min(9, capacity)] = len(kv) - 1 + ids[1, 10::97] = len(kv) + 17 + return q, kv, ids + + +def assert_reference(out, lse, q, kv, ids, *, q_scale=1.0, kv_scale=1.0, scale=1 / 16): + """Independent FP64 attention after the documented scalar Q encoding.""" + decoded_q = ( + q.double() * q_scale + if q.dtype == torch.float8_e4m3fn + else (q.float() / q_scale).to(torch.float8_e4m3fn).double() * q_scale + ) + assert out.shape == q.shape and out.dtype == torch.bfloat16 + assert torch.isfinite(out).all() + rows = sorted({0, min(1, len(q) - 1), len(q) // 2, len(q) - 1}) if len(q) else [] + worst = 0.0 + for row in rows: + selected = ids[row].long() + selected = selected[(selected >= 0) & (selected < len(kv))] + if selected.numel(): + values = kv.view(torch.uint8)[selected].view(kv.dtype).double() * kv_scale + scores = (decoded_q[row] @ values.T) * scale + expected = scores.softmax(-1) @ values + expected_lse = scores.logsumexp(-1) + else: + expected = torch.zeros((64, 512), device=q.device, dtype=torch.float64) + expected_lse = torch.full( + (64,), -math.inf, device=q.device, dtype=torch.float64 + ) + error = out[row].double() - expected + rrms = error.square().mean(-1).sqrt() / expected.square().mean( + -1 + ).sqrt().clamp_min(1e-3) + worst = max(worst, rrms.max().item()) + assert rrms.max() <= TOLERANCE + torch.testing.assert_close( + out[row].double(), expected, atol=TOLERANCE, rtol=TOLERANCE + ) + if lse is not None: + torch.testing.assert_close( + lse[row].double(), expected_lse, atol=TOLERANCE, rtol=TOLERANCE + ) + if not selected.numel(): + assert torch.count_nonzero(out[row]) == 0 + if lse is not None: + assert torch.isneginf(lse[row]).all() + return worst + + +@pytest.mark.parametrize("n", QUERIES) +def test_quantized_reference(n): + q, kv, ids = inputs(n) + out, lse = sparse_mla_fp8_fwd(q, kv, ids, 1 / 16, return_lse=True) + assert_reference(out, lse, q, kv, ids) + encoded = q.to(torch.float8_e4m3fn) + control, control_lse = sparse_mla_fp8_fwd(encoded, kv, ids, 1 / 16, return_lse=True) + assert torch.equal(out, control) + assert torch.equal(lse, control_lse) + + +@pytest.mark.parametrize("n", [1, 8, 13, 24, 65, 129]) +@pytest.mark.parametrize("capacity", CAPACITIES) +def test_capacity_stride_scale(n, capacity): + q, kv, ids = inputs(n, capacity, strided=True) + kwargs = dict(q_scale=0.5, kv_scale=1.25, return_lse=True) + out, lse = sparse_mla_fp8_fwd( + q, kv.view(-1, 1, 1, 512), ids[:, None, :], 1 / 16, **kwargs + ) + assert_reference(out, lse, q, kv, ids, q_scale=0.5, kv_scale=1.25) + control, control_lse = sparse_mla_fp8_fwd( + (q.float() / 0.5).to(torch.float8_e4m3fn), kv, ids, 1 / 16, **kwargs + ) + assert torch.equal(out, control) + assert torch.equal(lse, control_lse) + + +@pytest.mark.parametrize("n", [1, 6, 7, 8, 12, 13, 22, 24, 43, 65, 129, 2048]) +def test_changed_graph(n): + q, kv, ids = inputs(n) + workspace = allocate_sparse_mla_fp8_workspace(q) + out = torch.empty_like(q) + lse = torch.empty(q.shape[:2], dtype=torch.float32, device=q.device) + kwargs = dict(out=out, lse=lse, return_lse=True, workspace=workspace) + sparse_mla_fp8_fwd(q, kv, ids, 1 / 16, **kwargs) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured, captured_lse = sparse_mla_fp8_fwd(q, kv, ids, 1 / 16, **kwargs) + pointers = [t.data_ptr() for t in [q, kv, ids, out, lse, *workspace]] + for replay in range(4): + q.normal_() + fresh = torch.randn(kv.shape, device=kv.device, dtype=torch.bfloat16).to( + kv.dtype + ) + kv.copy_(fresh) + ids.random_(0, len(kv)) + ids[:, 2051:] = -1 + if replay == 1: + ids[:, :128] = len(kv) - 1 + elif replay == 2: + ids[:, 3::11] = -1 + elif replay == 3: + ids.fill_(-1) + out.fill_(math.nan) + lse.fill_(math.nan) + for t in workspace: + t.fill_(math.nan) + graph.replay() + torch.cuda.synchronize() + expected, expected_lse = sparse_mla_fp8_fwd(q, kv, ids, 1 / 16, return_lse=True) + assert torch.equal(captured, expected) + assert torch.equal(captured_lse, expected_lse) + assert_reference(out, lse, q, kv, ids) + assert pointers == [t.data_ptr() for t in [q, kv, ids, out, lse, *workspace]] + + +@pytest.mark.parametrize("scale", [0.0, -1.0, math.inf, math.nan]) +@pytest.mark.parametrize("argument", ["q_scale", "kv_scale", "softmax_scale"]) +def test_bad_scales(argument, scale): + q, kv, ids = inputs(8, 1) + kwargs = {"softmax_scale": 1 / 16, argument: scale} + with pytest.raises((ValueError, RuntimeError), match="scale"): + sparse_mla_fp8_fwd(q, kv, ids, **kwargs) + + +def test_metadata_and_storage(): + q, kv, ids = inputs(8, 1) + with pytest.raises((ValueError, RuntimeError), match="dtype|FP8|float8"): + sparse_mla_fp8_fwd(q, kv.to(torch.bfloat16), ids, 1 / 16) + with pytest.raises((ValueError, RuntimeError), match="dtype"): + sparse_mla_fp8_fwd(q, kv, ids.long(), 1 / 16) + with pytest.raises((ValueError, RuntimeError), match="overlap"): + sparse_mla_fp8_fwd(q, kv, ids, 1 / 16, out=q) + unaligned = torch.empty((q.numel() + 1,), device=q.device, dtype=q.dtype)[1:].view( + q.shape + ) + with pytest.raises((ValueError, RuntimeError), match="align"): + sparse_mla_fp8_fwd(unaligned, kv, ids, 1 / 16) + with pytest.raises((ValueError, RuntimeError), match="shape"): + sparse_mla_fp8_fwd( + q, kv, ids, 1 / 16, workspace=allocate_sparse_mla_fp8_workspace(q[:1]) + ) + oversized = torch.empty((8, 2177), device=q.device, dtype=torch.int32) + with pytest.raises((ValueError, RuntimeError), match="2176"): + sparse_mla_fp8_fwd(q, kv, oversized, 1 / 16) + empty, empty_lse = sparse_mla_fp8_fwd(q[:0], kv, ids[:0], 1 / 16, return_lse=True) + assert empty.shape == (0, 64, 512) and empty_lse.shape == (0, 64) + + +def test_torch_compile(): + q, kv, ids = inputs(8, 129) + + def run(q, kv, ids): + return sparse_mla_fp8_fwd(q, kv, ids, 1 / 16)[0] + + expected = run(q, kv, ids) + compiled = torch.compile(run, backend="eager", fullgraph=True) + assert torch.equal(compiled(q, kv, ids), expected) From ac78aad7b4a860a2e16beb03001af6235a73466a Mon Sep 17 00:00:00 2001 From: Sumin Hong Date: Fri, 16 Oct 2026 00:17:39 +0000 Subject: [PATCH 3/3] Allow sparse MLA correctness tests to run as a standalone script Signed-off-by: Sumin Hong --- op_tests/test_sparse_mla_fp8.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/op_tests/test_sparse_mla_fp8.py b/op_tests/test_sparse_mla_fp8.py index 9c07f470eb5..05df1e00bd0 100644 --- a/op_tests/test_sparse_mla_fp8.py +++ b/op_tests/test_sparse_mla_fp8.py @@ -189,3 +189,7 @@ def run(q, kv, ids): expected = run(q, kv, ids) compiled = torch.compile(run, backend="eager", fullgraph=True) assert torch.equal(compiled(q, kv, ids), expected) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__]))