diff --git a/aiter/__init__.py b/aiter/__init__.py index 5aac376c0d..f3aec12849 100644 --- a/aiter/__init__.py +++ b/aiter/__init__.py @@ -113,6 +113,7 @@ def getLogger(): from .ops.moe_sorting_opus import * from .ops.moe_mxfp4_aux import * from .ops.pa_sparse_prefill_opus import * + from .ops.minimax_m3_fused_qknorm_rope import * from .ops.pos_encoding import * from .ops.cache import * from .ops.rmsnorm import * diff --git a/aiter/jit/optCompilerConfig.json b/aiter/jit/optCompilerConfig.json index ee2e6347a9..20c37cd028 100644 --- a/aiter/jit/optCompilerConfig.json +++ b/aiter/jit/optCompilerConfig.json @@ -1200,6 +1200,23 @@ "verbose": "False", "blob_gen_cmd": "''" }, + "module_minimax_m3_fused_qknorm_rope_cache_shuffle": { + "srcs": [ + "f'{AITER_CSRC_DIR}/pybind/minimax_m3_fused_qknorm_rope_cache_shuffle_pybind.cu'", + "f'{AITER_CSRC_DIR}/kernels/minimax_m3_fused_qknorm_rope_cache_shuffle.cu'" + ], + "flags_extra_cc": [], + "flags_extra_hip": [ + "'-DENABLE_FP8'" + ], + "extra_ldflags": "None", + "extra_include": [ + "f'{AITER_CSRC_DIR}/include/ck_tile'", + "f'{AITER_CSRC_DIR}/include/opus'" + ], + "verbose": "False", + "blob_gen_cmd": "''" + }, "module_fused_qk_norm_rope_cache_quant_shuffle": { "srcs": [ "f'{AITER_CSRC_DIR}/pybind/fused_qk_norm_rope_cache_quant_pybind.cu'", diff --git a/aiter/ops/minimax_m3_fused_qknorm_rope.py b/aiter/ops/minimax_m3_fused_qknorm_rope.py new file mode 100644 index 0000000000..fd251e76e3 --- /dev/null +++ b/aiter/ops/minimax_m3_fused_qknorm_rope.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. + + +from torch import Tensor + +from ..jit.core import compile_ops + + +@compile_ops("module_minimax_m3_fused_qknorm_rope_cache_shuffle", develop=True) +def minimax_m3_qknorm_rope_cache_shuffle_insert( + qkv: Tensor, + q_norm_weight: Tensor, + k_norm_weight: Tensor, + cos_sin_cache: Tensor, + positions: Tensor, + num_heads: int, + num_kv_heads: int, + num_index_heads: int, + rotary_dim: int, + eps: float, + slot_mapping: Tensor, + k_cache: Tensor, + v_cache: Tensor, + q_out: Tensor, + index_q_norm_weight: Tensor | None = None, + index_k_norm_weight: Tensor | None = None, + index_slot_mapping: Tensor | None = None, + index_cache: Tensor | None = None, + index_q_out: Tensor | None = None, + kv_cache_dtype: str = "auto", + k_scale: Tensor | None = None, + v_scale: Tensor | None = None, + skip_index_branch: bool = False, +) -> None: + """Fused MiniMax-M3 QK-norm + partial NeoX RoPE + page-16 SHUFFLE KV insert. + + Consumes the sparse layer's packed projection row + ``[q | k | v | index_q | index_k]`` (all head_dim=128) and, in one pass: + + * writes Gemma-normed + roped ``q`` into ``q_out`` and ``index_q`` into + ``index_q_out``, + * scatters normed + roped ``k`` and verbatim ``v`` into the paged caches in + the layout ``pa_decode_gluon`` reads:: + + k_cache: [num_pages, num_kv_heads, head_dim // x, page_size, x] + v_cache: [num_pages, num_kv_heads, page_size // x, head_dim, x] + + where ``x = 16 // k_cache.element_size()`` and ``page_size`` is + ``k_cache.shape[3]`` (16 for the gluon decode path), + * scatters normed + roped ``index_k`` into ``index_cache`` + ``[num_pages, page_size, 128]``. + + ``k``/``v``/``index_k`` are *not* written back into ``qkv``; they are + consumed only by the cache scatter. ``skip_index_branch=True`` skips all + index work (the row still carries the index sub-blocks). + + ``k_scale``/``v_scale`` are optional single-element fp32 **device** tensors + applied as ``value / scale`` when ``kv_cache_dtype`` is quantized, matching + ``reshape_and_cache``. + """ diff --git a/csrc/include/minimax_m3_fused_qknorm_rope_cache_shuffle.h b/csrc/include/minimax_m3_fused_qknorm_rope_cache_shuffle.h new file mode 100644 index 0000000000..5455a2ef08 --- /dev/null +++ b/csrc/include/minimax_m3_fused_qknorm_rope_cache_shuffle.h @@ -0,0 +1,46 @@ +#pragma once +// SPDX-License-Identifier: MIT +// Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. + +#include +#include + +#include "aiter_tensor.h" + +namespace aiter { + +// Fused MiniMax-M3 attention pre-processing that writes the paged K/V caches +// directly in the page-16 SHUFFLE layout consumed by pa_decode_gluon. +// +// k_cache: [num_pages, num_kv_heads, head_dim/x, page_size, x] +// v_cache: [num_pages, num_kv_heads, page_size/x, head_dim, x] x = 16/elem_size +// +// Replaces the three-kernel sequence (fused QK-norm/RoPE -> reshape_and_cache +// with asm_layout=True -> index-cache scatter) with a single pass over the +// fused ``qkv`` row [q | k | v | index_q | index_k]. +void minimax_m3_qknorm_rope_cache_shuffle_insert( + aiter_tensor_t& qkv, // [num_tokens, qkv_row] fp16/bf16 + aiter_tensor_t& q_norm_weight, // [head_dim] + aiter_tensor_t& k_norm_weight, // [head_dim] + aiter_tensor_t& cos_sin_cache, // [max_pos, rotary_dim] + aiter_tensor_t& positions, // [num_tokens] i64 + int64_t num_heads, + int64_t num_kv_heads, + int64_t num_index_heads, + int64_t rotary_dim, + double eps, + aiter_tensor_t& slot_mapping, // [num_tokens] i64 + aiter_tensor_t& k_cache, + aiter_tensor_t& v_cache, + aiter_tensor_t& q_out, // [num_tokens, num_heads * head_dim] + std::optional index_q_norm_weight, + std::optional index_k_norm_weight, + std::optional index_slot_mapping, // [num_tokens] i64 + std::optional index_cache, // [pages, page_size, head_dim] + std::optional index_q_out, // [num_tokens, niq*head_dim] + const std::string& kv_cache_dtype, + std::optional k_scale, // fp32 device scalar + std::optional v_scale, // fp32 device scalar + bool skip_index_branch); + +} // namespace aiter diff --git a/csrc/include/rocm_ops.hpp b/csrc/include/rocm_ops.hpp index d5940d3e86..1d753b16fd 100644 --- a/csrc/include/rocm_ops.hpp +++ b/csrc/include/rocm_ops.hpp @@ -2154,6 +2154,33 @@ namespace py = pybind11; py::arg("quant_group_size") = 128, \ py::arg("scale_layout") = 0); +#define MINIMAX_M3_FUSED_QKNORM_ROPE_CACHE_SHUFFLE_PYBIND \ + m.def("minimax_m3_qknorm_rope_cache_shuffle_insert", \ + &aiter::minimax_m3_qknorm_rope_cache_shuffle_insert, \ + py::arg("qkv"), \ + py::arg("q_norm_weight"), \ + py::arg("k_norm_weight"), \ + py::arg("cos_sin_cache"), \ + py::arg("positions"), \ + py::arg("num_heads"), \ + py::arg("num_kv_heads"), \ + py::arg("num_index_heads"), \ + py::arg("rotary_dim"), \ + py::arg("eps"), \ + py::arg("slot_mapping"), \ + py::arg("k_cache"), \ + py::arg("v_cache"), \ + py::arg("q_out"), \ + py::arg("index_q_norm_weight") = std::nullopt, \ + py::arg("index_k_norm_weight") = std::nullopt, \ + py::arg("index_slot_mapping") = std::nullopt, \ + py::arg("index_cache") = std::nullopt, \ + py::arg("index_q_out") = std::nullopt, \ + py::arg("kv_cache_dtype") = "auto", \ + py::arg("k_scale") = std::nullopt, \ + py::arg("v_scale") = std::nullopt, \ + py::arg("skip_index_branch") = false); + #define SMOOTHQUANT_PYBIND \ m.def("smoothquant_fwd", &smoothquant_fwd); \ m.def("moe_smoothquant_fwd", &moe_smoothquant_fwd); diff --git a/csrc/kernels/minimax_m3_fused_qknorm_rope_cache_shuffle.cu b/csrc/kernels/minimax_m3_fused_qknorm_rope_cache_shuffle.cu new file mode 100644 index 0000000000..4fba98e049 --- /dev/null +++ b/csrc/kernels/minimax_m3_fused_qknorm_rope_cache_shuffle.cu @@ -0,0 +1,674 @@ +// SPDX-License-Identifier: MIT +// Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +/* + * Fused MiniMax-M3 attention pre-processing writing the page-16 SHUFFLE KV + * layout that pa_decode_gluon consumes. + * + * The sparse layer's fused projection emits one packed row per token: + * + * [ q (nq heads) | k (nkv) | v (nkv) | index_q (niq) | index_k (1) ] + * + * all with head_dim=128, all sharing the same partial-NeoX RoPE table. The + * four norms are Gemma-style RMSNorm, ``x * rsqrt(mean(x^2)+eps) * (1 + w)``, + * with independent weights. This kernel replaces the three-pass sequence + * + * fused QK-norm/RoPE -> reshape_and_cache(asm_layout=True) + * -> index-cache scatter + * + * with a single pass: q and index_q are gathered into contiguous output + * buffers, while k/v/index_k go straight from the projection output into their + * caches and are never written back to ``qkv``. + * + * One warp (32 lanes) owns one (token, slot) pair; each lane owns 4 contiguous + * head dims. Slot order matches the packed row exactly, so a slot index + * doubles as the head index inside the row: + * + * [0, nq) Q -> norm(q_w) + RoPE -> q_out + * [nq, nq+nkv) K -> norm(k_w) + RoPE -> k_cache (shuffle) + * [nq+nkv, nq+2*nkv) V -> verbatim -> v_cache (shuffle) + * [nq+2*nkv, +niq) IQ -> norm(iq_w) + RoPE -> index_q_out + * nq+2*nkv+niq IK -> norm(ik_w) + RoPE -> index_cache + * + * Cache index math (page_size = k_cache.size(3), x = 16/elem_size): + * + * page = slot / page_size, off = slot % page_size + * k: page*k_block_stride + h*head_dim*page_size + * + (d/x)*page_size*x + off*x + (d%x) + * v: page*v_block_stride + h*head_dim*page_size + * + (off/x)*head_dim*x + d*x + (off%x) + * + * K keeps head_dim contiguous inside a 16-byte vector: x is a multiple of 4, so + * a lane's 4 dims always land in one x group and the store stays a single + * vector. V transposes token against head_dim, which forces per-element stores + * strided by x. + */ + +#include +#include + +#include "aiter_dispatch.h" +#include "aiter_hip_common.h" +#include "aiter_stream.h" +#include "minimax_m3_fused_qknorm_rope_cache_shuffle.h" +#include "quant_utils.cuh" + +#include "attention_dtypes.h" +#include "aiter_opus_plus.h" +#include "opus/opus.hpp" + +#include + +namespace aiter { +namespace minimax_m3 { + +constexpr int kHeadDim = 128; +constexpr int kLanesPerHead = 32; +constexpr int kElemsPerLane = kHeadDim / kLanesPerHead; // 4 +constexpr float kFp8Max = 448.0f; + +// Sum across the 32-lane group owning one head (a wave64 holds two groups). +__device__ __forceinline__ float headReduceSum(float v) +{ +#pragma unroll + for(int mask = kLanesPerHead / 2; mask > 0; mask >>= 1) + { + v += __shfl_xor(v, mask, kLanesPerHead); + } + return v; +} + +// Gemma RMSNorm over the full head (skipped when ``weight == nullptr``) then +// partial NeoX RoPE over the leading ``rotary_dim`` dims. Lane L owns dims +// [4L, 4L+4); ``half`` is a multiple of 4, so a lane lies wholly in one half of +// the rotary range and its RoPE partner sits ``half/4`` lanes away. +template +__device__ __forceinline__ void normAndRope(float (&elems)[kElemsPerLane], + const int lane, + const float eps, + const scalar_t* __restrict__ weight, + const bool do_rope, + const int rotary_dim, + const scalar_t* __restrict__ cos_ptr) +{ + if(weight != nullptr) + { + float sumsq = 0.0f; +#pragma unroll + for(int i = 0; i < kElemsPerLane; i++) + sumsq += elems[i] * elems[i]; + sumsq = headReduceSum(sumsq); + const float rms_rcp = rsqrtf(sumsq / static_cast(kHeadDim) + eps); +#pragma unroll + for(int i = 0; i < kElemsPerLane; i++) + { + const float w = 1.0f + static_cast(weight[lane * kElemsPerLane + i]); + elems[i] = elems[i] * rms_rcp * w; + } + } + + if(!do_rope) + return; + + const int half = rotary_dim / 2; + const int dim0 = lane * kElemsPerLane; + const int lane_xor = half / kElemsPerLane; + float partner[kElemsPerLane]; +#pragma unroll + for(int i = 0; i < kElemsPerLane; i++) + partner[i] = __shfl_xor(elems[i], lane_xor, kLanesPerHead); + + if(dim0 >= rotary_dim) + return; // trailing dims pass through unrotated + + const bool first_half = dim0 < half; + const int i_base = first_half ? dim0 : (dim0 - half); + const scalar_t* __restrict__ s_ptr = cos_ptr + half; +#pragma unroll + for(int i = 0; i < kElemsPerLane; i++) + { + const float c = static_cast(cos_ptr[i_base + i]); + const float s = static_cast(s_ptr[i_base + i]); + elems[i] = first_half ? (elems[i] * c - partner[i] * s) + : (elems[i] * c + partner[i] * s); + } +} + +// ──────────────────────────────────────────────────────────────────────────── +// Kernel +// ──────────────────────────────────────────────────────────────────────────── +// idx_t is the index-K cache / index-Q output dtype: scalar_t, or fp8 e4m3 for +// the AITER fp8 score path. kProcessIndex compiles away the index branch for +// skip-index-topk reuse layers (whose rows still carry the index sub-blocks). +template +__global__ void minimaxM3QKNormRopeCacheShuffleInsertKernel( + const scalar_t* __restrict__ qkv, + scalar_t* __restrict__ q_out, + idx_t* __restrict__ index_q_out, + const scalar_t* __restrict__ q_norm_w, + const scalar_t* __restrict__ k_norm_w, + const scalar_t* __restrict__ iq_norm_w, + const scalar_t* __restrict__ ik_norm_w, + const scalar_t* __restrict__ cos_sin_cache, + const int64_t* __restrict__ positions, + const int64_t* __restrict__ slot_mapping, + const int64_t* __restrict__ index_slot_mapping, + cache_t* __restrict__ k_cache, + cache_t* __restrict__ v_cache, + idx_t* __restrict__ index_cache, + const float* __restrict__ k_scale, + const float* __restrict__ v_scale, + const float eps, + const int rotary_dim, + const int num_tokens, + const int nq, + const int nkv, + const int niq, + const int page_size, + const int x, + const int64_t k_block_stride, + const int64_t v_block_stride, + const int64_t idx_page_stride, + const int64_t idx_token_stride, + const int idx_page_size) +{ + const int warps_per_block = blockDim.x / kLanesPerHead; + const int lane = threadIdx.x % kLanesPerHead; + const int global_warp = blockIdx.x * warps_per_block + (threadIdx.x / kLanesPerHead); + + const int v_begin = nq + nkv; + const int iq_begin = nq + 2 * nkv; + const int ik_slot = iq_begin + niq; + const int slots = kProcessIndex ? (ik_slot + 1) : iq_begin; + + const int token = global_warp / slots; + const int slot = global_warp % slots; + if(token >= num_tokens) + return; + + const bool isQ = slot < nq; + const bool isK = slot >= nq && slot < v_begin; + const bool isV = slot >= v_begin && slot < iq_begin; + bool isIQ = false, isIK = false; + if constexpr(kProcessIndex) + { + isIQ = slot >= iq_begin && slot < ik_slot; + isIK = slot == ik_slot; + } + + // Slot order mirrors the packed row, so ``slot`` is also the head index + // inside the row -- one formula covers q/k/v/index_q/index_k. + const int qkv_row = (nq + 2 * nkv + niq + 1) * kHeadDim; + const int dim_base = lane * kElemsPerLane; + const scalar_t* __restrict__ row_ptr = + qkv + static_cast(token) * qkv_row + static_cast(slot) * kHeadDim; + + const scalar_t* norm_w = nullptr; + if(isQ) + norm_w = q_norm_w; + else if(isK) + norm_w = k_norm_w; + else if(isIQ) + norm_w = iq_norm_w; + else if(isIK) + norm_w = ik_norm_w; + + float elems[kElemsPerLane]; +#pragma unroll + for(int i = 0; i < kElemsPerLane; i++) + elems[i] = static_cast(row_ptr[dim_base + i]); + + if(!isV) + { + const int64_t pos = positions[token]; + normAndRope(elems, + lane, + eps, + norm_w, + /*do_rope=*/true, + rotary_dim, + cos_sin_cache + pos * rotary_dim); + } + + // Round to the model dtype first: bf16/fp16 -> fp32 -> the same bf16/fp16 is + // an exact round trip, so this reproduces the value the unfused path + // materialized in ``qkv`` before quantizing. + scalar_t rounded[kElemsPerLane]; +#pragma unroll + for(int i = 0; i < kElemsPerLane; i++) + rounded[i] = opus::cast(elems[i]); + + // ── Q / index_q gather into contiguous buffers ───────────────────────── + if(isQ) + { + scalar_t* dst = + q_out + static_cast(token) * nq * kHeadDim + slot * kHeadDim + dim_base; +#pragma unroll + for(int i = 0; i < kElemsPerLane; i++) + dst[i] = rounded[i]; + return; + } + if constexpr(kProcessIndex) + { + if(isIQ) + { + if(index_q_out != nullptr) + { + idx_t* dst = index_q_out + static_cast(token) * niq * kHeadDim + + (slot - iq_begin) * kHeadDim + dim_base; +#pragma unroll + for(int i = 0; i < kElemsPerLane; i++) + { + if constexpr(kFp8Idx) + { + // Straight from fp32, matching the unfused kernel's + // index_q store; going via bf16 would double-round. + const float v = fminf(fmaxf(elems[i], -kFp8Max), kFp8Max); + dst[i] = opus::cast(v); + } + else + { + dst[i] = rounded[i]; + } + } + } + return; + } + } + + // ── Cache inserts ────────────────────────────────────────────────────── + int64_t sm = -1; + if(isK || isV) + sm = slot_mapping[token]; + else if constexpr(kProcessIndex) + { + if(isIK) + sm = index_slot_mapping[token]; + } + if(sm < 0) + return; // padded / unscheduled token + + if constexpr(kProcessIndex) + { + if(isIK) + { + if(index_cache != nullptr) + { + idx_t* dst = index_cache + (sm / idx_page_size) * idx_page_stride + + (sm % idx_page_size) * idx_token_stride + dim_base; +#pragma unroll + for(int i = 0; i < kElemsPerLane; i++) + { + // Deliberately quantized from the rounded scalar_t value: + // the path this replaces materialized index_k in ``qkv`` as + // bf16 before a separate insert cast it to fp8, so rounding + // here keeps the fused op bit-identical to it. + if constexpr(kFp8Idx) + { + const float v = + fminf(fmaxf(static_cast(rounded[i]), -kFp8Max), kFp8Max); + dst[i] = opus::cast(v); + } + else + { + dst[i] = rounded[i]; + } + } + } + return; + } + } + + const int64_t page = sm / page_size; + const int off = static_cast(sm % page_size); + const int64_t head_stride = static_cast(kHeadDim) * page_size; + + if(isK) + { + const int head = slot - nq; + cache_t* dst = k_cache + page * k_block_stride + head * head_stride + + static_cast(dim_base / x) * page_size * x + + static_cast(off) * x + (dim_base % x); + const float inv = (k_scale == nullptr) ? 1.0f : 1.0f / (*k_scale); +#pragma unroll + for(int i = 0; i < kElemsPerLane; i++) + { + if constexpr(kv_dt == vllm::Fp8KVCacheDataType::kAuto) + dst[i] = rounded[i]; + else + dst[i] = opus::cast(static_cast(rounded[i]) * inv); + } + } + else // isV + { + const int head = slot - v_begin; + cache_t* dst = v_cache + page * v_block_stride + head * head_stride + + static_cast(off / x) * kHeadDim * x + + static_cast(dim_base) * x + (off % x); + const float inv = (v_scale == nullptr) ? 1.0f : 1.0f / (*v_scale); +#pragma unroll + for(int i = 0; i < kElemsPerLane; i++) + { + if constexpr(kv_dt == vllm::Fp8KVCacheDataType::kAuto) + dst[i * x] = rounded[i]; + else + dst[i * x] = opus::cast(static_cast(rounded[i]) * inv); + } + } +} + +template +void launch(const scalar_t* qkv, + scalar_t* q_out, + void* index_q_out, + const scalar_t* q_norm_w, + const scalar_t* k_norm_w, + const scalar_t* iq_norm_w, + const scalar_t* ik_norm_w, + const scalar_t* cos_sin_cache, + const int64_t* positions, + const int64_t* slot_mapping, + const int64_t* index_slot_mapping, + cache_t* k_cache, + cache_t* v_cache, + void* index_cache, + const float* k_scale, + const float* v_scale, + float eps, + int rotary_dim, + int num_tokens, + int nq, + int nkv, + int niq, + int page_size, + int x, + int64_t k_block_stride, + int64_t v_block_stride, + int64_t idx_page_stride, + int64_t idx_token_stride, + int idx_page_size, + bool process_index, + bool fp8_idx, + hipStream_t stream) +{ + constexpr int kBlockSize = 256; + const int slots = process_index ? (nq + 2 * nkv + niq + 1) : (nq + 2 * nkv); + const int64_t total_warps = static_cast(num_tokens) * slots; + const int warps_per_block = kBlockSize / kLanesPerHead; + const int grid = static_cast((total_warps + warps_per_block - 1) / warps_per_block); + if(grid == 0) + return; + +#define LAUNCH_M3_SHUFFLE(PROCESS_INDEX, FP8_IDX, IDX_T) \ + minimaxM3QKNormRopeCacheShuffleInsertKernel<<>>( \ + qkv, \ + q_out, \ + reinterpret_cast(index_q_out), \ + q_norm_w, \ + k_norm_w, \ + iq_norm_w, \ + ik_norm_w, \ + cos_sin_cache, \ + positions, \ + slot_mapping, \ + index_slot_mapping, \ + k_cache, \ + v_cache, \ + reinterpret_cast(index_cache), \ + k_scale, \ + v_scale, \ + eps, \ + rotary_dim, \ + num_tokens, \ + nq, \ + nkv, \ + niq, \ + page_size, \ + x, \ + k_block_stride, \ + v_block_stride, \ + idx_page_stride, \ + idx_token_stride, \ + idx_page_size) + + if(!process_index) + { + LAUNCH_M3_SHUFFLE(false, false, scalar_t); + } + else if(fp8_idx) + { + LAUNCH_M3_SHUFFLE(true, true, opus::fp8_t); + } + else + { + LAUNCH_M3_SHUFFLE(true, false, scalar_t); + } +#undef LAUNCH_M3_SHUFFLE +} + +} // namespace minimax_m3 + +namespace { +// Like is_contiguous() but ignoring dim 0, which permits an interleaved block +// stride (e.g. a vLLM [num_blocks, 2, ...] cache after unbind(1)). +inline bool contiguous_from_dim1(const aiter_tensor_t& t) +{ + if(t.numel() == 0) + return true; + int64_t expected = 1; + for(int d = t.dim() - 1; d >= 1; --d) + { + if(t.size(d) != 1 && t.stride(d) != expected) + return false; + expected *= t.size(d); + } + return true; +} + +// The scale is dereferenced on the device (as reshape_and_cache does), so a +// device scalar is required; nullptr means an identity scale. +inline const float* scale_ptr_of(const std::optional& scale, const char* name) +{ + if(!scale.has_value() || scale->numel() == 0) + return nullptr; + AITER_CHECK(scale->dtype() == AITER_DTYPE_fp32 && scale->numel() == 1 && scale->is_gpu(), + name, + " must be a single-element fp32 GPU tensor"); + return reinterpret_cast(scale->data_ptr()); +} +} // namespace + +void minimax_m3_qknorm_rope_cache_shuffle_insert(aiter_tensor_t& qkv, + aiter_tensor_t& q_norm_weight, + aiter_tensor_t& k_norm_weight, + aiter_tensor_t& cos_sin_cache, + aiter_tensor_t& positions, + int64_t num_heads, + int64_t num_kv_heads, + int64_t num_index_heads, + int64_t rotary_dim, + double eps, + aiter_tensor_t& slot_mapping, + aiter_tensor_t& k_cache, + aiter_tensor_t& v_cache, + aiter_tensor_t& q_out, + std::optional index_q_norm_weight, + std::optional index_k_norm_weight, + std::optional index_slot_mapping, + std::optional index_cache, + std::optional index_q_out, + const std::string& kv_cache_dtype, + std::optional k_scale, + std::optional v_scale, + bool skip_index_branch) +{ + using namespace minimax_m3; + + const int nq = static_cast(num_heads); + const int nkv = static_cast(num_kv_heads); + const int niq = static_cast(num_index_heads); + AITER_CHECK(nq > 0 && nkv > 0 && niq > 0, + "minimax_m3 fused insert is the sparse-layer path: " + "num_heads/num_kv_heads/num_index_heads must all be > 0"); + + AITER_CHECK(qkv.is_gpu() && qkv.is_contiguous() && qkv.dim() == 2, + "qkv must be a contiguous 2-D GPU tensor"); + AITER_CHECK(qkv.dtype() == AITER_DTYPE_bf16 || qkv.dtype() == AITER_DTYPE_fp16, + "qkv must be bf16 or fp16"); + const int num_tokens = static_cast(qkv.size(0)); + AITER_CHECK(qkv.size(1) == static_cast(nq + 2 * nkv + niq + 1) * kHeadDim, + "qkv row must be (num_heads + 2*num_kv_heads + num_index_heads + 1) * 128, got ", + qkv.size(1)); + + AITER_CHECK(rotary_dim > 0 && rotary_dim % (2 * kElemsPerLane) == 0 && rotary_dim <= kHeadDim, + "rotary_dim must be a positive multiple of 8 and <= 128, got ", + rotary_dim); + AITER_CHECK(cos_sin_cache.is_gpu() && cos_sin_cache.is_contiguous() && + cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == rotary_dim && + cos_sin_cache.dtype() == qkv.dtype(), + "cos_sin_cache must be contiguous [max_pos, rotary_dim] matching qkv dtype"); + AITER_CHECK(positions.dtype() == AITER_DTYPE_i64 && positions.numel() >= num_tokens, + "positions must be int64 with at least num_tokens elements"); + AITER_CHECK(slot_mapping.dtype() == AITER_DTYPE_i64 && slot_mapping.numel() >= num_tokens, + "slot_mapping must be int64 with at least num_tokens elements"); + + AITER_CHECK(q_norm_weight.numel() == kHeadDim && k_norm_weight.numel() == kHeadDim && + q_norm_weight.dtype() == qkv.dtype() && k_norm_weight.dtype() == qkv.dtype(), + "q/k norm weights must be 128 elements matching qkv dtype"); + + AITER_CHECK(q_out.is_gpu() && q_out.is_contiguous() && q_out.dtype() == qkv.dtype() && + q_out.numel() == static_cast(num_tokens) * nq * kHeadDim, + "q_out must be contiguous [num_tokens, num_heads*128] matching qkv dtype"); + + // Cache layout: k [pages, heads, D/x, page, x], v [pages, heads, page/x, D, x]. + AITER_CHECK(k_cache.dim() == 5 && v_cache.dim() == 5, + "shuffle-layout k_cache/v_cache must be 5-D"); + AITER_CHECK(k_cache.dtype() == v_cache.dtype(), "k_cache/v_cache dtype must match"); + AITER_CHECK(contiguous_from_dim1(k_cache) && contiguous_from_dim1(v_cache), + "k_cache/v_cache must be contiguous within a page (dims >= 1)"); + const int x = static_cast(k_cache.size(4)); + const int page_size = static_cast(k_cache.size(3)); + AITER_CHECK(x == static_cast(16 / k_cache.element_size()), + "k_cache last dim must be 16 / element_size, got ", + x); + AITER_CHECK(x % kElemsPerLane == 0, "x must be a multiple of 4, got ", x); + AITER_CHECK(page_size % x == 0, "page_size must be a multiple of x"); + AITER_CHECK(k_cache.size(1) == nkv && v_cache.size(1) == nkv, + "k_cache/v_cache must have num_kv_heads in dim 1"); + AITER_CHECK(k_cache.size(2) == kHeadDim / x, "k_cache dim 2 must be head_dim/x"); + AITER_CHECK(v_cache.size(2) == page_size / x && v_cache.size(3) == kHeadDim && + v_cache.size(4) == x, + "v_cache must be [pages, heads, page_size/x, head_dim, x]"); + + const bool process_index = !skip_index_branch; + const int64_t* index_slot_ptr = nullptr; + void* index_cache_ptr = nullptr; + void* index_q_out_ptr = nullptr; + int64_t idx_page_stride = 0, idx_token_stride = 0; + int idx_page_size = 1; + bool fp8_idx = false; + const void* iq_norm_ptr = nullptr; + const void* ik_norm_ptr = nullptr; + + if(process_index) + { + AITER_CHECK(index_q_norm_weight.has_value() && index_k_norm_weight.has_value(), + "index branch requires both index norm weights"); + AITER_CHECK(index_q_norm_weight->numel() == kHeadDim && + index_k_norm_weight->numel() == kHeadDim && + index_q_norm_weight->dtype() == qkv.dtype() && + index_k_norm_weight->dtype() == qkv.dtype(), + "index norm weights must be 128 elements matching qkv dtype"); + iq_norm_ptr = index_q_norm_weight->data_ptr(); + ik_norm_ptr = index_k_norm_weight->data_ptr(); + + AITER_CHECK(index_cache.has_value(), + "index branch requires index_cache (pass skip_index_branch=True to disable)"); + AITER_CHECK(index_cache->dim() == 3 && index_cache->stride(2) == 1 && + index_cache->size(2) == kHeadDim, + "index_cache must be [pages, page_size, 128] with contiguous head dim"); + AITER_CHECK(index_cache->dtype() == qkv.dtype() || index_cache->dtype() == AITER_DTYPE_fp8, + "index_cache must match qkv dtype or be fp8 e4m3"); + idx_page_stride = index_cache->stride(0); + idx_token_stride = index_cache->stride(1); + idx_page_size = static_cast(index_cache->size(1)); + index_cache_ptr = index_cache->data_ptr(); + fp8_idx = index_cache->dtype() == AITER_DTYPE_fp8; + + AITER_CHECK(index_slot_mapping.has_value() && + index_slot_mapping->dtype() == AITER_DTYPE_i64 && + index_slot_mapping->numel() >= num_tokens, + "index branch requires an int64 index_slot_mapping"); + index_slot_ptr = reinterpret_cast(index_slot_mapping->data_ptr()); + + if(index_q_out.has_value()) + { + AITER_CHECK(index_q_out->is_contiguous() && + index_q_out->numel() == + static_cast(num_tokens) * niq * kHeadDim, + "index_q_out must be contiguous [num_tokens, niq*128]"); + AITER_CHECK(index_q_out->dtype() == index_cache->dtype(), + "index_q_out dtype must match index_cache dtype"); + index_q_out_ptr = index_q_out->data_ptr(); + } + } + + const bool quantized = kv_cache_dtype != "auto"; + const float* k_scale_ptr = quantized ? scale_ptr_of(k_scale, "k_scale") : nullptr; + const float* v_scale_ptr = quantized ? scale_ptr_of(v_scale, "v_scale") : nullptr; + AITER_CHECK(quantized ? (k_cache.dtype() == AITER_DTYPE_fp8) + : (k_cache.dtype() == qkv.dtype()), + "k_cache must be fp8 for a quantized kv_cache_dtype and match qkv otherwise"); + + HipDeviceGuard device_guard(qkv.device_id); + const hipStream_t stream = aiter::getCurrentHIPStream(); + +#define CALL_M3_SHUFFLE_INSERT(SCALAR_T, CACHE_T, KV_DTYPE) \ + minimax_m3::launch( \ + reinterpret_cast(qkv.data_ptr()), \ + reinterpret_cast(q_out.data_ptr()), \ + index_q_out_ptr, \ + reinterpret_cast(q_norm_weight.data_ptr()), \ + reinterpret_cast(k_norm_weight.data_ptr()), \ + reinterpret_cast(iq_norm_ptr), \ + reinterpret_cast(ik_norm_ptr), \ + reinterpret_cast(cos_sin_cache.data_ptr()), \ + reinterpret_cast(positions.data_ptr()), \ + reinterpret_cast(slot_mapping.data_ptr()), \ + index_slot_ptr, \ + reinterpret_cast(k_cache.data_ptr()), \ + reinterpret_cast(v_cache.data_ptr()), \ + index_cache_ptr, \ + k_scale_ptr, \ + v_scale_ptr, \ + static_cast(eps), \ + static_cast(rotary_dim), \ + num_tokens, \ + nq, \ + nkv, \ + niq, \ + page_size, \ + x, \ + k_cache.stride(0), \ + v_cache.stride(0), \ + idx_page_stride, \ + idx_token_stride, \ + idx_page_size, \ + process_index, \ + fp8_idx, \ + stream) + + DISPATCH_BY_KV_CACHE_DTYPE_OPUS_rmTorch(qkv.dtype(), kv_cache_dtype, CALL_M3_SHUFFLE_INSERT) +#undef CALL_M3_SHUFFLE_INSERT +} + +} // namespace aiter diff --git a/csrc/pybind/minimax_m3_fused_qknorm_rope_cache_shuffle_pybind.cu b/csrc/pybind/minimax_m3_fused_qknorm_rope_cache_shuffle_pybind.cu new file mode 100644 index 0000000000..9cf8b78328 --- /dev/null +++ b/csrc/pybind/minimax_m3_fused_qknorm_rope_cache_shuffle_pybind.cu @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: MIT +// Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +#include "aiter_stream.h" +#include "minimax_m3_fused_qknorm_rope_cache_shuffle.h" +#include "rocm_ops.hpp" + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + AITER_SET_STREAM_PYBIND + MINIMAX_M3_FUSED_QKNORM_ROPE_CACHE_SHUFFLE_PYBIND; +} diff --git a/op_tests/test_minimax_m3_fused_qknorm_rope_cache_shuffle.py b/op_tests/test_minimax_m3_fused_qknorm_rope_cache_shuffle.py new file mode 100644 index 0000000000..595574f701 --- /dev/null +++ b/op_tests/test_minimax_m3_fused_qknorm_rope_cache_shuffle.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +"""Correctness tests for minimax_m3_qknorm_rope_cache_shuffle_insert. + +The reference reimplements Gemma QK-norm, partial NeoX RoPE and the page-16 +SHUFFLE cache index math in torch, so a layout mistake shows up as a gross +mismatch rather than a rounding difference. +""" + +import pytest +import torch + +from aiter import minimax_m3_qknorm_rope_cache_shuffle_insert +from aiter.utility.dtypes import get_dtype_fp8 + +HEAD_DIM = 128 +PAGE_SIZE = 16 + + +def gemma_rmsnorm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor: + """x * rsqrt(mean(x^2) + eps) * (1 + w), computed in fp32.""" + xf = x.float() + var = xf.pow(2).mean(-1, keepdim=True) + return xf * torch.rsqrt(var + eps) * (1.0 + w.float()) + + +def partial_neox_rope( + x: torch.Tensor, cos_sin: torch.Tensor, rotary_dim: int +) -> torch.Tensor: + """Rotate the leading rotary_dim dims; cos_sin is [cos(half) | sin(half)].""" + half = rotary_dim // 2 + cos, sin = cos_sin[..., :half].float(), cos_sin[..., half:].float() + out = x.clone() + x1, x2 = x[..., :half], x[..., half:rotary_dim] + out[..., :half] = x1 * cos - x2 * sin + out[..., half:rotary_dim] = x2 * cos + x1 * sin + return out + + +def make_inputs( + num_tokens, + nq, + nkv, + niq, + rotary_dim, + dtype, + cache_dtype, + idx_cache_dtype, + num_pages, + idx_num_pages, + seed=0, +): + torch.manual_seed(seed) + dev = "cuda" + row = (nq + 2 * nkv + niq + 1) * HEAD_DIM + qkv = torch.randn(num_tokens, row, device=dev, dtype=dtype) + + weights = { + name: torch.randn(HEAD_DIM, device=dev, dtype=dtype) * 0.1 + for name in ("q", "k", "iq", "ik") + } + max_pos = 4096 + cos_sin = torch.randn(max_pos, rotary_dim, device=dev, dtype=dtype) + positions = torch.randint(0, max_pos, (num_tokens,), device=dev, dtype=torch.int64) + + x = 16 // torch.empty((), dtype=cache_dtype).element_size() + k_cache = torch.zeros( + num_pages, nkv, HEAD_DIM // x, PAGE_SIZE, x, device=dev, dtype=cache_dtype + ) + v_cache = torch.zeros( + num_pages, nkv, PAGE_SIZE // x, HEAD_DIM, x, device=dev, dtype=cache_dtype + ) + index_cache = torch.zeros( + idx_num_pages, PAGE_SIZE, HEAD_DIM, device=dev, dtype=idx_cache_dtype + ) + + # Distinct slots so no two tokens collide, then punch a few padded tokens. + slot_mapping = torch.randperm(num_pages * PAGE_SIZE, device=dev)[:num_tokens].to( + torch.int64 + ) + index_slot_mapping = torch.randperm(idx_num_pages * PAGE_SIZE, device=dev)[ + :num_tokens + ].to(torch.int64) + if num_tokens >= 4: + slot_mapping[1] = -1 + index_slot_mapping[2] = -1 + + q_out = torch.zeros(num_tokens, nq * HEAD_DIM, device=dev, dtype=dtype) + index_q_out = torch.zeros( + num_tokens, niq * HEAD_DIM, device=dev, dtype=idx_cache_dtype + ) + return { + "qkv": qkv, + "weights": weights, + "cos_sin": cos_sin, + "positions": positions, + "k_cache": k_cache, + "v_cache": v_cache, + "index_cache": index_cache, + "slot_mapping": slot_mapping, + "index_slot_mapping": index_slot_mapping, + "q_out": q_out, + "index_q_out": index_q_out, + "x": x, + } + + +def reference(t, nq, nkv, niq, rotary_dim, eps, process_index, k_scale, v_scale): + """Returns (q_out, index_q_out, k_cache, v_cache, index_cache) references.""" + qkv, w = t["qkv"], t["weights"] + num_tokens = qkv.shape[0] + dtype = qkv.dtype + cache_dtype = t["k_cache"].dtype + idx_dtype = t["index_cache"].dtype + x = t["x"] + + heads = qkv.view(num_tokens, -1, HEAD_DIM) + cos_sin = t["cos_sin"][t["positions"]].unsqueeze(1) # [N, 1, rotary_dim] + + def norm_rope(sub, weight): + return partial_neox_rope(gemma_rmsnorm(sub, weight, eps), cos_sin, rotary_dim) + + q = norm_rope(heads[:, :nq], w["q"]).to(dtype) + k = norm_rope(heads[:, nq : nq + nkv], w["k"]).to(dtype) + v = heads[:, nq + nkv : nq + 2 * nkv] + iq_begin = nq + 2 * nkv + # index_q is quantized straight from fp32; index_k goes via the model dtype + # (see the kernel comment: it mirrors the unfused bf16 -> fp8 insert). + index_q_f32 = norm_rope(heads[:, iq_begin : iq_begin + niq], w["iq"]) + index_k = norm_rope(heads[:, iq_begin + niq : iq_begin + niq + 1], w["ik"]).to( + dtype + ) + + q_ref = q.reshape(num_tokens, nq * HEAD_DIM) + iq_src = index_q_f32 if idx_dtype != dtype else index_q_f32.to(dtype) + iq_ref = ( + iq_src.reshape(num_tokens, niq * HEAD_DIM).clamp(-448.0, 448.0).to(idx_dtype) + ) + + k_ref = torch.zeros_like(t["k_cache"]) + v_ref = torch.zeros_like(t["v_cache"]) + idx_ref = torch.zeros_like(t["index_cache"]) + + def quant(vals, scale): + if cache_dtype == dtype: + return vals + f = vals.float() + if scale is not None: + f = f / scale.item() + return f.to(cache_dtype) + + for tok in range(num_tokens): + slot = int(t["slot_mapping"][tok]) + if slot >= 0: + page, off = slot // PAGE_SIZE, slot % PAGE_SIZE + k_ref[page, :, :, off, :] = quant(k[tok], k_scale).view( + nkv, HEAD_DIM // x, x + ) + v_ref[page, :, off // x, :, off % x] = quant(v[tok], v_scale) + if not process_index: + continue + islot = int(t["index_slot_mapping"][tok]) + if islot >= 0: + ipage, ioff = islot // PAGE_SIZE, islot % PAGE_SIZE + idx_ref[ipage, ioff] = index_k[tok, 0].to(idx_dtype) + return q_ref, iq_ref, k_ref, v_ref, idx_ref + + +def compare(name, got, ref): + """bf16/fp8 tolerance; a layout bug misses by orders of magnitude.""" + g, r = got.float(), ref.float() + torch.testing.assert_close(g, r, rtol=2e-2, atol=2e-2, msg=lambda m: f"{name}: {m}") + exact = (g == r).float().mean().item() + assert exact > 0.9, f"{name}: only {exact:.1%} of elements bit-exact" + + +def run_op(t, nq, nkv, niq, rotary_dim, eps, **kwargs): + minimax_m3_qknorm_rope_cache_shuffle_insert( + t["qkv"], + t["weights"]["q"], + t["weights"]["k"], + t["cos_sin"], + t["positions"], + nq, + nkv, + niq, + rotary_dim, + eps, + t["slot_mapping"], + t["k_cache"], + t["v_cache"], + t["q_out"], + **kwargs, + ) + torch.cuda.synchronize() + + +@pytest.mark.parametrize("quantized", [False, True]) +@pytest.mark.parametrize("fp8_index", [False, True]) +@pytest.mark.parametrize("num_tokens", [1, 7, 64]) +def test_fused_insert(quantized, fp8_index, num_tokens): + nq, nkv, niq, rotary_dim, eps = 16, 1, 1, 64, 1e-6 + dtype = torch.bfloat16 + fp8 = get_dtype_fp8() + t = make_inputs( + num_tokens, + nq, + nkv, + niq, + rotary_dim, + dtype, + fp8 if quantized else dtype, + fp8 if fp8_index else dtype, + num_pages=8, + idx_num_pages=8, + ) + k_scale = torch.tensor([0.7], device="cuda") if quantized else None + v_scale = torch.tensor([1.3], device="cuda") if quantized else None + + run_op( + t, + nq, + nkv, + niq, + rotary_dim, + eps, + index_q_norm_weight=t["weights"]["iq"], + index_k_norm_weight=t["weights"]["ik"], + index_slot_mapping=t["index_slot_mapping"], + index_cache=t["index_cache"], + index_q_out=t["index_q_out"], + kv_cache_dtype="fp8" if quantized else "auto", + k_scale=k_scale, + v_scale=v_scale, + ) + + q_ref, iq_ref, k_ref, v_ref, idx_ref = reference( + t, nq, nkv, niq, rotary_dim, eps, True, k_scale, v_scale + ) + compare("q_out", t["q_out"], q_ref) + compare("index_q_out", t["index_q_out"], iq_ref) + compare("k_cache", t["k_cache"], k_ref) + compare("v_cache", t["v_cache"], v_ref) + compare("index_cache", t["index_cache"], idx_ref) + + +def test_skip_index_branch_leaves_index_outputs_untouched(): + nq, nkv, niq, rotary_dim, eps = 16, 1, 1, 64, 1e-6 + dtype = torch.bfloat16 + t = make_inputs( + 32, nq, nkv, niq, rotary_dim, dtype, dtype, dtype, num_pages=8, idx_num_pages=8 + ) + run_op(t, nq, nkv, niq, rotary_dim, eps, skip_index_branch=True) + + q_ref, _, k_ref, v_ref, _ = reference( + t, nq, nkv, niq, rotary_dim, eps, False, None, None + ) + compare("q_out", t["q_out"], q_ref) + compare("k_cache", t["k_cache"], k_ref) + compare("v_cache", t["v_cache"], v_ref) + assert t["index_cache"].float().abs().sum() == 0 + assert t["index_q_out"].float().abs().sum() == 0 + + +def test_multi_kv_head_and_full_rotary(): + nq, nkv, niq, rotary_dim, eps = 32, 4, 2, HEAD_DIM, 1e-6 + dtype = torch.bfloat16 + fp8 = get_dtype_fp8() + t = make_inputs( + 48, nq, nkv, niq, rotary_dim, dtype, fp8, fp8, num_pages=16, idx_num_pages=8 + ) + k_scale = torch.tensor([1.0], device="cuda") + v_scale = torch.tensor([1.0], device="cuda") + run_op( + t, + nq, + nkv, + niq, + rotary_dim, + eps, + index_q_norm_weight=t["weights"]["iq"], + index_k_norm_weight=t["weights"]["ik"], + index_slot_mapping=t["index_slot_mapping"], + index_cache=t["index_cache"], + index_q_out=t["index_q_out"], + kv_cache_dtype="fp8", + k_scale=k_scale, + v_scale=v_scale, + ) + q_ref, iq_ref, k_ref, v_ref, idx_ref = reference( + t, nq, nkv, niq, rotary_dim, eps, True, k_scale, v_scale + ) + compare("q_out", t["q_out"], q_ref) + compare("index_q_out", t["index_q_out"], iq_ref) + compare("k_cache", t["k_cache"], k_ref) + compare("v_cache", t["v_cache"], v_ref) + compare("index_cache", t["index_cache"], idx_ref) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/op_tests/test_minimax_m3_fused_shuffle_vs_unfused.py b/op_tests/test_minimax_m3_fused_shuffle_vs_unfused.py new file mode 100644 index 0000000000..b8aed62e05 --- /dev/null +++ b/op_tests/test_minimax_m3_fused_shuffle_vs_unfused.py @@ -0,0 +1,366 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +"""End-to-end equivalence: the fused MiniMax-M3 shuffle insert must reproduce +the three-kernel path vLLM's AMD sparse layer runs today. + +Unfused path (3 kernels): + 1. vLLM fused_minimax_m3_qknorm_rope_kv_insert (norm + RoPE only) + 2. aiter reshape_and_cache(..., asm_layout=True) + 3. vLLM minimax_m3_insert_index_cache (Triton) + +Fused path (1 kernel): + aiter minimax_m3_qknorm_rope_cache_shuffle_insert + +Also checks that pa_decode_gluon reads the fused cache to the same output as +the unfused one, i.e. the layout is what the decode kernel expects. +""" + +import pytest +import torch + +from aiter import ( + minimax_m3_qknorm_rope_cache_shuffle_insert, + reshape_and_cache, +) +from aiter.utility.dtypes import get_dtype_fp8 + +HEAD_DIM = 128 +PAGE_SIZE = 16 +PAGES_PER_BLOCK = 8 # a 128-token vLLM block is 8 physical page-16 pages + + +def _vllm_ops(): + pytest.importorskip("vllm") + import vllm._custom_ops as ops + + sparse_pa = pytest.importorskip("vllm.models.minimax_m3.amd.ops.sparse_pa") + return ops, sparse_pa.minimax_m3_insert_index_cache + + +def build_case(num_tokens, nq, nkv, niq, rotary_dim, cache_dtype, idx_dtype, seed=0): + torch.manual_seed(seed) + dev = "cuda" + dtype = torch.bfloat16 + row = (nq + 2 * nkv + niq + 1) * HEAD_DIM + qkv = torch.randn(num_tokens, row, device=dev, dtype=dtype) + + w = { + n: torch.randn(HEAD_DIM, device=dev, dtype=dtype) * 0.1 + for n in ("q", "k", "iq", "ik") + } + max_pos = 8192 + cos_sin = torch.randn(max_pos, rotary_dim, device=dev, dtype=dtype) + positions = torch.randint(0, max_pos, (num_tokens,), device=dev, dtype=torch.int64) + + num_blocks = 4 + num_pages = num_blocks * PAGES_PER_BLOCK + x = 16 // torch.empty((), dtype=cache_dtype).element_size() + + def caches(): + k = torch.zeros( + num_pages, nkv, HEAD_DIM // x, PAGE_SIZE, x, device=dev, dtype=cache_dtype + ) + v = torch.zeros( + num_pages, nkv, PAGE_SIZE // x, HEAD_DIM, x, device=dev, dtype=cache_dtype + ) + idx = torch.zeros(num_pages, PAGE_SIZE, HEAD_DIM, device=dev, dtype=idx_dtype) + return k, v, idx + + # vLLM slot_mapping is a global token index; page = slot // 16 works because + # a logical 128-token block's 8 page-16 pages are contiguous. + slots = torch.randperm(num_pages * PAGE_SIZE, device=dev)[:num_tokens] + slot_mapping = slots.to(torch.int64) + idx_slots = torch.randperm(num_pages * PAGE_SIZE, device=dev)[:num_tokens] + index_slot_mapping = idx_slots.to(torch.int64) + if num_tokens >= 4: + slot_mapping[1] = -1 # padded token + + return { + "qkv": qkv, + "w": w, + "cos_sin": cos_sin, + "positions": positions, + "slot_mapping": slot_mapping, + "index_slot_mapping": index_slot_mapping, + "caches": caches, + "x": x, + "dtype": dtype, + "num_tokens": num_tokens, + } + + +def run_unfused(c, nq, nkv, niq, rotary_dim, eps, kv_cache_dtype, k_scale, v_scale): + ops, insert_index_cache = _vllm_ops() + qkv = c["qkv"].clone() + n = c["num_tokens"] + q = qkv.new_empty((n, nq * HEAD_DIM)) + k_cache, v_cache, index_cache = c["caches"]() + index_q = torch.empty( + (n, niq * HEAD_DIM), device=qkv.device, dtype=index_cache.dtype + ) + + ops.fused_minimax_m3_qknorm_rope_kv_insert( + qkv, + c["w"]["q"], + c["w"]["k"], + c["cos_sin"], + c["positions"], + nq, + nkv, + rotary_dim, + eps, + c["w"]["iq"], + c["w"]["ik"], + niq, + q_out=q, + index_q_out=index_q, + kv_cache_dtype=kv_cache_dtype, + ) + + k_start = nq * HEAD_DIM + v_start = k_start + nkv * HEAD_DIM + ik_start = v_start + nkv * HEAD_DIM + niq * HEAD_DIM + k = qkv[:, k_start:v_start].view(n, nkv, HEAD_DIM) + v = qkv[:, v_start : v_start + nkv * HEAD_DIM].view(n, nkv, HEAD_DIM) + index_k = qkv[:, ik_start : ik_start + HEAD_DIM].view(n, HEAD_DIM) + + reshape_and_cache( + k.contiguous(), + v.contiguous(), + k_cache, + v_cache, + c["slot_mapping"], + kv_cache_dtype=kv_cache_dtype, + k_scale=k_scale, + v_scale=v_scale, + asm_layout=True, + ) + insert_index_cache(index_k, index_cache, c["index_slot_mapping"]) + torch.cuda.synchronize() + return q, index_q, k_cache, v_cache, index_cache + + +def run_fused(c, nq, nkv, niq, rotary_dim, eps, kv_cache_dtype, k_scale, v_scale): + qkv = c["qkv"].clone() + n = c["num_tokens"] + q = qkv.new_empty((n, nq * HEAD_DIM)) + k_cache, v_cache, index_cache = c["caches"]() + index_q = torch.empty( + (n, niq * HEAD_DIM), device=qkv.device, dtype=index_cache.dtype + ) + + minimax_m3_qknorm_rope_cache_shuffle_insert( + qkv, + c["w"]["q"], + c["w"]["k"], + c["cos_sin"], + c["positions"], + nq, + nkv, + niq, + rotary_dim, + eps, + c["slot_mapping"], + k_cache, + v_cache, + q, + index_q_norm_weight=c["w"]["iq"], + index_k_norm_weight=c["w"]["ik"], + index_slot_mapping=c["index_slot_mapping"], + index_cache=index_cache, + index_q_out=index_q, + kv_cache_dtype=kv_cache_dtype, + k_scale=k_scale, + v_scale=v_scale, + ) + torch.cuda.synchronize() + return q, index_q, k_cache, v_cache, index_cache + + +def assert_same(name, a, b, exact_frac=0.995): + fa, fb = a.float(), b.float() + torch.testing.assert_close( + fa, fb, rtol=8e-3, atol=8e-3, msg=lambda m: f"{name}: {m}" + ) + frac = (fa == fb).float().mean().item() + assert frac >= exact_frac, f"{name}: only {frac:.3%} bit-identical" + + +@pytest.mark.parametrize("quantized", [False, True]) +@pytest.mark.parametrize("fp8_index", [False, True]) +@pytest.mark.parametrize("num_tokens", [1, 5, 33]) +def test_fused_matches_unfused(quantized, fp8_index, num_tokens): + nq, nkv, niq, rotary_dim, eps = 16, 1, 1, 64, 1e-6 + _check_matches_unfused( + nq, nkv, niq, rotary_dim, eps, quantized, fp8_index, num_tokens + ) + + +@pytest.mark.parametrize("nkv", [2, 4]) +def test_fused_matches_unfused_multi_kv_head(nkv): + """Both writers must agree on page-major/head-minor placement. + + The AITER sparse-PA prototype only runs num_kv_heads == 1 per rank today, so + this pins the layout contract that lifting that restriction depends on. + """ + _check_matches_unfused(8 * nkv, nkv, 1, 64, 1e-6, True, True, 40) + + +def _check_matches_unfused( + nq, nkv, niq, rotary_dim, eps, quantized, fp8_index, num_tokens +): + fp8 = get_dtype_fp8() + cache_dtype = fp8 if quantized else torch.bfloat16 + idx_dtype = fp8 if fp8_index else torch.bfloat16 + kv_cache_dtype = "fp8" if quantized else "auto" + k_scale = torch.tensor([0.8], device="cuda") if quantized else None + v_scale = torch.tensor([1.25], device="cuda") if quantized else None + + c = build_case(num_tokens, nq, nkv, niq, rotary_dim, cache_dtype, idx_dtype) + args = (nq, nkv, niq, rotary_dim, eps, kv_cache_dtype, k_scale, v_scale) + ref = run_unfused(c, *args) + got = run_fused(c, *args) + + for name, a, b in zip( + ("q_out", "index_q_out", "k_cache", "v_cache", "index_cache"), ref, got + ): + assert_same(name, a, b) + + +@pytest.mark.parametrize("nkv", [2, 4]) +def test_gluon_head_folding_matches_per_head_decode(nkv): + """Reader-side check for num_kv_heads > 1. + + ``_run_gluon_decode`` folds (page, kv_head) into gluon's page dim with the + head minor, so a folded run with page ids ``page * nkv + head`` must equal + running each head separately against its own slice of the cache. This is the + other half of the layout contract needed to lift the one-KV-head limit. + """ + sparse_pa = pytest.importorskip("vllm.models.minimax_m3.amd.ops.sparse_pa") + + group, num_tokens, pages_used = 8, 6, 3 + nq = group * nkv + fp8 = get_dtype_fp8() + c = build_case(num_tokens, nq, nkv, 1, 64, fp8, fp8, seed=11) + k_scale = torch.tensor([1.0], device="cuda") + v_scale = torch.tensor([1.0], device="cuda") + q, _, k_cache, v_cache, _ = run_fused( + c, nq, nkv, 1, 64, 1e-6, "fp8", k_scale, v_scale + ) + q = q.view(num_tokens, nq, HEAD_DIM) + + num_pages = k_cache.shape[0] + ctx = pages_used * PAGE_SIZE + # Logical page choice per token, shared by every kv head (index_k is one head). + pages = ( + torch.arange(num_tokens * pages_used, device="cuda", dtype=torch.int32) + % num_pages + ).view(num_tokens, pages_used) + ctx_lens = torch.full((num_tokens,), ctx, device="cuda", dtype=torch.int32) + + # Folded: one row per (token, kv head), page ids scaled by nkv. + folded_bt = pages.repeat_interleave(nkv, dim=0) * nkv + torch.arange( + nkv, device="cuda", dtype=torch.int32 + ).repeat(num_tokens).unsqueeze(1) + folded_out = torch.empty( + num_tokens, nq, HEAD_DIM, device="cuda", dtype=torch.bfloat16 + ) + sparse_pa._run_gluon_decode( + q, + k_cache, + v_cache, + folded_bt, + ctx_lens.repeat_interleave(nkv), + nkv, + HEAD_DIM**-0.5, + folded_out, + k_scale, + v_scale, + ) + torch.cuda.synchronize() + + for h in range(nkv): + per_head_out = torch.empty( + num_tokens, group, HEAD_DIM, device="cuda", dtype=torch.bfloat16 + ) + sparse_pa._run_gluon_decode( + q[:, h * group : (h + 1) * group].contiguous(), + k_cache[:, h : h + 1].contiguous(), + v_cache[:, h : h + 1].contiguous(), + pages, + ctx_lens, + 1, + HEAD_DIM**-0.5, + per_head_out, + k_scale, + v_scale, + ) + torch.cuda.synchronize() + assert per_head_out.abs().sum().item() > 0, "per-head decode produced zeros" + assert torch.equal( + folded_out[:, h * group : (h + 1) * group], per_head_out + ), f"folded gluon decode differs from per-head decode for kv head {h}" + + +@pytest.mark.parametrize("quantized", [False, True]) +def test_fused_cache_feeds_gluon_decode_identically(quantized): + """The fused cache must drive the production gluon decode to the same output.""" + sparse_pa = pytest.importorskip("vllm.models.minimax_m3.amd.ops.sparse_pa") + + nq, nkv, niq, rotary_dim, eps = 16, 1, 1, 64, 1e-6 + num_tokens = 8 + fp8 = get_dtype_fp8() + cache_dtype = fp8 if quantized else torch.bfloat16 + k_scale = torch.tensor([1.0], device="cuda") if quantized else None + v_scale = torch.tensor([1.0], device="cuda") if quantized else None + c = build_case(num_tokens, nq, nkv, niq, rotary_dim, cache_dtype, fp8, seed=3) + + args = ( + nq, + nkv, + niq, + rotary_dim, + eps, + "fp8" if quantized else "auto", + k_scale, + v_scale, + ) + q_ref, _, k_ref, v_ref, _ = run_unfused(c, *args) + q_got, _, k_got, v_got, _ = run_fused(c, *args) + + # Every token attends to a fixed window of physical pages, so the decode + # touches slots both paths wrote. + num_pages = k_ref.shape[0] + pages_used = 4 + ctx = pages_used * PAGE_SIZE + sparse_bt = ( + torch.arange(num_tokens * pages_used, device="cuda", dtype=torch.int32) + % num_pages + ).view(num_tokens, pages_used) + sparse_ctx = torch.full((num_tokens,), ctx, device="cuda", dtype=torch.int32) + + def decode(q, k_cache, v_cache): + out = torch.empty(num_tokens, nq, HEAD_DIM, device="cuda", dtype=torch.bfloat16) + sparse_pa._run_gluon_decode( + q.view(num_tokens, nq, HEAD_DIM), + k_cache, + v_cache, + sparse_bt, + sparse_ctx, + nkv, + HEAD_DIM**-0.5, + out, + k_scale, + v_scale, + ) + torch.cuda.synchronize() + return out + + out_ref = decode(q_ref, k_ref, v_ref) + out_got = decode(q_got, k_got, v_got) + assert out_ref.abs().sum().item() > 0, "decode produced all zeros" + assert torch.equal(out_ref, out_got), "gluon decode differs on the fused cache" + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"]))