diff --git a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu index 2eaf4ba15fa6..607485edfb1f 100644 --- a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu +++ b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu @@ -88,6 +88,22 @@ __device__ __forceinline__ void storeHeadElements( } } +template +__device__ __forceinline__ void storeFp8HeadElements64( + __nv_fp8_e4m3* out, int64_t offsetThread, float const (&elements)[numElemsPerThread]) +{ + static_assert(numElemsPerThread == 4, "MiniMax-M3 FP8 store expects four elements per thread"); + static_assert(sizeof(__nv_fp8x2_storage_t) == 2, "MiniMax-M3 FP8 pair storage must be 16 bits"); + // Form the final pointer with 64-bit arithmetic before one aligned 32-bit + // store. Production coalesced paged-cache offsets can exceed INT32_MAX + // FP8 elements even though each individual head row is small. + auto* threadOut = out + offsetThread; + __nv_fp8x2_e4m3 const low(make_float2(elements[0], elements[1])); + __nv_fp8x2_e4m3 const high(make_float2(elements[2], elements[3])); + uint32_t const packed = static_cast(low.__x) | (static_cast(high.__x) << 16); + *reinterpret_cast(threadOut) = packed; +} + // Perform per-head QK Norm and RoPE in a single kernel, reading a BF16 input and // writing to a (possibly different-dtype) output buffer. // head_dim: the dimension of each head @@ -351,6 +367,296 @@ __global__ void fusedQKNormRopeKernel( storeHeadElements(qkv_out, offsetThread, elements); } +namespace +{ + +constexpr int kMinimaxM3HeadDim = 128; +constexpr int kMinimaxM3RotaryDim = 64; +constexpr int kMinimaxM3PageSize = 128; +static_assert((kMinimaxM3PageSize & (kMinimaxM3PageSize - 1)) == 0, "page size must be a power of two"); +constexpr int kMinimaxM3ElemsPerThread = kMinimaxM3HeadDim / 32; + +// MiniMax-M3-only direct-cache specialization for eager pure prefill. The +// general fused QK-norm/RoPE producer plus the #16755 Triton scatter remains +// the fallback for decode, mixed batches, BF16 caches, and unsupported layouts. +__global__ void minimaxM3Fp8QKNormRopeKVInsertKernel(__nv_bfloat16 const* qkvInput, __nv_fp8_e4m3* qOutput, + __nv_fp8_e4m3* kvCache, int const* outCacheLoc, int64_t pageStride, int64_t planeStride, int64_t headStride, + int64_t tokenStride, int64_t numPages, int numTokens, int numHeadsQ, int numHeadsK, int numHeadsV, float eps, + __nv_bfloat16 const* qWeight, __nv_bfloat16 const* kWeight, float base, int const* positionIds) +{ + int const warpsPerBlock = blockDim.x / 32; + int const warpId = threadIdx.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarp = blockIdx.x * warpsPerBlock + warpId; + int const totalHeads = numHeadsQ + numHeadsK + numHeadsV; + int const tokenIdx = globalWarp / totalHeads; + int const localHead = globalWarp % totalHeads; + if (tokenIdx >= numTokens) + { + return; + } + + int const totalQKHeads = numHeadsQ + numHeadsK; + bool const isQ = localHead < numHeadsQ; + bool const isV = localHead >= totalQKHeads; + int const headIdx = isQ ? localHead : (isV ? localHead - totalQKHeads : localHead - numHeadsQ); + int64_t const inputOffset = (static_cast(tokenIdx) * totalHeads + localHead) * kMinimaxM3HeadDim + + laneId * kMinimaxM3ElemsPerThread; + + float elements[kMinimaxM3ElemsPerThread]; + float sumSquares = 0.0F; + constexpr int kVecSize = kMinimaxM3ElemsPerThread * sizeof(__nv_bfloat16) / 4; + using VecT = typename tensorrt_llm::common::packed_as::type; + VecT const packedInput = *reinterpret_cast(qkvInput + inputOffset); +#pragma unroll + for (int i = 0; i < kVecSize; ++i) + { + float2 const values = __bfloat1622float2( + *reinterpret_cast<__nv_bfloat162 const*>(reinterpret_cast(&packedInput) + i)); + if (!isV) + { + sumSquares += values.x * values.x; + sumSquares += values.y * values.y; + } + elements[2 * i] = values.x; + elements[2 * i + 1] = values.y; + } + + __nv_fp8_e4m3* output = qOutput; + int64_t outputOffset; + if (isQ) + { + outputOffset = (static_cast(tokenIdx) * numHeadsQ + headIdx) * kMinimaxM3HeadDim + + laneId * kMinimaxM3ElemsPerThread; + } + else + { + int slot = laneId == 0 ? outCacheLoc[tokenIdx] : 0; + slot = __shfl_sync(0xffffffff, slot, 0); + // CUDA-graph padding uses -1 for non-live cache destinations. Q is + // still produced for the padded row, but K/V must not address it. + if (slot < 0) + { + return; + } + int const page = slot / kMinimaxM3PageSize; + if (page >= numPages) + { + return; + } + int const withinPage = slot & (kMinimaxM3PageSize - 1); + int const plane = isV ? 1 : 0; + output = kvCache; + outputOffset = static_cast(page) * pageStride + static_cast(plane) * planeStride + + static_cast(headIdx) * headStride + static_cast(withinPage) * tokenStride + + laneId * kMinimaxM3ElemsPerThread; + } + + // V is copy-cast only. + if (isV) + { + storeFp8HeadElements64(output, outputOffset, elements); + return; + } + + sumSquares = tensorrt_llm::common::warpReduceSum(sumSquares); + float const rmsReciprocal = rsqrtf(sumSquares / static_cast(kMinimaxM3HeadDim) + eps); +#pragma unroll + for (int i = 0; i < kMinimaxM3ElemsPerThread; ++i) + { + int const dim = laneId * kMinimaxM3ElemsPerThread + i; + float const weight = isQ ? __bfloat162float(qWeight[dim]) : __bfloat162float(kWeight[dim]); + elements[i] *= rmsReciprocal * (1.0F + weight); + } + + // MiniMax-M3 uses NeoX partial RoPE over the first 64 of 128 channels. + // Only lanes 0..7 calculate the 32 distinct angles; lanes 8..15 reuse + // them for the paired half, while lanes 16..31 bypass RoPE. + float pairedElements[kMinimaxM3ElemsPerThread]; + float cosineValues[kMinimaxM3ElemsPerThread] = {}; + float sineValues[kMinimaxM3ElemsPerThread] = {}; + __syncwarp(); + constexpr int kPairOffset = (kMinimaxM3RotaryDim / 2) / kMinimaxM3ElemsPerThread; + int positionId = laneId == 0 ? positionIds[tokenIdx] : 0; + positionId = __shfl_sync(0xffffffff, positionId, 0); +#pragma unroll + for (int i = 0; i < kMinimaxM3ElemsPerThread; ++i) + { + int const dim = laneId * kMinimaxM3ElemsPerThread + i; + pairedElements[i] = __shfl_xor_sync(0xffffffff, elements[i], kPairOffset); + if (laneId < kPairOffset) + { + pairedElements[i] = -pairedElements[i]; + } + + if (laneId < kPairOffset) + { + int const halfDim = dim; + float const frequency = powf(base, -2.0F * halfDim / static_cast(kMinimaxM3RotaryDim)); + __sincosf(static_cast(positionId) * frequency, &sineValues[i], &cosineValues[i]); + } + if (laneId < 2 * kPairOffset) + { + int const sourceLane = laneId % kPairOffset; + cosineValues[i] = __shfl_sync(0x0000ffff, cosineValues[i], sourceLane); + sineValues[i] = __shfl_sync(0x0000ffff, sineValues[i], sourceLane); + } + } + __syncwarp(); + +#pragma unroll + for (int i = 0; i < kMinimaxM3ElemsPerThread; ++i) + { + int const dim = laneId * kMinimaxM3ElemsPerThread + i; + if (dim < kMinimaxM3RotaryDim) + { + elements[i] = elements[i] * cosineValues[i] + pairedElements[i] * sineValues[i]; + } + } + + storeFp8HeadElements64(output, outputOffset, elements); +} + +// Horizontal sparse producer for a packed [Q|K|V|index-Q|index-K] row. +// One warp owns one (token, head slot). All four norm/RoPE branches share the +// model's precomputed FP32 RoPE table, eliminating per-head powf/sincos work. +__global__ void minimaxM3Fp8QKVIndexerNormRopeKVInsertKernel(__nv_bfloat16 const* packedInput, __nv_fp8_e4m3* qOutput, + __nv_fp8_e4m3* indexQOutput, __nv_fp8_e4m3* kvCache, __nv_fp8_e4m3* indexKCache, int const* outCacheLoc, + int64_t kvPageStride, int64_t kvPlaneStride, int64_t kvHeadStride, int64_t kvTokenStride, int64_t indexPageStride, + int64_t indexTokenStride, int64_t numPages, int numTokens, int numHeadsQ, int numHeadsKV, int numHeadsIndex, + float eps, __nv_bfloat16 const* qWeight, __nv_bfloat16 const* kWeight, __nv_bfloat16 const* indexQWeight, + __nv_bfloat16 const* indexKWeight, float const* rotaryCosSin, int const* positionIds) +{ + int const warpsPerBlock = blockDim.x / 32; + int const warpId = threadIdx.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarp = blockIdx.x * warpsPerBlock + warpId; + int const totalHeads = numHeadsQ + 2 * numHeadsKV + numHeadsIndex + 1; + int const tokenIdx = globalWarp / totalHeads; + int const localHead = globalWarp % totalHeads; + if (tokenIdx >= numTokens) + { + return; + } + + int const kBegin = numHeadsQ; + int const vBegin = kBegin + numHeadsKV; + int const indexQBegin = vBegin + numHeadsKV; + int const indexKHead = indexQBegin + numHeadsIndex; + bool const isQ = localHead < kBegin; + bool const isK = localHead >= kBegin && localHead < vBegin; + bool const isV = localHead >= vBegin && localHead < indexQBegin; + bool const isIndexQ = localHead >= indexQBegin && localHead < indexKHead; + bool const isIndexK = localHead == indexKHead; + + int64_t const inputOffset = (static_cast(tokenIdx) * totalHeads + localHead) * kMinimaxM3HeadDim + + laneId * kMinimaxM3ElemsPerThread; + constexpr int kVecSize = kMinimaxM3ElemsPerThread * sizeof(__nv_bfloat16) / 4; + using VecT = typename tensorrt_llm::common::packed_as::type; + VecT const packed = *reinterpret_cast(packedInput + inputOffset); + + float elements[kMinimaxM3ElemsPerThread]; + float sumSquares = 0.0F; +#pragma unroll + for (int pair = 0; pair < kVecSize; ++pair) + { + float2 const values = __bfloat1622float2( + *reinterpret_cast<__nv_bfloat162 const*>(reinterpret_cast(&packed) + pair)); + elements[2 * pair] = values.x; + elements[2 * pair + 1] = values.y; + if (!isV) + { + sumSquares += values.x * values.x + values.y * values.y; + } + } + + if (!isV) + { + auto const* normWeight = isQ ? qWeight : (isK ? kWeight : (isIndexQ ? indexQWeight : indexKWeight)); + sumSquares = tensorrt_llm::common::warpReduceSum(sumSquares); + float const rmsReciprocal = rsqrtf(sumSquares / static_cast(kMinimaxM3HeadDim) + eps); +#pragma unroll + for (int i = 0; i < kMinimaxM3ElemsPerThread; ++i) + { + int const dim = laneId * kMinimaxM3ElemsPerThread + i; + elements[i] *= rmsReciprocal * (1.0F + __bfloat162float(normWeight[dim])); + } + + __syncwarp(); + constexpr int kPairOffset = (kMinimaxM3RotaryDim / 2) / kMinimaxM3ElemsPerThread; + int positionId = laneId == 0 ? positionIds[tokenIdx] : 0; + positionId = __shfl_sync(0xffffffff, positionId, 0); + int64_t const ropeRow = static_cast(positionId) * kMinimaxM3RotaryDim; +#pragma unroll + for (int i = 0; i < kMinimaxM3ElemsPerThread; ++i) + { + int const dim = laneId * kMinimaxM3ElemsPerThread + i; + float paired = __shfl_xor_sync(0xffffffff, elements[i], kPairOffset); + if (dim < kMinimaxM3RotaryDim) + { + bool const firstHalf = dim < kMinimaxM3RotaryDim / 2; + if (firstHalf) + { + paired = -paired; + } + int const coefficient = firstHalf ? dim : dim - kMinimaxM3RotaryDim / 2; + float const cosine = rotaryCosSin[ropeRow + coefficient]; + float const sine = rotaryCosSin[ropeRow + kMinimaxM3RotaryDim / 2 + coefficient]; + elements[i] = elements[i] * cosine + paired * sine; + } + } + __syncwarp(); + } + + if (isQ) + { + int const head = localHead; + int64_t const outputOffset = (static_cast(tokenIdx) * numHeadsQ + head) * kMinimaxM3HeadDim + + laneId * kMinimaxM3ElemsPerThread; + storeFp8HeadElements64(qOutput, outputOffset, elements); + return; + } + if (isIndexQ) + { + int const head = localHead - indexQBegin; + int64_t const outputOffset = (static_cast(tokenIdx) * numHeadsIndex + head) * kMinimaxM3HeadDim + + laneId * kMinimaxM3ElemsPerThread; + // Match vLLM's CUDA path: normalized/RoPE FP32 registers convert + // directly to saturating E4M3, without an intermediate BF16 round. + storeFp8HeadElements64(indexQOutput, outputOffset, elements); + return; + } + + int slot = laneId == 0 ? outCacheLoc[tokenIdx] : 0; + slot = __shfl_sync(0xffffffff, slot, 0); + if (slot < 0) + { + return; + } + int const page = slot / kMinimaxM3PageSize; + if (page >= numPages) + { + return; + } + int const withinPage = slot & (kMinimaxM3PageSize - 1); + if (isIndexK) + { + int64_t const outputOffset = static_cast(page) * indexPageStride + + static_cast(withinPage) * indexTokenStride + laneId * kMinimaxM3ElemsPerThread; + storeFp8HeadElements64(indexKCache, outputOffset, elements); + return; + } + + int const head = isK ? localHead - kBegin : localHead - vBegin; + int const plane = isV ? 1 : 0; + int64_t const outputOffset = static_cast(page) * kvPageStride + static_cast(plane) * kvPlaneStride + + static_cast(head) * kvHeadStride + static_cast(withinPage) * kvTokenStride + + laneId * kMinimaxM3ElemsPerThread; + storeFp8HeadElements64(kvCache, outputOffset, elements); +} + +} // namespace + // Borrowed from // https://github.com/flashinfer-ai/flashinfer/blob/8125d079a43e9a0ba463a4ed1b639cefd084cec9/include/flashinfer/pos_enc.cuh#L568 #define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \ @@ -459,6 +765,63 @@ void launchFusedQKNormRopeToFp8(void const* qkv_in, void* qkv_out, int const num static_cast<__nv_bfloat16 const*>(k_weight), base, interleave, position_ids, factor, low, high, attention_factor, stream, is_qk_norm, use_gemma, use_mrope, mrope_section1, mrope_section2); } + +void launchMinimaxM3Fp8QKNormRopeKVInsert(void const* qkv_input, void* q_output, void* kv_cache, + int const* out_cache_loc, int64_t page_stride, int64_t plane_stride, int64_t head_stride, int64_t token_stride, + int64_t num_pages, int page_size, int num_tokens, int num_heads_q, int num_heads_k, int num_heads_v, int head_dim, + int rotary_dim, float eps, void const* q_weight, void const* k_weight, float base, int const* position_ids, + cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(head_dim == kMinimaxM3HeadDim, "MiniMax-M3 FP8 main Q/K/V producer requires head_dim=128"); + TLLM_CHECK_WITH_INFO( + rotary_dim == kMinimaxM3RotaryDim, "MiniMax-M3 FP8 main Q/K/V producer requires rotary_dim=64"); + TLLM_CHECK_WITH_INFO(num_heads_q > 0, "MiniMax-M3 FP8 main Q/K/V producer requires query heads"); + TLLM_CHECK_WITH_INFO( + num_heads_k > 0 && num_heads_v > 0, "MiniMax-M3 FP8 main Q/K/V producer requires K and V heads"); + TLLM_CHECK_WITH_INFO(page_size == kMinimaxM3PageSize, "MiniMax-M3 FP8 main Q/K/V producer requires page_size=128"); + + constexpr int kBlockSize = 256; + constexpr int kWarpsPerBlock = kBlockSize / 32; + int const totalWarps = num_tokens * (num_heads_q + num_heads_k + num_heads_v); + int const gridSize = common::divUp(totalWarps, kWarpsPerBlock); + minimaxM3Fp8QKNormRopeKVInsertKernel<<>>( + static_cast<__nv_bfloat16 const*>(qkv_input), static_cast<__nv_fp8_e4m3*>(q_output), + static_cast<__nv_fp8_e4m3*>(kv_cache), out_cache_loc, page_stride, plane_stride, head_stride, token_stride, + num_pages, num_tokens, num_heads_q, num_heads_k, num_heads_v, eps, static_cast<__nv_bfloat16 const*>(q_weight), + static_cast<__nv_bfloat16 const*>(k_weight), base, position_ids); + TLLM_CUDA_CHECK(cudaGetLastError()); +} + +void launchMinimaxM3Fp8QKVIndexerNormRopeKVInsert(void const* packed_input, void* q_output, void* index_q_output, + void* kv_cache, void* index_k_cache, int const* out_cache_loc, int64_t kv_page_stride, int64_t kv_plane_stride, + int64_t kv_head_stride, int64_t kv_token_stride, int64_t index_page_stride, int64_t index_token_stride, + int64_t num_pages, int page_size, int num_tokens, int num_heads_q, int num_heads_kv, int num_heads_index, + int head_dim, int rotary_dim, float eps, void const* q_weight, void const* k_weight, void const* index_q_weight, + void const* index_k_weight, float const* rotary_cos_sin, int const* position_ids, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(head_dim == kMinimaxM3HeadDim, "MiniMax-M3 horizontal producer requires head_dim=128"); + TLLM_CHECK_WITH_INFO(rotary_dim == kMinimaxM3RotaryDim, "MiniMax-M3 horizontal producer requires rotary_dim=64"); + TLLM_CHECK_WITH_INFO(page_size == kMinimaxM3PageSize, "MiniMax-M3 horizontal producer requires page_size=128"); + TLLM_CHECK_WITH_INFO(num_heads_q > 0 && num_heads_kv > 0 && num_heads_index > 0, + "MiniMax-M3 horizontal producer requires Q, KV, and index heads"); + TLLM_CHECK_WITH_INFO( + num_heads_index == num_heads_kv, "MiniMax-M3 horizontal producer requires index heads to equal KV heads"); + + constexpr int kBlockSize = 256; + constexpr int kWarpsPerBlock = kBlockSize / 32; + int const slotsPerToken = num_heads_q + 2 * num_heads_kv + num_heads_index + 1; + int const totalWarps = num_tokens * slotsPerToken; + int const gridSize = common::divUp(totalWarps, kWarpsPerBlock); + minimaxM3Fp8QKVIndexerNormRopeKVInsertKernel<<>>( + static_cast<__nv_bfloat16 const*>(packed_input), static_cast<__nv_fp8_e4m3*>(q_output), + static_cast<__nv_fp8_e4m3*>(index_q_output), static_cast<__nv_fp8_e4m3*>(kv_cache), + static_cast<__nv_fp8_e4m3*>(index_k_cache), out_cache_loc, kv_page_stride, kv_plane_stride, kv_head_stride, + kv_token_stride, index_page_stride, index_token_stride, num_pages, num_tokens, num_heads_q, num_heads_kv, + num_heads_index, eps, static_cast<__nv_bfloat16 const*>(q_weight), static_cast<__nv_bfloat16 const*>(k_weight), + static_cast<__nv_bfloat16 const*>(index_q_weight), static_cast<__nv_bfloat16 const*>(index_k_weight), + rotary_cos_sin, position_ids); + TLLM_CUDA_CHECK(cudaGetLastError()); +} } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h index fd2401f592f1..4b41467d457c 100644 --- a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h +++ b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -17,6 +17,8 @@ #pragma once #include "tensorrt_llm/common/config.h" + +#include #include TRTLLM_NAMESPACE_BEGIN @@ -61,6 +63,25 @@ void launchFusedQKNormRopeToFp8(void const* qkv_in, // BF16 input [num_tokens, t bool const interleave, int const* position_ids, float factor, float low, float high, float attention_factor, cudaStream_t stream, bool is_qk_norm, bool use_gemma, bool use_mrope, int mrope_section1, int mrope_section2); +// MiniMax-M3-specific main-branch producer. It returns contiguous +// FP8 Q and inserts normalized/RoPE'd FP8 K plus copy-cast FP8 V directly into +// a paged HND pool [num_pages, 2, num_heads, page_size, head_dim]. +void launchMinimaxM3Fp8QKNormRopeKVInsert(void const* qkv_input, void* q_output, void* kv_cache, + int const* out_cache_loc, int64_t page_stride, int64_t plane_stride, int64_t head_stride, int64_t token_stride, + int64_t num_pages, int page_size, int num_tokens, int num_heads_q, int num_heads_k, int num_heads_v, int head_dim, + int rotary_dim, float eps, void const* q_weight, void const* k_weight, float base, int const* position_ids, + cudaStream_t stream); + +// MiniMax-M3 sparse producer for the packed [Q|K|V|index-Q|index-K] +// projection. It uses a precomputed FP32 RoPE table, emits compact FP8 Q and +// index-Q, and inserts main K/V plus index-K into their paged FP8 HND caches. +void launchMinimaxM3Fp8QKVIndexerNormRopeKVInsert(void const* packed_input, void* q_output, void* index_q_output, + void* kv_cache, void* index_k_cache, int const* out_cache_loc, int64_t kv_page_stride, int64_t kv_plane_stride, + int64_t kv_head_stride, int64_t kv_token_stride, int64_t index_page_stride, int64_t index_token_stride, + int64_t num_pages, int page_size, int num_tokens, int num_heads_q, int num_heads_kv, int num_heads_index, + int head_dim, int rotary_dim, float eps, void const* q_weight, void const* k_weight, void const* index_q_weight, + void const* index_k_weight, float const* rotary_cos_sin, int const* position_ids, cudaStream_t stream); + } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp b/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp index 12c124b85aaf..933ed62432c2 100644 --- a/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp +++ b/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -18,9 +18,12 @@ #include "tensorrt_llm/thop/thUtils.h" #include +#include #include #include +#include + TRTLLM_NAMESPACE_BEGIN namespace torch_ext @@ -65,6 +68,35 @@ int64_t validateFusedQKNormRopeInputs(torch::Tensor const& qkv, torch::Tensor co return num_tokens; } +void checkMinimaxM3HndKVPool(torch::Tensor const& kvCache, int64_t numHeads, int64_t headDim) +{ + TORCH_CHECK(kvCache.is_cuda(), "kv_cache must be a CUDA tensor"); + TORCH_CHECK(kvCache.scalar_type() == at::ScalarType::Float8_e4m3fn, "kv_cache must use torch.float8_e4m3fn"); + TORCH_CHECK(kvCache.dim() == 5, "kv_cache must be HND [num_pages, 2, num_heads, page_size, head_dim]"); + TORCH_CHECK(kvCache.size(0) > 0 && kvCache.size(3) > 0, "kv_cache must have positive num_pages and page_size"); + TORCH_CHECK(kvCache.size(1) == 2, "kv_cache plane dimension must contain K and V"); + TORCH_CHECK(kvCache.size(2) == numHeads, "kv_cache num_heads mismatch"); + TORCH_CHECK(kvCache.size(4) == headDim, "kv_cache head_dim mismatch"); + TORCH_CHECK(kvCache.stride(4) == 1 && kvCache.stride(3) == headDim, + "kv_cache must have contiguous head_dim rows in HND layout"); + TORCH_CHECK(kvCache.stride(2) == kvCache.size(3) * kvCache.stride(3), + "kv_cache must have contiguous [page_size, head_dim] blocks in HND layout"); + TORCH_CHECK(kvCache.stride(1) >= kvCache.size(2) * kvCache.stride(2), "kv_cache K and V planes must not overlap"); + TORCH_CHECK(kvCache.stride(0) >= kvCache.size(1) * kvCache.stride(1), + "kv_cache page stride must not overlap adjacent HND pages"); + TORCH_CHECK(kvCache.stride(0) % 4 == 0 && kvCache.stride(1) % 4 == 0, + "kv_cache page and plane strides must preserve 32-bit FP8 store alignment"); +} + +void checkMinimaxM3Int32LaunchGeometry(int64_t numTokens, int64_t slotsPerToken) +{ + TORCH_CHECK(numTokens <= std::numeric_limits::max(), "MiniMax-M3 producer num_tokens exceeds int32"); + TORCH_CHECK(slotsPerToken > 0 && slotsPerToken <= std::numeric_limits::max(), + "MiniMax-M3 producer head geometry exceeds int32"); + TORCH_CHECK(numTokens == 0 || slotsPerToken <= std::numeric_limits::max() / numTokens, + "MiniMax-M3 producer launch geometry exceeds int32"); +} + } // namespace // Function for fused QK Norm and RoPE @@ -150,6 +182,189 @@ torch::Tensor fused_qk_norm_rope_to_fp8_meta(torch::Tensor const& qkv, int64_t n return torch::empty({num_tokens, total_heads * head_dim}, qkv.options().dtype(torch::kFloat8_e4m3fn)); } +torch::Tensor minimaxM3Fp8QKNormRopeKVInsert(torch::Tensor const& qkv, torch::Tensor& kvCache, + torch::Tensor const& outCacheLoc, int64_t numHeadsQ, int64_t numHeadsK, int64_t numHeadsV, int64_t headDim, + int64_t rotaryDim, double eps, torch::Tensor const& qWeight, torch::Tensor const& kWeight, double base, bool isNeox, + torch::Tensor const& positionIds) +{ + constexpr int64_t kHeadDim = 128; + constexpr int64_t kRotaryDim = 64; + constexpr int64_t kPageSize = 128; + TORCH_CHECK(numHeadsQ > 0, "MiniMax-M3 FP8 main Q/K/V producer requires num_heads_q > 0"); + TORCH_CHECK(numHeadsK > 0 && numHeadsV > 0, "MiniMax-M3 FP8 main Q/K/V producer requires K and V heads"); + TORCH_CHECK(numHeadsK == numHeadsV, "MiniMax-M3 FP8 main Q/K/V producer requires equal K and V head counts"); + TORCH_CHECK(headDim == kHeadDim, "MiniMax-M3 FP8 main Q/K/V producer requires head_dim=128"); + TORCH_CHECK(rotaryDim == kRotaryDim, "MiniMax-M3 FP8 main Q/K/V producer requires rotary_dim=64"); + TORCH_CHECK(isNeox, "MiniMax-M3 FP8 main Q/K/V producer requires NeoX RoPE"); + TORCH_CHECK(std::isfinite(eps) && eps > 0.0, "MiniMax-M3 FP8 main Q/K/V producer requires finite eps > 0"); + TORCH_CHECK(std::isfinite(base) && base > 0.0, "MiniMax-M3 FP8 main Q/K/V producer requires finite RoPE base > 0"); + auto const epsFloat = static_cast(eps); + auto const baseFloat = static_cast(base); + TORCH_CHECK(std::isfinite(epsFloat) && epsFloat > 0.0F, + "MiniMax-M3 FP8 main Q/K/V producer eps must remain finite and positive in float32"); + TORCH_CHECK(std::isfinite(baseFloat) && baseFloat > 0.0F, + "MiniMax-M3 FP8 main Q/K/V producer RoPE base must remain finite and positive in float32"); + + TORCH_CHECK(qkv.dim() == 2, "QKV tensor must be 2D: [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim]"); + TORCH_CHECK(outCacheLoc.dim() == 1, "out_cache_loc must be one-dimensional"); + TORCH_CHECK(positionIds.dim() == 1, "position_ids must be one-dimensional"); + TORCH_CHECK(qWeight.dim() == 1 && kWeight.dim() == 1, "Q/K norm weights must be one-dimensional"); + + CHECK_INPUT(qkv, torch::kBFloat16); + CHECK_INPUT(outCacheLoc, torch::kInt32); + CHECK_INPUT(positionIds, torch::kInt32); + CHECK_INPUT(qWeight, torch::kBFloat16); + CHECK_INPUT(kWeight, torch::kBFloat16); + checkMinimaxM3HndKVPool(kvCache, numHeadsK, headDim); + TORCH_CHECK(kvCache.size(3) == kPageSize, "MiniMax-M3 FP8 main Q/K/V producer requires page_size=128"); + + int64_t const numTokens = qkv.size(0); + int64_t const totalHeads = numHeadsQ + numHeadsK + numHeadsV; + checkMinimaxM3Int32LaunchGeometry(numTokens, totalHeads); + TORCH_CHECK(qkv.size(1) == totalHeads * headDim, + "QKV tensor width must equal (num_heads_q + num_heads_k + num_heads_v) * head_dim"); + TORCH_CHECK(outCacheLoc.numel() >= numTokens, "out_cache_loc is shorter than num_tokens"); + TORCH_CHECK(positionIds.numel() == numTokens, "position_ids length must equal num_tokens"); + TORCH_CHECK(qWeight.numel() == headDim && kWeight.numel() == headDim, "Q/K norm weight width must equal head_dim"); + TORCH_CHECK(reinterpret_cast(qkv.data_ptr()) % 8 == 0, + "QKV input must start at an 8-byte-aligned address for vectorized BF16 loads"); + TORCH_CHECK(reinterpret_cast(kvCache.data_ptr()) % 4 == 0, + "K/V cache must start at a 4-byte-aligned address for packed E4M3 stores"); + TORCH_CHECK(qkv.get_device() == kvCache.get_device() && qkv.get_device() == outCacheLoc.get_device() + && qkv.get_device() == positionIds.get_device() && qkv.get_device() == qWeight.get_device() + && qkv.get_device() == kWeight.get_device(), + "All MiniMax-M3 FP8 main Q/K/V producer tensors must be on the same CUDA device"); + + auto qOut = torch::empty({numTokens, numHeadsQ, headDim}, qkv.options().dtype(at::ScalarType::Float8_e4m3fn)); + if (numTokens == 0) + { + return qOut; + } + + auto stream = at::cuda::getCurrentCUDAStream(qkv.get_device()); + tensorrt_llm::kernels::launchMinimaxM3Fp8QKNormRopeKVInsert(qkv.data_ptr(), qOut.data_ptr(), kvCache.data_ptr(), + outCacheLoc.data_ptr(), kvCache.stride(0), kvCache.stride(1), kvCache.stride(2), kvCache.stride(3), + kvCache.size(0), static_cast(kvCache.size(3)), static_cast(numTokens), static_cast(numHeadsQ), + static_cast(numHeadsK), static_cast(numHeadsV), static_cast(headDim), + static_cast(rotaryDim), epsFloat, qWeight.data_ptr(), kWeight.data_ptr(), baseFloat, + positionIds.data_ptr(), stream); + return qOut; +} + +torch::Tensor minimaxM3Fp8QKNormRopeKVInsertMeta(torch::Tensor const& qkv, torch::Tensor& /*kvCache*/, + torch::Tensor const& /*outCacheLoc*/, int64_t numHeadsQ, int64_t /*numHeadsK*/, int64_t /*numHeadsV*/, + int64_t headDim, int64_t /*rotaryDim*/, double /*eps*/, torch::Tensor const& /*qWeight*/, + torch::Tensor const& /*kWeight*/, double /*base*/, bool /*isNeox*/, torch::Tensor const& /*positionIds*/) +{ + return torch::empty({qkv.size(0), numHeadsQ, headDim}, qkv.options().dtype(at::ScalarType::Float8_e4m3fn)); +} + +std::tuple minimaxM3Fp8QKVIndexerNormRopeKVInsert(torch::Tensor const& packed, + torch::Tensor& kvCache, torch::Tensor& indexKCache, torch::Tensor const& outCacheLoc, int64_t numHeadsQ, + int64_t numHeadsKV, int64_t numHeadsIndex, int64_t headDim, int64_t rotaryDim, double eps, + torch::Tensor const& qWeight, torch::Tensor const& kWeight, torch::Tensor const& indexQWeight, + torch::Tensor const& indexKWeight, torch::Tensor const& rotaryCosSin, torch::Tensor const& positionIds) +{ + constexpr int64_t kHeadDim = 128; + constexpr int64_t kRotaryDim = 64; + constexpr int64_t kPageSize = 128; + TORCH_CHECK(numHeadsQ > 0 && numHeadsKV > 0 && numHeadsIndex > 0, + "MiniMax-M3 horizontal producer requires Q, KV, and index heads"); + TORCH_CHECK(numHeadsKV == numHeadsIndex, "MiniMax-M3 horizontal producer requires index heads to equal KV heads"); + TORCH_CHECK(headDim == kHeadDim, "MiniMax-M3 horizontal producer requires head_dim=128"); + TORCH_CHECK(rotaryDim == kRotaryDim, "MiniMax-M3 horizontal producer requires rotary_dim=64"); + TORCH_CHECK(std::isfinite(eps) && eps > 0.0, "MiniMax-M3 horizontal producer requires finite eps > 0"); + auto const epsFloat = static_cast(eps); + TORCH_CHECK(std::isfinite(epsFloat) && epsFloat > 0.0F, + "MiniMax-M3 horizontal producer eps must remain finite and positive in float32"); + + TORCH_CHECK(packed.dim() == 2, "Packed QKV+index tensor must be two-dimensional"); + TORCH_CHECK(outCacheLoc.dim() == 1, "out_cache_loc must be one-dimensional"); + TORCH_CHECK(positionIds.dim() == 1, "position_ids must be one-dimensional"); + CHECK_INPUT(packed, torch::kBFloat16); + CHECK_INPUT(outCacheLoc, torch::kInt32); + CHECK_INPUT(positionIds, torch::kInt32); + CHECK_INPUT(qWeight, torch::kBFloat16); + CHECK_INPUT(kWeight, torch::kBFloat16); + CHECK_INPUT(indexQWeight, torch::kBFloat16); + CHECK_INPUT(indexKWeight, torch::kBFloat16); + CHECK_INPUT(rotaryCosSin, torch::kFloat32); + checkMinimaxM3HndKVPool(kvCache, numHeadsKV, headDim); + TORCH_CHECK(kvCache.size(3) == kPageSize, "MiniMax-M3 horizontal producer requires page_size=128"); + TORCH_CHECK(indexKCache.is_cuda() && indexKCache.scalar_type() == at::ScalarType::Float8_e4m3fn, + "Index-K cache must be CUDA torch.float8_e4m3fn"); + TORCH_CHECK(indexKCache.dim() == 4 && indexKCache.size(1) == 1 && indexKCache.size(2) == kPageSize + && indexKCache.size(3) == kHeadDim, + "Index-K cache must be HND [num_pages, 1, 128, 128]"); + TORCH_CHECK(indexKCache.stride(3) == 1 && indexKCache.stride(2) == kHeadDim, + "Index-K cache must have contiguous token rows"); + TORCH_CHECK(indexKCache.stride(1) >= indexKCache.size(2) * indexKCache.stride(2) + && indexKCache.stride(0) >= indexKCache.size(1) * indexKCache.stride(1), + "Index-K cache pages must not overlap"); + TORCH_CHECK(indexKCache.stride(0) % 4 == 0 && indexKCache.stride(1) % 4 == 0, + "Index-K cache page/head strides must preserve 32-bit FP8 store alignment"); + TORCH_CHECK( + indexKCache.size(0) == kvCache.size(0), "Main K/V and index-K caches must contain the same number of pages"); + TORCH_CHECK(rotaryCosSin.dim() == 3 && rotaryCosSin.size(1) == 2 && rotaryCosSin.size(2) == kRotaryDim / 2, + "rotary_cos_sin must be [max_positions, 2, rotary_dim/2]"); + + int64_t const numTokens = packed.size(0); + int64_t const totalHeads = numHeadsQ + 2 * numHeadsKV + numHeadsIndex + 1; + checkMinimaxM3Int32LaunchGeometry(numTokens, totalHeads); + TORCH_CHECK( + packed.size(1) == totalHeads * headDim, "Packed tensor width must equal (Q + 2*KV + index-Q + 1) * head_dim"); + TORCH_CHECK(outCacheLoc.numel() >= numTokens, "out_cache_loc is shorter than num_tokens"); + TORCH_CHECK(positionIds.numel() == numTokens, "position_ids length must equal num_tokens"); + TORCH_CHECK(qWeight.dim() == 1 && kWeight.dim() == 1 && indexQWeight.dim() == 1 && indexKWeight.dim() == 1, + "All norm weights must be one-dimensional"); + TORCH_CHECK(qWeight.numel() == headDim && kWeight.numel() == headDim && indexQWeight.numel() == headDim + && indexKWeight.numel() == headDim, + "All norm weights must contain head_dim elements"); + TORCH_CHECK(reinterpret_cast(packed.data_ptr()) % 8 == 0, + "Packed input must start at an 8-byte-aligned address for vectorized BF16 loads"); + TORCH_CHECK(reinterpret_cast(kvCache.data_ptr()) % 4 == 0 + && reinterpret_cast(indexKCache.data_ptr()) % 4 == 0, + "Paged caches must start at 4-byte-aligned addresses for packed E4M3 stores"); + TORCH_CHECK(packed.get_device() == kvCache.get_device() && packed.get_device() == indexKCache.get_device() + && packed.get_device() == outCacheLoc.get_device() && packed.get_device() == positionIds.get_device() + && packed.get_device() == qWeight.get_device() && packed.get_device() == kWeight.get_device() + && packed.get_device() == indexQWeight.get_device() && packed.get_device() == indexKWeight.get_device() + && packed.get_device() == rotaryCosSin.get_device(), + "All MiniMax-M3 horizontal producer tensors must be on the same CUDA device"); + + auto qOut = torch::empty({numTokens, numHeadsQ, headDim}, packed.options().dtype(at::ScalarType::Float8_e4m3fn)); + auto indexQOut + = torch::empty({numTokens, numHeadsIndex, headDim}, packed.options().dtype(at::ScalarType::Float8_e4m3fn)); + if (numTokens == 0) + { + return {qOut, indexQOut}; + } + + auto stream = at::cuda::getCurrentCUDAStream(packed.get_device()); + tensorrt_llm::kernels::launchMinimaxM3Fp8QKVIndexerNormRopeKVInsert(packed.data_ptr(), qOut.data_ptr(), + indexQOut.data_ptr(), kvCache.data_ptr(), indexKCache.data_ptr(), outCacheLoc.data_ptr(), + kvCache.stride(0), kvCache.stride(1), kvCache.stride(2), kvCache.stride(3), indexKCache.stride(0), + indexKCache.stride(2), kvCache.size(0), static_cast(kvCache.size(3)), static_cast(numTokens), + static_cast(numHeadsQ), static_cast(numHeadsKV), static_cast(numHeadsIndex), + static_cast(headDim), static_cast(rotaryDim), epsFloat, qWeight.data_ptr(), kWeight.data_ptr(), + indexQWeight.data_ptr(), indexKWeight.data_ptr(), rotaryCosSin.data_ptr(), positionIds.data_ptr(), + stream); + return {qOut, indexQOut}; +} + +std::tuple minimaxM3Fp8QKVIndexerNormRopeKVInsertMeta(torch::Tensor const& packed, + torch::Tensor& /*kvCache*/, torch::Tensor& /*indexKCache*/, torch::Tensor const& /*outCacheLoc*/, int64_t numHeadsQ, + int64_t /*numHeadsKV*/, int64_t numHeadsIndex, int64_t headDim, int64_t /*rotaryDim*/, double /*eps*/, + torch::Tensor const& /*qWeight*/, torch::Tensor const& /*kWeight*/, torch::Tensor const& /*indexQWeight*/, + torch::Tensor const& /*indexKWeight*/, torch::Tensor const& /*rotaryCosSin*/, torch::Tensor const& /*positionIds*/) +{ + auto options = packed.options().dtype(at::ScalarType::Float8_e4m3fn); + return { + torch::empty({packed.size(0), numHeadsQ, headDim}, options), + torch::empty({packed.size(0), numHeadsIndex, headDim}, options), + }; +} + // Register the PyTorch operators TORCH_LIBRARY_FRAGMENT(trtllm, m) { @@ -164,6 +379,15 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "rotary_dim, float eps, Tensor q_weight, Tensor k_weight, float base, bool is_neox, Tensor position_ids, float " "factor, float low, float high, float attention_factor, bool is_qk_norm, bool use_gemma, bool use_mrope, int " "mrope_section1, int mrope_section2) -> Tensor"); + m.def( + "minimax_m3_fp8_qk_norm_rope_kv_insert(Tensor qkv, Tensor(a!) kv_cache, Tensor out_cache_loc, int " + "num_heads_q, int num_heads_k, int num_heads_v, int head_dim, int rotary_dim, float eps, Tensor q_weight, " + "Tensor k_weight, float base, bool is_neox, Tensor position_ids) -> Tensor"); + m.def( + "minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert(Tensor packed, Tensor(a!) kv_cache, Tensor(b!) " + "index_k_cache, Tensor out_cache_loc, int num_heads_q, int num_heads_kv, int num_heads_index, int head_dim, " + "int rotary_dim, float eps, Tensor q_weight, Tensor k_weight, Tensor index_q_weight, Tensor index_k_weight, " + "Tensor rotary_cos_sin, Tensor position_ids) -> (Tensor, Tensor)"); } // Register the CUDA implementation @@ -171,12 +395,16 @@ TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { m.impl("fused_qk_norm_rope", &fused_qk_norm_rope); m.impl("fused_qk_norm_rope_to_fp8", &fused_qk_norm_rope_to_fp8); + m.impl("minimax_m3_fp8_qk_norm_rope_kv_insert", &minimaxM3Fp8QKNormRopeKVInsert); + m.impl("minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert", &minimaxM3Fp8QKVIndexerNormRopeKVInsert); } // Register the Meta implementation (shape/dtype inference for torch.compile). TORCH_LIBRARY_IMPL(trtllm, Meta, m) { m.impl("fused_qk_norm_rope_to_fp8", &fused_qk_norm_rope_to_fp8_meta); + m.impl("minimax_m3_fp8_qk_norm_rope_kv_insert", &minimaxM3Fp8QKNormRopeKVInsertMeta); + m.impl("minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert", &minimaxM3Fp8QKVIndexerNormRopeKVInsertMeta); } } // namespace torch_ext diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py index 5151fa2d4b6d..f288495eeff1 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py @@ -32,12 +32,39 @@ _LOCAL_SCORE = 1e29 +def index_head_range( + num_index_heads: int, num_kv_heads: int, mapping: Optional["Mapping"] = None +) -> Tuple[int, int]: + """Return this rank's global index-head interval, preserving KV groups. + + Whole KV groups (including all their index heads) replicate when TP + exceeds the KV-head count. Attention-DP retains every head on each rank. + """ + tp_size = 1 if mapping is None or mapping.enable_attention_dp else mapping.tp_size + if num_index_heads <= 0: + raise ValueError("MiniMax-M3 requires positive index heads") + # Metadata without model geometry is usable only for an unsharded view. + if num_kv_heads == 0 and tp_size == 1: + return 0, num_index_heads + if num_kv_heads <= 0 or num_index_heads % num_kv_heads != 0: + raise ValueError("MiniMax-M3 index heads must be divisible by global KV heads") + shard_count = min(tp_size, num_kv_heads) + if tp_size % shard_count != 0 or num_kv_heads % shard_count != 0: + raise ValueError("MiniMax-M3 TP and KV heads must divide one another") + rank = 0 if tp_size == 1 else mapping.tp_rank // (tp_size // shard_count) + count = num_index_heads // shard_count + return rank * count, (rank + 1) * count + + @dataclass(frozen=True) class MiniMaxM3SparseParams(SparseParams): """Lowered runtime parameters for the MiniMax-M3 sparse backend.""" algorithm: Literal["minimax_m3"] = field(init=False, default="minimax_m3") num_index_heads: int = 4 + # None keeps explicit rank-local backend construction supported. Model + # lowering supplies the global count so both backends preserve KV groups. + global_num_kv_heads: Optional[int] = None sparse_index_dim: int = 128 block_size: int = 128 topk: int = 16 @@ -47,6 +74,7 @@ class MiniMaxM3SparseParams(SparseParams): disable_index_value: bool = True implementation: Literal["triton", "msa"] = "triton" indexer_kv_dtype: Literal["bf16", "fp8"] = "bf16" + fuse_qkv_index_projection: bool = False @property def indices_block_size(self) -> int: @@ -83,6 +111,10 @@ def _shard(num_heads: int) -> int: return _shard(self.global_num_q_heads), _shard(self.global_num_kv_heads) + def sharded_index_head_count(self, mapping: Optional["Mapping"] = None) -> int: + start, end = index_head_range(self.num_index_heads, self.global_num_kv_heads, mapping) + return end - start + @dataclass(frozen=True) class MiniMaxM3SparseConfig: @@ -149,11 +181,17 @@ def from_sparse_params( """Build a kernel param bundle from lowered ``MiniMaxM3SparseParams`` and the per-rank model geometry. """ + num_index_heads = int(sparse_params.num_index_heads) + if sparse_params.global_num_kv_heads is not None: + global_kv_heads = int(sparse_params.global_num_kv_heads) + if global_kv_heads <= 0 or num_index_heads % global_kv_heads != 0: + raise ValueError("MiniMax-M3 index heads must be divisible by global KV heads") + num_index_heads = num_index_heads // global_kv_heads * int(num_kv_heads) return cls( num_q_heads=int(num_q_heads), num_kv_heads=int(num_kv_heads), head_dim=int(head_dim), - num_index_heads=int(sparse_params.num_index_heads), + num_index_heads=num_index_heads, sparse_index_dim=int(sparse_params.sparse_index_dim), block_size=int(sparse_params.block_size), topk=int(sparse_params.topk), diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py index 78944e3c0dab..3e45870d64a8 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py @@ -368,14 +368,15 @@ def _validate_decode_kernel_support(self) -> None: # num_index_heads * query tokens into one Q block, which bounds the # draft length it can verify. decode_query_len = self._msa_max_decode_query_len() + num_index_heads = params.sharded_index_head_count(self.mapping) if not self._cutedsl_indexer_supported( - num_index_heads=params.num_index_heads, + num_index_heads=num_index_heads, page_size=page_size, decode_query_len=decode_query_len, ): raise RuntimeError( "The MiniMax-M3 CuTe DSL indexer scorer does not support this " - f"configuration: {params.num_index_heads} index heads, page size " + f"configuration: {num_index_heads} index heads, page size " f"{page_size}, index dtype {self._msa_index_kv_dtype()}, up to " f"{decode_query_len} query tokens per generation request." ) @@ -475,13 +476,13 @@ def _create_msa_buffers(self) -> None: fmha_sm100 = require_msa_module() max_k_tiles = _worst_case_proxy_max_k_tiles( fmha_sm100, - num_index_heads=params.num_index_heads, + num_index_heads=params.sharded_index_head_count(self.mapping), kv_cache_manager=kv_cache_manager, max_batch=max_num_sequences, ) self._msa_worst_case_max_k_tiles = int(max_k_tiles) self._alloc_msa_proxy_scratch( - num_index_heads=params.num_index_heads, + num_index_heads=params.sharded_index_head_count(self.mapping), max_tokens=self._msa_max_decode_tokens(), max_k_tiles=max_k_tiles, capture_graph=capture_graph, @@ -863,7 +864,7 @@ def _build_step_plans(self) -> None: params = self._msa_params if params is None: return - num_index_heads = params.num_index_heads + num_index_heads = params.sharded_index_head_count(self.mapping) qo_lens_cpu = self.msa_qo_lens_cpu kv_lens_cpu = self.msa_kv_lens_cpu qo_offset_cpu = self.msa_qo_offset_cpu diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 59d2d40c1123..fb443da45e70 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -47,6 +47,7 @@ _gather_paged_batched, _write_main_kv_slots_to_pool, ) +from ..attention.backends.sparse.minimax_m3.common import index_head_range from ..attention.backends.sparse.params import SparseBackendForwardArgs from ..distributed import AllReduce, AllReduceFusionOp, AllReduceParams, MiniMaxAllReduceRMS from ..modules.decoder_layer import DecoderLayer @@ -84,6 +85,169 @@ def _moe_routed_output_is_global(experts: nn.Module) -> bool: return getattr(backend, "scheduler_kind", None) == MoESchedulerKind.FUSED_COMM +class MiniMaxM3QKVIndexerLinear(Linear): + """Five-way MiniMax-M3 projection with vLLM-compatible TP sharding. + + Each rank emits ``[Q | K | V | index-Q | index-K]``. Q follows normal + attention head sharding, K/V/index-Q follow KV-head sharding (including + replication when TP exceeds the KV-head count), and the single index-K + head is replicated. The underlying :class:`Linear` remains the standard + quantized implementation; only checkpoint packing is model-specific. + """ + + _SHARD_NAMES = ("q", "k", "v", "index_q", "index_k") + + def __init__( + self, + *, + hidden_size: int, + head_dim: int, + total_num_heads: int, + total_num_kv_heads: int, + total_num_index_heads: int, + index_head_dim: int, + dtype: torch.dtype, + mapping: Mapping, + quant_config: Optional[QuantConfig], + skip_create_weights_in_init: bool, + force_dynamic_quantization: bool, + disable_deep_gemm: bool, + use_custom_cublas_mm: bool, + use_cute_dsl_bf16_gemm: bool, + use_cute_dsl_blockscaling_mm: bool, + ) -> None: + if total_num_index_heads != total_num_kv_heads: + raise ValueError( + "MiniMax-M3 fused QKV+index projection requires index heads " + f"({total_num_index_heads}) to equal KV heads ({total_num_kv_heads})." + ) + if index_head_dim != head_dim: + raise ValueError( + "MiniMax-M3 fused QKV+index projection requires index_head_dim " + f"({index_head_dim}) to equal head_dim ({head_dim})." + ) + + tp_size = int(mapping.tp_size) + if total_num_heads % tp_size != 0: + raise ValueError(f"Q heads ({total_num_heads}) must be divisible by TP ({tp_size}).") + if total_num_kv_heads >= tp_size: + if total_num_kv_heads % tp_size != 0: + raise ValueError( + f"KV heads ({total_num_kv_heads}) must be divisible by TP ({tp_size})." + ) + local_num_kv_heads = total_num_kv_heads // tp_size + else: + if tp_size % total_num_kv_heads != 0: + raise ValueError( + f"TP ({tp_size}) must be divisible by KV heads " + f"({total_num_kv_heads}) for replication." + ) + local_num_kv_heads = 1 + + self.total_num_heads = int(total_num_heads) + self.total_num_kv_heads = int(total_num_kv_heads) + self.total_num_index_heads = int(total_num_index_heads) + self.head_dim = int(head_dim) + self.index_head_dim = int(index_head_dim) + self.local_num_heads = total_num_heads // tp_size + self.local_num_kv_heads = local_num_kv_heads + self.local_num_index_heads = local_num_kv_heads + self.local_output_sizes = ( + self.local_num_heads * head_dim, + local_num_kv_heads * head_dim, + local_num_kv_heads * head_dim, + local_num_kv_heads * index_head_dim, + index_head_dim, + ) + local_out_features = sum(self.local_output_sizes) + + super().__init__( + hidden_size, + tp_size * local_out_features, + bias=False, + dtype=dtype, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + quant_config=quant_config, + weights_loading_config=WeightsLoadingConfig(weight_mode=WeightMode.FUSED_QKV_LINEAR), + reduce_output=False, + skip_create_weights_in_init=skip_create_weights_in_init, + force_dynamic_quantization=force_dynamic_quantization, + disable_deep_gemm=disable_deep_gemm, + use_custom_cublas_mm=use_custom_cublas_mm, + use_cute_dsl_bf16_gemm=use_cute_dsl_bf16_gemm, + use_cute_dsl_blockscaling_mm=use_cute_dsl_blockscaling_mm, + ) + + def _shard_geometry(self, shard_name: str) -> Tuple[int, int]: + """Return effective (world size, rank) for one checkpoint shard.""" + if shard_name == "q": + return self.tp_size, self.tp_rank + if shard_name == "index_k": + return 1, 0 + + total_heads = ( + self.total_num_index_heads if shard_name == "index_q" else self.total_num_kv_heads + ) + if self.tp_size <= total_heads: + return self.tp_size, self.tp_rank + replicas = self.tp_size // total_heads + return total_heads, self.tp_rank // replicas + + def load_five_way_weights(self, shards: Dict[str, Dict]) -> None: + """Load five checkpoint projections into this rank's packed MXFP8 matrix.""" + local_shards: Dict[str, Dict[str, torch.Tensor]] = {} + for shard_name in self._SHARD_NAMES: + shard = shards[shard_name] + if "weight" not in shard: + raise KeyError(f"Missing {shard_name} projection weight.") + shard_world, shard_rank = self._shard_geometry(shard_name) + local = {} + for key in ("weight", "weight_scale_inv", "weight_scale", "bias"): + if key in shard: + local[key] = load_weight_shard( + shard[key], + shard_world, + shard_rank, + TensorParallelMode.COLUMN, + device=torch.device("cuda"), + ) + local_shards[shard_name] = local + + combined: Dict[str, torch.Tensor] = { + "weight": torch.cat( + [local_shards[name]["weight"] for name in self._SHARD_NAMES], dim=0 + ).contiguous() + } + for key in ("weight_scale_inv", "weight_scale", "bias"): + present = [key in local_shards[name] for name in self._SHARD_NAMES] + if any(present): + if not all(present): + raise KeyError(f"Incomplete {key} across fused QKV+index shards.") + combined[key] = torch.cat( + [local_shards[name][key] for name in self._SHARD_NAMES], dim=0 + ).contiguous() + + # The checkpoint tensors above are already rank-local and packed. + # Temporarily select vanilla loading so Linear copies them without a + # second TP split or the three-shard QKV loader. + saved_tp_size = self.tp_size + saved_tp_rank = self.tp_rank + saved_tp_mode = self.tp_mode + saved_loading_config = self.weights_loading_config + try: + self.tp_size = 1 + self.tp_rank = 0 + self.tp_mode = None + self.weights_loading_config = WeightsLoadingConfig(weight_mode=WeightMode.VANILLA) + self.load_weights([combined]) + finally: + self.tp_size = saved_tp_size + self.tp_rank = saved_tp_rank + self.tp_mode = saved_tp_mode + self.weights_loading_config = saved_loading_config + + # --------------------------------------------------------------------------- # Config normalization helpers # --------------------------------------------------------------------------- @@ -176,6 +340,24 @@ def _validate_sparse_attention_runtime_config( "Set the following in the LLM API configuration:\n" "sparse_attention_config:\n algorithm: minimax_m3" ) + if getattr(sparse_config, "fuse_qkv_index_projection", False): + if getattr(sparse_config, "implementation", None) != "msa": + raise ValueError( + "MiniMax-M3 fuse_qkv_index_projection=True requires the 'msa' implementation." + ) + if getattr(sparse_config, "indexer_kv_dtype", None) != "fp8": + raise ValueError( + "MiniMax-M3 fuse_qkv_index_projection=True requires indexer_kv_dtype='fp8'." + ) + quant_config = model_config.quant_config + if ( + quant_config is None + or quant_config.quant_mode is None + or not quant_config.quant_mode.has_fp8_kv_cache() + ): + raise ValueError( + "MiniMax-M3 fuse_qkv_index_projection=True requires an FP8 main KV cache." + ) def get_sparse_layer_ids(text_config: PretrainedConfig) -> Tuple[List[int], List[int]]: @@ -622,23 +804,82 @@ def _extract_minimax_m3_attention_extra_attrs(layer_idx: str): return metadata, attn_layer +@torch.library.custom_op("trtllm::minimax_m3_qkv_index_proj", mutates_args=()) +def minimax_m3_qkv_index_proj( + hidden_states: torch.Tensor, + position_ids: Optional[torch.Tensor], + layer_idx: str, +) -> torch.Tensor: + """Run the five-way projection as one CUDA-graph-capturable custom op.""" + del position_ids # Symbolic token-shape carrier; projection itself is position agnostic. + _, attn_layer = _extract_minimax_m3_attention_extra_attrs(layer_idx) + return attn_layer.qkv_proj(hidden_states) + + +@minimax_m3_qkv_index_proj.register_fake +def _minimax_m3_qkv_index_proj_fake( + hidden_states: torch.Tensor, + position_ids: Optional[torch.Tensor], + layer_idx: str, +) -> torch.Tensor: + """Preserve the symbolic token dimension across the opaque projection.""" + _, attn_layer = _extract_minimax_m3_attention_extra_attrs(layer_idx) + qkv_proj = attn_layer.qkv_proj + if not isinstance(qkv_proj, MiniMaxM3QKVIndexerLinear): + raise RuntimeError( + "MiniMax-M3 fused QKV/index projection custom op requires " + f"MiniMaxM3QKVIndexerLinear, got {type(qkv_proj).__name__}." + ) + # The real projection is token-major over ``hidden_states``. Keep that + # exact output contract here so generation CUDA-graph batch sizes can + # share Dynamo's dynamic-shape specialization. ``position_ids`` remains + # an explicit custom-op input to carry the unpadded token symbol into + # piecewise context segments; the attention output below uses that symbol. + return hidden_states.new_empty((hidden_states.shape[0], sum(qkv_proj.local_output_sizes))) + + @torch.library.custom_op("trtllm::minimax_m3_attn_custom_op_inplace", mutates_args=("output",)) def minimax_m3_attn_custom_op_inplace( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, + q: Optional[torch.Tensor], + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], idx_q: Optional[torch.Tensor], idx_k: Optional[torch.Tensor], + packed_qkv_index: Optional[torch.Tensor], + position_ids: Optional[torch.Tensor], layer_idx: str, output: torch.Tensor, ) -> None: - """Run MiniMax-M3 cache and attention work behind a compile boundary.""" + """Run MiniMax-M3 cache and attention work behind a compile boundary. + + The horizontal producer needs live paged-cache tensors and cache-slot + metadata, which are intentionally resolved inside this opaque attention + boundary rather than traced through Dynamo. Projection remains in the + captured segment; only the cache-writing producer and MSA attention stay + on the eager side of the existing piecewise boundary. + """ attn_metadata, attn_layer = _extract_minimax_m3_attention_extra_attrs(layer_idx) num_tokens = attn_metadata.num_tokens + if packed_qkv_index is not None: + horizontal = attn_layer._fused_fp8_qkv_indexer_norm_rope_kv_insert( + packed_qkv_index[:num_tokens], + position_ids[..., :num_tokens] if position_ids is not None else None, + attn_metadata, + ) + if horizontal is None: + raise RuntimeError( + "MiniMax-M3 fused horizontal producer is required inside the " + f"piecewise attention boundary for layer {layer_idx}, but its " + "runtime geometry or cache layout was unsupported." + ) + q, idx_q = horizontal + k = v = idx_k = None + if q is None: + raise RuntimeError(f"MiniMax-M3 attention layer {layer_idx} received no query tensor.") attn_layer._dispatch_attention_backend( q[:num_tokens], - k[:num_tokens], - v[:num_tokens], + k[:num_tokens] if k is not None else None, + v[:num_tokens] if v is not None else None, idx_q[:num_tokens] if idx_q is not None else None, idx_k[:num_tokens] if idx_k is not None else None, attn_metadata, @@ -737,9 +978,16 @@ def __init__( self.is_sparse_attention_layer = bool(is_sparse_attention_layer) self.disable_index_value = bool(disable_index_value) + sparse_runtime_cfg = getattr(model_config, "sparse_attention_config", None) + self.enable_fused_qkv_index_projection = bool( + self.is_sparse_attention_layer + and sparse_runtime_cfg is not None + and getattr(sparse_runtime_cfg, "fuse_qkv_index_projection", False) + ) if self.is_sparse_attention_layer: sparse_cfg = getattr(config, "sparse_attention_config", None) or {} - self.sparse_num_index_heads = int(sparse_cfg.get("sparse_num_index_heads", 4)) + total_num_index_heads = int(sparse_cfg.get("sparse_num_index_heads", 4)) + self.sparse_num_index_heads = total_num_index_heads self.sparse_index_dim = int(sparse_cfg.get("sparse_index_dim", 128)) self.sparse_block_size = int(sparse_cfg.get("sparse_block_size", 128)) self.sparse_topk_blocks = int(sparse_cfg.get("sparse_topk_blocks", 16)) @@ -747,25 +995,57 @@ def __init__( self.sparse_local_block = int(sparse_cfg.get("sparse_local_block", 1)) self.sparse_score_type = str(sparse_cfg.get("sparse_score_type", "max")) - # Index Q and K are both replicated and project the same - # hidden_states, so fuse them into one GEMM with output - # [idx_q | idx_k]. idx_q holds all index heads; idx_k is a single K - # per token, broadcast across heads when scoring. + if self.enable_fused_qkv_index_projection: + old_qkv_proj = self.qkv_proj + self.qkv_proj = MiniMaxM3QKVIndexerLinear( + hidden_size=config.hidden_size, + head_dim=self.head_dim, + total_num_heads=config.num_attention_heads, + total_num_kv_heads=config.num_key_value_heads, + total_num_index_heads=total_num_index_heads, + index_head_dim=self.sparse_index_dim, + dtype=config.torch_dtype, + mapping=old_qkv_proj.mapping, + quant_config=old_qkv_proj.quant_config, + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + force_dynamic_quantization=old_qkv_proj.force_dynamic_quantization, + disable_deep_gemm=old_qkv_proj.disable_deep_gemm, + use_custom_cublas_mm=old_qkv_proj.use_custom_cublas_mm, + use_cute_dsl_bf16_gemm=old_qkv_proj.use_cute_dsl_bf16_gemm, + use_cute_dsl_blockscaling_mm=old_qkv_proj.use_cute_dsl_blockscaling_mm, + ) + self.sparse_num_index_heads = self.qkv_proj.local_num_index_heads + else: + # Use the same KV-group ownership as the five-way projection. + # Linear's existing per-shard override handles index-Q slicing + # and index-K replication without a custom checkpoint loader. + index_start, index_end = index_head_range( + total_num_index_heads, config.num_key_value_heads, self.qkv_proj.mapping + ) + self.sparse_num_index_heads = index_end - index_start + self.index_qk_proj = Linear( + config.hidden_size, + total_num_index_heads * self.sparse_index_dim + self.sparse_index_dim, + bias=False, + dtype=config.torch_dtype, + mapping=self.qkv_proj.mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + quant_config=None, + weights_loading_config=WeightsLoadingConfig( + weight_mode=WeightMode.FUSED_GATE_UP_LINEAR + ), + override_tp_sharding={ + "gate": ( + index_start * self.sparse_index_dim, + index_end * self.sparse_index_dim, + ), + "up": (0, self.sparse_index_dim), + }, + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + ) + self.index_q_size = self.sparse_num_index_heads * self.sparse_index_dim self.index_k_size = self.sparse_index_dim - self.index_qk_proj = Linear( - config.hidden_size, - self.index_q_size + self.index_k_size, - bias=False, - dtype=config.torch_dtype, - mapping=model_config.mapping, - tensor_parallel_mode=None, - quant_config=None, - weights_loading_config=WeightsLoadingConfig( - weight_mode=WeightMode.FUSED_GATE_UP_LINEAR - ), - skip_create_weights_in_init=model_config.skip_create_weights_in_init, - ) # Per-head Gemma RMSNorm of width ``sparse_index_dim``; # applied to the projected index Q/K before partial RoPE in # the sparse forward path. @@ -820,7 +1100,7 @@ def apply_index_qk_norm( Mirrors :meth:`apply_qk_norm` for the index branch: reshapes ``idx_q`` (shape ``[..., num_index_heads * sparse_index_dim]``, - possibly TP-sharded along the head axis) and ``idx_k`` (shape + sharded with the KV heads) and ``idx_k`` (shape ``[..., sparse_index_dim]`` — single replicated head, not per index head) to per-head rows of width ``sparse_index_dim``, applies :attr:`index_q_norm` / :attr:`index_k_norm`, and reshapes @@ -1007,6 +1287,131 @@ def _fused_fp8_index_qk_norm_rope( position_ids.reshape(-1).contiguous().to(torch.int32), ).flatten(1) + def _fused_fp8_qkv_indexer_norm_rope_kv_insert( + self, + packed: torch.Tensor, + position_ids: Optional[torch.Tensor], + attn_metadata: AttentionMetadata, + ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """Run the vLLM-style horizontal producer for every sparse batch. + + The CUDA kernel is token-major and batch-type agnostic: per-token + positions and cache slots cover pure prefill, mixed aggregate batches, + and CUDA-graph decode. + """ + if ( + not self.enable_fused_qkv_index_projection + or not isinstance(self.attn, MiniMaxM3MsaSparseAttention) + or not self._emit_fp8_main_qkv() + or self.attn.indexer_kv_dtype != "fp8" + ): + return None + rope = self.pos_embd_params.rope if self.pos_embd_params is not None else None + rotary_dim = int(rope.dim) if rope is not None else 0 + if ( + packed.dtype != torch.bfloat16 + or position_ids is None + or self.head_dim != 128 + or self.sparse_index_dim != 128 + or self.sparse_num_index_heads != self.num_key_value_heads + or not self.use_gemma_norm + or self.pos_embd_params is None + or not self.pos_embd_params.is_neox + or self.rotary_emb is None + or rope is None + or rotary_dim != 64 + ): + return None + norm_eps = self.q_norm.variance_epsilon + if any( + module.variance_epsilon != norm_eps + for module in (self.k_norm, self.index_q_norm, self.index_k_norm) + ): + return None + norm_weights = ( + self.q_norm.weight, + self.k_norm.weight, + self.index_q_norm.weight, + self.index_k_norm.weight, + ) + if any(weight.dtype != torch.bfloat16 or not weight.is_cuda for weight in norm_weights): + return None + + kv_cache_manager = getattr(attn_metadata, "kv_cache_manager", None) + if kv_cache_manager is None: + return None + buffers = kv_cache_manager.get_buffers(self.layer_idx, kv_layout="HND") + index_k_cache = attn_metadata.msa_idx_k_cache(self.layer_idx) + out_cache_loc = getattr(attn_metadata, "msa_out_cache_loc", None) + num_tokens = int(packed.shape[0]) + supported_main_cache = ( + buffers is not None + and buffers.is_cuda + and buffers.dtype == torch.float8_e4m3fn + and buffers.dim() == 5 + and tuple(buffers.shape[1:]) == (2, self.num_key_value_heads, 128, 128) + and buffers.stride(4) == 1 + and buffers.stride(3) == 128 + and buffers.stride(2) == 128 * 128 + and buffers.stride(1) >= self.num_key_value_heads * buffers.stride(2) + and buffers.stride(0) >= 2 * buffers.stride(1) + and buffers.stride(0) % 4 == 0 + and buffers.stride(1) % 4 == 0 + ) + supported_index_cache = ( + index_k_cache.is_cuda + and index_k_cache.dtype == torch.float8_e4m3fn + and index_k_cache.dim() == 4 + and tuple(index_k_cache.shape[1:]) == (1, 128, 128) + and index_k_cache.stride(3) == 1 + and index_k_cache.stride(2) == 128 + and index_k_cache.stride(1) >= 128 * 128 + and index_k_cache.stride(0) >= index_k_cache.stride(1) + and index_k_cache.stride(0) % 4 == 0 + and index_k_cache.stride(1) % 4 == 0 + and buffers is not None + and index_k_cache.shape[0] == buffers.shape[0] + ) + rotary_cos_sin = self.rotary_emb.rotary_cos_sin + supported_rope_cache = ( + rotary_cos_sin.is_cuda + and rotary_cos_sin.dtype == torch.float32 + and rotary_cos_sin.is_contiguous() + and rotary_cos_sin.dim() == 3 + and tuple(rotary_cos_sin.shape[1:]) == (2, 32) + ) + if ( + not supported_main_cache + or not supported_index_cache + or not supported_rope_cache + or out_cache_loc is None + or not out_cache_loc.is_cuda + or out_cache_loc.dtype != torch.int32 + or not out_cache_loc.is_contiguous() + or out_cache_loc.numel() < num_tokens + ): + return None + + q, index_q = torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + packed.contiguous(), + buffers, + index_k_cache, + out_cache_loc[:num_tokens], + self.num_heads, + self.num_key_value_heads, + self.sparse_num_index_heads, + self.head_dim, + rotary_dim, + norm_eps, + self.q_norm.weight, + self.k_norm.weight, + self.index_q_norm.weight, + self.index_k_norm.weight, + rotary_cos_sin, + position_ids.reshape(-1).contiguous().to(torch.int32), + ) + return q.flatten(1), index_q.flatten(1) + def _expect_fused_qk_norm_rope(self, position_ids: Optional[torch.Tensor]) -> bool: """Whether the fused kernel is expected to run instead of the fallback. @@ -1361,8 +1766,8 @@ def _sdpa_dense_attention_core( def _forward_attention_core( self, q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], idx_q: Optional[torch.Tensor], idx_k: Optional[torch.Tensor], attn_metadata: AttentionMetadata, @@ -1380,6 +1785,8 @@ def _forward_attention_core( v, idx_q, idx_k, + None, + None, self.layer_idx_str, output, ) @@ -1390,8 +1797,8 @@ def _forward_attention_core( def _dispatch_attention_backend( self, q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], idx_q: Optional[torch.Tensor], idx_k: Optional[torch.Tensor], attn_metadata: AttentionMetadata, @@ -1409,6 +1816,7 @@ def _dispatch_attention_backend( """ if isinstance(self.attn, MiniMaxM3MsaSparseAttention): return self._msa_attention_core(q, k, v, idx_q, idx_k, attn_metadata, output) + assert k is not None and v is not None if self.is_sparse_attention_layer: assert idx_q is not None and idx_k is not None return self._triton_sparse_attention_core(q, k, v, idx_q, idx_k, attn_metadata, output) @@ -1418,8 +1826,8 @@ def _dispatch_attention_backend( def _msa_attention_core( self, q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], idx_q: Optional[torch.Tensor], idx_k: Optional[torch.Tensor], attn_metadata: AttentionMetadata, @@ -1437,11 +1845,16 @@ def _msa_attention_core( then receives k=v=None, which is the backend's contract for "K/V are already resident", so neither FMHA phase writes them again. """ + assert (k is None) == (v is None) if self.is_sparse_attention_layer: assert idx_q is not None # On the FP8 indexer path idx_k is None: the fused producer already # inserted E4M3 index-K into the side cache, so only K/V are written. - self.attn.write_layer_caches(k, v, idx_k, attn_metadata) + if k is not None: + self.attn.write_layer_caches(k, v, idx_k, attn_metadata) + else: + # The horizontal producer has already written both caches. + assert idx_k is None # Publish the selected blocks so the FMHA runs the sparse path. # idx_k_prewritten: index-K is already in the cache (written above # on bf16, or by the FP8 producer), so run_indexer must not write it. @@ -1454,7 +1867,8 @@ def _msa_attention_core( ) else: assert idx_q is None and idx_k is None - self.attn.write_layer_caches(k, v, None, attn_metadata) + if k is not None: + self.attn.write_layer_caches(k, v, None, attn_metadata) # No top-k selection means the FMHA attends the full page table. forward_args = AttentionForwardArgs(output=output) self.attn.forward(q, None, None, attn_metadata, forward_args=forward_args) @@ -1508,11 +1922,58 @@ def _sparse_forward( "attn_metadata; received None." ) - # Project, norm, and apply RoPE for the main and index branches. Both - # read only hidden_states and write disjoint outputs, so they overlap on - # the aux stream and join before the attention core. + # The opt-in projection emits [Q|K|V|index-Q|index-K] with one GEMM. + # Its horizontal producer writes both paged caches directly and returns + # only compact Q/index-Q. Unsupported compile/geometry/layout cases + # split the packed output and retain the existing producer paths. + packed_qkv = None + packed_idx_qk = None + if self.enable_fused_qkv_index_projection: + if self.register_to_config and (is_torch_compiling() or is_in_breakable_cuda_graph()): + # Keep the projection in the captured segment while hiding + # its shape-specializing MXFP8 internals behind a symbolic + # fake implementation. Cache insertion and MSA remain in the + # existing eager attention boundary below. + packed = torch.ops.trtllm.minimax_m3_qkv_index_proj( + hidden_states, position_ids, self.layer_idx_str + ) + num_tokens = ( + position_ids.shape[-1] if position_ids is not None else hidden_states.shape[0] + ) + output = packed.new_empty( + (num_tokens, self.num_heads * self.head_dim), + dtype=self.attn_activation_dtype, + ) + maybe_bcg_minimax_m3_attn_custom_op_inplace( + None, + None, + None, + None, + None, + packed, + position_ids, + self.layer_idx_str, + output, + ) + return self.o_proj(output, all_reduce_params=all_reduce_params) + packed = self.qkv_proj(hidden_states) + horizontal = self._fused_fp8_qkv_indexer_norm_rope_kv_insert( + packed, position_ids, attn_metadata + ) + if horizontal is not None: + q, idx_q = horizontal + o = self._forward_attention_core(q, None, None, idx_q, None, attn_metadata) + return self.o_proj(o, all_reduce_params=all_reduce_params) + main_size = self.q_size + 2 * self.kv_size + packed_qkv, packed_idx_qk = packed.split( + [main_size, self.index_q_size + self.index_k_size], dim=-1 + ) + + # Project, norm, and apply RoPE for the compatibility/fallback main and + # index branches. Both read only hidden_states and write disjoint + # outputs, so they overlap on the aux stream and join before attention. def _main_norm_rope(): - qkv = self.qkv_proj(hidden_states) + qkv = packed_qkv if packed_qkv is not None else self.qkv_proj(hidden_states) fused_qkv = self._fused_qk_norm_rope( qkv, position_ids, @@ -1539,7 +2000,9 @@ def _main_norm_rope(): return q, k, v def _index_norm_rope(): - idx_qk = self.index_qk_proj(hidden_states) + idx_qk = ( + packed_idx_qk if packed_idx_qk is not None else self.index_qk_proj(hidden_states) + ) fp8_idx_q = self._fused_fp8_index_qk_norm_rope(idx_qk, position_ids, attn_metadata) if fp8_idx_q is not None: # Index-K was inserted directly into the paged side cache. @@ -2077,6 +2540,32 @@ def _load_index_qk_proj_weights(model: nn.Module, weights) -> None: del weights[key] +def _load_qkv_index_proj_weights(model: nn.Module, weights) -> List[str]: + """Pack five checkpoint shards and return generic-loader module skips.""" + checkpoint_names = ("q_proj", "k_proj", "v_proj", "index_q_proj", "index_k_proj") + shard_names = MiniMaxM3QKVIndexerLinear._SHARD_NAMES + loaded_modules = [] + for name, module in model.named_modules(): + if not isinstance(module, MiniMaxM3QKVIndexerLinear): + continue + parent = name.rsplit(".", 1)[0] + shards = { + shard_name: filter_weights(f"{parent}.{checkpoint_name}", weights) + for shard_name, checkpoint_name in zip(shard_names, checkpoint_names, strict=True) + } + module.load_five_way_weights(shards) + loaded_modules.append(name) + for checkpoint_name in checkpoint_names: + prefix = f"{parent}.{checkpoint_name}" + if hasattr(weights, "mark_consumed"): + weights.mark_consumed(prefix) + else: + for key in list(weights.keys()): + if key.startswith(f"{prefix}."): + del weights[key] + return loaded_modules + + # Layer-boundary RMSNorms whose Gemma (1 + weight) scaling is folded into the # stored weight at load time so the runtime norm is a plain RMSNorm (see # MiniMaxM3DecoderLayer.__init__ / MiniMaxM3Model.__init__). These are exactly @@ -2152,8 +2641,9 @@ def load_weights( params_map: Optional[Dict[str, str]] = None, allow_partial_loading: bool = False, ) -> None: - # The generic loader has no rule for this fusion. The VL subclass routes - # its text weights through here, so both paths are covered. + # Pack the opt-in five-way projection with MiniMax-specific TP + # sharding before the generic mapper handles the compatibility path. + packed_projection_modules = _load_qkv_index_proj_weights(self, weights) _load_index_qk_proj_weights(self, weights) # Fold Gemma (1 + weight) into the layer-boundary RMSNorm weights so the # runtime norms can be plain (non-Gemma) and drive the fused @@ -2163,6 +2653,9 @@ def load_weights( weights = _fold_gemma_boundary_norm_weights(weights) if weight_mapper is None: weight_mapper = MiniMaxM3HfWeightMapper() + # The generic mapper understands three-way QKV fusion, not the + # already-loaded five-way Q/K/V/index-Q/index-K modules. + weight_mapper.add_skip_modules(packed_projection_modules) weight_mapper.init_model_and_config(self, self.model_config) merged_params_map = {**MINIMAX_M3_PARAMS_MAP, **(params_map or {})} super().load_weights( @@ -2180,6 +2673,7 @@ def setup_aliases(self) -> None: layer's input_layernorm; the last layer chains the final model norm so its output AllReduce folds the final normalization too. """ + super().setup_aliases() layers = self.model.layers num_layers = len(layers) for idx, layer in enumerate(layers): diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 27dfc35fcafc..0b7dc7bfab57 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -859,7 +859,9 @@ class MiniMaxM3SparseAttentionConfig(BaseSparseAttentionConfig): algorithm: Literal["minimax_m3"] = "minimax_m3" sparse_num_index_heads: PositiveInt = Field( default=4, - description="Number of index-attention heads (per TP rank's view).", + description= + "Global checkpoint index-attention head count. Index heads shard with " + "their KV-head groups in both separate and fused projections.", ) sparse_index_dim: int = Field( default=128, @@ -899,6 +901,17 @@ class MiniMaxM3SparseAttentionConfig(BaseSparseAttentionConfig): "by the MSA implementation.", status="prototype", ) + fuse_qkv_index_projection: bool = Field( + default=False, + description= + "Fuse Q/K/V and index-Q/index-K into one quantized projection. Index-Q " + "is sharded with the KV heads and index-K is replicated. MSA batches " + "also use a horizontal norm/RoPE/cache-insertion producer for prefill, " + "mixed, and CUDA-graph decode execution. The MiniMax-M3-specific path " + "requires the MSA implementation, indexer_kv_dtype='fp8', and an FP8 " + "main KV cache.", + status="prototype", + ) num_attention_heads: Optional[int] = Field( default=None, description= @@ -933,6 +946,14 @@ def _validate_msa_configuration(self): if self.indexer_kv_dtype == "fp8" and not self.sparse_disable_index_value: raise ValueError("MiniMax-M3 indexer_kv_dtype='fp8' requires " "sparse_disable_index_value=True.") + if self.fuse_qkv_index_projection and self.implementation != "msa": + raise ValueError( + "MiniMax-M3 fuse_qkv_index_projection=True currently requires " + "the 'msa' implementation.") + if self.fuse_qkv_index_projection and self.indexer_kv_dtype != "fp8": + raise ValueError( + "MiniMax-M3 fuse_qkv_index_projection=True currently requires " + "indexer_kv_dtype='fp8'.") return self def supports_backend(self, backend: str) -> bool: @@ -947,6 +968,9 @@ def to_sparse_params(self, **kwargs): return MiniMaxM3SparseParams( num_index_heads=self.sparse_num_index_heads, + global_num_kv_heads=( + self.to_sparse_metadata_params(**kwargs).global_num_kv_heads + or None), sparse_index_dim=self.sparse_index_dim, block_size=self.sparse_block_size, topk=self.sparse_topk_blocks, @@ -956,6 +980,7 @@ def to_sparse_params(self, **kwargs): disable_index_value=self.sparse_disable_index_value, implementation=self.implementation, indexer_kv_dtype=self.indexer_kv_dtype, + fuse_qkv_index_projection=self.fuse_qkv_index_projection, ) def to_sparse_metadata_params(self, **kwargs): diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 943cbfb3d0f6..469601b2ece8 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1277,6 +1277,11 @@ "kind": "value", "path": "sparse_attention_config.enable_heuristic_topk" }, + { + "capture_policy": "bool", + "kind": "value", + "path": "sparse_attention_config.fuse_qkv_index_projection" + }, { "allowed_values": [ "triton", diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 7c90aae98df9..ff6459cc211d 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -7358,11 +7358,13 @@ def test_nvfp4(self, use_msa): @pytest.mark.skip_less_device(4) @pytest.mark.skip_less_device_memory(140000) @parametrize_with_ids("eval_mode", ["default", "inferencex"]) + @parametrize_with_ids("fuse_qkv_index_projection", [False, True]) @parametrize_with_ids("overlap_scheduler", [False, True]) @parametrize_with_ids("attention_dp", [False, True]) @parametrize_with_ids("tp_size,ep_size", [(4, 4)]) def test_nvfp4_eagle3(self, tp_size, ep_size, attention_dp, - overlap_scheduler, eval_mode): + overlap_scheduler, fuse_qkv_index_projection, + eval_mode): # One-model Eagle3 on the MSA backend with an FP8 KV cache and CUDA # graphs; the GQA drafter shares the target KV cache. MMLU + GSM8K, or # InferenceX GSM8K, plus a chat-GSM8K acceptance probe, since accuracy @@ -7397,7 +7399,9 @@ def test_nvfp4_eagle3(self, tp_size, ep_size, attention_dp, moe_expert_parallel_size=ep_size, kv_cache_config=kv_cache_config, sparse_attention_config=MiniMaxM3SparseAttentionConfig( - implementation="msa", indexer_kv_dtype="fp8"), + implementation="msa", + indexer_kv_dtype="fp8", + fuse_qkv_index_projection=fuse_qkv_index_projection), moe_config=MoeConfig(backend="CUTLASS"), max_seq_len=max_seq_len, max_batch_size=max_batch_size, diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 24818d8a8ef4..4cad54969aa3 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -462,7 +462,9 @@ accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] TIMEOUT (180) accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True] accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True] TIMEOUT (60) -accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True-eval_mode=inferencex] +accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True-fuse_qkv_index_projection=False-eval_mode=inferencex] +accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True-fuse_qkv_index_projection=True-eval_mode=inferencex] +accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=True-overlap_scheduler=True-fuse_qkv_index_projection=True-eval_mode=inferencex] accuracy/test_llm_api_pytorch.py::TestMinistral8BInstruct::test_auto_dtype accuracy/test_llm_api_pytorch.py::TestMinistral8BInstruct::test_fp8 accuracy/test_llm_api_pytorch.py::TestMistralLarge3_675B::test_fp8[latency_moe_deepgemm] diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 55c319f1d36b..58d74f586db6 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -31,6 +31,10 @@ l0_b200: backend: pytorch tests: # ------------- PyTorch tests --------------- + - unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_horizontal_producer.py + - unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py + - unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_main_kv_insert.py + - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_unfused_index_projection_tp_construction - unittest/others/test_kv_cache_transceiver.py::test_cpp_nixl_sync_transfer_stress - unittest/others/test_lora_manager.py - accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_nvfp4 diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 0e7e5e66f7f1..38247c9ad825 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -47,6 +47,20 @@ l0_cpu: - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_fp8_indexer_rejects_different_qk_norm_epsilons - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_moe_reduces_only_local_terms - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_decoder_layer_sets_post_fusion_from_moe_scheduler + - unittest/_torch/models/test_minimax_m3.py::test_validate_fused_projection_requires_fp8_main_kv_cache + - unittest/_torch/models/test_minimax_m3.py::test_fused_qkv_index_projection_preserves_index_head_groups + - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_five_way_projection_shards_index_rows + - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_index_tp_matches_unsharded_reference + - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_index_head_range_rejects_invalid_geometry + - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_uses_one_engine_speculative_base + - unittest/_torch/models/test_minimax_m3.py::test_setup_aliases_preserves_one_engine_draft_weight_loading + - unittest/_torch/models/test_minimax_m3.py::test_eagle_capture_precedes_next_layer_norm + - unittest/_torch/models/test_minimax_m3.py::test_piecewise_attention_boundary_runs_horizontal_producer + - unittest/_torch/models/test_minimax_m3.py::test_piecewise_projection_fake_preserves_padded_hidden_rows + - unittest/_torch/models/test_minimax_m3.py::test_piecewise_fused_projection_preserves_input_token_dimension + - unittest/_torch/models/test_minimax_m3.py::test_msa_attention_core_routes_compact_q_to_attention_dispatcher + - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_five_way_projection_shard_geometry + - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_five_way_loader_returns_exact_generic_skip - unittest/_torch/models/checkpoints - unittest/_torch/modules - unittest/_torch/moe diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 22652d73f8d1..61a9b0881f18 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -54,7 +54,9 @@ l0_dgx_b200: - disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt_oss/gpt-oss-120b] - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[latency_adp_lmtp_tp4] - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True] TIMEOUT (60) - - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True-eval_mode=inferencex] + - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True-fuse_qkv_index_projection=False-eval_mode=inferencex] + - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True-fuse_qkv_index_projection=True-eval_mode=inferencex] + - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=True-overlap_scheduler=True-fuse_qkv_index_projection=True-eval_mode=inferencex] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] TIMEOUT (60) - unittest/_torch/modeling/test_modeling_deepseekv4.py - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_auto_dtype TIMEOUT (60) diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index e5df14afd095..59564a029292 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -31,7 +31,15 @@ from transformers import AutoConfig from utils.llm_data import llm_models_root +import tensorrt_llm._torch.models.modeling_minimaxm3 as modeling_minimaxm3 from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.common import ( + MiniMaxM3SparseConfig, + MiniMaxM3SparseMetadataParams, + MiniMaxM3SparseParams, + index_head_range, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_indexer import _group_max_reduce from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.checkpoints.hf.minimaxm3_weight_mapper import ( MiniMaxM3HfWeightMapper, @@ -39,9 +47,12 @@ from tensorrt_llm._torch.models.modeling_minimaxm3 import ( MiniMaxM3Attention, MiniMaxM3DecoderLayer, + MiniMaxM3ForCausalLM, MiniMaxM3Model, MiniMaxM3MoE, + MiniMaxM3QKVIndexerLinear, _build_swiglu_oai_dense_mlp, + _load_qkv_index_proj_weights, _minimax_m3_swiglu_oai, _moe_routed_output_is_global, _strip_language_model_prefix, @@ -53,6 +64,7 @@ get_text_config, is_minimax_m3_vl_config, ) +from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM from tensorrt_llm._torch.models.modeling_utils import _load_weights_impl_v2 from tensorrt_llm._torch.modules.rms_norm import RMSNorm from tensorrt_llm._torch.moe.fused_moe.interface import MoESchedulerKind @@ -63,6 +75,8 @@ from tensorrt_llm._torch.utils import AuxStreamType, EventType from tensorrt_llm.llmapi import MiniMaxM3SparseAttentionConfig, RocketSparseAttentionConfig from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo # --------------------------------------------------------------------------- # Fixtures @@ -239,6 +253,311 @@ def test_validate_sparse_attention_runtime_config_accepts_minimax_m3() -> None: _validate_sparse_attention_runtime_config(model_config) +@pytest.mark.cpu_only +def test_validate_fused_projection_requires_fp8_main_kv_cache() -> None: + sparse_config = MiniMaxM3SparseAttentionConfig( + implementation="msa", + indexer_kv_dtype="fp8", + fuse_qkv_index_projection=True, + ) + model_config = ModelConfig( + pretrained_config=_make_text_config(), + sparse_attention_config=sparse_config, + ) + with pytest.raises(ValueError, match="requires an FP8 main KV cache"): + _validate_sparse_attention_runtime_config(model_config) + + model_config.quant_config = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) + _validate_sparse_attention_runtime_config(model_config) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) +@pytest.mark.parametrize("attention_dp", [False, True]) +def test_fused_qkv_index_projection_preserves_index_head_groups( + tp_size: int, attention_dp: bool +) -> None: + cfg = MiniMaxM3SparseAttentionConfig( + implementation="msa", + indexer_kv_dtype="fp8", + fuse_qkv_index_projection=True, + num_attention_heads=64, + num_key_value_heads=4, + ) + sparse_params = cfg.to_sparse_params() + metadata_params = cfg.to_sparse_metadata_params() + checkpoint_cfg = cfg.model_copy( + update={"num_attention_heads": None, "num_key_value_heads": None} + ) + checkpoint = SimpleNamespace(num_attention_heads=64, num_key_value_heads=4) + assert checkpoint_cfg.to_sparse_params(pretrained_config=checkpoint) == sparse_params + assert checkpoint_cfg.to_sparse_metadata_params(pretrained_config=checkpoint) == metadata_params + mapping = SimpleNamespace(tp_size=tp_size, tp_rank=0, enable_attention_dp=attention_dp) + local_q_heads = 64 if attention_dp else 64 // tp_size + local_kv_heads = 4 if attention_dp else max(4 // tp_size, 1) + + assert sparse_params.fuse_qkv_index_projection is True + assert metadata_params.sharded_head_counts(mapping) == (local_q_heads, local_kv_heads) + assert metadata_params.num_index_heads == 4 + assert metadata_params.sharded_index_head_count(mapping) == local_kv_heads + kernel_cfg = MiniMaxM3SparseConfig.from_sparse_params( + sparse_params, num_q_heads=local_q_heads, num_kv_heads=local_kv_heads, head_dim=128 + ) + assert kernel_cfg.num_index_heads == local_kv_heads + + compatibility_cfg = cfg.model_copy(update={"fuse_qkv_index_projection": False}) + assert compatibility_cfg.to_sparse_metadata_params() == metadata_params + compatibility_kernel_cfg = MiniMaxM3SparseConfig.from_sparse_params( + compatibility_cfg.to_sparse_params(), + num_q_heads=local_q_heads, + num_kv_heads=local_kv_heads, + head_dim=128, + ) + # Rank zero keeps its own index/KV pairs, not a max over other ranks' + # index heads. Each head favors a distinct block to expose misgrouping. + scores = torch.diag(torch.tensor([1.0, 2.0, 3.0, 4.0])).unsqueeze(-1) + fused_scores = _group_max_reduce(scores[:local_kv_heads], kernel_cfg) + reference_scores = _group_max_reduce(scores[:local_kv_heads], compatibility_kernel_cfg) + torch.testing.assert_close(fused_scores, reference_scores) + expected_blocks = torch.arange(local_kv_heads) + torch.testing.assert_close(fused_scores.argmax(dim=1).flatten(), expected_blocks) + with pytest.raises(ValueError, match=r"requires the 'msa' implementation"): + MiniMaxM3SparseAttentionConfig(implementation="triton", fuse_qkv_index_projection=True) + with pytest.raises(ValueError, match=r"requires indexer_kv_dtype='fp8'"): + MiniMaxM3SparseAttentionConfig(implementation="msa", fuse_qkv_index_projection=True) + + +@pytest.mark.cpu_only +def test_minimax_m3_uses_one_engine_speculative_base() -> None: + assert issubclass(MiniMaxM3ForCausalLM, SpecDecOneEngineForCausalLM) + + +@pytest.mark.cpu_only +def test_setup_aliases_preserves_one_engine_draft_weight_loading() -> None: + loaded = [] + + class DraftModel: + shares_target_kv_cache = True + + def load_weights_from_target_model(self, target) -> None: + loaded.append(target) + + target = MiniMaxM3ForCausalLM.__new__(MiniMaxM3ForCausalLM) + layers = [ + SimpleNamespace(input_layernorm=object()), + SimpleNamespace(input_layernorm=object()), + ] + final_norm = object() + object.__setattr__(target, "draft_model", DraftModel()) + object.__setattr__(target, "model", SimpleNamespace(layers=layers, norm=final_norm)) + + target.setup_aliases() + + assert loaded == [target] + assert layers[0].next_layer_layernorm is layers[1].input_layernorm + assert layers[1].next_layer_layernorm is final_norm + + +@pytest.mark.cpu_only +def test_eagle_capture_precedes_next_layer_norm() -> None: + class CaptureMetadata: + def __init__(self) -> None: + self.captured = None + + def is_layer_capture(self, layer_idx: int) -> bool: + return layer_idx == 25 + + def maybe_capture_hidden_states(self, layer_idx, hidden_states, residual) -> None: + self.captured = (layer_idx, hidden_states.clone(), residual.clone()) + + layer = SimpleNamespace( + layer_idx=25, + _apply_pre_feed_forward_norm=lambda hidden, residual: (hidden + 1, residual + 2), + block_sparse_moe=lambda hidden, unused_metadata, **unused_kwargs: hidden + 3, + _feed_forward_all_reduce_params=lambda: None, + _apply_next_layer_layernorm=lambda hidden, residual: (hidden + 10, residual + 20), + ) + spec_metadata = CaptureMetadata() + hidden_states = torch.tensor([1.0]) + residual = torch.tensor([2.0]) + + output, output_residual = MiniMaxM3DecoderLayer.forward_MoE( + layer, + hidden_states, + SimpleNamespace(), + residual, + spec_metadata, + ) + + assert spec_metadata.captured is not None + layer_idx, captured_hidden, captured_residual = spec_metadata.captured + assert layer_idx == 25 + torch.testing.assert_close(captured_hidden, torch.tensor([5.0])) + torch.testing.assert_close(captured_residual, torch.tensor([4.0])) + torch.testing.assert_close(output, torch.tensor([15.0])) + torch.testing.assert_close(output_residual, torch.tensor([24.0])) + + +@pytest.mark.cpu_only +def test_piecewise_attention_boundary_runs_horizontal_producer(monkeypatch) -> None: + class FakeAttentionLayer: + def __init__(self) -> None: + self.producer_shapes = None + + def _fused_fp8_qkv_indexer_norm_rope_kv_insert(self, packed, position_ids, attn_metadata): + self.producer_shapes = ( + tuple(packed.shape), + tuple(position_ids.shape), + attn_metadata.num_tokens, + ) + return packed[:, :3].clone(), packed[:, :1].clone() + + def _dispatch_attention_backend(self, q, k, v, idx_q, idx_k, attn_metadata, output) -> None: + assert k is None and v is None and idx_k is None + assert idx_q.shape == (attn_metadata.num_tokens, 1) + output.copy_(q) + + metadata = SimpleNamespace(num_tokens=2) + layer = FakeAttentionLayer() + monkeypatch.setattr( + modeling_minimaxm3, + "_extract_minimax_m3_attention_extra_attrs", + lambda layer_idx: (metadata, layer), + ) + packed = torch.arange(20, dtype=torch.float32).reshape(4, 5) + position_ids = torch.arange(4, dtype=torch.int32).reshape(1, 4) + output = torch.full((4, 3), -1.0) + + modeling_minimaxm3.minimax_m3_attn_custom_op_inplace( + None, + None, + None, + None, + None, + packed, + position_ids, + "3", + output, + ) + + assert layer.producer_shapes == ((2, 5), (1, 2), 2) + torch.testing.assert_close(output[:2], packed[:2, :3]) + torch.testing.assert_close(output[2:], torch.full((2, 3), -1.0)) + + +@pytest.mark.cpu_only +def test_piecewise_projection_fake_preserves_padded_hidden_rows(monkeypatch) -> None: + projection = object.__new__(MiniMaxM3QKVIndexerLinear) + nn.Module.__init__(projection) + projection.local_output_sizes = (3, 4) + layer = SimpleNamespace(qkv_proj=projection) + monkeypatch.setattr( + modeling_minimaxm3, + "_extract_minimax_m3_attention_extra_attrs", + lambda layer_idx: (SimpleNamespace(), layer), + ) + hidden_states = torch.randn(256, 5) + position_ids = torch.arange(6).reshape(1, 6) + + packed = modeling_minimaxm3._minimax_m3_qkv_index_proj_fake(hidden_states, position_ids, "3") + + # The real GEMM projects every padded hidden row. Position IDs remain an + # input solely to carry the unpadded token symbol to the piecewise segment. + assert packed.shape == (hidden_states.shape[0], 7) + + +@pytest.mark.cpu_only +def test_piecewise_fused_projection_preserves_input_token_dimension(monkeypatch) -> None: + """Do not inherit a bucket-specialized token dimension from the GEMM output.""" + packed = torch.randn(2, 7) + captured = {} + + def fake_boundary(q, k, v, idx_q, idx_k, packed_arg, position_ids, layer_idx, output): + assert q is None and k is None and v is None + assert idx_q is None and idx_k is None + captured["packed"] = packed_arg + captured["position_ids"] = position_ids + captured["output_shape"] = tuple(output.shape) + output.zero_() + + layer = SimpleNamespace( + enable_fused_qkv_index_projection=True, + qkv_proj=lambda hidden_states: packed, + register_to_config=True, + num_heads=1, + head_dim=3, + attn_activation_dtype=torch.float32, + layer_idx_str="3", + o_proj=lambda output, all_reduce_params: output, + ) + monkeypatch.setattr( + modeling_minimaxm3, + "_extract_minimax_m3_attention_extra_attrs", + lambda layer_idx: (SimpleNamespace(), layer), + ) + monkeypatch.setattr(modeling_minimaxm3, "is_torch_compiling", lambda: True) + monkeypatch.setattr( + modeling_minimaxm3, + "maybe_bcg_minimax_m3_attn_custom_op_inplace", + fake_boundary, + ) + hidden_states = torch.randn(4, 5) + position_ids = torch.arange(6).reshape(1, 6) + + result = MiniMaxM3Attention._sparse_forward( + layer, + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=SimpleNamespace(), + ) + + assert captured["packed"] is packed + assert captured["position_ids"] is position_ids + assert captured["output_shape"] == (position_ids.shape[-1], 3) + assert result.shape == (position_ids.shape[-1], 3) + + +@pytest.mark.cpu_only +def test_msa_attention_core_routes_compact_q_to_attention_dispatcher() -> None: + selected_blocks = torch.zeros(2, 1, 16, dtype=torch.int32) + + class FakeMsaBackend: + def __init__(self) -> None: + self.prepopulated_call = None + + def write_layer_caches(self, k, v, idx_k, metadata) -> None: + pytest.fail("The horizontal producer has already written both caches") + + def run_indexer(self, idx_q, idx_k, metadata, *, idx_k_prewritten): + assert idx_k is None + assert idx_k_prewritten + assert metadata is attn_metadata + return selected_blocks + + def forward(self, q, k, v, metadata, forward_args) -> None: + assert k is None and v is None + self.prepopulated_call = (q, metadata, forward_args) + + layer = MiniMaxM3Attention.__new__(MiniMaxM3Attention) + backend = FakeMsaBackend() + layer.attn = backend + layer.is_sparse_attention_layer = True + q = torch.randn(2, 8) + idx_q = torch.randn(2, 4) + attn_metadata = SimpleNamespace() + output = torch.empty_like(q) + + result = layer._msa_attention_core(q, None, None, idx_q, None, attn_metadata, output) + + assert result is output + assert backend.prepopulated_call is not None + called_q, called_metadata, forward_args = backend.prepopulated_call + assert called_q is q + assert called_metadata is attn_metadata + assert forward_args.output is output + assert forward_args.sparse_backend_args.topk_indices is selected_blocks + + def test_model_init_validates_sparse_attention_runtime_config() -> None: model_config = ModelConfig( pretrained_config=_make_text_config(), @@ -464,7 +783,7 @@ def test_get_moe_layer_ids_length_mismatch_raises(): # ``use_gemma=True`` and ``hidden_size=head_dim``; the # :meth:`apply_qk_norm` reshape matches an independent hand-written # reference. -# * Sparse index branch: fused replicated (tp_mode None) index_qk_proj with +# * Sparse index branch: KV-group-sharded index_qk_proj with # output [idx_q | idx_k] = num_index_heads * sparse_index_dim + sparse_index_dim # (idx_k is one K per token). # * Dense layers do not expose any index branch attributes (negative @@ -663,9 +982,9 @@ def test_minimax_m3_attention_apply_qk_norm_matches_reference(): def test_minimax_m3_attention_sparse_construction_matches_config(): """Sparse layer adds the index branch with the fused index projection. - * index_qk_proj is replicated (tp_mode None), out = + * At TP1, index_qk_proj retains all index heads, out = num_index_heads * sparse_index_dim (idx_q) + sparse_index_dim (idx_k), - where idx_k is one K per token (SGLang ReplicatedLinear contract). + where idx_k is one replicated K per token. * index_q_norm / index_k_norm are per-head Gemma RMSNorm of width sparse_index_dim. """ @@ -684,16 +1003,12 @@ def test_minimax_m3_attention_sparse_construction_matches_config(): assert attn.is_sparse_attention_layer is True assert attn.disable_index_value is True - # Replication keeps the idx_q -> [num_tokens, num_index_heads, - # sparse_index_dim] reshape valid at any TP geometry, whereas a - # column-parallel split would slice the head dimension. + # Explicit shard ranges retain whole index heads and replicate index-K. assert attn.index_q_size == num_index_heads * sparse_index_dim assert attn.index_k_size == sparse_index_dim assert attn.index_qk_proj.in_features == hidden assert attn.index_qk_proj.out_features == num_index_heads * sparse_index_dim + sparse_index_dim - assert attn.index_qk_proj.tp_mode is None, ( - f"index_qk_proj must be replicated, got {attn.index_qk_proj.tp_mode!r}" - ) + assert attn.index_qk_proj.tp_mode == modeling_minimaxm3.TensorParallelMode.COLUMN # Only the fused projection exists. assert not hasattr(attn, "index_q_proj") assert not hasattr(attn, "index_k_proj") @@ -718,6 +1033,243 @@ def test_minimax_m3_attention_sparse_construction_matches_config(): raise AssertionError("sparse forward must raise RuntimeError when attn_metadata is None") +@pytest.mark.cpu_only +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) +def test_minimax_m3_five_way_projection_shard_geometry(tp_size: int) -> None: + shard_geometry = MiniMaxM3QKVIndexerLinear._shard_geometry + for tp_rank in range(tp_size): + module = SimpleNamespace( + tp_size=tp_size, tp_rank=tp_rank, total_num_kv_heads=4, total_num_index_heads=4 + ) + kv_world = min(tp_size, 4) + kv_rank = tp_rank // max(tp_size // 4, 1) + assert shard_geometry(module, "q") == (tp_size, tp_rank) + assert shard_geometry(module, "k") == (kv_world, kv_rank) + assert shard_geometry(module, "v") == (kv_world, kv_rank) + assert shard_geometry(module, "index_q") == (kv_world, kv_rank) + assert shard_geometry(module, "index_k") == (1, 0) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) +def test_minimax_m3_five_way_projection_shards_index_rows(monkeypatch, tp_size: int) -> None: + def init_linear(self, in_features, out_features, **kwargs) -> None: + nn.Module.__init__(self) + self.tp_size = kwargs["mapping"].tp_size + self.tp_rank = kwargs["mapping"].tp_rank + self.tp_mode = kwargs["tensor_parallel_mode"] + self.weights_loading_config = kwargs["weights_loading_config"] + self.out_features = out_features // self.tp_size + + monkeypatch.setattr(modeling_minimaxm3.Linear, "__init__", init_linear) + # Exercise the real checkpoint loader on CPU; only its destination device + # is substituted, leaving the TP slicing and five-way packing intact. + load_weight_shard = modeling_minimaxm3.load_weight_shard + + def load_cpu_shard(weight, world_size, rank, mode, *, device): + return load_weight_shard(weight, world_size, rank, mode, device=torch.device("cpu")) + + monkeypatch.setattr(modeling_minimaxm3, "load_weight_shard", load_cpu_shard) + for tp_rank in range(tp_size): + projection = MiniMaxM3QKVIndexerLinear( + hidden_size=2, + head_dim=128, + total_num_heads=64, + total_num_kv_heads=4, + total_num_index_heads=4, + index_head_dim=128, + dtype=torch.bfloat16, + mapping=SimpleNamespace(tp_size=tp_size, tp_rank=tp_rank), + quant_config=None, + skip_create_weights_in_init=True, + force_dynamic_quantization=False, + disable_deep_gemm=False, + use_custom_cublas_mm=False, + use_cute_dsl_bf16_gemm=False, + use_cute_dsl_blockscaling_mm=False, + ) + kv_heads = max(4 // tp_size, 1) + assert projection.local_output_sizes == ( + 64 // tp_size * 128, + kv_heads * 128, + kv_heads * 128, + kv_heads * 128, + 128, + ) + assert projection.local_num_index_heads == kv_heads + assert projection.out_features == sum(projection.local_output_sizes) + shards = { + name: {"weight": torch.arange(heads * 128 * 2).reshape(heads * 128, 2)} + for name, heads in zip(projection._SHARD_NAMES, (64, 4, 4, 4, 1), strict=True) + } + loaded = [] + projection.load_weights = loaded.extend + projection.load_five_way_weights(shards) + packed = loaded[0]["weight"].split(projection.local_output_sizes) + torch.testing.assert_close(packed[4], shards["index_k"]["weight"]) + kv_rank = tp_rank // max(tp_size // 4, 1) + torch.testing.assert_close( + packed[3], shards["index_q"]["weight"].chunk(min(tp_size, 4))[kv_rank] + ) + torch.testing.assert_close(packed[1], shards["k"]["weight"].chunk(min(tp_size, 4))[kv_rank]) + assert projection.tp_size == tp_size and projection.tp_rank == tp_rank + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) +@pytest.mark.parametrize("attention_dp", [False, True]) +@pytest.mark.parametrize("num_index_heads", [4, 8]) +def test_minimax_m3_index_tp_matches_unsharded_reference( + tp_size: int, attention_dp: bool, num_index_heads: int +) -> None: + """Checkpoint row ownership and selected blocks must be TP invariant.""" + generator = torch.Generator().manual_seed(2718) + head_dim, hidden_dim, global_kv_heads = 2, 8, 4 + q_weight = torch.randn(num_index_heads * head_dim, hidden_dim, generator=generator) + k_weight = torch.randn(head_dim, hidden_dim, generator=generator) + hidden = torch.randn(5, hidden_dim, generator=generator) + key_hidden = torch.randn(11, 3, hidden_dim, generator=generator) + global_q = torch.nn.functional.linear(hidden, q_weight).reshape(5, num_index_heads, head_dim) + keys = torch.nn.functional.linear(key_hidden, k_weight) + global_scores = torch.einsum("thd,bkd->hbkt", global_q, keys).amax(dim=2) + params = MiniMaxM3SparseParams(num_index_heads=num_index_heads, global_num_kv_heads=4) + reference_cfg = MiniMaxM3SparseConfig.from_sparse_params( + params, num_q_heads=64, num_kv_heads=4, head_dim=head_dim + ) + reference_scores = _group_max_reduce(global_scores, reference_cfg) + metadata = MiniMaxM3SparseMetadataParams( + global_num_q_heads=64, global_num_kv_heads=4, num_index_heads=num_index_heads + ) + effective_tp = 1 if attention_dp else tp_size + local_kv_heads = max(global_kv_heads // effective_tp, 1) + for rank in range(tp_size): + mapping = SimpleNamespace(tp_size=tp_size, tp_rank=rank, enable_attention_dp=attention_dp) + start, end = index_head_range(num_index_heads, global_kv_heads, mapping) + kv_start = 0 if attention_dp else rank // max(tp_size // 4, 1) * local_kv_heads + assert (start, end) == ( + kv_start * (num_index_heads // 4), + (kv_start + local_kv_heads) * (num_index_heads // 4), + ) + assert metadata.sharded_index_head_count(mapping) == end - start + + # Exercise Linear's real per-shard loader, as used by index_qk_proj: + # index-Q uses the KV group's rows; index-K always uses all its rows. + projection = object.__new__(modeling_minimaxm3.Linear) + nn.Module.__init__(projection) + projection.tp_size = effective_tp + projection.tp_rank = 0 if attention_dp else rank + projection.tp_mode = modeling_minimaxm3.TensorParallelMode.COLUMN + projection.tp_sharding = {"gate": (start * head_dim, end * head_dim), "up": (0, head_dim)} + projection.weights_loading_config = modeling_minimaxm3.WeightsLoadingConfig( + weight_mode=modeling_minimaxm3.WeightMode.FUSED_GATE_UP_LINEAR + ) + local_q_weight = projection.load_shard({"weight": q_weight}, "weight", name="gate") + local_k_weight = projection.load_shard({"weight": k_weight}, "weight", name="up") + assert ( + projection.calculate_local_out_features((num_index_heads + 1) * head_dim) + == (end - start + 1) * head_dim + ) + torch.testing.assert_close(local_k_weight, k_weight) + local_q = torch.nn.functional.linear(hidden, local_q_weight).reshape( + 5, end - start, head_dim + ) + torch.testing.assert_close(local_q, global_q[:, start:end]) + local_scores = torch.einsum("thd,bkd->hbkt", local_q, keys).amax(dim=2) + local_cfg = MiniMaxM3SparseConfig.from_sparse_params( + params, num_q_heads=64 // effective_tp, num_kv_heads=local_kv_heads, head_dim=head_dim + ) + assert local_cfg.num_index_heads == end - start + grouped_scores = _group_max_reduce(local_scores, local_cfg) + expected = reference_scores[kv_start : kv_start + local_kv_heads] + torch.testing.assert_close(grouped_scores, expected) + torch.testing.assert_close( + grouped_scores.topk(3, dim=1).indices, expected.topk(3, dim=1).indices + ) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "num_index_heads,num_kv_heads,tp_size", [(0, 4, 1), (6, 4, 2), (4, 4, 3), (4, 0, 2)] +) +def test_minimax_m3_index_head_range_rejects_invalid_geometry( + num_index_heads: int, num_kv_heads: int, tp_size: int +) -> None: + with pytest.raises(ValueError): + index_head_range( + num_index_heads, + num_kv_heads, + SimpleNamespace(tp_size=tp_size, tp_rank=0, enable_attention_dp=False), + ) + + +@pytest.mark.gpu +@pytest.mark.skipif(not _has_cuda(), reason="MiniMax-M3 attention construction needs CUDA") +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) +@pytest.mark.parametrize("attention_dp", [False, True]) +def test_minimax_m3_unfused_index_projection_tp_construction( + tp_size: int, attention_dp: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + """Verify the model wires Linear's shard overrides to effective attention TP.""" + import tensorrt_llm._torch.distributed as distributed + + # Simulate each rank on one GPU without creating distributed workspaces. + # No forward runs here; projection construction and shard sizing stay real. + monkeypatch.setattr(distributed, "AllReduce", lambda **_kwargs: nn.Identity()) + for rank in range(tp_size): + text_cfg, model_cfg = _make_attention_test_config() + text_cfg.hidden_size = 256 + text_cfg.num_attention_heads = 8 + text_cfg.num_key_value_heads = 4 + text_cfg.sparse_attention_config["sparse_num_index_heads"] = 4 + model_cfg.mapping = Mapping( + world_size=tp_size, rank=rank, tp_size=tp_size, enable_attention_dp=attention_dp + ) + attn = MiniMaxM3Attention( + model_config=model_cfg, + layer_idx=3, + is_sparse_attention_layer=True, + disable_index_value=True, + ) + kv_heads = 4 if attention_dp else max(4 // tp_size, 1) + kv_start = 0 if attention_dp else rank // max(tp_size // 4, 1) * kv_heads + assert attn.sparse_num_index_heads == kv_heads + assert attn.index_qk_proj.out_features == (kv_heads + 1) * 32 + assert attn.index_qk_proj.tp_sharding == { + "gate": (kv_start * 32, (kv_start + kv_heads) * 32), + "up": (0, 32), + } + assert attn.index_qk_proj.mapping == attn.qkv_proj.mapping + + +@pytest.mark.cpu_only +def test_minimax_m3_five_way_loader_returns_exact_generic_skip() -> None: + projection = object.__new__(MiniMaxM3QKVIndexerLinear) + nn.Module.__init__(projection) + captured = {} + projection.load_five_way_weights = lambda shards: captured.update(shards) + + model = nn.Module() + model.sparse = nn.Module() + model.sparse.qkv_proj = projection + weights = { + f"sparse.{name}_proj.weight": torch.empty(1) + for name in ("q", "k", "v", "index_q", "index_k") + } + + loaded_modules = _load_qkv_index_proj_weights(model, weights) + + assert loaded_modules == ["sparse.qkv_proj"] + assert set(captured) == {"q", "k", "v", "index_q", "index_k"} + assert all(set(shard) == {"weight"} for shard in captured.values()) + assert weights == {} + + mapper = MiniMaxM3HfWeightMapper() + mapper.add_skip_modules(loaded_modules) + mapper._model = SimpleNamespace(config=SimpleNamespace(tie_word_embeddings=False)) + assert mapper.should_skip_module("sparse.qkv_proj") + assert not mapper.should_skip_module("dense.qkv_proj") + + @pytest.mark.gpu @pytest.mark.skipif(not _has_cuda(), reason="MiniMax-M3 attention construction needs CUDA") def test_minimax_m3_attention_apply_index_qk_norm_matches_reference(): @@ -1090,7 +1642,7 @@ def test_minimax_m3_attention_real_config_index_branch_shapes(): assert attn.index_k_size == sparse_index_dim assert attn.index_qk_proj.in_features == int(text_cfg.hidden_size) assert attn.index_qk_proj.out_features == num_index_heads * sparse_index_dim + sparse_index_dim - assert attn.index_qk_proj.tp_mode is None + assert attn.index_qk_proj.tp_mode == modeling_minimaxm3.TensorParallelMode.COLUMN assert not hasattr(attn, "index_q_proj") assert not hasattr(attn, "index_k_proj") diff --git a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_horizontal_producer.py b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_horizontal_producer.py new file mode 100644 index 000000000000..6bec43157ebe --- /dev/null +++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_horizontal_producer.py @@ -0,0 +1,357 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _rope_cache(max_positions, rotary_dim=64, base=5_000_000.0): + positions = torch.arange(max_positions, dtype=torch.float32, device="cuda") + inverse_frequency = 1.0 / ( + base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32, device="cuda") / rotary_dim) + ) + frequency = torch.outer(positions, inverse_frequency) + return torch.stack((frequency.cos(), frequency.sin()), dim=1).contiguous() + + +def _main_cache(num_pages, num_kv_heads, stride_scale=3): + backing = torch.zeros( + num_pages * stride_scale, + 2, + num_kv_heads, + 128, + 128, + dtype=torch.float8_e4m3fn, + device="cuda", + ) + return backing[::stride_scale] + + +def _index_cache(num_pages, stride_scale=5): + backing = torch.zeros( + num_pages * stride_scale, + 1, + 128, + 128, + dtype=torch.float8_e4m3fn, + device="cuda", + ) + return backing[::stride_scale] + + +def _assert_fp8_within_one_ulp(actual: torch.Tensor, expected: torch.Tensor) -> None: + """Allow adjacent finite E4M3 values, including subnormals and signed zero.""" + assert actual.dtype == expected.dtype == torch.float8_e4m3fn + assert torch.isfinite(actual.float()).all() + assert torch.isfinite(expected.float()).all() + # E4M3 encodes sign/magnitude. Map both signs to monotonically ordered + # integers, collapsing +0 and -0, so atol=1 means exactly one FP8 step. + actual_bits = actual.view(torch.uint8).to(torch.int16) + expected_bits = expected.view(torch.uint8).to(torch.int16) + actual_ordered = torch.where(actual_bits < 128, actual_bits, 128 - actual_bits) + expected_ordered = torch.where(expected_bits < 128, expected_bits, 128 - expected_bits) + torch.testing.assert_close(actual_ordered, expected_ordered, rtol=0, atol=1) + + +def test_minimax_m3_fp8_one_ulp_comparison() -> None: + # Cover every adjacent finite pair, including exponent boundaries and + # subnormals, for both signs. Two-step errors and NaNs must still fail. + for sign in (0, 128): + values = (torch.arange(127, dtype=torch.uint8) + sign).view(torch.float8_e4m3fn) + _assert_fp8_within_one_ulp(values[:-1], values[1:]) + with pytest.raises(AssertionError): + _assert_fp8_within_one_ulp(values[:-2], values[2:]) + _assert_fp8_within_one_ulp( + torch.tensor([0, 128, 1, 129], dtype=torch.uint8).view(torch.float8_e4m3fn), + torch.tensor([128, 0, 128, 0], dtype=torch.uint8).view(torch.float8_e4m3fn), + ) + with pytest.raises(AssertionError): + _assert_fp8_within_one_ulp( + torch.tensor([1], dtype=torch.uint8).view(torch.float8_e4m3fn), + torch.tensor([129], dtype=torch.uint8).view(torch.float8_e4m3fn), + ) + for nan_bits in (127, 255): + nan = torch.tensor([nan_bits], dtype=torch.uint8).view(torch.float8_e4m3fn) + with pytest.raises(AssertionError): + _assert_fp8_within_one_ulp(nan, nan) + + +@pytest.mark.parametrize("num_tokens", [1, 16, 129]) +@pytest.mark.parametrize("num_kv_heads,num_index_heads", [(4, 4), (2, 2), (1, 1)]) +def test_minimax_m3_horizontal_producer_matches_separate_producers( + num_tokens, num_kv_heads, num_index_heads +): + torch.manual_seed(1234) + num_heads_q = 8 + num_pages = max(4, (num_tokens + 127) // 128 + 2) + total_heads = num_heads_q + 2 * num_kv_heads + num_index_heads + 1 + packed = torch.randn( + num_tokens, + total_heads * 128, + dtype=torch.bfloat16, + device="cuda", + ) + q_weight = torch.randn(128, dtype=torch.bfloat16, device="cuda") + k_weight = torch.randn(128, dtype=torch.bfloat16, device="cuda") + index_q_weight = torch.randn(128, dtype=torch.bfloat16, device="cuda") + index_k_weight = torch.randn(128, dtype=torch.bfloat16, device="cuda") + position_ids = torch.arange(num_tokens, dtype=torch.int32, device="cuda") + slots = (torch.arange(num_tokens, dtype=torch.int32, device="cuda") * 37) % ( + (num_pages - 1) * 128 + ) + # Keep the parity reference slots valid: the legacy separate main-K/V + # producer does not support negative slots. Negative-slot handling is + # exercised below using horizontal eager execution versus graph replay. + rope_cache = _rope_cache(max(256, num_tokens)) + + main_width = (num_heads_q + 2 * num_kv_heads) * 128 + main_input = packed[:, :main_width].contiguous() + index_input = packed[:, main_width:].contiguous() + reference_main_cache = _main_cache(num_pages, num_kv_heads) + reference_index_cache = _index_cache(num_pages) + q_reference = torch.ops.trtllm.minimax_m3_fp8_qk_norm_rope_kv_insert( + main_input, + reference_main_cache, + slots, + num_heads_q, + num_kv_heads, + num_kv_heads, + 128, + 64, + 1e-5, + q_weight, + k_weight, + 5_000_000.0, + True, + position_ids, + ) + index_q_reference = torch.ops.trtllm.minimax_m3_fp8_indexer_qk_norm_rope( + index_input, + reference_index_cache, + slots, + num_index_heads, + 128, + 64, + 1e-5, + index_q_weight, + index_k_weight, + 5_000_000.0, + position_ids, + ) + + main_cache = _main_cache(num_pages, num_kv_heads) + index_cache = _index_cache(num_pages) + q, index_q = torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + packed, + main_cache, + index_cache, + slots, + num_heads_q, + num_kv_heads, + num_index_heads, + 128, + 64, + 1e-5, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rope_cache, + position_ids, + ) + + valid = slots >= 0 + pages = slots[valid].long() // 128 + within = slots[valid].long() % 128 + # The horizontal producer reads precomputed FP32 RoPE coefficients; the + # separate main producer computes powf/__sincosf per head. Near an FP8 + # midpoint, their small FP32 differences can round to adjacent E4M3 bins. + # Bound Q/K differences to one actual FP8 step, not a broad atol/rtol. + _assert_fp8_within_one_ulp(q, q_reference) + # The horizontal producer follows vLLM's CUDA contract and converts its + # normalized/RoPE FP32 registers directly to E4M3. The existing separate + # TRT-LLM index producer first materializes BF16, so retain its numerical + # tolerance across the different rounding orders. V is copy-cast only and + # retains byte-exact parity below. + torch.testing.assert_close( + index_q.float(), + index_q_reference.float(), + rtol=0.13, + atol=0.05, + ) + _assert_fp8_within_one_ulp( + main_cache[pages, 0, :, within, :], + reference_main_cache[pages, 0, :, within, :], + ) + assert torch.equal( + main_cache[pages, 1, :, within, :].view(torch.uint8), + reference_main_cache[pages, 1, :, within, :].view(torch.uint8), + ) + torch.testing.assert_close( + index_cache[pages, :, within, :].float(), + reference_index_cache[pages, :, within, :].float(), + rtol=0.13, + atol=0.05, + ) + + # Aggregate decode captures this producer in a CUDA graph. The operator + # allocates compact Q/index-Q outputs while writing graph-stable paged + # caches through a graph-stable slot mapping, so exercise both capture and + # replay for decode-sized (1) and larger mixed/prefill token counts. + graph_packed = packed.clone() + graph_positions = position_ids.clone() + graph_slots = slots.clone() + graph_main_cache = _main_cache(num_pages, num_kv_heads) + graph_index_cache = _index_cache(num_pages) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_q, graph_index_q = torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + graph_packed, + graph_main_cache, + graph_index_cache, + graph_slots, + num_heads_q, + num_kv_heads, + num_index_heads, + 128, + 64, + 1e-5, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rope_cache, + graph_positions, + ) + + # Replay with different projection values, nonuniform positions, and new + # cache destinations. This proves replay reads the refreshed graph buffers + # rather than retaining capture-time values or slots. + replay_packed = torch.randn_like(packed) + replay_positions = ( + torch.arange(num_tokens, dtype=torch.int32, device="cuda") * 7 + 3 + ) % rope_cache.shape[0] + replay_slots = (torch.arange(num_tokens, dtype=torch.int32, device="cuda") * 53 + 11) % ( + (num_pages - 1) * 128 + ) + if num_tokens > 1: + replay_slots[-1] = -1 + graph_packed.copy_(replay_packed) + graph_positions.copy_(replay_positions) + graph_slots.copy_(replay_slots) + graph_main_cache.zero_() + graph_index_cache.zero_() + + replay_main_cache = _main_cache(num_pages, num_kv_heads) + replay_index_cache = _index_cache(num_pages) + replay_q, replay_index_q = torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + replay_packed, + replay_main_cache, + replay_index_cache, + replay_slots, + num_heads_q, + num_kv_heads, + num_index_heads, + 128, + 64, + 1e-5, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rope_cache, + replay_positions, + ) + graph.replay() + torch.cuda.synchronize() + + replay_valid = replay_slots >= 0 + replay_pages = replay_slots[replay_valid].long() // 128 + replay_within = replay_slots[replay_valid].long() % 128 + assert torch.equal(graph_q.view(torch.uint8), replay_q.view(torch.uint8)) + torch.testing.assert_close( + graph_index_q.float(), + replay_index_q.float(), + rtol=0.0, + atol=0.0, + ) + assert torch.equal( + graph_main_cache[replay_pages, :, :, replay_within, :].view(torch.uint8), + replay_main_cache[replay_pages, :, :, replay_within, :].view(torch.uint8), + ) + torch.testing.assert_close( + graph_index_cache[replay_pages, :, replay_within, :].float(), + replay_index_cache[replay_pages, :, replay_within, :].float(), + rtol=0.0, + atol=0.0, + ) + + +def test_minimax_m3_horizontal_producer_ignores_out_of_range_cache_slot(): + num_heads_q = 8 + num_kv_heads = 2 + total_heads = num_heads_q + 3 * num_kv_heads + 1 + packed = torch.randn(1, total_heads * 128, dtype=torch.bfloat16, device="cuda") + weights = [torch.randn(128, dtype=torch.bfloat16, device="cuda") for _ in range(4)] + positions = torch.zeros(1, dtype=torch.int32, device="cuda") + slots = torch.tensor([2 * 128], dtype=torch.int32, device="cuda") + main_cache = _main_cache(2, num_kv_heads) + index_cache = _index_cache(2) + main_before = main_cache.clone() + index_before = index_cache.clone() + + q, index_q = torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + packed, + main_cache, + index_cache, + slots, + num_heads_q, + num_kv_heads, + num_kv_heads, + 128, + 64, + 1e-5, + *weights, + _rope_cache(1), + positions, + ) + + assert q.shape == (1, num_heads_q, 128) + assert index_q.shape == (1, num_kv_heads, 128) + assert torch.equal(main_cache.view(torch.uint8), main_before.view(torch.uint8)) + assert torch.equal(index_cache.view(torch.uint8), index_before.view(torch.uint8)) + + +@pytest.mark.parametrize("invalid_weight_index", range(4)) +def test_minimax_m3_horizontal_producer_rejects_non_vector_norm_weight( + invalid_weight_index: int, +) -> None: + num_heads_q = 8 + num_kv_heads = 2 + num_index_heads = num_kv_heads + total_heads = num_heads_q + 2 * num_kv_heads + num_index_heads + 1 + packed = torch.randn(1, total_heads * 128, dtype=torch.bfloat16, device="cuda") + weights = [torch.randn(128, dtype=torch.bfloat16, device="cuda") for _ in range(4)] + weights[invalid_weight_index] = weights[invalid_weight_index].reshape(1, 128) + positions = torch.zeros(1, dtype=torch.int32, device="cuda") + slots = torch.zeros(1, dtype=torch.int32, device="cuda") + + with pytest.raises(RuntimeError, match="norm weights must be one-dimensional"): + torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + packed, + _main_cache(1, num_kv_heads), + _index_cache(1), + slots, + num_heads_q, + num_kv_heads, + num_index_heads, + 128, + 64, + 1e-5, + *weights, + _rope_cache(1), + positions, + ) diff --git a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_main_kv_insert.py b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_main_kv_insert.py new file mode 100644 index 000000000000..5d17670704d8 --- /dev/null +++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_main_kv_insert.py @@ -0,0 +1,189 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _reference(qkv, num_heads_q, num_kv_heads, q_weight, k_weight, position_ids): + output = torch.ops.trtllm.fused_qk_norm_rope_to_fp8( + qkv, + num_heads_q, + num_kv_heads, + num_kv_heads, + 128, + 64, + 1e-5, + q_weight, + k_weight, + 5_000_000.0, + True, + position_ids, + 1.0, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + return output.view(qkv.shape[0], num_heads_q + 2 * num_kv_heads, 128).split( + [num_heads_q, num_kv_heads, num_kv_heads], dim=1 + ) + + +def _strided_kv_cache(num_pages, num_kv_heads, page_size=128, stride_scale=3): + backing = torch.zeros( + num_pages * stride_scale, + 2, + num_kv_heads, + page_size, + 128, + dtype=torch.float8_e4m3fn, + device="cuda", + ) + return backing[::stride_scale] + + +def _inputs(num_tokens, num_heads_q, num_kv_heads): + qkv = torch.randn( + num_tokens, + (num_heads_q + 2 * num_kv_heads) * 128, + dtype=torch.bfloat16, + device="cuda", + ) + q_weight = torch.randn(128, dtype=torch.bfloat16, device="cuda") + k_weight = torch.randn(128, dtype=torch.bfloat16, device="cuda") + position_ids = torch.arange(num_tokens, dtype=torch.int32, device="cuda") + 8192 + return qkv, q_weight, k_weight, position_ids + + +def _run(qkv, kv_cache, slots, q_weight, k_weight, position_ids, num_heads_q, num_kv_heads): + return torch.ops.trtllm.minimax_m3_fp8_qk_norm_rope_kv_insert( + qkv, + kv_cache, + slots, + num_heads_q, + num_kv_heads, + num_kv_heads, + 128, + 64, + 1e-5, + q_weight, + k_weight, + 5_000_000.0, + True, + position_ids, + ) + + +@pytest.mark.parametrize(("num_heads_q", "num_kv_heads"), [(8, 1), (8, 8), (64, 4)]) +@pytest.mark.parametrize("num_tokens", [1, 16, 129]) +def test_minimax_m3_fp8_main_kv_insert_matches_materialize_then_scatter( + num_tokens, num_heads_q, num_kv_heads +): + torch.manual_seed(1234) + page_size = 128 + num_pages = max(4, (num_tokens + page_size - 1) // page_size + 2) + qkv, q_weight, k_weight, position_ids = _inputs(num_tokens, num_heads_q, num_kv_heads) + slots = (torch.arange(num_tokens, dtype=torch.int32, device="cuda") * 37) % ( + (num_pages - 1) * page_size + ) + kv_cache = _strided_kv_cache(num_pages, num_kv_heads, page_size) + guard_page = kv_cache[-1].clone() + + q_out = _run( + qkv, + kv_cache, + slots, + q_weight, + k_weight, + position_ids, + num_heads_q, + num_kv_heads, + ) + q_ref, k_ref, v_ref = _reference( + qkv, num_heads_q, num_kv_heads, q_weight, k_weight, position_ids + ) + pages = slots.long() // page_size + within = slots.long() % page_size + + # The specialized kernel uses powf while fused_qk_norm_rope_to_fp8 uses + # the exp2f/log2f equivalent, so values at an FP8 boundary can round to + # adjacent E4M3 values. + torch.testing.assert_close(q_out.float(), q_ref.float(), rtol=0.13, atol=0.05) + torch.testing.assert_close( + kv_cache[:, 0][pages, :, within, :].float(), + k_ref.float(), + rtol=0.13, + atol=0.05, + ) + assert torch.equal( + kv_cache[:, 1][pages, :, within, :].view(torch.uint8), + v_ref.contiguous().view(torch.uint8), + ) + assert torch.equal(kv_cache[-1].view(torch.uint8), guard_page.view(torch.uint8)) + + +def test_minimax_m3_fp8_main_kv_insert_uses_64bit_cache_offsets(): + """Exercise a real paged-cache address beyond INT32_MAX elements. + + The smallest contiguous HND pool whose page-65536 K row starts at + 2**31 FP8 elements is about 2 GiB. The former implicit int conversion in + the store helper wrapped this address negative; this test writes and reads + that real allocation so arithmetic-only tests cannot mask the bug. + """ + required_bytes = (65537 * 2 * 128 * 128) + (1 << 30) + free_bytes, _ = torch.cuda.mem_get_info() + if free_bytes < required_bytes: + pytest.skip("64-bit cache-offset test requires about 3 GiB free GPU memory") + + torch.manual_seed(4321) + page = 65536 + kv_cache = torch.empty( + page + 1, + 2, + 1, + 128, + 128, + dtype=torch.float8_e4m3fn, + device="cuda", + ) + qkv, q_weight, k_weight, position_ids = _inputs(1, 8, 1) + slots = torch.tensor([page * 128], dtype=torch.int32, device="cuda") + + q_out = _run(qkv, kv_cache, slots, q_weight, k_weight, position_ids, 8, 1) + q_ref, k_ref, v_ref = _reference(qkv, 8, 1, q_weight, k_weight, position_ids) + torch.cuda.synchronize() + + torch.testing.assert_close(q_out.float(), q_ref.float(), rtol=0.13, atol=0.05) + torch.testing.assert_close( + kv_cache[page, 0, :, 0, :].float(), + k_ref[0].float(), + rtol=0.13, + atol=0.05, + ) + assert torch.equal( + kv_cache[page, 1, :, 0, :].view(torch.uint8), + v_ref[0].contiguous().view(torch.uint8), + ) + + +@pytest.mark.parametrize("invalid_slot", [-1, 2 * 128]) +def test_minimax_m3_fp8_main_kv_insert_ignores_invalid_slot(invalid_slot): + torch.manual_seed(5678) + qkv, q_weight, k_weight, position_ids = _inputs(1, 8, 1) + kv_cache = _strided_kv_cache(2, 1) + kv_cache.fill_(1.0) + before = kv_cache.clone() + slots = torch.tensor([invalid_slot], dtype=torch.int32, device="cuda") + + q_out = _run(qkv, kv_cache, slots, q_weight, k_weight, position_ids, 8, 1) + q_ref, _, _ = _reference(qkv, 8, 1, q_weight, k_weight, position_ids) + + torch.testing.assert_close(q_out.float(), q_ref.float(), rtol=0.13, atol=0.05) + assert torch.equal(kv_cache.view(torch.uint8), before.view(torch.uint8))