diff --git a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu index 8987234060fb..47db29b36a47 100644 --- a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu +++ b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu @@ -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. @@ -88,6 +88,21 @@ __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"); + // 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 the result to a (possibly different-dtype) output buffer. // head_dim: the dimension of each head @@ -343,6 +358,279 @@ __global__ void fusedQKNormRopeKernel( storeHeadElements(qkv_out, offsetThread, elements); } +namespace +{ + +constexpr int kMinimaxM3HeadDim = 128; +constexpr int kMinimaxM3RotaryDim = 64; +constexpr int kMinimaxM3PageSize = 128; +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, 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); + int const inputOffset = (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); + int const page = slot >> 7; + 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, 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; + + int const inputOffset = (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 >> 7; + 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 +747,62 @@ void launchFusedQKNormRopeOut(void const* qkv_in, void* qkv_out, bool out_fp8, b 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, + 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_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, + 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_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 4eaaf77cd9f0..24dc796ff9d0 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 @@ -71,6 +73,24 @@ void launchFusedQKNormRopeOut(void const* qkv_in, // BF16 input [num_tokens, tot 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 producer for eager pure prefill. 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, + 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, + 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/kernels/minimaxM3SelectBlocks.cu b/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cu index 0490844cd82c..6e28c8f4a358 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.cu @@ -37,9 +37,403 @@ constexpr int kTopK = 16; constexpr int kWarpSize = 32; constexpr int kThreadsPerBlock = 256; constexpr int kWarpsPerBlock = kThreadsPerBlock / kWarpSize; +constexpr int kSmallMaxBlocks = 128; constexpr float kInitScore = 1.0e30F; constexpr float kLocalScore = 1.0e29F; +constexpr uint32_t kFullWarpMask = 0xFFFFFFFFU; +__forceinline__ __device__ bool candidateGreater( + uint32_t lhsScoreKey, int32_t lhsBlockId, uint32_t rhsScoreKey, int32_t rhsBlockId) +{ + return lhsScoreKey > rhsScoreKey || (lhsScoreKey == rhsScoreKey && lhsBlockId < rhsBlockId); +} + +__forceinline__ __device__ void warpBitonicSortDesc64( + uint32_t& scoreKey0, int32_t& blockId0, uint32_t& scoreKey1, int32_t& blockId1, int32_t lane) +{ +#pragma unroll + for (int32_t size = 2; size <= 2 * kWarpSize; size *= 2) + { +#pragma unroll + for (int32_t stride = size / 2; stride > 0; stride /= 2) + { + if (stride == kWarpSize) + { + bool const firstGreater = candidateGreater(scoreKey0, blockId0, scoreKey1, blockId1); + uint32_t const greaterScoreKey = firstGreater ? scoreKey0 : scoreKey1; + int32_t const greaterBlockId = firstGreater ? blockId0 : blockId1; + uint32_t const lesserScoreKey = firstGreater ? scoreKey1 : scoreKey0; + int32_t const lesserBlockId = firstGreater ? blockId1 : blockId0; + scoreKey0 = greaterScoreKey; + blockId0 = greaterBlockId; + scoreKey1 = lesserScoreKey; + blockId1 = lesserBlockId; + } + else + { + uint32_t const partnerScoreKey0 = __shfl_xor_sync(kFullWarpMask, scoreKey0, stride); + int32_t const partnerBlockId0 = __shfl_xor_sync(kFullWarpMask, blockId0, stride); + uint32_t const partnerScoreKey1 = __shfl_xor_sync(kFullWarpMask, scoreKey1, stride); + int32_t const partnerBlockId1 = __shfl_xor_sync(kFullWarpMask, blockId1, stride); + bool const takeGreater0 = ((lane & size) == 0) == ((lane & stride) == 0); + int32_t const secondIndex = lane + kWarpSize; + bool const takeGreater1 = ((secondIndex & size) == 0) == ((secondIndex & stride) == 0); + bool const partnerGreater0 = candidateGreater(partnerScoreKey0, partnerBlockId0, scoreKey0, blockId0); + bool const currentGreater0 = candidateGreater(scoreKey0, blockId0, partnerScoreKey0, partnerBlockId0); + bool const partnerGreater1 = candidateGreater(partnerScoreKey1, partnerBlockId1, scoreKey1, blockId1); + bool const currentGreater1 = candidateGreater(scoreKey1, blockId1, partnerScoreKey1, partnerBlockId1); + if ((takeGreater0 && partnerGreater0) || (!takeGreater0 && currentGreater0)) + { + scoreKey0 = partnerScoreKey0; + blockId0 = partnerBlockId0; + } + if ((takeGreater1 && partnerGreater1) || (!takeGreater1 && currentGreater1)) + { + scoreKey1 = partnerScoreKey1; + blockId1 = partnerBlockId1; + } + } + } + } +} + +__forceinline__ __device__ void warpBitonicSortDesc128(uint32_t (&scoreKeys)[4], int32_t (&blockIds)[4], int32_t lane) +{ + // Virtual item slot * 32 + lane stays in registers. Short strides exchange + // lanes with shuffles; strides 32 and 64 exchange slots within each lane. +#pragma unroll + for (int32_t size = 2; size <= kSmallMaxBlocks; size *= 2) + { +#pragma unroll + for (int32_t stride = size / 2; stride > 0; stride /= 2) + { + uint32_t previousScoreKeys[4]; + int32_t previousBlockIds[4]; +#pragma unroll + for (int32_t slot = 0; slot < 4; ++slot) + { + previousScoreKeys[slot] = scoreKeys[slot]; + previousBlockIds[slot] = blockIds[slot]; + } + +#pragma unroll + for (int32_t slot = 0; slot < 4; ++slot) + { + uint32_t partnerScoreKey; + int32_t partnerBlockId; + if (stride < kWarpSize) + { + partnerScoreKey = __shfl_xor_sync(kFullWarpMask, previousScoreKeys[slot], stride); + partnerBlockId = __shfl_xor_sync(kFullWarpMask, previousBlockIds[slot], stride); + } + else + { + int32_t const partnerSlot = slot ^ (stride / kWarpSize); + partnerScoreKey = previousScoreKeys[partnerSlot]; + partnerBlockId = previousBlockIds[partnerSlot]; + } + + int32_t const item = lane + slot * kWarpSize; + bool const takeGreater = ((item & size) == 0) == ((item & stride) == 0); + bool const partnerGreater = candidateGreater( + partnerScoreKey, partnerBlockId, previousScoreKeys[slot], previousBlockIds[slot]); + bool const currentGreater = candidateGreater( + previousScoreKeys[slot], previousBlockIds[slot], partnerScoreKey, partnerBlockId); + if ((takeGreater && partnerGreater) || (!takeGreater && currentGreater)) + { + scoreKeys[slot] = partnerScoreKey; + blockIds[slot] = partnerBlockId; + } + } + } + } +} + +template +__forceinline__ __device__ int64_t outputOffset( + int32_t query, int32_t kvHead, int32_t totalQueries, int32_t numKvHeads, int32_t rank) +{ + int64_t const outputRow = HeadMajorOutput ? static_cast(kvHead) * totalQueries + query + : static_cast(query) * numKvHeads + kvHead; + return outputRow * kTopK + rank; +} + +template +__forceinline__ __device__ void selectFromCandidates(cg::thread_block_tile const& warp, + float (&localScores)[NumCandidates], int32_t (&localIndices)[NumCandidates], int32_t* output, int32_t query, + int32_t kvHead, int32_t totalQueries, int32_t numKvHeads, int32_t numBlocks) +{ + float selectedScores[kTopK]; + int32_t selectedIndices[kTopK]; + reduce_topk::reduceTopK(warp, selectedScores, selectedIndices, localScores, localIndices, -INFINITY); + + if (warp.thread_rank() == 0) + { +#pragma unroll + for (int32_t rank = 0; rank < kTopK; ++rank) + { + if (selectedScores[rank] == -INFINITY) + { + selectedIndices[rank] = -1; + } + } + + // MSA consumes block IDs in ascending order. Sort the sixteen selected + // IDs in registers, treating -1 padding as greater than every valid ID. +#pragma unroll + for (int32_t rank = 1; rank < kTopK; ++rank) + { + int32_t const candidate = selectedIndices[rank]; + int32_t const candidateKey = candidate < 0 ? numBlocks : candidate; + int32_t insertion = rank; + while (insertion > 0) + { + int32_t const previous = selectedIndices[insertion - 1]; + int32_t const previousKey = previous < 0 ? numBlocks : previous; + if (previousKey <= candidateKey) + { + break; + } + selectedIndices[insertion] = previous; + --insertion; + } + selectedIndices[insertion] = candidate; + } + +#pragma unroll + for (int32_t rank = 0; rank < kTopK; ++rank) + { + output[outputOffset(query, kvHead, totalQueries, numKvHeads, rank)] + = selectedIndices[rank]; + } + } +} + +template +__global__ void minimaxM3SelectBlocksSmallKernel(float const* __restrict__ scores, int64_t headStride, + int64_t blockStride, int64_t queryStride, int32_t const* __restrict__ nValidBlocks, int32_t* __restrict__ output, + int32_t numKvHeads, int32_t numBlocks, int32_t totalQueries, int32_t initBlocks, int32_t localBlocks) +{ + static_assert(NumCandidates * kWarpSize <= kSmallMaxBlocks); + auto const warp = cg::tiled_partition(cg::this_thread_block()); + int32_t const warpInBlock = threadIdx.x / kWarpSize; + int32_t const outputRow = blockIdx.x * kWarpsPerBlock + warpInBlock; + int32_t const numOutputRows = totalQueries * numKvHeads; + if (outputRow >= numOutputRows) + { + return; + } + + int32_t const query = outputRow / numKvHeads; + int32_t const kvHead = outputRow % numKvHeads; + int32_t const rawValidBlocks = nValidBlocks[query]; + int32_t const validBlocks = max(0, min(rawValidBlocks, numBlocks)); + int64_t const localStart + = max(static_cast(rawValidBlocks) - static_cast(localBlocks), static_cast(0)); + + using RedType = reduce_topk::TopKRedType; + float localScores[NumCandidates]; + int32_t localIndices[NumCandidates]; +#pragma unroll + for (int32_t slot = 0; slot < NumCandidates; ++slot) + { + int32_t const block = warp.thread_rank() + slot * kWarpSize; + RedType candidate{-INFINITY, RedType::kMaxIdx}; + if (block < validBlocks) + { + int64_t const offset = static_cast(kvHead) * headStride + static_cast(block) * blockStride + + static_cast(query) * queryStride; + float score = scores[offset]; + if (block < initBlocks) + { + score = kInitScore; + } + // Match the PyTorch reference's second torch.where: local forcing + // overwrites init forcing if the two ranges overlap. + if (block >= localStart) + { + score = kLocalScore; + } + candidate = RedType{score, block}; + } + RedType::unpack(localScores[slot], localIndices[slot], candidate.compValIdx); + } + + selectFromCandidates( + warp, localScores, localIndices, output, query, kvHead, totalQueries, numKvHeads, numBlocks); +} + +template +__global__ void minimaxM3SelectBlocks64Kernel(float const* __restrict__ scores, int64_t headStride, int64_t blockStride, + int64_t queryStride, int32_t const* __restrict__ nValidBlocks, int32_t* __restrict__ output, int32_t numKvHeads, + int32_t numBlocks, int32_t totalQueries, int32_t initBlocks, int32_t localBlocks) +{ + int32_t const warpInBlock = threadIdx.x / kWarpSize; + int32_t const outputRow = blockIdx.x * kWarpsPerBlock + warpInBlock; + int32_t const lane = threadIdx.x % kWarpSize; + int32_t const numOutputRows = totalQueries * numKvHeads; + if (outputRow >= numOutputRows) + { + return; + } + + int32_t const query = outputRow / numKvHeads; + int32_t const kvHead = outputRow % numKvHeads; + int32_t const rawValidBlocks = nValidBlocks[query]; + int32_t const validBlocks = max(0, min(rawValidBlocks, numBlocks)); + int64_t const localStart + = max(static_cast(rawValidBlocks) - static_cast(localBlocks), static_cast(0)); + + using RedType = reduce_topk::TopKRedType; + RedType candidates[2]{{-INFINITY, RedType::kMaxIdx}, {-INFINITY, RedType::kMaxIdx}}; +#pragma unroll + for (int32_t slot = 0; slot < 2; ++slot) + { + int32_t const block = lane + slot * kWarpSize; + if (block < validBlocks) + { + int64_t const offset = static_cast(kvHead) * headStride + static_cast(block) * blockStride + + static_cast(query) * queryStride; + float score = scores[offset]; + if (block < initBlocks) + { + score = kInitScore; + } + // Match the PyTorch reference's second torch.where: local forcing + // overwrites init forcing if the two ranges overlap. + if (block >= localStart) + { + score = kLocalScore; + } + candidates[slot] = RedType{score, block}; + } + } + + uint32_t scoreKey0 = static_cast(candidates[0].compValIdx >> RedType::kMoveBits); + int32_t blockId0 + = RedType::kMaxIdx - static_cast(static_cast(candidates[0].compValIdx) & 0xFFFFU); + uint32_t scoreKey1 = static_cast(candidates[1].compValIdx >> RedType::kMoveBits); + int32_t blockId1 + = RedType::kMaxIdx - static_cast(static_cast(candidates[1].compValIdx) & 0xFFFFU); + warpBitonicSortDesc64(scoreKey0, blockId0, scoreKey1, blockId1, lane); + + if (lane < kTopK) + { + RedType const negativeInfinity{-INFINITY, 0}; + uint32_t const negativeInfinityScoreKey + = static_cast(negativeInfinity.compValIdx >> RedType::kMoveBits); + if (scoreKey0 == negativeInfinityScoreKey) + { + blockId0 = numBlocks; + } + + // MSA consumes block IDs in ascending order. numBlocks is the sentinel + // so padding naturally follows every valid ID. +#pragma unroll + for (int32_t size = 2; size <= kTopK; size *= 2) + { +#pragma unroll + for (int32_t stride = size / 2; stride > 0; stride /= 2) + { + int32_t const partnerBlockId = __shfl_xor_sync(0xFFFFU, blockId0, stride); + bool const takeMin = ((lane & size) == 0) == ((lane & stride) == 0); + blockId0 = takeMin ? min(blockId0, partnerBlockId) : max(blockId0, partnerBlockId); + } + } + + output[outputOffset(query, kvHead, totalQueries, numKvHeads, lane)] + = blockId0 == numBlocks ? -1 : blockId0; + } +} + +template +__global__ void minimaxM3SelectBlocks128Kernel(float const* __restrict__ scores, int64_t headStride, + int64_t blockStride, int64_t queryStride, int32_t const* __restrict__ nValidBlocks, int32_t* __restrict__ output, + int32_t numKvHeads, int32_t numBlocks, int32_t totalQueries, int32_t initBlocks, int32_t localBlocks) +{ + int32_t const warpInBlock = threadIdx.x / kWarpSize; + int32_t const outputRow = blockIdx.x * kWarpsPerBlock + warpInBlock; + int32_t const lane = threadIdx.x % kWarpSize; + int32_t const numOutputRows = totalQueries * numKvHeads; + if (outputRow >= numOutputRows) + { + return; + } + + int32_t const query = outputRow / numKvHeads; + int32_t const kvHead = outputRow % numKvHeads; + int32_t const rawValidBlocks = nValidBlocks[query]; + int32_t const validBlocks = max(0, min(rawValidBlocks, numBlocks)); + int64_t const localStart + = max(static_cast(rawValidBlocks) - static_cast(localBlocks), static_cast(0)); + + using RedType = reduce_topk::TopKRedType; + RedType candidates[4]; +#pragma unroll + for (int32_t slot = 0; slot < 4; ++slot) + { + candidates[slot] = RedType{-INFINITY, RedType::kMaxIdx}; + int32_t const block = lane + slot * kWarpSize; + if (block < validBlocks) + { + int64_t const offset = static_cast(kvHead) * headStride + static_cast(block) * blockStride + + static_cast(query) * queryStride; + float score = scores[offset]; + if (block < initBlocks) + { + score = kInitScore; + } + // Match the PyTorch reference's second torch.where: local forcing + // overwrites init forcing if the two ranges overlap. + if (block >= localStart) + { + score = kLocalScore; + } + candidates[slot] = RedType{score, block}; + } + } + + uint32_t scoreKeys[4]; + int32_t blockIds[4]; +#pragma unroll + for (int32_t slot = 0; slot < 4; ++slot) + { + scoreKeys[slot] = static_cast(candidates[slot].compValIdx >> RedType::kMoveBits); + blockIds[slot] + = RedType::kMaxIdx - static_cast(static_cast(candidates[slot].compValIdx) & 0xFFFFU); + } + warpBitonicSortDesc128(scoreKeys, blockIds, lane); + + if (lane < kTopK) + { + RedType const negativeInfinity{-INFINITY, 0}; + uint32_t const negativeInfinityScoreKey + = static_cast(negativeInfinity.compValIdx >> RedType::kMoveBits); + if (scoreKeys[0] == negativeInfinityScoreKey) + { + blockIds[0] = numBlocks; + } + + // MSA consumes block IDs in ascending order. numBlocks is the sentinel + // so padding naturally follows every valid ID. +#pragma unroll + for (int32_t size = 2; size <= kTopK; size *= 2) + { +#pragma unroll + for (int32_t stride = size / 2; stride > 0; stride /= 2) + { + int32_t const partnerBlockId = __shfl_xor_sync(0xFFFFU, blockIds[0], stride); + bool const takeMin = ((lane & size) == 0) == ((lane & stride) == 0); + blockIds[0] = takeMin ? min(blockIds[0], partnerBlockId) : max(blockIds[0], partnerBlockId); + } + } + + output[outputOffset(query, kvHead, totalQueries, numKvHeads, lane)] + = blockIds[0] == numBlocks ? -1 : blockIds[0]; + } +} + +template __global__ void minimaxM3SelectBlocksKernel(float const* __restrict__ scores, int64_t headStride, int64_t blockStride, int64_t queryStride, int32_t const* __restrict__ nValidBlocks, int32_t* __restrict__ output, int32_t numKvHeads, int32_t numBlocks, int32_t totalQueries, int32_t initBlocks, int32_t localBlocks) @@ -109,48 +503,44 @@ __global__ void minimaxM3SelectBlocksKernel(float const* __restrict__ scores, in RedType::unpack(localScores[rank], localIndices[rank], localTopK[rank].compValIdx); } - float selectedScores[kTopK]; - int32_t selectedIndices[kTopK]; - reduce_topk::reduceTopK(warp, selectedScores, selectedIndices, localScores, localIndices, -INFINITY); + selectFromCandidates( + warp, localScores, localIndices, output, query, kvHead, totalQueries, numKvHeads, numBlocks); +} - if (warp.thread_rank() == 0) +template +void launchMinimaxM3SelectBlocks(float const* scores, int64_t headStride, int64_t blockStride, int64_t queryStride, + int32_t const* nValidBlocks, int32_t* output, int32_t numKvHeads, int32_t numBlocks, int32_t totalQueries, + int32_t initBlocks, int32_t localBlocks, cudaStream_t stream) +{ + int32_t const numOutputRows = totalQueries * numKvHeads; + if (numOutputRows == 0) { -#pragma unroll - for (int32_t rank = 0; rank < kTopK; ++rank) - { - if (selectedScores[rank] == -INFINITY) - { - selectedIndices[rank] = -1; - } - } - - // MSA consumes block IDs in ascending order. Sort the sixteen selected - // IDs in registers, treating -1 padding as greater than every valid ID. -#pragma unroll - for (int32_t rank = 1; rank < kTopK; ++rank) - { - int32_t const candidate = selectedIndices[rank]; - int32_t const candidateKey = candidate < 0 ? numBlocks : candidate; - int32_t insertion = rank; - while (insertion > 0) - { - int32_t const previous = selectedIndices[insertion - 1]; - int32_t const previousKey = previous < 0 ? numBlocks : previous; - if (previousKey <= candidateKey) - { - break; - } - selectedIndices[insertion] = previous; - --insertion; - } - selectedIndices[insertion] = candidate; - } - -#pragma unroll - for (int32_t rank = 0; rank < kTopK; ++rank) - { - output[static_cast(outputRow) * kTopK + rank] = selectedIndices[rank]; - } + return; + } + int32_t const gridSize = (numOutputRows + kWarpsPerBlock - 1) / kWarpsPerBlock; + if (numBlocks <= kWarpSize) + { + minimaxM3SelectBlocksSmallKernel<1, HeadMajorOutput><<>>(scores, + headStride, blockStride, queryStride, nValidBlocks, output, numKvHeads, numBlocks, totalQueries, initBlocks, + localBlocks); + } + else if (numBlocks <= 2 * kWarpSize) + { + minimaxM3SelectBlocks64Kernel<<>>(scores, headStride, + blockStride, queryStride, nValidBlocks, output, numKvHeads, numBlocks, totalQueries, initBlocks, + localBlocks); + } + else if (numBlocks <= kSmallMaxBlocks) + { + minimaxM3SelectBlocks128Kernel<<>>(scores, headStride, + blockStride, queryStride, nValidBlocks, output, numKvHeads, numBlocks, totalQueries, initBlocks, + localBlocks); + } + else + { + minimaxM3SelectBlocksKernel<<>>(scores, headStride, + blockStride, queryStride, nValidBlocks, output, numKvHeads, numBlocks, totalQueries, initBlocks, + localBlocks); } } @@ -158,16 +548,18 @@ __global__ void minimaxM3SelectBlocksKernel(float const* __restrict__ scores, in void invokeMinimaxM3SelectBlocks(float const* scores, int64_t headStride, int64_t blockStride, int64_t queryStride, int32_t const* nValidBlocks, int32_t* output, int32_t numKvHeads, int32_t numBlocks, int32_t totalQueries, - int32_t initBlocks, int32_t localBlocks, cudaStream_t stream) + int32_t initBlocks, int32_t localBlocks, bool headMajorOutput, cudaStream_t stream) { - int32_t const numOutputRows = totalQueries * numKvHeads; - if (numOutputRows == 0) + if (headMajorOutput) { - return; + launchMinimaxM3SelectBlocks(scores, headStride, blockStride, queryStride, nValidBlocks, output, + numKvHeads, numBlocks, totalQueries, initBlocks, localBlocks, stream); + } + else + { + launchMinimaxM3SelectBlocks(scores, headStride, blockStride, queryStride, nValidBlocks, output, + numKvHeads, numBlocks, totalQueries, initBlocks, localBlocks, stream); } - int32_t const gridSize = (numOutputRows + kWarpsPerBlock - 1) / kWarpsPerBlock; - minimaxM3SelectBlocksKernel<<>>(scores, headStride, blockStride, queryStride, - nValidBlocks, output, numKvHeads, numBlocks, totalQueries, initBlocks, localBlocks); } } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.h b/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.h index 23aeaa0e1730..45d068c303ab 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.h @@ -29,7 +29,7 @@ namespace kernels void invokeMinimaxM3SelectBlocks(float const* scores, int64_t headStride, int64_t blockStride, int64_t queryStride, int32_t const* nValidBlocks, int32_t* output, int32_t numKvHeads, int32_t numBlocks, int32_t totalQueries, - int32_t initBlocks, int32_t localBlocks, cudaStream_t stream); + int32_t initBlocks, int32_t localBlocks, bool headMajorOutput, cudaStream_t stream); } // namespace kernels diff --git a/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp b/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp index 5d7d84f43ce6..1986e3abb09a 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. @@ -20,11 +20,38 @@ #include #include +#include + TRTLLM_NAMESPACE_BEGIN namespace torch_ext { +namespace +{ + +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"); +} + +} // namespace + // Function for fused QK Norm and RoPE // This operator applies RMS normalization and RoPE to Q and K tensors in a single CUDA kernel. // The OP performs operations in-place on the input qkv tensor. @@ -149,6 +176,160 @@ 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(eps >= 0.0, "MiniMax-M3 FP8 main Q/K/V producer requires eps >= 0"); + TORCH_CHECK(base > 0.0, "MiniMax-M3 FP8 main Q/K/V producer requires base > 0"); + + 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; + 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(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), + static_cast(kvCache.size(3)), static_cast(numTokens), static_cast(numHeadsQ), + static_cast(numHeadsK), static_cast(numHeadsV), static_cast(headDim), + static_cast(rotaryDim), static_cast(eps), qWeight.data_ptr(), kWeight.data_ptr(), + static_cast(base), 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(eps >= 0.0, "MiniMax-M3 horizontal producer requires eps >= 0"); + + 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(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; + 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.numel() == headDim && kWeight.numel() == headDim && indexQWeight.numel() == headDim + && indexKWeight.numel() == headDim, + "All norm weights must contain head_dim elements"); + 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), static_cast(kvCache.size(3)), static_cast(numTokens), + static_cast(numHeadsQ), static_cast(numHeadsKV), static_cast(numHeadsIndex), + static_cast(headDim), static_cast(rotaryDim), static_cast(eps), 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) { @@ -163,6 +344,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 @@ -170,12 +360,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/cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp index 6b4c3e3a9a85..be2d4aebd917 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp @@ -27,7 +27,7 @@ namespace torch_ext { torch::Tensor minimaxM3SelectBlocks(torch::Tensor const& scores, torch::Tensor const& nValidBlocks, int64_t topK, - int64_t initBlocks, int64_t localBlocks) + int64_t initBlocks, int64_t localBlocks, bool headMajorOutput) { constexpr int64_t kRequiredTopK = 16; constexpr int64_t kMaxBlockIndex = 65'535; @@ -61,13 +61,16 @@ torch::Tensor minimaxM3SelectBlocks(torch::Tensor const& scores, torch::Tensor c TORCH_CHECK(initBlocks <= std::numeric_limits::max() && localBlocks <= std::numeric_limits::max(), "minimax_m3_select_blocks forcing ranges exceed int32 range"); - auto output = torch::empty({scores.size(2), scores.size(0), topK}, scores.options().dtype(torch::kInt32)); + auto const outputOptions = scores.options().dtype(torch::kInt32); + auto output = headMajorOutput + ? torch::empty({scores.size(0), scores.size(2), topK}, outputOptions).permute({1, 0, 2}) + : torch::empty({scores.size(2), scores.size(0), topK}, outputOptions); auto const stream = at::cuda::getCurrentCUDAStream(scores.get_device()); tensorrt_llm::kernels::invokeMinimaxM3SelectBlocks(scores.data_ptr(), scores.stride(0), scores.stride(1), scores.stride(2), nValidBlocks.data_ptr(), output.data_ptr(), static_cast(scores.size(0)), static_cast(scores.size(1)), static_cast(scores.size(2)), static_cast(initBlocks), static_cast(localBlocks), - stream); + headMajorOutput, stream); return output; } @@ -79,7 +82,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def( "minimax_m3_select_blocks(Tensor scores, Tensor n_valid_blocks, int topk, int init_blocks, int " - "local_blocks) -> Tensor"); + "local_blocks, bool head_major_output=False) -> Tensor"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py index c508563bb720..1311a38f54e4 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py @@ -41,6 +41,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: @@ -60,6 +61,7 @@ class MiniMaxM3SparseMetadataParams(SparseMetadataParams): global_num_kv_heads: int = 0 num_index_heads: int = 4 topk: int = 16 + fuse_qkv_index_projection: bool = False def sharded_head_counts(self, mapping: Optional["Mapping"] = None) -> Tuple[int, int]: """Return per-rank (num_q_heads, num_kv_heads) for mapping. @@ -77,6 +79,13 @@ 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: + """Return the index-head count used by this rank's proxy attention.""" + if not self.fuse_qkv_index_projection: + return int(self.num_index_heads) + _, num_kv_heads = self.sharded_head_counts(mapping) + return num_kv_heads + @dataclass(frozen=True) class MiniMaxM3SparseConfig: @@ -147,7 +156,11 @@ def from_sparse_params( 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=( + int(num_kv_heads) + if sparse_params.fuse_qkv_index_projection + else int(sparse_params.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_backend/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py index e30a71b832ec..f3f8c73f8095 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py @@ -474,14 +474,15 @@ def _create_msa_buffers(self) -> None: params = self._msa_params if params is not None: fmha_sm100 = require_msa_module() + num_index_heads = params.sharded_index_head_count(self.mapping) max_k_tiles = _worst_case_proxy_max_k_tiles( fmha_sm100, - num_index_heads=params.num_index_heads, + num_index_heads=num_index_heads, kv_cache_manager=kv_cache_manager, max_batch=max_num_sequences, ) self._alloc_msa_proxy_scratch( - num_index_heads=params.num_index_heads, + num_index_heads=num_index_heads, max_tokens=self._msa_max_decode_tokens(), max_k_tiles=max_k_tiles, capture_graph=capture_graph, @@ -700,7 +701,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) num_q_heads, num_kv_heads = params.sharded_head_counts(self.mapping) topk = params.topk @@ -1082,6 +1083,9 @@ def run_indexer( config = self.m3_config idx_sm_scale = idx_sm_scale if idx_sm_scale is not None else config.sparse_index_dim**-0.5 num_tokens = int(idx_q.shape[0]) + head_major_output = ( + int(metadata.num_contexts or 0) > 0 and int(metadata.num_generations or 0) == 0 + ) # idx_q and idx_k may be strided column-views of a fused buffer, so # reshape to keep them zero-copy. The proxy fmha_sm100 and the index-K # scatter below both honor the source strides. @@ -1095,6 +1099,9 @@ def run_indexer( if idx_k is not None and not idx_k_prewritten: idx_k_view = idx_k.reshape(num_tokens, 1, config.sparse_index_dim) metadata.msa_write_idx_k(self.layer_idx, idx_k_view) + # Lightweight metadata implementations may install their cache on + # first write, so refresh the handle before the proxy reads it. + idx_k_cache = metadata.msa_idx_k_cache(self.layer_idx) elif idx_k is None and ( idx_k_cache.dtype != torch.float8_e4m3fn or idx_q_view.dtype != torch.float8_e4m3fn ): @@ -1131,12 +1138,18 @@ def run_indexer( max_score = None # The empty case was resolved on the host, so short-circuit here. if metadata.msa_eager_all_blocks_empty: - return torch.full( - (num_tokens, config.num_kv_heads, MSA_REQUIRED_TOPK), + output_shape = ( + (config.num_kv_heads, num_tokens, MSA_REQUIRED_TOPK) + if head_major_output + else (num_tokens, config.num_kv_heads, MSA_REQUIRED_TOPK) + ) + output = torch.full( + output_shape, -1, dtype=torch.int32, device=idx_q.device, ) + return output.permute(1, 0, 2) if head_major_output else output n_valid_blocks = metadata.msa_eager_n_valid_blocks if n_valid_blocks is not None: n_valid_blocks = n_valid_blocks[:num_tokens] @@ -1151,6 +1164,7 @@ def run_indexer( proxy_plan=proxy_plan, max_score=max_score, n_valid_blocks=n_valid_blocks, + head_major_output=head_major_output, ) def sparse_attn_predict( @@ -1174,6 +1188,48 @@ def sparse_kv_predict( ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: return None, None + def forward_prepopulated_kv( + self, + q: torch.Tensor, + metadata: MiniMaxM3MsaSparseAttentionMetadata, + forward_args: "AttentionForwardArgs", + ) -> None: + """Run MSA after the eager-prefill producer inserted main K/V. + + ``TrtllmAttention.forward`` interprets ``k=None`` as a fused QKV + buffer, so it cannot represent compact Q with prewritten paged K/V. + Dispatch the same MSA paged-GQA helper directly; the #16755 prewritten + marker is consumed there exactly as on its general scatter path. + """ + output = forward_args.output + if output is None: + raise RuntimeError( + f"{type(self).__name__}.forward_prepopulated_kv requires an output buffer." + ) + + kv_block_indexes = forward_args.topk_indices + if kv_block_indexes is not None: + plan = metadata.msa_decode_gqa_plan + if plan is None: + plan = metadata.msa_eager_gqa_plan + else: + plan = metadata.msa_decode_dense_plan + if plan is None: + plan = metadata.msa_eager_dense_plan + + from tensorrt_llm._torch.attention_backend.fmha.msa_sparse_gqa import run_msa_paged_gqa + + run_msa_paged_gqa( + self, + q, + None, + None, + metadata, + output, + kv_block_indexes=kv_block_indexes, + plan=plan, + ) + __all__ = [ "MiniMaxM3MsaSparseAttention", diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py index 61893c6ee96a..482ddda18d38 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py @@ -131,6 +131,7 @@ def select_blocks( proxy_plan: Optional[tuple] = None, max_score: Optional[torch.Tensor] = None, n_valid_blocks: Optional[torch.Tensor] = None, + head_major_output: bool = False, ) -> torch.Tensor: """Return [total_q, num_kv_heads, topk] selected block indices. @@ -180,18 +181,25 @@ def select_blocks( # Empty-selection guard. n_valid_blocks is a host tensor on # this path, so the .item() read does not sync the device. if n_valid_blocks.numel() == 0 or int(n_valid_blocks.max().item()) <= 0: - return torch.full( - (idx_q.shape[0], config.num_kv_heads, MSA_REQUIRED_TOPK), + output_shape = ( + (config.num_kv_heads, idx_q.shape[0], MSA_REQUIRED_TOPK) + if head_major_output + else (idx_q.shape[0], config.num_kv_heads, MSA_REQUIRED_TOPK) + ) + output = torch.full( + output_shape, -1, dtype=torch.int32, device=idx_q.device, ) + return output.permute(1, 0, 2) if head_major_output else output return select_blocks_from_maxscore( max_score_kv, topk=MSA_REQUIRED_TOPK, n_valid_blocks=n_valid_blocks, init_blocks=config.init_blocks, local_blocks=config.local_blocks, + head_major_output=head_major_output, ) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py index d06087b84376..f0616489de41 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py @@ -215,17 +215,25 @@ def select_blocks_from_maxscore( n_valid_blocks: torch.Tensor, init_blocks: int, local_blocks: int, + head_major_output: bool = False, ) -> torch.Tensor: """Select per-query top-k blocks from per-KV-head block scores. Applies init and local forced blocks and per-query valid-block masking on the amax-reduced scores [num_kv_heads, n_blocks, total_q]. Returns [total_q, num_kv_heads, topk] int32 ascending block ids with -1 tail - padding. + padding. When ``head_major_output`` is set, the logical result uses a + head-major backing so ``result.permute(1, 0, 2)`` is contiguous without a + copy. """ nvb = n_valid_blocks.to(device=max_score_kv.device, dtype=torch.int32).contiguous() return torch.ops.trtllm.minimax_m3_select_blocks( - max_score_kv, nvb, topk, init_blocks, local_blocks + max_score_kv, + nvb, + topk, + init_blocks, + local_blocks, + head_major_output, ) diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index ce3e410c2d8d..b31a859843bb 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -277,8 +277,18 @@ def _(logits, pass @torch.library.register_fake("trtllm::minimax_m3_select_blocks") - def _(scores, n_valid_blocks, topk, init_blocks, local_blocks): + def _( + scores, + n_valid_blocks, + topk, + init_blocks, + local_blocks, + head_major_output=False, + ): del n_valid_blocks, init_blocks, local_blocks + if head_major_output: + return scores.new_empty((scores.shape[0], scores.shape[2], topk), + dtype=torch.int32).permute(1, 0, 2) return scores.new_empty((scores.shape[2], scores.shape[0], topk), dtype=torch.int32) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 563471e0ab8e..b486c1f7eb53 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -81,6 +81,171 @@ # and flash SDPA does not accept attn_mask. _DENSE_SDPA_BACKENDS = [SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH] +# Experimental A/B gate for producing compact FP8 Q while inserting main K/V +# directly into the MSA paged cache during supported eager pure-prefill steps. +_FUSED_MAIN_KV_WRITE_ENV = "TRTLLM_MINIMAX_M3_FUSED_MAIN_KV_WRITE" + + +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_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_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 @@ -723,6 +888,7 @@ def __init__( and quant_config.quant_mode is not None and quant_config.quant_mode.has_fp8_kv_cache() ) + self.enable_fused_main_kv_write = os.environ.get(_FUSED_MAIN_KV_WRITE_ENV, "0") == "1" # Per-head Gemma RMSNorm — one set of weights shared across heads. self.q_norm = RMSNorm( @@ -749,9 +915,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)) @@ -759,25 +932,44 @@ 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 (no head sharding) and project - # the same hidden_states, so fuse them into one index_qk_proj GEMM - # with output [idx_q | idx_k]. idx_q holds all num_index_heads 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_blockscaling_mm=old_qkv_proj.use_cute_dsl_blockscaling_mm, + ) + self.sparse_num_index_heads = self.qkv_proj.local_num_index_heads + else: + # The compatibility path keeps index Q and K replicated and + # emits them from a separate GEMM. + 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=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, + ) + 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. @@ -1019,6 +1211,207 @@ def _fused_fp8_index_qk_norm_rope( position_ids.reshape(-1).contiguous().to(torch.int32), ).flatten(1) + def _fused_fp8_main_qk_norm_rope_kv_insert( + self, + qkv: torch.Tensor, + position_ids: Optional[torch.Tensor], + attn_metadata: AttentionMetadata, + ) -> Optional[torch.Tensor]: + """Produce compact FP8 Q and insert main K/V into their final cache. + + This path is intentionally narrower than the #16755 fused Triton + scatter. Unsupported decode, mixed, CUDA-graph, BF16-cache, or layout + cases return ``None`` and retain that general producer-plus-scatter + path unchanged. Sparse index-K remains owned by the existing FP8 index + producer and is not handled by this main-QKV kernel. + """ + if not self.enable_fused_main_kv_write or not self._emit_fp8_main_qkv(): + return None + if ( + is_torch_compiling() + or getattr(attn_metadata, "is_cuda_graph", False) + or int(getattr(attn_metadata, "num_generations", 0)) != 0 + or int(getattr(attn_metadata, "num_contexts", 0)) == 0 + or int(qkv.shape[0]) != int(attn_metadata.num_tokens) + ): + return None + if ( + qkv.dtype != torch.bfloat16 + or position_ids is None + or self.head_dim != 128 + 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 self.pos_embd_params.rope is None + or int(self.pos_embd_params.rope.dim) != 64 + or self.q_norm.weight.dtype != torch.bfloat16 + or self.k_norm.weight.dtype != torch.bfloat16 + ): + return None + # Sparse layers must leave index-K ownership with the #16742 direct + # FP8 producer. Do not create a second write path for BF16 index-K. + if self.is_sparse_attention_layer and self.attn.indexer_kv_dtype != "fp8": + 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") + if buffers is None: + return None + expected_head_stride = 128 * 128 + supported_hnd = ( + buffers.is_cuda + and buffers.dtype == torch.float8_e4m3fn + and buffers.dim() == 5 + and int(buffers.shape[0]) > 0 + 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) == expected_head_stride + and buffers.stride(1) >= self.num_key_value_heads * expected_head_stride + and buffers.stride(0) >= 2 * buffers.stride(1) + and buffers.stride(0) % 4 == 0 + and buffers.stride(1) % 4 == 0 + ) + out_cache_loc = getattr(attn_metadata, "msa_out_cache_loc", None) + num_tokens = int(qkv.shape[0]) + if ( + not supported_hnd + or out_cache_loc is None + or not out_cache_loc.is_cuda + or out_cache_loc.dtype != torch.int32 + or out_cache_loc.dim() != 1 + or not out_cache_loc.is_contiguous() + or out_cache_loc.numel() < num_tokens + ): + return None + + q = torch.ops.trtllm.minimax_m3_fp8_qk_norm_rope_kv_insert( + qkv.contiguous(), + buffers, + out_cache_loc[:num_tokens], + self.num_heads, + self.num_key_value_heads, + self.num_key_value_heads, + self.head_dim, + 64, + self.q_norm.variance_epsilon, + self.q_norm.weight, + self.k_norm.weight, + self.pos_embd_params.rope.theta, + True, + position_ids.reshape(-1).contiguous().to(torch.int32), + ).flatten(1) + # #16755 consumes this marker in run_msa_paged_gqa. It also lets the + # model core distinguish the direct producer from the general path and + # skip only this layer's fused Triton K/V scatter. + attn_metadata._msa_prewritten_layer = self.attn.layer_idx + return q + + 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 sparse producer for pure prefill.""" + 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 + if ( + is_torch_compiling() + or getattr(attn_metadata, "is_cuda_graph", False) + or int(getattr(attn_metadata, "num_generations", 0)) != 0 + or int(getattr(attn_metadata, "num_contexts", 0)) == 0 + or int(packed.shape[0]) != int(attn_metadata.num_tokens) + ): + return None + 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 self.pos_embd_params.rope is None + or int(self.pos_embd_params.rope.dim) != 64 + ): + 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 + ) + 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 + ) + 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, + 64, + self.q_norm.variance_epsilon, + 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), + ) + attn_metadata._msa_prewritten_layer = self.attn.layer_idx + return q.flatten(1), index_q.flatten(1) + def _expect_fused_qk_norm_rope(self, position_ids: Optional[torch.Tensor]) -> bool: """Whether the fused path must run rather than fall back. @@ -1179,30 +1572,34 @@ def _dense_forward( # Per-head Gemma RMSNorm and partial RoPE on Q/K. The bf16 fast # path fuses both into one kernel over the fused qkv; otherwise fall # back to separate norm and RoPE. - fused_qkv = self._fused_qk_norm_rope( - qkv, - position_ids, - num_heads_q=self.num_heads, - num_heads_k=self.num_key_value_heads, - num_heads_v=self.num_key_value_heads, - head_dim=self.head_dim, - q_norm=self.q_norm, - k_norm=self.k_norm, - out_fp8=self._emit_fp8_main_qkv(), - ) - if fused_qkv is not None: - q, k, v = self._split_main_qkv(fused_qkv) + fused_q = self._fused_fp8_main_qk_norm_rope_kv_insert(qkv, position_ids, attn_metadata) + if fused_q is not None: + q, k, v = fused_q, None, None else: - assert not self._expect_fused_qk_norm_rope(position_ids), ( - f"MiniMax-M3 dense attention (layer {self.layer_idx}) expected the " - f"fused QK-norm+RoPE kernel (bf16 activations, head_dim=" - f"{self.head_dim}) but fell back to the separate path; qkv dtype " - f"is {qkv.dtype} (expected {self.attn_activation_dtype})." + fused_qkv = self._fused_qk_norm_rope( + qkv, + position_ids, + num_heads_q=self.num_heads, + num_heads_k=self.num_key_value_heads, + num_heads_v=self.num_key_value_heads, + head_dim=self.head_dim, + q_norm=self.q_norm, + k_norm=self.k_norm, + out_fp8=self._emit_fp8_main_qkv(), ) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q, k = self.apply_qk_norm(q, k) - if self.rotary_emb is not None and position_ids is not None: - q, k = self.rotary_emb(position_ids, [q, k]) + if fused_qkv is not None: + q, k, v = self._split_main_qkv(fused_qkv) + else: + assert not self._expect_fused_qk_norm_rope(position_ids), ( + f"MiniMax-M3 dense attention (layer {self.layer_idx}) expected the " + f"fused QK-norm+RoPE kernel (bf16 activations, head_dim=" + f"{self.head_dim}) but fell back to the separate path; qkv dtype " + f"is {qkv.dtype} (expected {self.attn_activation_dtype})." + ) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.apply_qk_norm(q, k) + if self.rotary_emb is not None and position_ids is not None: + q, k = self.rotary_emb(position_ids, [q, k]) # Keep token-wise projections and the output projection visible to # torch.compile. Only the metadata/cache-dependent attention core is @@ -1416,8 +1813,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, @@ -1445,8 +1842,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, @@ -1465,6 +1862,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) @@ -1474,8 +1872,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, @@ -1487,29 +1885,36 @@ def _msa_attention_core( inherited FMHA forward; this layer selects the top-k blocks (sparse only) and builds the ``forward_args`` the FMHA reads. """ + prewritten_main_kv = k is None or v is None + assert (k is None) == (v is None) + if prewritten_main_kv: + assert attn_metadata._msa_prewritten_layer == self.attn.layer_idx + if self.is_sparse_attention_layer: assert idx_q is not None - # One launch writes this layer's K/V and, on the bf16 path, index-K, - # ahead of the proxy pass that reads the index-K cache. On the FP8 - # indexer path idx_k is None because the fused producer already - # inserted FP8 index-K into the side cache; the fused write then - # stores K/V only. - attn_metadata.msa_write_layer_caches(self.attn.layer_idx, k, v, idx_k) + if not prewritten_main_kv: + # General #16755 path: one Triton launch writes this layer's + # K/V and, on the BF16 path, index-K. + attn_metadata.msa_write_layer_caches(self.attn.layer_idx, k, v, idx_k) # Publish the selected blocks so the FMHA runs the sparse path. - # idx_k_prewritten marks that index-K is already in the cache (via - # the fused write above on bf16, or the FP8 producer when idx_k is - # None), so run_indexer must not write it again. kv_block_indexes = self.attn.run_indexer( - idx_q, idx_k, attn_metadata, idx_k_prewritten=True + idx_q, + idx_k, + attn_metadata, + idx_k_prewritten=(not prewritten_main_kv or idx_k is None), ) forward_args = AttentionForwardArgs(output=output, topk_indices=kv_block_indexes) else: assert idx_q is None and idx_k is None - # Dense layers get the same fused K/V write. - attn_metadata.msa_write_layer_caches(self.attn.layer_idx, k, v) + if not prewritten_main_kv: + # Dense layers retain the same general #16755 K/V scatter. + attn_metadata.msa_write_layer_caches(self.attn.layer_idx, k, v) # No top-k selection means the FMHA attends the full page table. forward_args = AttentionForwardArgs(output=output) - self.attn.forward(q, k, v, attn_metadata, forward_args=forward_args) + if prewritten_main_kv: + self.attn.forward_prepopulated_kv(q, attn_metadata, forward_args) + else: + self.attn.forward(q, k, v, attn_metadata, forward_args=forward_args) return output def _sparse_forward( @@ -1567,8 +1972,27 @@ def _sparse_forward( # only hidden_states and write disjoint outputs, so each runs its # projection, norm, and RoPE concurrently on the aux stream when # multi-stream is enabled, joining before the attention core. + packed_qkv = None + packed_idx_qk = None + if self.enable_fused_qkv_index_projection: + 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 + ) + 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_q = self._fused_fp8_main_qk_norm_rope_kv_insert(qkv, position_ids, attn_metadata) + if fused_q is not None: + return fused_q, None, None fused_qkv = self._fused_qk_norm_rope( qkv, position_ids, @@ -1595,7 +2019,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. @@ -2123,6 +2549,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) + } + 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 @@ -2177,8 +2629,10 @@ def load_weights( params_map: Optional[Dict[str, str]] = None, allow_partial_loading: bool = False, ) -> None: - # Fuse index_q/index_k into each index_qk_proj module (also covers - # the VL path, which routes text weights through here). + # The opt-in vLLM-style path packs all five projections with + # MiniMax-specific TP sharding. The compatibility path fuses only the + # replicated index Q/K pair. + 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 @@ -2188,6 +2642,11 @@ def load_weights( weights = _fold_gemma_boundary_norm_weights(weights) if weight_mapper is None: weight_mapper = MiniMaxM3HfWeightMapper() + # These five-way modules were loaded above with model-specific + # Q/K/V/index-Q/index-K sharding. The generic HF mapper only knows + # three-way Q/K/V fusion; after the source keys are marked consumed it + # would still invoke that callback with three empty dictionaries. + 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( diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index e60e923edd51..35d44a74027d 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -736,6 +736,16 @@ 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. This prototype " + "currently targets disaggregated prefill workers; leave it disabled on " + "decode or mixed workers. The MiniMax-M3-specific path requires the MSA " + "implementation.", + status="prototype", + ) num_attention_heads: Optional[int] = Field( default=None, description= @@ -770,6 +780,10 @@ def _validate_msa_block_size(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.") return self def supports_backend(self, backend: str) -> bool: @@ -793,6 +807,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): @@ -822,6 +837,7 @@ def _value(name: str, default=None): global_num_kv_heads=num_kv_heads, num_index_heads=self.sparse_num_index_heads, topk=self.sparse_topk_blocks, + fuse_qkv_index_projection=self.fuse_qkv_index_projection, ) diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py index 90639786c8cf..bfd9f2aa3d6b 100644 --- a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py @@ -8,9 +8,12 @@ reference is covered by the SM100 integration accuracy test. """ +from types import SimpleNamespace + import pytest import torch +from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs from tensorrt_llm._torch.attention_backend.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.common import write_kv_slots from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_scatter import ( @@ -54,6 +57,81 @@ def test_msa_fp8_indexer_config_is_explicit_and_lowered(): ) +def test_fused_qkv_index_projection_is_explicit_and_shards_index_heads(): + cfg = MiniMaxM3SparseAttentionConfig( + implementation="msa", + fuse_qkv_index_projection=True, + ) + sparse_params = cfg.to_sparse_params() + metadata_params = cfg.to_sparse_metadata_params( + pretrained_config=SimpleNamespace(num_attention_heads=64, num_key_value_heads=4) + ) + mapping = SimpleNamespace(tp_size=2, enable_attention_dp=False) + + assert sparse_params.fuse_qkv_index_projection is True + assert metadata_params.sharded_head_counts(mapping) == (32, 2) + assert metadata_params.sharded_index_head_count(mapping) == 2 + + compatibility = MiniMaxM3SparseAttentionConfig(implementation="msa").to_sparse_metadata_params( + pretrained_config=SimpleNamespace(num_attention_heads=64, num_key_value_heads=4) + ) + assert compatibility.sharded_index_head_count(mapping) == 4 + + with pytest.raises(ValueError, match=r"requires the 'msa' implementation"): + MiniMaxM3SparseAttentionConfig( + implementation="triton", + fuse_qkv_index_projection=True, + ) + + +def test_prepopulated_kv_dispatches_compact_q_and_consumes_marker(monkeypatch): + from tensorrt_llm._torch.attention_backend.fmha import msa_sparse_gqa + + attention = MiniMaxM3MsaSparseAttention.__new__(MiniMaxM3MsaSparseAttention) + q = torch.empty(3, 8 * 128) + output = torch.empty_like(q) + topk = torch.zeros(3, 1, 16, dtype=torch.int32) + eager_plan = object() + + class FakeMetadata: + msa_decode_gqa_plan = None + msa_eager_gqa_plan = eager_plan + msa_decode_dense_plan = None + msa_eager_dense_plan = object() + _msa_prewritten_layer = 7 + + captured = {} + + def fake_run(attn, q_arg, k, v, metadata, output_arg, **kwargs): + captured.update( + attn=attn, + q=q_arg, + k=k, + v=v, + metadata=metadata, + output=output_arg, + kwargs=kwargs, + ) + metadata._msa_prewritten_layer = None + + monkeypatch.setattr(msa_sparse_gqa, "run_msa_paged_gqa", fake_run) + metadata = FakeMetadata() + attention.forward_prepopulated_kv( + q, + metadata, + AttentionForwardArgs(output=output, topk_indices=topk), + ) + + assert captured["attn"] is attention + assert captured["q"] is q + assert captured["k"] is None and captured["v"] is None + assert captured["metadata"] is metadata + assert captured["output"] is output + assert captured["kwargs"]["kv_block_indexes"] is topk + assert captured["kwargs"]["plan"] is eager_plan + assert metadata._msa_prewritten_layer is None + + def test_msa_metadata_rejects_undersized_max_score_buffer(): metadata_cls = MiniMaxM3MsaSparseAttention.Metadata metadata = metadata_cls.__new__(metadata_cls) @@ -220,6 +298,8 @@ class FakeMetadata: msa_qo_lens_cpu = torch.ones(2, dtype=torch.int32) msa_kv_lens_cpu = torch.full((2,), 128, dtype=torch.int32) msa_qo_offset_cpu = torch.full((2,), 127, dtype=torch.int32) + num_contexts = 0 + num_generations = 2 def __init__(self): self.cache = torch.empty(2, 1, 128, 128, dtype=torch.float8_e4m3fn, device="cuda") @@ -252,6 +332,67 @@ def msa_idx_k_cache(self, layer_idx): assert "write" not in captured +@pytest.mark.parametrize( + ("num_contexts", "num_generations", "expected_head_major"), + [(2, 0, True), (1, 1, False), (0, 2, False)], +) +def test_run_indexer_routes_head_major_output_by_batch_mode( + num_contexts, num_generations, expected_head_major +): + num_tokens, num_index_heads, sparse_index_dim = 3, 4, 128 + captured = {} + + class FakeIndexer: + def select_blocks(self, *args, **kwargs): + del args + captured["head_major_output"] = kwargs["head_major_output"] + return torch.zeros(num_tokens, 1, 16, dtype=torch.int32) + + class FakeMetadata: + msa_decode_proxy_plan = None + msa_eager_proxy_plan = ("eager",) + msa_eager_all_blocks_empty = False + msa_eager_n_valid_blocks = torch.ones(num_tokens, dtype=torch.int32) + msa_kv_indices = torch.arange(num_tokens, dtype=torch.int32) + msa_qo_lens_cpu = torch.tensor([num_tokens], dtype=torch.int32) + msa_kv_lens_cpu = torch.tensor([num_tokens], dtype=torch.int32) + msa_qo_offset_cpu = torch.tensor([0], dtype=torch.int32) + + def __init__(self): + self.num_contexts = num_contexts + self.num_generations = num_generations + self.idx_k_cache = None + + def msa_write_idx_k(self, layer_idx, idx_k): + del layer_idx + self.idx_k_cache = idx_k + + def msa_idx_k_cache(self, layer_idx): + del layer_idx + return self.idx_k_cache + + attention = SimpleNamespace( + layer_idx=0, + m3_config=SimpleNamespace( + sparse_index_dim=sparse_index_dim, + num_index_heads=num_index_heads, + num_kv_heads=1, + ), + indexer=FakeIndexer(), + ) + metadata = FakeMetadata() + + result = MiniMaxM3MsaSparseAttention.run_indexer( + attention, + torch.zeros(num_tokens, num_index_heads * sparse_index_dim), + torch.zeros(num_tokens, sparse_index_dim), + metadata, + ) + + assert result.shape == (num_tokens, 1, 16) + assert captured["head_major_output"] is expected_head_major + + def test_msa_proxy_max_score_strided_index_k_matches_packed(): if not torch.cuda.is_available(): pytest.skip("CUDA required") diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py index d8617108fe8a..44e6a9b20aff 100644 --- a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py @@ -57,7 +57,9 @@ def _reference_select_blocks( @pytest.mark.parametrize("num_kv_heads", [1, 4]) -@pytest.mark.parametrize("num_blocks", [1, 8, 16, 17, 127, 1024, 1537]) +@pytest.mark.parametrize( + "num_blocks", [1, 8, 16, 17, 32, 33, 64, 65, 96, 127, 128, 129, 1024, 1537] +) def test_fused_selector_matches_reference_random(num_kv_heads, num_blocks): total_q = 19 generator = torch.Generator(device="cuda").manual_seed(num_blocks) @@ -95,7 +97,82 @@ def test_fused_selector_matches_reference_random(num_kv_heads, num_blocks): assert actual.dtype == torch.int32 assert actual.shape == (total_q, num_kv_heads, 16) + assert actual.stride() == (num_kv_heads * 16, 16, 1) + assert actual.is_contiguous() + assert torch.equal(actual, expected) + + +@pytest.mark.parametrize("num_blocks", [65, 96, 128]) +def test_fused_selector_128_path_matches_reference_ties_and_forcing(num_blocks): + scores = torch.zeros((2, num_blocks, 5), device="cuda", dtype=torch.float32) + n_valid_blocks = torch.tensor( + [0, 8, 16, num_blocks - 1, num_blocks], device="cuda", dtype=torch.int32 + ) + + expected = _reference_select_blocks( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=8, + local_blocks=12, + ) + actual = select_blocks_from_maxscore( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=8, + local_blocks=12, + ) + + assert torch.equal(actual, expected) + + +@pytest.mark.parametrize("num_blocks", [16, 48, 96, 129]) +def test_fused_selector_head_major_output_is_zero_copy_q2k(num_blocks): + total_q, num_kv_heads = 7, 4 + generator = torch.Generator(device="cuda").manual_seed(num_blocks) + scores = torch.randn( + num_kv_heads, + num_blocks, + total_q, + generator=generator, + device="cuda", + dtype=torch.float32, + ) + n_valid_blocks = torch.randint( + 0, + num_blocks + 1, + (total_q,), + generator=generator, + device="cuda", + dtype=torch.int32, + ) + + expected = _reference_select_blocks( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=8, + local_blocks=12, + ) + actual = select_blocks_from_maxscore( + scores, + topk=16, + n_valid_blocks=n_valid_blocks, + init_blocks=8, + local_blocks=12, + head_major_output=True, + ) + q2k = actual.permute(1, 0, 2).contiguous().to(torch.int32) + assert torch.equal(actual, expected) + assert actual.shape == (total_q, num_kv_heads, 16) + assert actual.stride() == (16, total_q * 16, 1) + assert not actual.is_contiguous() + assert q2k.shape == (num_kv_heads, total_q, 16) + assert q2k.is_contiguous() + assert q2k.data_ptr() == actual.data_ptr() + assert q2k.untyped_storage().data_ptr() == actual.untyped_storage().data_ptr() @pytest.mark.parametrize( diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index ea97c53985c1..f5c0048992fb 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -37,7 +37,9 @@ ) from tensorrt_llm._torch.models.modeling_minimaxm3 import ( MiniMaxM3Attention, + MiniMaxM3QKVIndexerLinear, _build_swiglu_oai_dense_mlp, + _load_qkv_index_proj_weights, _minimax_m3_swiglu_oai, _strip_language_model_prefix, _wrap_dict_as_config, @@ -407,6 +409,28 @@ def test_minimax_m3_attention_dense_construction_matches_config(): assert not hasattr(attn, name), f"dense layer should not declare {name!r}" +@pytest.mark.gpu +@pytest.mark.skipif(not _has_cuda(), reason="MiniMax-M3 attention construction needs CUDA") +def test_minimax_m3_fused_main_kv_write_is_opt_in(monkeypatch): + _, model_cfg = _make_attention_test_config() + + monkeypatch.delenv("TRTLLM_MINIMAX_M3_FUSED_MAIN_KV_WRITE", raising=False) + baseline = MiniMaxM3Attention( + model_config=model_cfg, + layer_idx=0, + is_sparse_attention_layer=False, + ) + assert baseline.enable_fused_main_kv_write is False + + monkeypatch.setenv("TRTLLM_MINIMAX_M3_FUSED_MAIN_KV_WRITE", "1") + candidate = MiniMaxM3Attention( + model_config=model_cfg, + layer_idx=0, + is_sparse_attention_layer=False, + ) + assert candidate.enable_fused_main_kv_write is True + + @pytest.mark.gpu @pytest.mark.skipif(not _has_cuda(), reason="MiniMax-M3 attention construction needs CUDA") def test_minimax_m3_attention_partial_rope_dim_is_rotary_dim(): @@ -536,6 +560,56 @@ def test_minimax_m3_attention_sparse_construction_matches_config(): raise AssertionError("sparse forward must raise RuntimeError when attn_metadata is None") +def test_minimax_m3_five_way_projection_shard_geometry(): + module = SimpleNamespace( + tp_size=2, + tp_rank=1, + total_num_kv_heads=4, + total_num_index_heads=4, + ) + shard_geometry = MiniMaxM3QKVIndexerLinear._shard_geometry + assert shard_geometry(module, "q") == (2, 1) + assert shard_geometry(module, "k") == (2, 1) + assert shard_geometry(module, "v") == (2, 1) + assert shard_geometry(module, "index_q") == (2, 1) + assert shard_geometry(module, "index_k") == (1, 0) + + # At TP8, each of four KV/index-Q heads is replicated on two ranks. + module.tp_size = 8 + module.tp_rank = 5 + assert shard_geometry(module, "k") == (4, 2) + assert shard_geometry(module, "index_q") == (4, 2) + assert shard_geometry(module, "index_k") == (1, 0) + + +def test_minimax_m3_five_way_loader_returns_exact_generic_skip(): + 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(): 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..c85d2008bee4 --- /dev/null +++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_horizontal_producer.py @@ -0,0 +1,147 @@ +# 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] + + +@pytest.mark.parametrize("num_tokens", [1, 16, 129]) +def test_minimax_m3_horizontal_producer_matches_separate_producers(num_tokens): + torch.manual_seed(1234) + num_heads_q = 8 + num_kv_heads = 2 + num_index_heads = num_kv_heads + 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 + ) + 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, + ) + + pages = slots.long() // 128 + within = slots.long() % 128 + assert torch.equal(q.view(torch.uint8), q_reference.view(torch.uint8)) + # 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 compare within one FP8 + # ULP rather than requiring byte identity across the different rounding + # orders. Main Q/K/V retain byte-exact parity above and below. + torch.testing.assert_close( + index_q.float(), + index_q_reference.float(), + rtol=0.13, + atol=0.05, + ) + assert torch.equal( + main_cache[pages, :, :, within, :].view(torch.uint8), + reference_main_cache[pages, :, :, 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, + ) 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..65282bcd859a --- /dev/null +++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_main_kv_insert.py @@ -0,0 +1,166 @@ +# 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 + + assert torch.equal(q_out.view(torch.uint8), q_ref.contiguous().view(torch.uint8)) + assert torch.equal( + kv_cache[:, 0][pages, :, within, :].view(torch.uint8), + k_ref.contiguous().view(torch.uint8), + ) + 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() + + assert torch.equal(q_out.view(torch.uint8), q_ref.contiguous().view(torch.uint8)) + assert torch.equal( + kv_cache[page, 0, :, 0, :].view(torch.uint8), + k_ref[0].contiguous().view(torch.uint8), + ) + assert torch.equal( + kv_cache[page, 1, :, 0, :].view(torch.uint8), + v_ref[0].contiguous().view(torch.uint8), + )