forked from Karylab-cklius/vllm
[Model] Add fmha_sm100 MSA indexer backend + fp8 index cache for MiniMax M3
Add an SM100/Blackwell lightning-indexer impl that computes the per-128-block QK max-scores with fmha_sm100's score-only (OnlyScore) path and selects the top-k blocks with the existing Triton minimax_m3_index_topk kernel, mirroring how the main MSA attention pairs the SM100 attend with Triton. Decode and prefill requests are split manually (decode-first batch) and each side gets its own _fmha_sm100_plan / _fmha_sm100 call. Auto-selected on SM100 when topk_blocks in (4, 8, 16, 32) for both bf16 and fp8 index caches; falls back to the Triton indexer otherwise. The builder declares AttentionCGSupport.NEVER (eager; the attention is broken out of the graph by _run_attention). Extend the fused qknorm+rope+kv-insert kernel to optionally emit fp8 (e4m3) for the index-K cache and index-Q via a direct cast with no scale tensors (RMSNorm outputs are O(1) and scalar scales do not change top-k ordering). Only the index outputs go fp8; q/k/v and q_out stay bf16 and bit-identical to the existing path. MiniMaxM3IndexerCache now accepts fp8 caches and the model allocates index_q in the cache dtype. Tests: test_fmha_sm100_indexer_matches_reference (bf16/fp8 x prefill/decode) and test_msa_indexer_impl_matches_triton (full impl parity vs the Triton indexer through the real metadata builders); fp8 fused-kernel parity is covered in test_fused_minimax_m3_qknorm_rope_kv_insert. AI assistance (Claude Code) was used for this change. Signed-off-by: Yongye Zhu <yongye@inferact.ai> Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
This commit is contained in:
Yongye Zhu
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
4c62663315
commit
bb4844dba3
@@ -67,6 +67,13 @@
|
||||
#include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
|
||||
#endif
|
||||
|
||||
// Direct float -> E4M3 FP8 conversion for the indexer Q / index-K outputs.
|
||||
#ifndef USE_ROCM
|
||||
#include <cuda_fp8.h>
|
||||
#else
|
||||
#include <hip/hip_fp8.h>
|
||||
#endif
|
||||
|
||||
#ifndef FINAL_MASK
|
||||
#ifdef USE_ROCM
|
||||
#define FINAL_MASK 0xffffffffffffffffULL
|
||||
@@ -75,6 +82,19 @@
|
||||
#endif
|
||||
#endif
|
||||
|
||||
#ifdef USE_ROCM
|
||||
// ROCm-compatible direct float -> E4M3 FP8 conversion (mirrors the DeepSeek V4
|
||||
// fused kernel).
|
||||
__device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) {
|
||||
#if defined(HIP_FP8_TYPE_OCP)
|
||||
__hip_fp8_e4m3 fp8_val(val);
|
||||
#else
|
||||
__hip_fp8_e4m3_fnuz fp8_val(val);
|
||||
#endif
|
||||
return reinterpret_cast<uint8_t&>(fp8_val);
|
||||
}
|
||||
#endif
|
||||
|
||||
namespace vllm {
|
||||
namespace minimax_m3_fused_ops {
|
||||
|
||||
@@ -193,6 +213,8 @@ __device__ __forceinline__ void storeElems(
|
||||
*reinterpret_cast<uint2*>(dst) = v;
|
||||
}
|
||||
|
||||
// Main K/V cache store. kAuto = unquantized (cache_t == scalar_t); fp8 cache
|
||||
// dtypes use the scaled-convert path with identity scale.
|
||||
template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
|
||||
__device__ __forceinline__ void storeCacheElems(
|
||||
cache_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
|
||||
@@ -208,6 +230,32 @@ __device__ __forceinline__ void storeCacheElems(
|
||||
}
|
||||
}
|
||||
|
||||
// Store 4 fp32 registers -> 4 contiguous E4M3 FP8 bytes (direct cast,
|
||||
// saturating to ±448). Used for the fp8 indexer-Q / index-K outputs; no scale
|
||||
// (RMSNorm outputs are O(1) and the score path only needs relative block
|
||||
// ordering).
|
||||
__device__ __forceinline__ void storeElemsFp8(
|
||||
uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
|
||||
constexpr float kFp8Max = 448.0f;
|
||||
#ifndef USE_ROCM
|
||||
__nv_fp8x2_storage_t out2[kElemsPerLane / 2];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane / 2; i++) {
|
||||
float2 vv = make_float2(elems[2 * i], elems[2 * i + 1]);
|
||||
vv.x = fminf(fmaxf(vv.x, -kFp8Max), kFp8Max);
|
||||
vv.y = fminf(fmaxf(vv.y, -kFp8Max), kFp8Max);
|
||||
out2[i] = __nv_cvt_float2_to_fp8x2(vv, __NV_SATFINITE, __NV_E4M3);
|
||||
}
|
||||
*reinterpret_cast<uint32_t*>(dst) = *reinterpret_cast<uint32_t const*>(out2);
|
||||
#else
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
float vv = fminf(fmaxf(elems[i], -kFp8Max), kFp8Max);
|
||||
dst[i] = rocm_cvt_float_to_fp8_e4m3(vv);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Kernel
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
@@ -224,12 +272,14 @@ __device__ __forceinline__ void storeCacheElems(
|
||||
// V : nkv only if kInsertKV (V-cache insert; no warps in dense)
|
||||
// IQ: niq only if kIsSparse (norm+RoPE)
|
||||
// IK: 1 only if kIsSparse (norm+RoPE; +index-cache insert)
|
||||
// cache_t/kv_dt: main attention KV-cache dtype (auto/fp8). out_idx_t/kFp8Idx:
|
||||
// indexer index-K cache + index-Q output dtype (scalar_t or e4m3 byte).
|
||||
template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
|
||||
bool kIsSparse, bool kInsertKV>
|
||||
typename out_idx_t, bool kIsSparse, bool kInsertKV, bool kFp8Idx>
|
||||
__global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
|
||||
scalar_t* __restrict__ qkv, // [N, qkv_row] in/out (packs index if sparse)
|
||||
scalar_t* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr
|
||||
scalar_t* __restrict__ index_q_out, // [N, niq*128] contiguous, or nullptr
|
||||
scalar_t* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr
|
||||
out_idx_t* __restrict__ index_q_out, // [N, niq*128]; scalar_t or e4m3 byte
|
||||
scalar_t const* __restrict__ q_norm_w,
|
||||
scalar_t const* __restrict__ k_norm_w,
|
||||
scalar_t const* __restrict__ iq_norm_w,
|
||||
@@ -238,8 +288,8 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
|
||||
int64_t const* __restrict__ positions, // [N] i64
|
||||
int64_t const* __restrict__ slot_mapping, // main K/V slots or nullptr
|
||||
int64_t const* __restrict__ index_slot_mapping, // index K slots/nullptr
|
||||
cache_t* __restrict__ kv_cache, // [nb,2,bs,nkv,128] or nullptr
|
||||
scalar_t* __restrict__ index_cache, // [nb*bs, 128] or nullptr
|
||||
cache_t* __restrict__ kv_cache, // [nb,2,bs,nkv,128] or nullptr
|
||||
out_idx_t* __restrict__ index_cache, // [nb*bs, 128]; scalar_t or e4m3 byte
|
||||
float const eps, int const rotary_dim, int const num_tokens, int const nq,
|
||||
int const nkv, int const niq, int const block_size,
|
||||
// kv_cache strides (in elements) for logical shape [nb, 2, bs, nkv, 128].
|
||||
@@ -334,9 +384,12 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
|
||||
store_ptr = q_out + static_cast<int64_t>(tokenIdx) * nq * kHeadDim +
|
||||
slot * kHeadDim;
|
||||
} else if (isIQ && index_q_out != nullptr) {
|
||||
store_ptr = index_q_out +
|
||||
static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
|
||||
(slot - iq_begin) * kHeadDim;
|
||||
// bf16 index_q_out: gather here. fp8: written by the explicit fp8 store.
|
||||
if constexpr (!kFp8Idx) {
|
||||
store_ptr = index_q_out +
|
||||
static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
|
||||
(slot - iq_begin) * kHeadDim;
|
||||
}
|
||||
}
|
||||
|
||||
// PDL: wait for the predecessor kernel (the qkv-projection GEMM that
|
||||
@@ -356,7 +409,19 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
|
||||
scalar_t const* cos_ptr = cos_sin_cache + pos * rotary_dim;
|
||||
normAndRope<scalar_t>(elems, laneId, eps, norm_w, do_rope, rotary_dim,
|
||||
cos_ptr, /*apply_norm=*/norm_w != nullptr);
|
||||
storeElems<scalar_t>(store_ptr + dim_base, elems);
|
||||
if constexpr (kFp8Idx) {
|
||||
// index_q is e4m3 bytes; Q/K (and in-place index_k) stay scalar_t.
|
||||
if (isIQ && index_q_out != nullptr) {
|
||||
storeElemsFp8(index_q_out +
|
||||
static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
|
||||
(slot - iq_begin) * kHeadDim + dim_base,
|
||||
elems);
|
||||
} else {
|
||||
storeElems<scalar_t>(store_ptr + dim_base, elems);
|
||||
}
|
||||
} else {
|
||||
storeElems<scalar_t>(store_ptr + dim_base, elems);
|
||||
}
|
||||
}
|
||||
|
||||
// ── Cache inserts (sparse serving only). ───────────────────────────────
|
||||
@@ -367,8 +432,11 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
|
||||
: (isIK ? index_slot_mapping[tokenIdx] : -1);
|
||||
if (sm >= 0) { // skip padded / unscheduled tokens
|
||||
if (isIK) {
|
||||
scalar_t* dst = index_cache + sm * kHeadDim + dim_base;
|
||||
storeElems<scalar_t>(dst, elems);
|
||||
if constexpr (kFp8Idx) {
|
||||
storeElemsFp8(index_cache + sm * kHeadDim + dim_base, elems);
|
||||
} else {
|
||||
storeElems<scalar_t>(index_cache + sm * kHeadDim + dim_base, elems);
|
||||
}
|
||||
} else if (isK || isV) {
|
||||
// kv_cache logical shape [num_blocks, 2, block_size, nkv, head_dim].
|
||||
// Paging is logical (block = sm/block_size, token = sm%block_size);
|
||||
@@ -398,19 +466,19 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
|
||||
// Launch wrapper
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
|
||||
void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
|
||||
scalar_t const* q_norm_w, scalar_t const* k_norm_w,
|
||||
scalar_t const* iq_norm_w, scalar_t const* ik_norm_w,
|
||||
scalar_t const* cos_sin_cache,
|
||||
int64_t const* positions, int64_t const* slot_mapping,
|
||||
int64_t const* index_slot_mapping, cache_t* kv_cache,
|
||||
scalar_t* index_cache, float const eps,
|
||||
int const rotary_dim, int const num_tokens,
|
||||
int const nq, int const nkv, int const niq,
|
||||
int const block_size, int64_t const kv_s_block,
|
||||
int64_t const kv_s_kv, int64_t const kv_s_token,
|
||||
int64_t const kv_s_head, bool const has_index,
|
||||
bool const insert_kv, cudaStream_t stream) {
|
||||
void launchFusedMiniMaxM3(
|
||||
scalar_t* qkv, scalar_t* q_out, void* index_q_out, scalar_t const* q_norm_w,
|
||||
scalar_t const* k_norm_w, scalar_t const* iq_norm_w,
|
||||
scalar_t const* ik_norm_w, scalar_t const* cos_sin_cache,
|
||||
int64_t const* positions, int64_t const* slot_mapping,
|
||||
int64_t const* index_slot_mapping, cache_t* kv_cache, void* index_cache,
|
||||
float const eps, int const rotary_dim, int const num_tokens, int const nq,
|
||||
int const nkv, int const niq, int const block_size,
|
||||
int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token,
|
||||
int64_t const kv_s_head, bool const has_index, bool const insert_kv,
|
||||
bool const fp8_idx, cudaStream_t stream) {
|
||||
// Index outputs are scalar_t (bf16) or e4m3 bytes (uint8_t); reinterpret the
|
||||
// void* pointers per instantiation in the LAUNCH macro.
|
||||
// Slot count must match the kernel's compile-time gating.
|
||||
int const v_slots = insert_kv ? nkv : 0;
|
||||
int const idx_slots = has_index ? niq + 1 : 0;
|
||||
@@ -440,25 +508,27 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
|
||||
config.attrs = attrs;
|
||||
config.numAttrs = (sm_version >= 90) ? 1 : 0;
|
||||
|
||||
#define LAUNCH(IS_SPARSE, INSERT) \
|
||||
cudaLaunchKernelEx( \
|
||||
&config, \
|
||||
fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, \
|
||||
IS_SPARSE, INSERT>, \
|
||||
qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, \
|
||||
cos_sin_cache, positions, slot_mapping, index_slot_mapping, kv_cache, \
|
||||
index_cache, eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, \
|
||||
kv_s_block, kv_s_kv, kv_s_token, kv_s_head)
|
||||
#define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T) \
|
||||
cudaLaunchKernelEx( \
|
||||
&config, \
|
||||
fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T, \
|
||||
IS_SPARSE, INSERT, FP8>, \
|
||||
qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, k_norm_w, \
|
||||
iq_norm_w, ik_norm_w, cos_sin_cache, positions, slot_mapping, \
|
||||
index_slot_mapping, kv_cache, reinterpret_cast<OUT_T*>(index_cache), \
|
||||
eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \
|
||||
kv_s_kv, kv_s_token, kv_s_head)
|
||||
#else
|
||||
// ROCm: standard kernel launch syntax (no PDL/stream serialization).
|
||||
// clang-format off
|
||||
#define LAUNCH(IS_SPARSE, INSERT) \
|
||||
fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, \
|
||||
IS_SPARSE, INSERT> \
|
||||
#define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T) \
|
||||
fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T, \
|
||||
IS_SPARSE, INSERT, FP8> \
|
||||
<<<grid, kBlockSize, 0, stream>>>( \
|
||||
qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, \
|
||||
ik_norm_w, cos_sin_cache, positions, slot_mapping, \
|
||||
index_slot_mapping, kv_cache, index_cache, eps, rotary_dim, \
|
||||
qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, \
|
||||
k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
|
||||
slot_mapping, index_slot_mapping, kv_cache, \
|
||||
reinterpret_cast<OUT_T*>(index_cache), eps, rotary_dim, \
|
||||
num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_kv, \
|
||||
kv_s_token, kv_s_head)
|
||||
// clang-format on
|
||||
@@ -466,14 +536,22 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
|
||||
|
||||
if (has_index) {
|
||||
if (insert_kv) {
|
||||
LAUNCH(true, true); // sparse serving
|
||||
if (fp8_idx) {
|
||||
LAUNCH(true, true, true, uint8_t); // sparse serving, fp8 index outputs
|
||||
} else {
|
||||
LAUNCH(true, true, false, scalar_t); // sparse serving, bf16
|
||||
}
|
||||
} else {
|
||||
LAUNCH(true, false); // sparse profiling
|
||||
if (fp8_idx) {
|
||||
LAUNCH(true, false, true, uint8_t); // sparse profiling, fp8 index_q
|
||||
} else {
|
||||
LAUNCH(true, false, false, scalar_t); // sparse profiling, bf16
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Dense layer: never has an index branch and never inserts here (the
|
||||
// generic Attention layer owns the KV insert).
|
||||
LAUNCH(false, false);
|
||||
LAUNCH(false, false, false, scalar_t);
|
||||
}
|
||||
#undef LAUNCH
|
||||
}
|
||||
@@ -485,8 +563,9 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
|
||||
vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st, CACHE_T, KV_DTYPE>( \
|
||||
reinterpret_cast<st*>(qkv.data_ptr()), \
|
||||
q_out.has_value() ? reinterpret_cast<st*>(q_out->data_ptr()) : nullptr, \
|
||||
index_q_out.has_value() ? reinterpret_cast<st*>(index_q_out->data_ptr()) \
|
||||
: nullptr, \
|
||||
index_q_out.has_value() \
|
||||
? reinterpret_cast<void*>(index_q_out->data_ptr()) \
|
||||
: nullptr, \
|
||||
reinterpret_cast<st const*>(q_norm_weight.data_ptr()), \
|
||||
reinterpret_cast<st const*>(k_norm_weight.data_ptr()), \
|
||||
has_index ? reinterpret_cast<st const*>(index_q_norm_weight->data_ptr()) \
|
||||
@@ -502,11 +581,11 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
|
||||
: nullptr, \
|
||||
insert_kv ? reinterpret_cast<CACHE_T*>(kv_cache->data_ptr()) : nullptr, \
|
||||
(insert_kv && has_index) \
|
||||
? reinterpret_cast<st*>(index_cache->data_ptr()) \
|
||||
? reinterpret_cast<void*>(index_cache->data_ptr()) \
|
||||
: nullptr, \
|
||||
static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens, nq, \
|
||||
nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_kv, kv_s_token, \
|
||||
kv_s_head, has_index, insert_kv, stream)
|
||||
kv_s_head, has_index, insert_kv, fp8_idx, stream)
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Torch op wrapper
|
||||
@@ -612,6 +691,7 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
|
||||
torch::headeronly::ScalarType::Long &&
|
||||
index_slot_mapping->numel() == slot_mapping->numel()),
|
||||
"index_slot_mapping must be int64 CUDA with slot_mapping length");
|
||||
// Main attention KV cache: auto matches qkv, fp8 uses uint8 storage.
|
||||
if (kv_dt == vllm::Fp8KVCacheDataType::kAuto) {
|
||||
STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
|
||||
"auto kv_cache dtype must match qkv");
|
||||
@@ -620,9 +700,13 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
|
||||
kv_cache->scalar_type() == torch::headeronly::ScalarType::Byte,
|
||||
"fp8 kv_cache must use uint8 storage");
|
||||
}
|
||||
STD_TORCH_CHECK(index_cache.has_value() &&
|
||||
index_cache->scalar_type() == qkv.scalar_type(),
|
||||
"insert mode requires matching index_cache");
|
||||
// Indexer index-K cache: independent dtype -- qkv dtype or fp8 e4m3.
|
||||
STD_TORCH_CHECK(
|
||||
index_cache.has_value() &&
|
||||
(index_cache->scalar_type() == qkv.scalar_type() ||
|
||||
index_cache->scalar_type() ==
|
||||
torch::headeronly::ScalarType::Float8_e4m3fn),
|
||||
"insert mode requires index_cache matching qkv dtype or fp8 e4m3");
|
||||
STD_TORCH_CHECK(kv_cache->dim() == 5 && kv_cache->stride(4) == 1,
|
||||
"kv_cache must be [nb,2,bs,nkv,head_dim] with contiguous "
|
||||
"head_dim (stride(4)==1)");
|
||||
@@ -652,14 +736,31 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
|
||||
"index_q_out requires the index branch (num_index_heads > 0)");
|
||||
STD_TORCH_CHECK(
|
||||
index_q_out->is_cuda() && index_q_out->is_contiguous() &&
|
||||
index_q_out->scalar_type() == qkv.scalar_type(),
|
||||
"index_q_out must be a contiguous CUDA tensor matching qkv dtype");
|
||||
(index_q_out->scalar_type() == qkv.scalar_type() ||
|
||||
index_q_out->scalar_type() ==
|
||||
torch::headeronly::ScalarType::Float8_e4m3fn),
|
||||
"index_q_out must be contiguous CUDA, qkv dtype or fp8 e4m3");
|
||||
STD_TORCH_CHECK(index_q_out->numel() ==
|
||||
static_cast<int64_t>(num_tokens) * niq * kHeadDim,
|
||||
"index_q_out must have num_tokens * num_index_heads * 128 "
|
||||
"elements");
|
||||
}
|
||||
|
||||
// fp8 index path: the index-K cache and index-Q outputs are e4m3 bytes while
|
||||
// q/k/v + q_out stay qkv dtype. Both index outputs must agree.
|
||||
auto const kFp8 = torch::headeronly::ScalarType::Float8_e4m3fn;
|
||||
bool const fp8_idx =
|
||||
(index_cache.has_value() && index_cache->scalar_type() == kFp8) ||
|
||||
(index_q_out.has_value() && index_q_out->scalar_type() == kFp8);
|
||||
if (fp8_idx) {
|
||||
STD_TORCH_CHECK(
|
||||
!index_cache.has_value() || index_cache->scalar_type() == kFp8,
|
||||
"fp8 index path: index_cache must be fp8 e4m3");
|
||||
STD_TORCH_CHECK(
|
||||
!index_q_out.has_value() || index_q_out->scalar_type() == kFp8,
|
||||
"fp8 index path: index_q_out must be fp8 e4m3");
|
||||
}
|
||||
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
qkv.get_device_index());
|
||||
auto stream = get_current_cuda_stream(qkv.get_device_index());
|
||||
|
||||
@@ -244,6 +244,251 @@ def test_prefill_index_topk_correctness():
|
||||
_assert_topk_indices_equal_unordered(actual, expected)
|
||||
|
||||
|
||||
# MSA indexer (SM100): fmha_sm100 OnlyScore for the per-block scores, then the
|
||||
# Triton minimax_m3_index_topk for selection (no sparse_topk_select). Uses a
|
||||
# deterministic construction (idx_q == 1, distinct e4m3-exact per-block values)
|
||||
# so scores are strictly monotonic in the block id -> exact top-k agreement.
|
||||
def _fmha_indexer_topk(
|
||||
idx_q: torch.Tensor, # [total_q, H, 128] bf16/e4m3
|
||||
index_cache: torch.Tensor, # [num_pages, 128, 128] bf16/e4m3
|
||||
block_table: torch.Tensor,
|
||||
q_lens: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
prefix_lens: torch.Tensor,
|
||||
sm_scale: float,
|
||||
topk: int,
|
||||
) -> torch.Tensor:
|
||||
"""Replicate MiniMaxM3IndexerMSAImpl's score path (single decode/prefill side)."""
|
||||
from vllm.third_party.fmha_sm100.api import _fmha_sm100, _fmha_sm100_plan
|
||||
|
||||
num_idx_heads, head_dim = idx_q.shape[1], idx_q.shape[2]
|
||||
nvp = [(s + 127) // 128 for s in seq_lens.tolist()]
|
||||
kv_indices = torch.cat([block_table[r, : nvp[r]] for r in range(len(nvp))]).to(
|
||||
torch.int32
|
||||
)
|
||||
|
||||
qo = q_lens.cpu().to(torch.int32)
|
||||
kv = seq_lens.cpu().to(torch.int32)
|
||||
plan = _fmha_sm100_plan(
|
||||
qo,
|
||||
kv,
|
||||
num_idx_heads,
|
||||
num_kv_heads=1,
|
||||
qo_offset=kv - qo,
|
||||
page_size=128,
|
||||
output_maxscore=True,
|
||||
causal=True,
|
||||
num_kv_splits=1,
|
||||
)
|
||||
k_pages = index_cache.view(index_cache.shape[0], 1, 128, head_dim)
|
||||
_, max_score = _fmha_sm100(
|
||||
idx_q,
|
||||
k_pages,
|
||||
k_pages,
|
||||
plan,
|
||||
kv_indices=kv_indices,
|
||||
output_o=False,
|
||||
output_maxscore=True,
|
||||
sm_scale=sm_scale,
|
||||
)
|
||||
|
||||
batch = q_lens.numel()
|
||||
cu = torch.zeros(batch + 1, dtype=torch.int32, device=idx_q.device)
|
||||
cu[1:] = q_lens.to(torch.int32).cumsum(0)
|
||||
# max_score [H, k_tiles, total_q] -> transpose to [H, total_q, k_tiles].
|
||||
return minimax_m3_index_topk(
|
||||
max_score.transpose(1, 2),
|
||||
cu,
|
||||
prefix_lens.to(torch.int32),
|
||||
int(q_lens.max()),
|
||||
topk,
|
||||
0, # init_blocks
|
||||
0, # local_blocks
|
||||
)
|
||||
|
||||
|
||||
# e4m3-exact, strictly-increasing per-block values: with idx_q == 1 (also exact)
|
||||
# the per-block scores are exact and distinct in BOTH bf16 and e4m3, so the fp8
|
||||
# score path selects the same top-k as the reference (no quantization ties).
|
||||
_E4M3_EXACT_VALUES = [
|
||||
*range(1, 17), # 1..16 (step 1)
|
||||
*range(18, 33, 2), # 18..32 (step 2)
|
||||
*range(36, 65, 4), # 36..64 (step 4)
|
||||
*range(72, 129, 8), # 72..128 (step 8)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_device_capability_family(100),
|
||||
reason="fmha_sm100 indexer requires SM100 (Blackwell).",
|
||||
)
|
||||
@pytest.mark.parametrize("index_dtype", [torch.bfloat16, torch.float8_e4m3fn])
|
||||
@pytest.mark.parametrize(
|
||||
("q_lens", "prefix_lens"),
|
||||
[
|
||||
((4, 3), (2048, 2560)), # prefill: every token sees >= 16 causal blocks
|
||||
((1, 1, 1), (2048, 3000, 4096)), # decode: one query token per request
|
||||
],
|
||||
)
|
||||
def test_fmha_sm100_indexer_matches_reference(q_lens, prefix_lens, index_dtype):
|
||||
torch.manual_seed(0)
|
||||
num_idx_heads, head_dim = 4, HEAD_DIM
|
||||
device = "cuda"
|
||||
|
||||
q_lens_t = torch.tensor(q_lens, device=device, dtype=torch.int32)
|
||||
prefix_lens_t = torch.tensor(prefix_lens, device=device, dtype=torch.int32)
|
||||
seq_lens = prefix_lens_t + q_lens_t
|
||||
batch = len(q_lens)
|
||||
max_blocks = (int(seq_lens.max()) + BLOCK_SIZE - 1) // BLOCK_SIZE
|
||||
assert max_blocks <= len(_E4M3_EXACT_VALUES)
|
||||
num_pages = batch * max_blocks
|
||||
block_table = torch.randperm(num_pages, device=device, dtype=torch.int32).reshape(
|
||||
batch, max_blocks
|
||||
)
|
||||
|
||||
idx_q = torch.ones(
|
||||
int(q_lens_t.sum()), num_idx_heads, head_dim, device=device, dtype=index_dtype
|
||||
)
|
||||
index_cache = torch.empty(
|
||||
num_pages, BLOCK_SIZE, head_dim, device=device, dtype=index_dtype
|
||||
)
|
||||
for r in range(batch):
|
||||
for b in range(max_blocks):
|
||||
index_cache[block_table[r, b]] = float(_E4M3_EXACT_VALUES[b])
|
||||
|
||||
sm_scale = head_dim**-0.5
|
||||
actual = _fmha_indexer_topk(
|
||||
idx_q,
|
||||
index_cache,
|
||||
block_table,
|
||||
q_lens_t,
|
||||
seq_lens,
|
||||
prefix_lens_t,
|
||||
sm_scale,
|
||||
TOPK,
|
||||
)
|
||||
expected = _reference_index_topk(
|
||||
idx_q,
|
||||
index_cache,
|
||||
block_table,
|
||||
q_lens_t,
|
||||
seq_lens,
|
||||
prefix_lens_t,
|
||||
TOPK,
|
||||
init_blocks=0,
|
||||
local_blocks=0,
|
||||
sm_scale=sm_scale,
|
||||
)
|
||||
_assert_topk_indices_equal_unordered(actual, expected)
|
||||
|
||||
|
||||
# Full impl-level parity: drive both MiniMaxM3IndexerMSAImpl (fmha_sm100 score +
|
||||
# Triton top-k) and MiniMaxM3IndexerTritonImpl through their real metadata
|
||||
# builders on the SAME CommonAttentionMetadata + index cache, and assert the
|
||||
# selected blocks agree. This exercises all the metadata the impl/kernels consume
|
||||
# (decode/prefill split, cu_seqlens_q rebasing, prefix_lens, kv_indices gather,
|
||||
# decode_pages split) -- a metadata bug on either side shifts the causal window
|
||||
# or the block->page mapping and breaks the comparison.
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_device_capability_family(100),
|
||||
reason="fmha_sm100 indexer requires SM100 (Blackwell).",
|
||||
)
|
||||
@pytest.mark.parametrize("topk", [8, 16])
|
||||
def test_msa_indexer_impl_matches_triton(topk, monkeypatch):
|
||||
import vllm.models.minimax_m3.common.indexer as indexer_mod
|
||||
from tests.v1.attention.utils import (
|
||||
BatchSpec,
|
||||
create_common_attn_metadata,
|
||||
create_vllm_config,
|
||||
)
|
||||
from vllm.config import set_current_vllm_config
|
||||
from vllm.forward_context import set_forward_context
|
||||
from vllm.models.minimax_m3.common.indexer import (
|
||||
MiniMaxM3IndexerTritonImpl,
|
||||
MiniMaxM3IndexerTritonMetadataBuilder,
|
||||
)
|
||||
from vllm.models.minimax_m3.nvidia.indexer_msa import (
|
||||
MiniMaxM3IndexerMSAImpl,
|
||||
MiniMaxM3IndexerMSAMetadataBuilder,
|
||||
)
|
||||
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda")
|
||||
num_idx_heads, head_dim = 4, HEAD_DIM
|
||||
# TP=1: avoid requiring an initialized distributed group in a unit test.
|
||||
monkeypatch.setattr(indexer_mod, "get_tensor_model_parallel_world_size", lambda: 1)
|
||||
|
||||
vllm_config = create_vllm_config(
|
||||
block_size=BLOCK_SIZE, max_model_len=8192, max_num_batched_tokens=8192
|
||||
)
|
||||
vllm_config.model_config.hf_config.sparse_attention_config = {
|
||||
"sparse_num_index_heads": num_idx_heads
|
||||
}
|
||||
|
||||
# Decode-first mixed batch: 2 decode reqs (q_len 1) then 2 prefill reqs. Long
|
||||
# prefixes so every token sees > TOPK causal blocks (non-trivial selection).
|
||||
batch = BatchSpec(seq_lens=[2305, 2561, 2624, 2720], query_lens=[1, 1, 64, 96])
|
||||
common = create_common_attn_metadata(
|
||||
batch, BLOCK_SIZE, device, arange_block_indices=True
|
||||
)
|
||||
num_tokens = batch.compute_num_tokens()
|
||||
|
||||
# Deterministic index cache: distinct, monotonic per-logical-block values so
|
||||
# the top-k is unambiguous (both kernels pick the same blocks, no fp ties).
|
||||
block_table = common.block_table_tensor
|
||||
num_pages = int(block_table.max().item()) + 1
|
||||
index_cache = torch.zeros(
|
||||
num_pages, BLOCK_SIZE, head_dim, device=device, dtype=DTYPE
|
||||
)
|
||||
for r, seq_len in enumerate(batch.seq_lens):
|
||||
for b in range((seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE):
|
||||
index_cache[block_table[r, b]] = float(b + 1)
|
||||
index_q = torch.ones(
|
||||
num_tokens, num_idx_heads * head_dim, device=device, dtype=DTYPE
|
||||
)
|
||||
|
||||
spec = MLAAttentionSpec(
|
||||
block_size=BLOCK_SIZE, num_kv_heads=1, head_size=head_dim, dtype=DTYPE
|
||||
)
|
||||
impl_kwargs = dict(
|
||||
num_kv_heads=num_idx_heads,
|
||||
scale=head_dim**-0.5,
|
||||
topk_blocks=topk,
|
||||
sparse_block_size=BLOCK_SIZE,
|
||||
num_index_heads=num_idx_heads,
|
||||
index_head_dim=head_dim,
|
||||
init_blocks=0,
|
||||
local_blocks=0,
|
||||
)
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
msa_impl = MiniMaxM3IndexerMSAImpl(prefix="idx_msa", **impl_kwargs)
|
||||
triton_impl = MiniMaxM3IndexerTritonImpl(prefix="idx_triton", **impl_kwargs)
|
||||
msa_builder = MiniMaxM3IndexerMSAMetadataBuilder(
|
||||
spec, [msa_impl.index_cache.prefix], vllm_config, device
|
||||
)
|
||||
triton_builder = MiniMaxM3IndexerTritonMetadataBuilder(
|
||||
spec, [triton_impl.index_cache.prefix], vllm_config, device
|
||||
)
|
||||
|
||||
# Both impls score against the same index keys.
|
||||
msa_impl.index_cache.kv_cache = index_cache
|
||||
triton_impl.index_cache.kv_cache = index_cache
|
||||
|
||||
attn_metadata = {
|
||||
msa_impl.index_cache.prefix: msa_builder.build(0, common),
|
||||
triton_impl.index_cache.prefix: triton_builder.build(0, common),
|
||||
}
|
||||
with set_forward_context(attn_metadata, vllm_config):
|
||||
msa_decode, msa_prefill = msa_impl(index_q)
|
||||
tri_decode, tri_prefill = triton_impl(index_q)
|
||||
|
||||
assert msa_decode is not None and tri_decode is not None
|
||||
assert msa_prefill is not None and tri_prefill is not None
|
||||
_assert_topk_indices_equal_unordered(msa_decode, tri_decode)
|
||||
_assert_topk_indices_equal_unordered(msa_prefill, tri_prefill)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("decode_query_len", "max_decode_query_len"),
|
||||
[
|
||||
|
||||
@@ -278,3 +278,99 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
|
||||
torch.testing.assert_close(
|
||||
index_cache.view(-1, HEAD_DIM), expected_index_cache, rtol=0, atol=0
|
||||
)
|
||||
|
||||
|
||||
# ── Test 3: fp8 (e4m3) index outputs ─────────────────────────────────────────
|
||||
# The fp8 score path stores index_q and the index-K cache as e4m3 while q/k/v +
|
||||
# q_out stay bf16. Asserts: (1) q/k/v/q_out are bit-identical to the bf16 run
|
||||
# (the index dtype must not perturb the main branch), and (2) the e4m3 index
|
||||
# outputs dequantize close to the bf16 reference.
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9),
|
||||
reason="e4m3 conversion requires CUDA SM89+.",
|
||||
)
|
||||
@pytest.mark.parametrize("num_tokens", [1, 7, 64, 513])
|
||||
@pytest.mark.parametrize("block_size", [16, 64])
|
||||
def test_sparse_full_fp8_index(num_tokens, block_size):
|
||||
torch.manual_seed(1)
|
||||
device, dtype, eps = "cuda", torch.bfloat16, 1e-6
|
||||
base, max_pos = 5_000_000.0, 4096
|
||||
num_heads, num_kv_heads, num_idx_heads = 16, 4, 4
|
||||
|
||||
q_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
|
||||
k_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
|
||||
iq_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
|
||||
ik_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
|
||||
cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device)
|
||||
positions = torch.randint(
|
||||
0, max_pos, (num_tokens,), dtype=torch.int64, device=device
|
||||
)
|
||||
|
||||
qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM
|
||||
iqsz, iksz = num_idx_heads * HEAD_DIM, HEAD_DIM
|
||||
qkv0 = torch.randn(
|
||||
num_tokens, qsz + 2 * kvsz + iqsz + iksz, dtype=dtype, device=device
|
||||
)
|
||||
|
||||
num_blocks = (num_tokens + block_size - 1) // block_size + 1
|
||||
slot_mapping = torch.randperm(
|
||||
num_blocks * block_size, dtype=torch.int64, device=device
|
||||
)[:num_tokens]
|
||||
index_slot_mapping = torch.roll(slot_mapping, shifts=1)
|
||||
|
||||
def run(index_dtype):
|
||||
qkv = qkv0.clone()
|
||||
kv_cache = torch.zeros(
|
||||
num_blocks,
|
||||
2,
|
||||
block_size,
|
||||
num_kv_heads,
|
||||
HEAD_DIM,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
index_cache = torch.zeros(
|
||||
num_blocks, block_size, HEAD_DIM, dtype=index_dtype, device=device
|
||||
)
|
||||
q_out = torch.empty(num_tokens, qsz, dtype=dtype, device=device)
|
||||
index_q = torch.empty(num_tokens, iqsz, dtype=index_dtype, device=device)
|
||||
ops.fused_minimax_m3_qknorm_rope_kv_insert(
|
||||
qkv,
|
||||
q_w,
|
||||
k_w,
|
||||
cos_sin,
|
||||
positions,
|
||||
num_heads,
|
||||
num_kv_heads,
|
||||
ROTARY_DIM,
|
||||
eps,
|
||||
iq_w,
|
||||
ik_w,
|
||||
num_idx_heads,
|
||||
slot_mapping,
|
||||
index_slot_mapping,
|
||||
kv_cache,
|
||||
index_cache,
|
||||
block_size,
|
||||
q_out,
|
||||
index_q,
|
||||
)
|
||||
return qkv, kv_cache, index_cache, q_out, index_q
|
||||
|
||||
qkv_bf, kvc_bf, idxc_bf, qo_bf, iq_bf = run(torch.bfloat16)
|
||||
qkv_fp, kvc_fp, idxc_fp, qo_fp, iq_fp = run(torch.float8_e4m3fn)
|
||||
|
||||
assert iq_fp.dtype == torch.float8_e4m3fn
|
||||
assert idxc_fp.dtype == torch.float8_e4m3fn
|
||||
|
||||
# (1) The main branch (q/k/v in qkv, q_out, kv cache) must be bit-identical:
|
||||
# the index output dtype must not perturb anything else.
|
||||
torch.testing.assert_close(qo_fp, qo_bf, rtol=0, atol=0)
|
||||
torch.testing.assert_close(qkv_fp, qkv_bf, rtol=0, atol=0)
|
||||
torch.testing.assert_close(kvc_fp, kvc_bf, rtol=0, atol=0)
|
||||
|
||||
# (2) Dequantized e4m3 index outputs match the bf16 reference within fp8 ulp.
|
||||
torch.testing.assert_close(iq_fp.float(), iq_bf.float(), rtol=0.13, atol=0.05)
|
||||
torch.testing.assert_close(idxc_fp.float(), idxc_bf.float(), rtol=0.13, atol=0.05)
|
||||
|
||||
@@ -31,6 +31,7 @@ from vllm.models.minimax_m3.common.ops.index_topk import (
|
||||
minimax_m3_index_score,
|
||||
minimax_m3_index_topk,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.attention.backend import (
|
||||
AttentionBackend,
|
||||
AttentionCGSupport,
|
||||
@@ -120,16 +121,20 @@ class MiniMaxM3IndexerCache(nn.Module, AttentionLayerBase):
|
||||
backend_cls: type[AttentionBackend] = MiniMaxM3IndexerBackend,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if indexer_kv_dtype != "bf16":
|
||||
if indexer_kv_dtype in ("fp8", "fp8_e4m3"):
|
||||
cache_dtype = torch.float8_e4m3fn
|
||||
elif indexer_kv_dtype == "bf16":
|
||||
cache_dtype = torch.bfloat16
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported yet "
|
||||
"for the MiniMax M3 indexer cache (only 'bf16')."
|
||||
f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported by the "
|
||||
"MiniMax M3 indexer cache (only 'bf16' or 'fp8'/'fp8_e4m3')."
|
||||
)
|
||||
self.kv_cache = torch.tensor([])
|
||||
self.head_dim = head_dim
|
||||
self.indexer_kv_dtype = indexer_kv_dtype
|
||||
# Storage dtype for the side cache (bf16 today; quantized layouts later).
|
||||
self.dtype = torch.bfloat16
|
||||
# Side-cache storage dtype: bf16, or e4m3 for the fp8 score path.
|
||||
self.dtype = cache_dtype
|
||||
self.prefix = prefix
|
||||
self.cache_config = cache_config
|
||||
# Impl-chosen backend -> each impl gets its own builder (get_attn_backend).
|
||||
@@ -438,18 +443,33 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
|
||||
|
||||
def select_indexer_impl_cls(
|
||||
*,
|
||||
topk_blocks: int,
|
||||
indexer_kv_dtype: IndexerKVDType = "bf16",
|
||||
) -> type[MiniMaxM3IndexerImpl]:
|
||||
"""Pick the indexer impl off the index-cache dtype.
|
||||
"""Pick the indexer impl off the platform, top-k count, and cache dtype.
|
||||
|
||||
The SM100 MSA indexer score path is disabled for now; use the local Triton
|
||||
indexer. If re-enabled, add a NVIDIA-specific ``MiniMaxM3IndexerImpl`` here.
|
||||
On Blackwell (SM100) with ``topk_blocks`` in ``(4, 8, 16, 32)`` (matching the
|
||||
main MSA attend), the fmha_sm100 score path + Triton top-k is used for both
|
||||
bf16 and fp8 index caches. Everything else falls back to the Triton indexer
|
||||
(bf16 only).
|
||||
"""
|
||||
if indexer_kv_dtype in ("mxfp4", "nvfp4"):
|
||||
raise NotImplementedError(
|
||||
f"indexer_kv_dtype={indexer_kv_dtype!r} needs the (not-yet-added) "
|
||||
"CuteDSL indexer impl."
|
||||
)
|
||||
if (
|
||||
current_platform.is_cuda()
|
||||
and current_platform.is_device_capability_family(100)
|
||||
and topk_blocks in (4, 8, 16, 32)
|
||||
and indexer_kv_dtype in ("bf16", "fp8", "fp8_e4m3")
|
||||
):
|
||||
# Lazy import so AMD / non-SM100 never import fmha_sm100.
|
||||
from vllm.models.minimax_m3.nvidia.indexer_msa import (
|
||||
MiniMaxM3IndexerMSAImpl,
|
||||
)
|
||||
|
||||
return MiniMaxM3IndexerMSAImpl
|
||||
if indexer_kv_dtype != "bf16":
|
||||
raise NotImplementedError(
|
||||
f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported by the "
|
||||
@@ -483,6 +503,7 @@ class MiniMaxM3Indexer(nn.Module):
|
||||
) -> None:
|
||||
super().__init__()
|
||||
impl_cls = select_indexer_impl_cls(
|
||||
topk_blocks=topk_blocks,
|
||||
indexer_kv_dtype=indexer_kv_dtype,
|
||||
)
|
||||
self.impl = impl_cls(
|
||||
|
||||
@@ -0,0 +1,248 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""MSA (SM100/Blackwell) indexer impl for MiniMax M3.
|
||||
|
||||
The lightning indexer's per-128-block QK max-score is computed with
|
||||
``fmha_sm100``'s score-only (``OnlyScore``) path; the top-k block selection
|
||||
reuses the existing Triton ``minimax_m3_index_topk`` kernel (it handles the
|
||||
per-token causal window + forced init/local blocks for any ``topk``). This
|
||||
mirrors how ``MiniMaxM3SparseMSAImpl`` pairs the SM100 attend with Triton.
|
||||
|
||||
Decode and prefill requests are split manually (the batch is decode-first) and
|
||||
each side gets its own ``_fmha_sm100_plan`` / ``_fmha_sm100`` call -- the public
|
||||
``fmha_sm100`` wrapper's mixed-batch split + max-score merge is bypassed.
|
||||
|
||||
The plan is built eagerly at metadata-build time (it ``.tolist()``s the segment
|
||||
lengths and allocates fresh workspaces), so this builder is not
|
||||
cudagraph-replay-safe -- it declares ``AttentionCGSupport.NEVER`` and the
|
||||
attention runs eager (broken out of the graph by ``_run_attention``).
|
||||
``fmha_sm100`` imports are function-local so this module is import-safe on
|
||||
AMD / non-SM100.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import ClassVar
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.forward_context import get_forward_context
|
||||
from vllm.models.minimax_m3.common.indexer import (
|
||||
MiniMaxM3IndexerBackend,
|
||||
MiniMaxM3IndexerImpl,
|
||||
MiniMaxM3IndexerMetadata,
|
||||
MiniMaxM3IndexerMetadataBuilder,
|
||||
)
|
||||
from vllm.models.minimax_m3.common.ops.index_topk import minimax_m3_index_topk
|
||||
from vllm.v1.attention.backend import (
|
||||
AttentionBackend,
|
||||
AttentionCGSupport,
|
||||
CommonAttentionMetadata,
|
||||
)
|
||||
from vllm.v1.attention.backends.utils import split_decodes_and_prefills
|
||||
|
||||
# Page size == sparse block size == index-K block; fmha tile id == M3 block id.
|
||||
PAGE_SIZE = 128
|
||||
|
||||
|
||||
class MiniMaxM3IndexerMSABackend(MiniMaxM3IndexerBackend):
|
||||
"""Indexer side-cache backend selecting the MSA builder."""
|
||||
|
||||
@staticmethod
|
||||
def get_builder_cls() -> type["MiniMaxM3IndexerMSAMetadataBuilder"]:
|
||||
return MiniMaxM3IndexerMSAMetadataBuilder
|
||||
|
||||
|
||||
@dataclass
|
||||
class MiniMaxM3IndexerMSASubMetadata:
|
||||
"""Per-side (decode or prefill) fmha score plan + Triton top-k inputs."""
|
||||
|
||||
plan: dict # _fmha_sm100_plan PlanInfo
|
||||
cu_seqlens_q: torch.Tensor # [n + 1] int32, rebased to 0
|
||||
prefix_lens: torch.Tensor # [n] int32, context tokens before this side
|
||||
max_query_len: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class MiniMaxM3IndexerMSAMetadata(MiniMaxM3IndexerMetadata):
|
||||
"""Indexer metadata with separate decode/prefill fmha score plans."""
|
||||
|
||||
# Whole-batch flattened physical page table [sum_pages] int32, request-major
|
||||
# (decode reqs first); split at ``decode_pages`` for the two _fmha_sm100 runs.
|
||||
kv_indices: torch.Tensor | None = None
|
||||
decode_pages: int = 0
|
||||
decode_metadata: MiniMaxM3IndexerMSASubMetadata | None = None
|
||||
prefill_metadata: MiniMaxM3IndexerMSASubMetadata | None = None
|
||||
|
||||
|
||||
class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):
|
||||
"""Builds separate decode/prefill fmha_sm100 score plans (eager only)."""
|
||||
|
||||
# Plans/workspaces are allocated fresh each build() and the run reads
|
||||
# per-request metadata off the plan, so this is not cudagraph-replay-safe.
|
||||
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER
|
||||
|
||||
def build(
|
||||
self,
|
||||
common_prefix_len: int,
|
||||
common_attn_metadata: CommonAttentionMetadata,
|
||||
fast_build: bool = False,
|
||||
) -> MiniMaxM3IndexerMSAMetadata:
|
||||
from vllm.third_party.fmha_sm100.api import _fmha_sm100_plan
|
||||
|
||||
num_reqs = common_attn_metadata.num_reqs
|
||||
num_tokens = common_attn_metadata.num_actual_tokens
|
||||
seq_lens = common_attn_metadata.seq_lens
|
||||
block_table = common_attn_metadata.block_table_tensor
|
||||
query_start_loc = common_attn_metadata.query_start_loc
|
||||
|
||||
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
|
||||
split_decodes_and_prefills(
|
||||
common_attn_metadata,
|
||||
decode_threshold=self.reorder_batch_threshold,
|
||||
require_uniform=True,
|
||||
)
|
||||
)
|
||||
assert num_decodes + num_prefills == num_reqs
|
||||
assert num_decode_tokens + num_prefill_tokens == num_tokens
|
||||
|
||||
# Context (prefix) lengths into the stable buffer; sliced per side below.
|
||||
context_lens = self.context_len_buffer[:num_reqs]
|
||||
context_lens.copy_(
|
||||
common_attn_metadata.compute_num_computed_tokens(), non_blocking=True
|
||||
)
|
||||
|
||||
# Exact per-request lengths (host): the plan .tolist()s these and uses the
|
||||
# exact KV lengths for run-time causal masking.
|
||||
qsl_cpu = common_attn_metadata.query_start_loc_cpu[: num_reqs + 1]
|
||||
qo_lens_cpu = (qsl_cpu[1:] - qsl_cpu[:-1]).to(torch.int32)
|
||||
kv_lens_cpu = seq_lens[:num_reqs].cpu().to(torch.int32)
|
||||
nvp = (kv_lens_cpu + PAGE_SIZE - 1) // PAGE_SIZE
|
||||
|
||||
# Whole-batch request-major flat page table; decode pages come first.
|
||||
max_blocks = block_table.shape[1]
|
||||
cols = torch.arange(max_blocks, device=block_table.device)
|
||||
valid = cols[None, :] < nvp.to(block_table.device)[:, None]
|
||||
kv_indices = block_table[valid].to(torch.int32)
|
||||
decode_pages = int(nvp[:num_decodes].sum())
|
||||
|
||||
def build_side_metadata(
|
||||
lo: int, hi: int
|
||||
) -> MiniMaxM3IndexerMSASubMetadata | None:
|
||||
if hi <= lo:
|
||||
return None
|
||||
side_qo = qo_lens_cpu[lo:hi]
|
||||
side_kv = kv_lens_cpu[lo:hi]
|
||||
plan = _fmha_sm100_plan(
|
||||
side_qo,
|
||||
side_kv,
|
||||
self.num_index_heads,
|
||||
num_kv_heads=1,
|
||||
qo_offset=side_kv - side_qo, # bottom-right causal
|
||||
page_size=PAGE_SIZE,
|
||||
output_maxscore=True,
|
||||
causal=True,
|
||||
num_kv_splits=1,
|
||||
)
|
||||
cu = query_start_loc[lo : hi + 1] - query_start_loc[lo]
|
||||
return MiniMaxM3IndexerMSASubMetadata(
|
||||
plan=plan,
|
||||
cu_seqlens_q=cu.to(torch.int32),
|
||||
prefix_lens=context_lens[lo:hi],
|
||||
max_query_len=int(side_qo.max()),
|
||||
)
|
||||
|
||||
# Decode requests precede prefill requests in the reordered batch.
|
||||
decode_metadata = build_side_metadata(0, num_decodes)
|
||||
prefill_metadata = build_side_metadata(num_decodes, num_reqs)
|
||||
|
||||
return MiniMaxM3IndexerMSAMetadata(
|
||||
seq_lens=seq_lens,
|
||||
max_seq_len=common_attn_metadata.max_seq_len,
|
||||
slot_mapping=common_attn_metadata.slot_mapping,
|
||||
num_actual_tokens=num_tokens,
|
||||
num_decodes=num_decodes,
|
||||
num_decode_tokens=num_decode_tokens,
|
||||
num_prefills=num_prefills,
|
||||
num_prefill_tokens=num_prefill_tokens,
|
||||
kv_indices=kv_indices,
|
||||
decode_pages=decode_pages,
|
||||
decode_metadata=decode_metadata,
|
||||
prefill_metadata=prefill_metadata,
|
||||
)
|
||||
|
||||
|
||||
class MiniMaxM3IndexerMSAImpl(MiniMaxM3IndexerImpl):
|
||||
"""fmha_sm100 OnlyScore for the per-block scores; Triton top-k selection."""
|
||||
|
||||
indexer_backend_cls: ClassVar[type[AttentionBackend]] = MiniMaxM3IndexerMSABackend
|
||||
|
||||
def forward(
|
||||
self,
|
||||
index_query: torch.Tensor,
|
||||
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
|
||||
from vllm.third_party.fmha_sm100.api import _fmha_sm100
|
||||
|
||||
attn_metadata = get_forward_context().attn_metadata
|
||||
if not isinstance(attn_metadata, dict):
|
||||
return None, None # profiling run; caches unbound
|
||||
md = attn_metadata[self.index_cache.prefix]
|
||||
assert isinstance(md, MiniMaxM3IndexerMSAMetadata)
|
||||
|
||||
num_tokens = md.num_actual_tokens
|
||||
nd = md.num_decode_tokens
|
||||
index_q = index_query[:num_tokens].view(
|
||||
-1, self.num_index_heads, self.index_head_dim
|
||||
)
|
||||
# Index-K cache (num_blocks, 128, D) -> paged MQA (num_blocks, 1, 128, D).
|
||||
kv = self.index_cache.kv_cache
|
||||
k_pages = kv.view(kv.shape[0], 1, PAGE_SIZE, self.index_head_dim)
|
||||
kv_indices = md.kv_indices
|
||||
|
||||
def score_topk(
|
||||
meta: MiniMaxM3IndexerMSASubMetadata,
|
||||
query: torch.Tensor,
|
||||
page_table: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
# OnlyScore -> max_score [num_index_heads, max_k_tiles, num_tokens].
|
||||
_, max_score = _fmha_sm100(
|
||||
query,
|
||||
k_pages,
|
||||
k_pages, # V placeholder; not read in OnlyScore
|
||||
meta.plan,
|
||||
kv_indices=page_table,
|
||||
output_o=False,
|
||||
output_maxscore=True,
|
||||
sm_scale=self.scale,
|
||||
)
|
||||
# Triton top-k wants [num_index_heads, num_tokens, max_block]; the
|
||||
# transpose is a strided view (the kernel reads via strides). One
|
||||
# 128-token KV tile == one M3 sparse block.
|
||||
return minimax_m3_index_topk(
|
||||
max_score.transpose(1, 2),
|
||||
meta.cu_seqlens_q,
|
||||
meta.prefix_lens,
|
||||
meta.max_query_len,
|
||||
self.topk_blocks,
|
||||
self.init_blocks,
|
||||
self.local_blocks,
|
||||
)
|
||||
|
||||
def run_decode() -> torch.Tensor | None:
|
||||
decode_metadata = md.decode_metadata
|
||||
if decode_metadata is None:
|
||||
return None
|
||||
decode_pages = (
|
||||
kv_indices[: md.decode_pages] if kv_indices is not None else None
|
||||
)
|
||||
return score_topk(decode_metadata, index_q[:nd], decode_pages)
|
||||
|
||||
def run_prefill() -> torch.Tensor | None:
|
||||
prefill_metadata = md.prefill_metadata
|
||||
if prefill_metadata is None:
|
||||
return None
|
||||
prefill_pages = (
|
||||
kv_indices[md.decode_pages :] if kv_indices is not None else None
|
||||
)
|
||||
return score_topk(prefill_metadata, index_q[nd:], prefill_pages)
|
||||
|
||||
return run_decode(), run_prefill()
|
||||
@@ -576,7 +576,12 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
|
||||
main_slot_mapping = fwd_slot_mapping[self.layer_name]
|
||||
index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
|
||||
q = qkv.new_empty((num_tokens, self.q_size))
|
||||
index_q = qkv.new_empty((num_tokens, self.index_q_size))
|
||||
# index_q matches the index-K cache dtype (e4m3 for the fp8 score path);
|
||||
# the fused kernel emits fp8 directly when this buffer is e4m3.
|
||||
index_q = qkv.new_empty(
|
||||
(num_tokens, self.index_q_size),
|
||||
dtype=self.indexer.index_cache.dtype,
|
||||
)
|
||||
ops.fused_minimax_m3_qknorm_rope_kv_insert(
|
||||
qkv,
|
||||
self.q_norm.weight,
|
||||
|
||||
Reference in New Issue
Block a user